How Uber Made GPU Training Faster and Reproducible

Peter Bubenik · Uber Data / ML · · Source
How Uber Made GPU Training Faster and Reproducible

After studying this material, you should be able to:

  1. Identify data pipeline bottlenecks that cause GPU underutilization in deep learning training
  2. Explain how push-down transformations and local disk caching improve throughput
  3. Describe hidden sources of non-determinism in multi-worker data loaders
  4. Apply systematic profiling methodology to diagnose training pipeline problems
  5. Understand the engineering tradeoffs involved in each optimization decision

Step-by-Step Teaching

Step 1: Understanding the Problem Space

Why Data Loading Matters

Before writing a single line of optimization code, you need to understand the producer-consumer relationship in GPU training.

[Data Storage] → [CPU Data Loader] → [GPU Training]
   Producer           Processor          Consumer

The GPU is your most expensive resource. It should be computing as close to 100% of the time as possible.

When the GPU sits idle waiting for data, you are:

  • Wasting money on expensive hardware
  • Slowing down your experiment iteration cycle
  • Scaling costs linearly without scaling performance

The Specific Scenario

Uber's team had:

  • Dataset size: tens of terabytes, tens of billions of rows
  • GPU utilization: 10–15% (extremely poor)
  • Training time: 22 hours per run

The question was: why are GPUs idle 85–90% of the time?


Step 2: Systematic Profiling — The Detective Work

Key Principle: Never assume you know the bottleneck. Prove it experimentally.

The team followed a structured elimination process. This is the most important skill demonstrated in the article.

Hypothesis 1: Is the model too small for the GPU?

Test: Run with a much larger synthetic model (deeper, wider layers)

Result: GPU utilization jumped to ~80%

Conclusion: The GPU can run at high utilization. The problem is that the real model finishes each batch so quickly that it has to wait for the next one. The data pipeline is too slow.

Small model:  [===compute===][........waiting........][===compute===]
Large model:  [==========compute==========][wait][==========compute==========]

Hypothesis 2: Is the network (HDFS) too slow?

Test: Force a small subset of data into RAM (in-memory cache)

Result: GPU utilization spiked to 60%+

Conclusion: When data is instantly available, GPUs stay busy. Network I/O is a bottleneck.

Hypothesis 3: Will a local disk cache solve it?

Test: Cache raw Parquet files on local disk instead of reading from HDFS

Result: Utilization remained low. Nothing changed.

Why did this fail? This is the critical insight:

Cache TypeWhat is storedCPU work still needed?
RAM cachePre-transformed NumPy arrays❌ No
Disk cache (raw files)Raw Parquet/Arrow data✅ Yes, every batch

The RAM cache had accidentally hidden a second bottleneck: the CPU transformation from PyArrow format → NumPy format. When data came from RAM, it was already transformed. When it came from disk, the CPU had to convert it on every single batch, in the main thread, blocking the GPU.

The Two-Bottleneck Model

Bottleneck 1: Network I/O     (reading from HDFS every epoch)
Bottleneck 2: CPU Transform   (PyArrow → NumPy in the main thread)

You must solve both to see improvement.


Step 3: The Solutions

Solution A — Push-Down Transformations

Original architecture:

Worker Thread → [reads raw Arrow data] → Queue → Main Thread → [transforms to NumPy] → GPU
                                                                      ↑
                                                              BOTTLENECK HERE

Optimized architecture:

Worker Thread → [reads raw Arrow data] → [transforms to NumPy] → Queue → Main Thread → GPU
                        ↑
               Transformation moved here

Why this works:

  • The worker pool has multiple threads running in parallel
  • Transformation is now parallelized across all workers
  • The main thread does nothing except pass ready-to-use arrays to the GPU

Tradeoff to understand:

PyArrow is more memory-efficient than NumPy. Storing transformed data in the queue increases RAM usage.

This is an acceptable tradeoff because throughput gains outweigh the memory cost.


Solution B — Local Disk Cache with FanoutCache

The insight: You read the same row groups every epoch. Why pay the network cost repeatedly?

Optimized worker flow:

Worker receives task: "read row group #4521"
        ↓
Check local FanoutCache
        ↓
   [Cache HIT]              [Cache MISS]
       ↓                         ↓
Read pre-transformed       Fetch from HDFS
data from local disk            ↓
       ↓                  Apply transformation
   Return array                 ↓
                          Write to cache
                                ↓
                           Return array

Crucially: Because push-down transformation is already in place, the cache stores pre-transformed NumPy arrays. On a cache hit, zero CPU transformation is needed.

Handling Datasets Larger Than Local Disk

Standard cache eviction policies (like LRU — Least Recently Used) work poorly here.

Why LRU fails for training:

  • In web services, some data is accessed rarely → evict it
  • In training, every data point is accessed exactly once per epoch → LRU has no useful signal

Solution: Quota management

Cache row groups until disk quota is reached
        ↓
For remaining row groups: fall back to HDFS

This doesn't eliminate the bottleneck for the largest datasets, but significantly reduces it.


Solution C — Thread Pool Instead of Process Pool

Process Pool:

Worker Process A → [serialize/pickle data] → [send over IPC] → Main Process
                         ↑
                   CPU overhead here

