Data Loading at Scale — infra.bea.ai

At 1000+ GPU scale, data loading becomes a distributed systems problem: thousands of workers reading from shared storage simultaneously, needing global shuffling without communication, and maintaining determinism for reproducibility.

1. The Scale Challenge

Scale	GPUs	Workers (8/GPU)	Data throughput needed	Storage pressure
Small (1 node)	8	64	~50 MB/s	Local NVMe handles it easily
Medium (8 nodes)	64	512	~400 MB/s	Parallel FS or cloud streaming
Large (64 nodes)	512	4,096	~3 GB/s	Shared storage under load
Mega (512 nodes)	4,096	32,768	~25 GB/s	Distributed design critical

How these numbers are derived: At ~3000 tokens/sec per GPU (typical LLM training throughput), each GPU consumes ~6 KB/s of pre-tokenized data (3000 tokens × 2 bytes). Even 4,096 GPUs only need ~25 MB/s of raw token data. The actual I/O demand is higher due to metadata, shuffle reads, and reading ahead — but still far below raw storage bandwidth. The bottleneck at scale is metadata operations and concurrent file opens, not raw read bandwidth.

At mega-scale, thousands of workers opening files simultaneously creates metadata storms — file system metadata servers overwhelmed, directory listings stall, training hangs.

2. Shard-Based Data Distribution

Split the dataset into thousands of shards, assign disjoint shard sets to each worker:

Dataset: 100,000 shards (256 MB each = 25 TB total)

Node 0 (8 GPUs, 64 workers):  shards 0-99 (worker 0 reads shards 0-1, etc.)
Node 1 (8 GPUs, 64 workers):  shards 100-199
...
Node 99 (8 GPUs, 64 workers): shards 9900-9999

→ No two workers read the same shard simultaneously
→ No I/O contention on any single file

Deterministic Shard Assignment

# Each rank gets a deterministic subset of shards
def get_shards_for_rank(total_shards, rank, world_size, seed=42):
    rng = np.random.RandomState(seed)
    all_shards = rng.permutation(total_shards)
    return all_shards[rank::world_size]  # interleaved assignment

Deterministic: Same seed + same world_size → same shard assignment. Enables reproducible training.

3. Global Shuffle Without Communication

True global shuffle across all shards would require all-to-all communication — impractical at scale. Instead, use hierarchical shuffling:

Level 1: Shuffle shard order per worker (each epoch, different order)
Level 2: Shuffle samples within each shard (in-memory)
Level 3: Shuffle within the local batch buffer

Result: approximate global shuffle with zero cross-node communication

Quality: This isn’t perfectly uniform, but empirically, shard-level + within-shard shuffle produces similar training quality to true global shuffle.

4. Avoiding I/O Hotspots

Problem	Cause	Fix
Metadata storm	32K workers listing files simultaneously	Pre-compute file lists, distribute as config
Sequential reads of same file	Multiple workers assigned same shard	Disjoint shard assignment
Bursty reads	All workers start new shards at same time	Stagger shard boundaries (random offset)
Cache thrashing	Local cache smaller than working set	Increase cache or use streaming with eviction
Network congestion	Too many concurrent S3 requests	Rate-limit requests, use caching layer (Alluxio)

5. Checkpointing Data State

When training restarts from a checkpoint, the data loader must resume from the exact position:

# Save data state with training checkpoint
data_state = {
    'epoch': current_epoch,
    'shard_index': current_shard,
    'sample_index': current_sample_in_shard,
    'rng_state': dataloader.rng_state,
}
torch.save({**model_state, **data_state}, checkpoint_path)

# Resume: skip to saved position
dataloader.load_state_dict(data_state)

Without data state checkpointing: After restart, the model may re-train on data it already saw, or skip data — breaking the intended curriculum and wasting compute.

For the storage systems that serve this data, see Parallel File Systems and Cloud Object Storage.

Self-Check

Calculate total read bandwidth needed for a given GPU count and throughput
Design shard-based data distribution for 512 GPUs
Explain hierarchical shuffling and why it approximates global shuffle
Identify and fix 3 common I/O hotspot causes
Implement data state checkpointing for resumable training