Vol. XV · Running a Training Cluster Plate 5 of 9data loader, shuffling, determinism, resumption, checkpoint, epochs, repeated data, reproducibilityDE

Feeding the Machine.

After a crash the model goes back to its checkpoint. Does the data go back with it?

The bandwidth of feeding a model is modest — Llama 3 405B’s 16 million tokens a step every few seconds is a few megabytes a second of token ids — but the bookkeeping is not. The loader shuffles documents across thousands of shards, packs them into sequences (plate IX.9), and hands each data-parallel rank its own slice. All of it must be deterministic: the order should be a function of a seed and the step number, so a run can be reproduced, debugged and replayed (plate XV.4).

Resumption is where it goes wrong. When a job restarts from a checkpoint, the model rewinds to step c, and the loader must rewind to exactly the same place. The workshop runs two epochs over 240 documents, fails at a chosen step, restarts from the last checkpoint, and counts how often each document reached the final model under three loaders: one whose position is saved in the checkpoint, one that restarts its epoch, and one that reshuffles.

Below: left, how many documents the final model saw 0, 1, 2, 3 or 4 times — every one should be 2. Right, the count for each document.

What you will learn
Workshop
240 documents in batches of 8, so 30 steps an epoch and 60 in all; each epoch’s order is a permutation seeded by the run seed and the epoch. Checkpoints every 10 steps; the job fails at the chosen step and the model restarts from the last checkpoint. Loaders: restore (position saved), epoch (restart the current epoch), reshuffle (a new random order from the same position).

Rewind the data with the model

A failure at step 47 sends the model back to step 40. A loader that restarts its epoch replays the epoch’s first ten batches and then runs out of steps before the epoch’s last ten: 80 document views repeated and 80 missing, and only 67% of steps match the uninterrupted run. A loader that reshuffles on restart does no better: 56 extra and 56 missing. A loader whose position is in the checkpoint reproduces the clean run exactly: 0 extra, 0 missing, 100% identical.

At this toy scale the damage is a few dozen views; on a real run it is millions of documents seen twice or never, with no error anywhere — and with a failure every few hours (plate XV.1), every restart would add more. The bandwidth, by contrast, is easy: Llama 3 405B at 400 TFLOPs per GPU takes about 6.2 seconds a step, so its 16 million tokens arrive at 2.7 million tokens a second — 10.8 MB a second as 4-byte ids. The loader’s job is not speed but exactness.

Determinism also needs the per-rank split to be a function of the step: if the data-parallel degree changes on restart — after losing a node, say — the mapping from global batch to ranks must change without changing the global batch itself.

Terms

Data loader
The component that reads, shuffles, packs and delivers training batches.
Deterministic order
A data order that is a function of the seed and the step alone.
Loader state
The position and random state a checkpoint must save to resume exactly.
Epoch
One pass over the training data.
A checkpoint must save where the data loader is; otherwise every restart silently repeats some data and skips other data.

Sources

  1. Llama 3 paper, arXiv:2407.21783, sections 3.3 and 3.4 (16M-token batches; 400 TFLOPs per GPU at 16 384 GPUs).
  2. PyTorch documentation, “torch.utils.data” and stateful data loaders.