Step-by-Step Study Material

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 deep learning training:

[Data Pipeline] → produces batches → [GPU] → consumes batches → trains model

The GPU is your most expensive resource. It should never wait. When it does, you have a producer bottleneck — your data pipeline cannot produce batches fast enough.

The Case Study Numbers

Key Insight: The GPU was idle ~85–90% of the time. This means the hardware was paid for but not used. Optimization here is more impactful than improving the model architecture itself.

---

Step 2: The Original Data Flow

Understanding what existed before helps you understand why it was slow.

HDFS (remote storage)
       ↓  [Network I/O — slow]
Worker Pool reads raw Parquet row groups
       ↓  [CPU transformation — slow, in main thread]
PyArrow → NumPy conversion
       ↓
Batch Dataloader batches the data
       ↓
Async Dataloader feeds GPU
       ↓
GPU trains

Two Bottlenecks Hidden Here

Critical Point: Both bottlenecks existed simultaneously. Fixing only one would not solve the problem — as Uber discovered when disk caching alone failed.

---

Step 3: The Profiling Methodology (How to Find Bottlenecks)

Uber did not guess. They used systematic elimination. This is the most transferable skill in this article.

The Three-Experiment Diagnostic

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

• Action: Run a synthetic benchmark with a much larger model
• Result: GPU utilization jumped to ~80%
• Conclusion: GPUs are capable. The real model's batches arrive too slowly

Larger model → more compute per batch → GPU stays busy longer → masks data delay

Experiment 2: Is the network the bottleneck?

• Action: Force a small dataset subset into RAM (in-memory cache)
• Result: Utilization spiked to 60%+
• Conclusion: When data is instantly available, GPUs stay busy → Network I/O is a bottleneck

Experiment 3: Is disk caching enough?

• Action: Cache raw files on local disk (skip HDFS)
• Result: Utilization remained low
• Conclusion: Disk caching alone failed because CPU transformation was a second bottleneck

The Lesson from Experiment 3

RAM Cache experiment:
  Read once → Transform once → Store as NumPy → Reuse instantly ✓

Disk Cache (naive) experiment:
  Read from disk (fast) → Transform EVERY TIME in main thread (slow) ✗

The RAM cache had accidentally hidden the CPU bottleneck. The disk cache exposed it.

Rule: When you fix one bottleneck, the next one becomes visible. Always profile again after each fix.

---

Step 4: The Two-Pillar Solution

Pillar 1: Push-Down Transformations

Before:

Worker Thread → reads raw row group → puts raw data in queue
Main Thread   → pulls raw data → converts PyArrow→NumPy → sends to GPU

After:

Worker Thread → reads raw row group → converts PyArrow→NumPy → puts NumPy array in queue
Main Thread   → pulls ready NumPy array → sends directly to GPU

Why this works:

• The CPU-heavy transformation is now parallelized across the worker pool
• The main thread does nothing but pass data — it is never blocked
• Multiple workers transform simultaneously instead of one main thread doing it sequentially

The tradeoff:

• NumPy arrays use more RAM than PyArrow tables
• The queue now holds larger objects
• This was deemed an acceptable trade-off for the throughput gain

Pillar 2: Local Disk Caching with FanoutCache

The insight: The same row groups are read from HDFS every single epoch. You pay the network cost repeatedly for identical data.

The solution: Cache the pre-transformed data on local disk after the first read.

Worker Thread receives task: "read row group #42"
       ↓
Check local FanoutCache
       ↓
   [Cache HIT]              [Cache MISS]
       ↓                         ↓
Load NumPy array          Fetch from HDFS
from local disk           Apply transformation
(very fast)               Write to cache
       ↓                         ↓
       └──────────┬──────────────┘
                  ↓
         Return ready NumPy array

Why "pre-transformed" matters:

• On a cache hit, zero CPU transformation is needed
• The worker simply loads a ready-to-use array
• This eliminates both bottlenecks simultaneously on subsequent epochs

---

Step 5: Engineering Challenges and Solutions

Challenge 1: Dataset Larger Than Local Disk

Problem: Datasets are tens of terabytes. Local disk on a worker node cannot hold everything.

Naive solution (wrong): Use LRU eviction — discard least recently used data when disk is full.

Why LRU fails here:

Training epoch: every row group must be read exactly once
LRU evicts row group #5 to make room for row group #100
Next epoch: row group #5 must be read again → cache miss → back to HDFS

LRU creates a cycle of eviction and re-fetching for training workloads.

