PyTorch 为 TorchTitan 引入 Zephon 弹性确定性数据加载

Why Zephon?

The data that a model sees during training is one of the most powerful levers we have for influencing how well it performs. Much of the day-to-day work of improving a model happens on the data side: we change the mixture, or the filters that we apply, or how we tokenize and pack the samples, and then we run an ablation experiment to see whether the change helped. The first problem with this is that each of these changes often requires us to do a full pass over our training corpus and run expensive offline operations to prepare the data and store it on disk before our ablation experiment can run. Since a typical training run reads through that copy exactly once, and because most experiments don’t end up producing a better model, we spend a lot of time and storage creating datasets that will never be read again.

This matters even more now that experiments can be designed, launched, and interpreted by agents instead of people. Agents are great for running lots of experiments in parallel, but they aren’t especially good at noticing when a result actually depends on something other than the change they made, like the number of GPUs that an experiment happened to run on (in fairness to the agents, humans aren’t great at noticing this either). And with most data loaders, the number of GPUs does change the result, because it changes the data order.

To illustrate this, we train a 1B-parameter model in TorchTitan on 16 GPUs, take a checkpoint at step 150, and then resume from that checkpoint on 8, 32, and 64 GPUs alongside the original run, once with a standard data loader and once with Zephon (think of it as the run hitting a hardware fault and coming back on whatever set of GPUs happened to be free). As seen in Figure 1, with the standard data loader, changing the GPU count means that the data loader hands the model a different sequence of samples after the resume, and so each one ends up with 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.

[Figure 1: Loss curves after resuming on 8, 32, and 64 GPUs from a 16-GPU checkpoint at step 150, standard loader vs. Zephon, with a barcode comparing the per-step SHA-256 of each global batch against the original run.] 

To get a sense for how much this matters, we train the same 1B-parameter model from scratch for 20B tokens on 8, 16, 32, and 64 GPUs using a standard data loader. We show the scores per GPU count in Figure 2. The model’s DCLM Core v1 score varies by up to 0.82 points depending on how many GPUs it was trained on, and its score based on the FineWeb evaluation suite by up to 0.64 points. That is roughly four times the relevant thresholds that the DCLM and FineWeb authors used when evaluating data curation methods, which means that an ablation comparing two mixtures on two different allocations could easily be measuring the GPU count instead of the data. With Zephon, the spread across the different GPU counts was only 0.05 and 0.011 points, respectively.

[Figure 2: Downstream evaluations at 8–64 GPUs on DCLM Core v1 and FineWeb with non-elastically deterministic (standard) data loading and Zephon. Grey bands mark the eval noise floor.]

Note that Zephon does all of this while doing the data prep work on the fly: it tokenizes, packs, shuffles, and mixes the data inside of the training job, taking advantage of all of those CPUs sitting around with nothing to do while the model trains, which means there’s no preprocessed copy of the corpus to create (or to throw away when the experiment doesn’t pan out).

A Zephon Pipeline

Zephon asks you to describe three things: the Datasets you want to train on, a WorkSource that decides how those datasets are mixed and ordered into a training curriculum, and a Pipeline of operators that turn the samples in their stored representations into the tensors your training loop consumes. Here’s an example that mixes two text datasets directly from Hugging Face, then tokenizes, packs, shuffles, and batches the result:

from zephon import Pipeline
from zephon.io import Dataset
from zephon.work import StaticMixtureWorkSource, TokenEstimation

source = StaticMixtureWorkSource(
    datasets=[
        Dataset.from_path("wiki", "hf://Salesforce/wikitext/wikitext-2-raw-v1/train"),
        Dataset.from_path("news", "hf://fancyzhx/ag_news/default/train"),
    ],
    mixture={"wiki": 0.5, "news": 0.5},   # a token ratio, not a sample ratio
    token_estimation=TokenEstimation(),
    seed=42,
)
pipeline = (
    Pipeline(source)
    .options(io_options={"cache": {"enabled": True}})
    .tokenize(tokenizer_id="gpt2", field="text", parallelism=4)
    .pack_flat(max_length=2048, algorithm="wrap")
    .shuffle(buffer_size=1024, seed=42)
    .batch(microbatch_size=8)
)
for batch in pipeline:
    train_step(batch.to_training())

The WorkSource plans the curriculum from each Dataset’s metadata and is kept separate from the operators that actually process those samples so you can change the mixture without touching the processing pipeline and vice-versa. In order to ensure we get the mixture percentages correct on a token basis, we include a TokenEstimation utility that measures the token/byte ratio for a small sample of each dataset and adjusts the plan to target the mixing proportions on a token-level even though the samples are stored as raw strings.

Each of the individual Pipeline operators can have its own level of parallelism; here, tokenization gets four workers because it’s the most expensive step. Parallelizing the preprocessing operators is how Zephon keeps the GPUs fed, even on more complex pipelines: on a full vision-language pipeline that decodes images, tokenizes, packs, and batches its samples, Zephon keeps pace with Google Grain, the fastest loader we tested, while offering the stronger elastic determinism guarantees.

[Figure 3: Usable (non-padding) tokens per second per GPU on a full vision-language pipeline as parallelism increases, for Zephon, Grain, and a quickly hacked-together internal loader we used internally before building Zephon. The horizontal line is the throughput needed to keep the GPUs from stalling.]

