Zephon keeps data order stable across GPU counts

DatologyAI has introduced Zephon, a data loader that preserves training batches when a distributed run resumes with more or fewer GPUs. Its TorchTitan integration produces the same bytes in the same order for every global batch, isolating data order from cluster size during training and controlled experiments.

GPU count was moving the score

Many loaders assign samples according to the number of data-parallel ranks, usually one process per GPU. Changing that count can alter which samples appear together and when the model sees them. A global batch, the combined batch processed across all ranks for one optimization step, can therefore change even when the dataset, seed, and training recipe remain fixed.

In a test cited by the Zephon authors, a 1B-parameter model trained from scratch for 20B tokens on 8, 16, 32, and 64 GPUs. With a standard loader, evaluation scores varied substantially across otherwise equivalent runs:

The standard-loader spread was about four times the decision thresholds used in the DCLM and FineWeb data-curation studies. Cluster allocation could therefore overwhelm the effect attributed to a new mixture or preprocessing method.

Zephon also passed a direct resume test. The team checkpointed a 16-GPU run at step 150, resumed it on 8, 32, and 64 GPUs, and compared each result with the continued 16-GPU baseline. The SHA-256 hash of every global batch matched.

On-the-fly packing turns order into state

A loader that tokenizes documents, combines them into fixed-length sequences, and mixes datasets during training cannot retrieve the final sample order from a static index. The contents of each sequence depend on buffered documents, partially filled packing bins, shuffle state, and the timing of earlier operations.

Mosaic Streaming and Google Grain can preserve global batches across GPU counts when an index determines sample order in advance. That requirement excludes pipelines that pack records during loading. Megatron-Energon can reproduce batches after a GPU-count change when the total worker count remains fixed.

Three separations make elasticity deterministic

• Stateful decisions are separated from parallel work. Each operator has a single-threaded accumulator that observes records in order and manages state. A stateless function handles expensive work across a configurable worker pool. Zephon assigns sequence numbers to record groups and restores output order after workers finish, allowing worker parallelism to change without changing the stream.
• Logical lanes replace rank-owned state. Zephon divides data into a fixed number of ordered lanes. Each lane maintains operator state, including partially filled packing bins. When the data-parallel world size changes, Zephon redistributes lanes among ranks while preserving lane contents and order.
• Checkpoints use flushing and replay. Zephon periodically injects a sentinel that moves through the pipeline, causes stateful operators to emit buffered records, and establishes a clean boundary. A checkpoint records the input chunks still in flight. Restoration replays input from the latest flush, keeping checkpoint state to a few MiB regardless of the run’s duration.

The pipeline stays declarative

Zephon pipelines define the datasets, a work source that controls mixing, and the operators applied to each record. The blog provides this minimal example:

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},
    token_estimation=TokenEstimation(),
    seed=42,
)

pipeline = (
    Pipeline(source)
    .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)
)

TokenEstimation samples each dataset to estimate its token-to-byte ratio, then adjusts the work plan so raw-text datasets reach the requested mixture by token count. In the team’s vision-language benchmark, a pipeline that decoded images, tokenized text, packed sequences, and formed batches matched Google Grain’s throughput while retaining deterministic stateful operations across GPU counts.

TorchTitan enforces the invariants

The TorchTitan adapter replaces the default loader with a Zephon pipeline and stores loader state in TorchTitan’s existing checkpoint. Startup validation requires the lane count to divide evenly across data-parallel replicas and each optimizer-step batch to split evenly across lanes. Invalid configurations fail before training instead of producing a silently different data stream.

Constraints to plan around

• Choose the lane count before training. Use a common multiple of the data-parallel sizes the run may use. For allocations of 8, 16, and 32 data-parallel ranks, 32 lanes supports all three. Additional lanes can increase restoration time.
• Keep the training recipe fixed. Resumed runs must retain the same data, tokenizer, random seed, mixture, context length, and global batch size. Performance settings such as worker parallelism may change.
• Expect small numerical differences. Zephon preserves global-batch contents, but losses may differ slightly because floating-point reduction order in the training loop changes with the number of GPUs.

Free-threaded Python could cut loader overhead

DatologyAI is also testing Python 3.14t, whose free-threaded build can execute Python threads without the global interpreter lock. In the team’s benchmark, a thread-based runner fetched vision-language samples about 3.7 times faster than a process-based runner on standard Python. The comparison covers sample fetching rather than end-to-end training throughput.

The accompanying research paper describes the stateful n:m pipeline model behind Zephon. For schedulers that resume jobs on available hardware, the design decouples data order from GPU allocation while leaving performance parallelism adjustable.