Thread Pool:

Worker Thread A → [shared memory] → Main Thread
                       ↑
                  No serialization needed

New problem introduced: Python threads blocked on I/O (slow HDFS reads) are hard to kill cleanly.

Fix: Implement sentinel logic — a special signal that tells threads to stop gracefully, rather than relying on the interpreter to force-kill them. Also enforce strict HDFS timeouts.


Step 4: Preserving Training Randomness While Caching

This is a subtle but important engineering concern.

The risk: If you cache pre-batched data, you might feed the model the exact same sequence every epoch. Deep learning requires stochastic (random) data ordering to generalize well.

The solution: Cache row groups, not batches.

Row Group (smaller unit) → cached
        ↓
Multiple worker threads fetch row groups in parallel
        ↓
Arrival order in queue naturally varies (provides randomness)
        ↓
Final shuffle step applied to each batch in memory
        ↓
GPU receives randomly ordered data

This preserves the stochastic nature of training while still getting cache benefits.


Step 5: Fixing Reproducibility — Eliminating Hidden Randomness

The Problem

Even with identical seeds and configurations, two training runs produced different results. This is dangerous because:

  • You cannot reliably compare model architectures
  • You cannot tell if an improvement is real or just lucky variance
  • Production jobs are inconsistent

Where Was the Randomness Hiding?

The team ran controlled experiments varying: number of workers, shuffle settings, and seeds.

Finding: Most randomness came from race conditions, not from shuffle operations.

Understanding the Race Condition

Original architecture used shared queues:

Ventilator Queue (shared) → Worker 1
                          → Worker 2
                          → Worker 3
                          
Results Queue (shared) ← Worker 1 (finishes first sometimes)
                       ← Worker 2 (finishes first other times)
                       ← Worker 3

The problem: Which worker picks up which row group depends on:

  • Thread scheduling by the OS
  • Network I/O timing
  • CPU execution speed

None of these are controlled by your random seed. So even with seed=42, the data order changes run to run.

As you add more workers to improve throughput, race conditions become more pronounced. Faster training = less reproducible training (without a fix).

The Fix — Dedicated Queues + Round-Robin

New architecture:

Ventilator Queue 1 → Worker 1 → Results Queue 1 ↘
Ventilator Queue 2 → Worker 2 → Results Queue 2 → Merge (round-robin) → GPU
Ventilator Queue 3 → Worker 3 → Results Queue 3 ↗

Two changes:

  1. Each worker gets its own dedicated ventilator and results queues
  2. Row groups are assigned and collected in fixed round-robin order

Result: Data order is now fully deterministic regardless of which worker finishes first.

Important: This does not hurt throughput. Workers still run in parallel. The only change is how results are merged.

Also Fixed: Legacy Random API

Replaced np.random.RandomState (legacy, inconsistent behavior when reseeded across distributed components) with np.random.default_rng (modern, consistent behavior).


Step 6: Results and Mental Model

Quantitative Results

MetricBeforeAfter
Training time22 hours3 hours
GPU utilization10–15%60%+
Compute costbaseline~80% reduction
Run-to-run MAP variance~0.5% shift~0.13% shift

The Complete Mental Model

PROBLEM 1: GPU Starvation
├── Root Cause A: Network I/O (HDFS reads every epoch)
│   └── Fix: Local FanoutCache (cache pre-transformed data)
├── Root Cause B: CPU bottleneck (PyArrow→NumPy in main thread)
│   └── Fix: Push-down transformation to worker pool
└── Root Cause C: Process pool serialization overhead
    └── Fix: Thread pool with sentinel-based teardown

PROBLEM 2: Non-Reproducibility
├── Root Cause A: Race conditions in shared queues
│   └── Fix: Dedicated per-worker queues + round-robin merge
└── Root Cause B: Legacy np.random.RandomState
    └── Fix: np.random.default_rng

Step 7: Generalizable Principles

These lessons apply beyond Uber's specific stack:

PrincipleApplication
Profile before optimizingNever assume the bottleneck. Eliminate hypotheses experimentally.
One fix can reveal another bottleneckSolving I/O exposed the CPU transform problem. Expect this.
Cache transformed data, not raw dataRaw data caching only removes one of two bottlenecks.
Shared state causes non-determinismRace conditions in queues are invisible until you look for them.
Throughput and reproducibility are not oppositesDedicated queues with round-robin preserve both.
Understand your eviction policyLRU is wrong for training workloads. Match policy to access pattern.

Quick Self-Check Questions

  1. Why did caching raw Parquet files on local disk fail to improve GPU utilization?
  2. What is the tradeoff of push-down transformation, and why is it acceptable?
  3. Why does adding more workers make reproducibility worse in the original architecture?
  4. Why is LRU eviction a poor choice for a training data cache?
  5. What two changes together eliminate race conditions in the data loader?

Answers: (1) The CPU transformation bottleneck was still present in the main thread. (2) Higher RAM usage, acceptable because throughput gains outweigh it. (3) More workers = more race conditions in shared queues = more non-deterministic ordering. (4) Every data point is accessed exactly once per epoch, so LRU has no useful signal to evict on. (5) Dedicated per-worker queues + deterministic round-robin assignment and merging.

More to study