How Zephon Keeps the Batches Identical

Packing and mixing data on the fly is what makes elastic determinism hard. When a data loader tokenizes documents, packs them together into fixed-length sequences, and mixes several datasets together as it goes, the order of the training data is no longer something that you can look up at an index on disk (we call this stateful n:m pipelines in our paper). In this case, it depends on the state of the pipeline when and how a sample is used, and in most data loaders that state can’t be transferred to a different number of GPUs without changing the output. Existing loaders get part of the way there: Mosaic Streaming and Google Grain can keep the global batches the same across GPU counts, but only when the order of samples can be worked out ahead of time from an index, which rules out things like packing on the fly. Megatron-Energon can reproduce its batches on a different number of GPUs, but only as long as the total number of data loader workers across all of the GPUs stays the same. 

Zephon allows for elastic determinism without these constraints with three design decisions that build on each other.

Accumulators make the decisions; workers do the work. Every Zephon operator is split into two parts: a single-threaded accumulator that sees every record in order and makes all of the stateful decisions (like which sequences get packed together), and a stateless function that does the expensive work across however many parallel workers you give it. Since the workers can finish in any order, Zephon numbers each group of records before handing it out and puts the results back in order before passing them along. Unlike other data loaders, this allows that in Zephon the degree of parallelism is a pure performance knob: you can double the number of tokenization workers and the output won’t change.

Lanes instead of ranks. Instead of dividing the data up among the available GPUs directly, Zephon divides it into a fixed number of logical streams that are called lanes. Each lane has its own ordered sequence of samples and its own copy of any state that the operators need to keep (like the partially filled bins inside of a packing operator), so a lane produces exactly the same records regardless of which GPU ends up being assigned to it. When the number of data-parallel ranks changes, Zephon simply redistributes the lanes across the ranks, and the global batches come out the same. Accumulators make each lane deterministic; the lanes make Zephon elastically deterministic. The catch (you knew there had to be one) is that you need to choose the number of lanes up front, which we’ll come back to below.

Flushing and replaying instead of saving state. A Zephon checkpoint doesn’t contain the contents of any packing bins or shuffle buffers; it only records which chunks of the input were still being worked on when the checkpoint was taken. In order to keep that list short, Zephon periodically sends a flush sentinel through the pipeline that tells every stateful operator to emit whatever it’s holding and start over from a clean slate. When we restore from a checkpoint, Zephon replays the input from the most recent flush (which rebuilds the operators’ state exactly) and throws away anything the training loop has already consumed, which keeps checkpoints small (a few MiB in our experiments) no matter how far into the run you are.

Plugging Into TorchTitan

Our TorchTitan integration replaces TorchTitan’s default data loader with a Zephon Pipeline built from a data recipe (i.e., your datasets, their mixture weights, and whatever packing, shuffling, and parallelism settings you want to use), and everything else about the training run stays the same. Zephon’s state goes into TorchTitan’s checkpoint alongside the model, and since that state covers every lane, you can resume the run on a different number of GPUs from the same checkpoint. When the run launches, the integration checks that the lanes divide evenly across the data-parallel replicas and that each optimizer step’s batch splits evenly across the lanes, so a misconfigured run fails right away instead of silently changing the data several hours in. 

The repo also comes with a quick CPU-only demo that checkpoints a data stream, resumes it on a different number of data-parallel replicas, and checks that every global step matches an uninterrupted run, along with a GPU smoke test that does the same thing with a real (albeit tiny) model. The README has the commands for running each of these, plus the setup instructions, the full set of configuration options, and some example recipes to get you started.

Things to Know Before You Start

Choose your lane count. Pick a number that’s a common multiple of the data-parallel sizes you expect to run on. If you might run on 8, 16, or 32 GPUs, 32 lanes is a cap that covers all three. Don’t go overboard though (e.g., multiple thousands when your max is 16), since every extra lane per GPU adds overhead to the checkpoint’s restore time.

Keep the recipe fixed when you resume on more/fewer ranks. The data, tokenizer, seed, mixture, context length, and global batch size all need to stay the same. Of course, you may still adjust performance settings such as parallelism.

Zephon guarantees the data, not the arithmetic. The global batches will be identical across GPU counts, and the losses will be very close, but might not be identical since floating-point reduction order in the training loop still depends on how many GPUs are sharing the work.

Looking Ahead

There are a couple of directions that we’re especially excited about. The first is GIL-free Python: Zephon allows users to pick a runner, i.e., how to parallelize the operators. One runner we implemented is a thread-based runner. On Python 3.14t, it fetches vision-language samples about 3.7x faster than Zephon’s process runner does on standard Python. We think that’s where data loading is headed, and Zephon is ready for it. The second is dynamic curricula: because the curriculum is kept separate from the processing pipeline, a WorkSource can change the mixture weights over the course of a run without touching anything downstream. We’re working on a configurable curriculum component that does exactly this while keeping the determinism guarantees that we described above.

Beyond that, we would love to hear which parts of your data pipeline you would like to move into the training job and how we can help; issues and PRs are most welcome!

Learn more: Launch post · Paper · Zephon · Zephon on GitHub · TorchTitan Integration

The post Bringing Elastically Deterministic Data Loading to TorchTitan with Zephon appeared first on PyTorch.