Actual solution: Quota management

• Cache row groups until the disk quota is reached
• Fall back to HDFS for the remainder
• No eviction during an epoch
• Partial speedup is better than no speedup

Challenge 2: Preserving Training Randomness

Problem: If you cache pre-batched data, you replay the same sequence every epoch. Models need random data order to generalize.

Solution: Cache row groups, not batches.

Row group = smaller unit of data (not a full batch)

Multiple worker threads fetch row groups in parallel
→ Arrival order naturally varies (thread scheduling, I/O timing)
→ Final shuffle step applied to each batch in memory before GPU

This gives you:

• High throughput (cached data, fast reads)
• Stochastic training (varied order, in-memory shuffle)

Challenge 3: Thread Pool Stability

Why switch from Process Pool to Thread Pool?

New problem introduced: Python threads blocked on I/O (slow HDFS reads) cannot be easily killed. This caused "zombie threads" that prevented clean job shutdown.

Solutions:

1. Sentinel logic in the ventilator — explicitly signal threads to stop rather than relying on the interpreter
2. Stricter HDFS timeouts — ensure stuck network calls eventually fail rather than hanging forever

---

Step 6: Reproducibility — Eliminating Hidden Randomness

Why Reproducibility Matters

If two identical training runs produce different results, you cannot answer:

• Is this new model actually better, or did we get lucky?
• Did a code change improve performance or hurt it?

Where Randomness Was Hiding

Petastorm already controlled some randomness:

• ✅ Row group shuffling — seed-controlled
• ✅ Data partitioning across GPU workers — seed-controlled

But runs were still non-deterministic. Why?

The answer: Race conditions in shared queues

Original Architecture:
                    ┌─ Worker 1 ─┐
Ventilator Queue ──┤─ Worker 2 ─├──→ Shared Results Queue → GPU
(shared)           └─ Worker 3 ─┘         (shared)

• Which worker picks up which row group depends on thread scheduling
• Which result arrives first depends on I/O timing and execution speed
• These are not controlled by the random seed
• More workers = more race conditions = more variance

Secondary issue: Legacy np.random.RandomState API behaves inconsistently when reseeded across distributed components.

The Fix

Change 1: Replace the random API

# Before (legacy, inconsistent)
rng = np.random.RandomState(seed)

# After (modern, consistent)
rng = np.random.default_rng(seed)

Change 2: Dedicated queues with round-robin scheduling

New Architecture:
                    ┌─ Worker 1 ─→ Queue 1 ─┐
Ventilator ────────┤─ Worker 2 ─→ Queue 2 ─├──→ Round-Robin Merge → GPU
(round-robin)      └─ Worker 3 ─→ Queue 3 ─┘

• Each worker has its own dedicated queue — no shared queue race conditions
• Row groups are assigned in fixed round-robin order — deterministic assignment
• Results are merged in the same round-robin order — deterministic collection

Result: Data order is now fully determined by the seed, not by which thread finishes first.

Important: Dedicated queues do not reduce throughput. Workers still run in parallel. Only the collection order is made deterministic.

---

Step 7: Connecting Everything — The Full Picture

PROBLEM                    ROOT CAUSE               SOLUTION
─────────────────────────────────────────────────────────────────
Low GPU utilization   →  Network I/O latency    →  Local disk cache
(10-15%)              +  CPU transformation     +  Push-down transformation
                         in main thread            to worker pool

Non-deterministic     →  Race conditions in     →  Dedicated per-worker queues
training results         shared queues          +  Round-robin scheduling
                      +  Legacy random API      +  np.random.default_rng

---

Step 8: Key Principles to Remember

Principle 1: Profile Before Optimizing

Never assume where the bottleneck is. Use controlled experiments to isolate variables one at a time.

Principle 2: Bottlenecks Are Sequential

Fixing one bottleneck reveals the next. Uber fixed network I/O and immediately hit the CPU transformation wall.

Principle 3: Cache the Right Thing

Caching raw data and caching transformed data have very different performance implications. Cache the output of your most expensive operation.

Principle 4: Parallelism Introduces Non-Determinism

More workers = more throughput = more race conditions. Design for determinism explicitly, not as an afterthought.

Principle 5: Trade-offs Are Real

• Thread pool vs. process pool: lower overhead but harder to kill
• NumPy in queue vs. PyArrow: faster access but more RAM
• Quota management vs. LRU: simpler for training but no eviction

Every optimization has a cost. The goal is to choose trade-offs that align with your constraints.

---

Quick Reference Summary