Data Modelling & Storage advanced 8 min read 7 flashcards

Storage Layout for ML Training Reads

Why the access pattern of a training loop is the one pattern the lakehouse was not designed for, and how sharded sequential formats trade shuffle quality for two orders of magnitude fewer requests.

Analytical queries read a few columns from many rows and skip most files. A training loop reads every row, all columns, in a different random order every epoch, forever. That pattern defeats every optimisation a lakehouse table is built around: predicate pushdown has nothing to prune, column projection saves nothing, and the randomness destroys locality.

The consequence is measurable. Across the configurations studied by Mohan and colleagues, DNN training spent 10% to 70% of epoch time blocked on I/O despite prefetching and pipelining, and on a spinning-disk configuration ResNet-50 on ImageNet was stalled on I/O for 75% of its epoch time (Mohan et al., 2021, Analyzing and Mitigating Data Stalls in DNN Training, VLDB 14(5), arXiv:2007.06775). Google's fleet analysis found 20% of training jobs spent more than a third of their compute time ingesting data (Murray et al., 2021, tf.data: A Machine Learning Data Processing Framework, VLDB 14(12), arXiv:2101.12127).

One object per sample is the failure mode

The obvious layout stores each sample as its own object: one JPEG, one WAV, one JSON document per key. It is easy to inspect, easy to append to, and it makes the data loader issue one GET per sample.

At training rates that is fatal. Feeding eight H100s at 4,000 samples per second means 4,000 GET requests per second per node. Against a single S3 prefix capped at 5,500 GET/s, one node nearly saturates the prefix and four nodes cannot run at all. The bytes are trivial, a few hundred megabytes per second, well inside a 100 Gb/s NIC. The request rate is the wall.

Little's law makes the concurrency requirement concrete. Sustaining 4,000 samples per second against 20 ms first-byte latency needs \(4000 \times 0.020 = 80\) requests in flight at all times. Eight synchronous dataloader workers give you eight, so the pipeline delivers 400 samples per second and the accelerators idle at 10% utilisation on a storage system that is not remotely busy.

Sharding converts random reads into sequential ones

The standard fix packs samples into shards: contiguous archives of roughly 100 MB to 1 GB holding a thousand or more samples each, read start to finish. WebDataset uses plain POSIX tar, which is deliberate, because tar is sequentially readable, needs no index, and every tool already understands it (Aizman, Maltby & Breuel, 2019, High Performance I/O For Large Scale Deep Learning, IEEE Big Data, arXiv:2001.01858). Mosaic Streaming's MDS, TFRecord and Ray Data's formats make the same structural choice with different framing.

The arithmetic changes completely. One thousand samples per shard turns 4,000 GET/s into 4 shard opens per second. Request-rate pressure drops by three orders of magnitude, per-request cost drops with it, and each open streams at whatever the connection sustains rather than paying a round trip per sample.

Shuffling is what you pay with

A sequential shard read returns samples in the order they were written. Real shuffling of a billion-sample dataset is not available at any price; the standard approximation is two-level:

  1. Shuffle the shard list each epoch, and assign shards to workers.
  2. Shuffle within a buffer of the last \(k\) samples read, emitting a random element and refilling from the stream.

The resulting sample order is correlated: two samples from the same shard are far more likely to land in the same batch than two samples chosen uniformly. Whether that matters depends on how the shards were written. If shards were built from data already grouped by class, source or date, the buffer cannot rescue it and training sees a batch composition that is systematically non-random. Shuffle once at write time, thoroughly, then rely on cheap shuffling at read time. A shuffle buffer of 10,000 samples over shards of 1,000 samples mixes roughly ten shards at a time, which is adequate only if the shards themselves are unordered.

Resumption and elasticity

A dataloader that cannot resume mid-epoch turns every preemption into lost work. Doing it properly means sample order is a deterministic function of (seed, epoch, global step) rather than of worker count, so a run stopped on 24 GPUs can resume on 16 and finish on 48 with the same loss curve. Mosaic Streaming implements this through a fixed number of canonical sample buckets, num_canonical_nodes, which holds the global order stable as the physical topology changes (Mosaic Streaming, elastic determinism). Changing it changes the order, so pin it in the run config rather than tuning it.

When it breaks

Shards fix the request rate and create a rewrite problem. Samples are now immutable inside archives. Deleting one record for a takedown request, or fixing a mislabelled batch, means rewriting whole shards or maintaining a deny-list that the loader consults. Neither is as clean as deleting an object.

Shard size interacts with node count. With 900 shards and 128 workers, seven workers get an extra shard and the epoch's straggler is a full shard behind. Shard count should be a comfortable multiple of the maximum worker count, which means choosing it against a cluster size you do not yet know.

The last epoch is not like the others. Once the working set fits in page cache or a local NVMe cache, throughput jumps and the bottleneck moves to preprocessing. Benchmarks that measure epoch one and benchmarks that measure epoch five report different systems.

Preprocessing, not storage, is often the real bottleneck. JPEG decode and augmentation are CPU-bound, and a node with eight accelerators and too few host cores starves regardless of layout. Moving decode onto the accelerator (DALI, nvJPEG) shifts the balance, and GPUDirect Storage removes the CPU bounce buffer from the path entirely, with NVIDIA reporting roughly 3x latency reductions as typical for small transfers (NVIDIA GPUDirect Storage overview). Diagnose which stage is starving before changing the format.

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track