TorchTitan: Reproducible data with Zephon, preprocessing during training

1. Processing during training instead of making copies in advance

Traditional workflows often preprocess the entire dataset and create a disk copy whenever data mixing or processing changes. If an experiment yields no improvement, that copy may be read only once, wasting preparation time and storage.

Zephon tokenizes text for the model, packs it into fixed lengths, and shuffles it during training. Crucially, it preserves the global batch—the data all GPUs collectively receive in one training step—even when the GPU count changes.

2. Differences observed in resumed runs and evaluations

The source trained a 1 billion-parameter model on 16 GPUs, then resumed on 8, 32, and 64 GPUs from the checkpoint at step 150. The standard loader's loss curves diverged, while Zephon's nearly overlapped. Each global batch's SHA-256 hash also matched the original run. These hashes served as comparison values to verify identical batch contents.

A separate evaluation trained models of the same size from scratch for 20 billion tokens on 8, 16, 32, and 64 GPUs. The maximum DCLM Core v1 score gap across GPU counts was 0.82 points with the standard loader and 0.05 with Zephon. The FineWeb evaluation gaps were 0.64 and 0.011 points, respectively.

These results show reduced evaluation variation across GPU counts, not higher average scores. The source says the standard loader's gaps were about four times the relevant thresholds used by the two evaluations' authors to assess data selection methods. Comparing data mixtures across different GPU counts can thus mistake hardware configuration effects for data improvements.

3. Separating data planning from processing

Users specify datasets, a plan defining mixing ratios and order, and processing steps that turn samples into training tensors—numeric arrays used in model computation. Separating planning from processing lets users change either mixing ratios or processing methods independently.

Mixing ratios are based on tokens, not sample counts, and the plan is calibrated by measuring the token-to-byte ratio of a small sample.

The following is an unchanged excerpt from the source example, executed after configuring the data sources. Tokenization uses four workers, and parallelism can be set separately for each processing stage.

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())

4. How order is preserved and the integration works

Existing loaders can also reproduce batches after GPU counts change, under certain conditions. Mosaic Streaming and Google Grain require sample order to be predefined by index, ruling out processing such as packing during training. Megatron-Energon requires the total number of data loader workers across all GPUs to remain unchanged.

Zephon separates the accumulators that determine order and grouping from parallel workers, reordering outputs even when completion order differs.

Data is divided into fixed logical streams called lanes, which are reassigned when the GPU count changes.

Checkpoints record the input ranges being processed rather than buffer contents. During recovery, inputs are replayed from the most recent point where buffers were empty to reconstruct the state.

The TorchTitan integration replaces the default data loader with a Zephon pipeline built from data and processing settings, leaving other training settings unchanged. Data processing state is saved alongside the model. At startup, it checks whether lanes are evenly distributed across data-parallel replicas—processing units that share model training. It also checks whether each training step's batch splits evenly across lanes, catching invalid settings early.

5. Preparation savings and remaining constraints

Instead of creating preprocessed copies, CPUs within the training job handle processing.

The lane count must be chosen in advance as a common multiple of the expected data-parallel sizes. For 8, 16, and 32 GPUs, 32 lanes can be used. More lanes per GPU increase checkpoint restoration time, so arbitrarily large counts should be avoided.

When resuming with a different GPU count, the data, tokenizer, random seed, mixing ratios, context length, and global batch size must remain fixed. Performance settings such as parallelism can be adjusted.

The guarantee covers data, not all computation results. Changing GPU counts can alter the order of floating-point reductions, so loss values may be very close without being identical.

The value of Zephon integration lies less in faster training itself than in providing a basis for comparing data experiments unaffected by hardware configuration.