Swap the GPU count, keep the batches: Zephon makes TorchTitan data loading elastically deterministic
Synopsis
Changing the GPU count during training usually changes the order in which a data loader feeds samples to the model, so an ablation can end up measuring hardware allocation instead of a data change. Zephon splits each loading operator into a single-threaded accumulator that makes stateful decisions and stateless parallel workers, and divides data into a fixed number of logical lanes; plugged into TorchTitan in place of the default loader, it resumed a 16-GPU checkpoint on 8, 32 and 64 GPUs with every global batch matching the original run's SHA-256 hash.
Interpretation
In TorchTitan, a run resumed from a 16-GPU checkpoint on 8, 32 or 64 GPUs hands the model global batches bit-identical to the original run. A standard loader redistributes samples after the resume according to the new GPU count, so the same run sees a different sample sequence and each allocation ends on a distinct loss curve. A 1B-parameter model trained on 16 GPUs was checkpointed at step 150 and resumed on 8, 32 and 64 GPUs; with Zephon the loss curves lie on top of each other and the SHA-256 hash of every global batch matches the baseline.
The effect of data order on downstream evaluation can rival a data-recipe change, and Zephon shrinks the GPU-count spread to near the evaluation noise floor. Training the same 1B model from scratch for 20B tokens with a standard loader, the DCLM Core v1 score varies by up to 0.82 points across 8 to 64 GPUs and the FineWeb suite by up to 0.64 points, roughly four times the thresholds the DCLM and FineWeb authors used when evaluating data curation methods. Across the same 8, 16, 32 and 64 GPU counts, the Zephon spread was only 0.05 and 0.011 points respectively.
Zephon tokenizes, packs, shuffles and mixes data inside the training job, so no preprocessed copy of the corpus has to be created on disk. Previously each change to mixture, filters or tokenization required a full offline pass over the training corpus stored to disk, a copy a typical run reads exactly once and most experiments never improve on. These preprocessing steps run on CPUs left idle during training, with parallelism set per operator; on a full vision-language pipeline that decodes images, tokenizes, packs and batches, Zephon keeps pace with Google Grain, the fastest loader tested.
Perspective
The demonstration covers a 1B-parameter model trained and resumed on 8, 16, 32 and 64 GPUs, plus a throughput comparison on one vision-language pipeline; integration replaces TorchTitan's default loader with a Zephon pipeline while the rest of the training run stays the same. The repository includes a CPU-only demo and a GPU smoke test with a real but tiny model that check every global step against an uninterrupted run after resuming on a different number of data-parallel replicas.
Zephon guarantees the data, not the arithmetic: global batches are identical across GPU counts and losses will be very close, but may not be identical because floating-point reduction order in the training loop still depends on how many GPUs share the work. The lane count must be chosen up front, extra lanes per GPU add checkpoint restore overhead, and resuming requires the data, tokenizer, seed, mixture, context length and global batch size to remain unchanged.