Zephon brings elastically deterministic data loading to PyTorch's TorchTitan

The problem: GPU count changes the data order

The PyTorch team argues that data is one of the strongest levers for model quality, and that much daily work involves changing mixtures, filters, tokenization or packing and then running an ablation. Each change often requires a full pass over the corpus and expensive offline steps to store a prepared copy on disk. Because a typical run reads that copy once, and most experiments do not yield a better model, the team says a lot of time and storage goes to datasets that are never read again.

The post adds that this matters more as agents design, launch and interpret experiments. Agents run many experiments in parallel, but they are not good at spotting when a result depends on something other than the intended change, such as the number of GPUs. With most data loaders, the GPU count changes the data order and therefore the result.

To show this, the team trained a 1B-parameter model in TorchTitan on 16 GPUs, took a checkpoint at step 150, and resumed on 8, 32 and 64 GPUs. With a standard loader, each resumed run followed a distinct loss curve. With Zephon, the curves lie on top of each other, and the SHA-256 hash of every global batch matches the baseline run, according to the post.

How much it matters for benchmarks

The team also trained the same 1B-parameter model from scratch for 20B tokens on 8, 16, 32 and 64 GPUs using a standard loader. The DCLM Core v1 score varied by up to 0.82 points depending on GPU count, and the score on the FineWeb evaluation suite by up to 0.64 points. The authors say this is roughly four times the relevant thresholds that the DCLM and FineWeb authors used when evaluating data curation methods. An ablation comparing two mixtures on two different GPU allocations could therefore be measuring the GPU count instead of the data.

With Zephon, the spread across GPU counts was only 0.05 and 0.011 points, respectively.

Because the data preparation happens on the fly, using CPUs that would otherwise sit idle during training, there is no preprocessed copy of the corpus to create or discard.

Datasets, a WorkSource and a Pipeline

A Zephon setup has three parts: the Datasets to train on, a WorkSource that decides how they are mixed and ordered into a curriculum, and a Pipeline of operators that turn stored samples into tensors for the training loop. The example in the post mixes two text datasets from Hugging Face, then tokenizes, packs, shuffles and batches them.

The WorkSource plans the curriculum from each Dataset's metadata and is separate from the processing operators, so the mixture can change without touching the pipeline, and the reverse. A TokenEstimation utility measures the token to byte ratio on a small sample of each dataset so that mixing proportions target tokens rather than raw samples, even though samples are stored as raw strings.

Each operator can have its own parallelism. In the example, tokenization gets four workers because it is the most expensive step. On a full vision-language pipeline that decodes images, tokenizes, packs and batches, Zephon keeps pace with Google Grain, which the authors call the fastest loader they tested, while offering stronger elastic determinism guarantees.

How the batches stay identical

On-the-fly packing and mixing make elastic determinism hard, because the data order can no longer be looked up by index on disk. It depends on pipeline state, which most loaders cannot transfer to a different number of GPUs without changing the output.

The post says Mosaic Streaming and Google Grain can keep global batches the same across GPU counts only when sample order can be worked out ahead of time from an index, which rules out on-the-fly packing. Megatron-Energon can reproduce batches on a different number of GPUs only if the total number of data loader workers stays the same.

Zephon relies on design decisions that build on each other:

• Accumulators and workers: each operator has a single-threaded accumulator that sees every record in order and makes stateful decisions, such as which sequences are packed together, plus a stateless function that does the heavy work across parallel workers. Results are numbered and put back in order, so the degree of parallelism is a pure performance knob.
• Lanes instead of ranks: data is divided into a fixed number of logical streams called lanes, each with its own ordered samples and operator state. When the number of data-parallel ranks changes, lanes are redistributed and the global batches stay the same.
• Flushing and replaying: a checkpoint does not contain the contents of packing bins or shuffle buffers.

The authors note one catch: the number of lanes must be chosen up front.

What it means for practitioners

For teams running data ablations, the pitch is that results can be compared across jobs that landed on different GPU allocations, and that a run resumed after a hardware fault on a different set of GPUs sees the same batches. The post presents this as a way to avoid both storing throwaway datasets and misattributing score differences to data changes.