Bringing Elastically Deterministic Data Loading to TorchTitan with Zephon
Zephon ensures consistent training data order across varying GPU counts for reproducible machine learning experiments.
1The Reproducibility Challenge
Standard data loaders often change the sequence of training samples when the number of GPUs varies. This creates inconsistent loss curves, making it difficult to isolate whether model performance changes result from data mixture adjustments or simply the hardware configuration. Zephon addresses this by ensuring identical global batch ordering across different scales.
2Impact of GPU Scaling
Experiments using standard loaders show significant score variance when scaling from 8 to 64 GPUs. Evaluations on DCLM Core v1 and FineWeb suites revealed performance fluctuations up to 0.82 and 0.64 points, respectively. These margins exceed typical data curation thresholds, potentially leading researchers to misinterpret GPU-induced noise as actual model improvements.
3On-the-Fly Data Processing
Zephon performs tokenization, packing, and shuffling directly within the training job, utilizing idle CPU resources. This eliminates the need for expensive offline preprocessing and storing massive datasets on disk. By moving these operations into the training pipeline, teams avoid creating redundant copies of the corpus that are never read again.
4Defining the Zephon Pipeline
The system requires three components: datasets, a WorkSource for curriculum planning, and a pipeline of operators. By separating the curriculum logic from processing operators, users can modify mixture weights without altering the underlying pipeline. A TokenEstimation utility ensures that mixing proportions are accurately targeted on a token-level rather than sample-level.
5Parallelism as Performance Knob
Zephon splits operators into single-threaded accumulators and stateless workers. The accumulators handle stateful logic, while workers perform heavy tasks like tokenization in parallel. Because Zephon numbers records before distribution, users can increase the number of workers to improve throughput without affecting the deterministic output of the training data sequence.
6The Role of Lanes
Instead of assigning data directly to GPUs, Zephon divides the workload into fixed logical streams called lanes. Each lane maintains its own ordered sequence and operator state. When the number of data-parallel ranks changes, Zephon simply redistributes these lanes across the new ranks, maintaining global batch consistency regardless of the hardware scale.
7Efficient Checkpoint Management
Zephon checkpoints do not store bulky packing bins or shuffle buffers. Instead, they record input chunks being processed and use flush sentinels to reset stateful operators. Upon restoration, the system replays the input from the last flush, keeping checkpoint sizes to a few MiB regardless of the training duration or progress.
8TorchTitan Integration Details
The integration replaces the default TorchTitan loader with a Zephon pipeline. It stores Zephon’s state alongside the model checkpoint, allowing seamless resumption on different GPU counts. The system performs validation at launch to ensure lanes divide evenly across replicas, preventing silent data errors that could otherwise persist for many hours.
9Strategic Configuration Advice
Users must select a lane count that is a common multiple of their expected data-parallel sizes. While 32 lanes might cover 8, 16, and 32 GPU setups, excessive lane counts increase checkpoint restore times. Consistency in the recipe, including seeds and context length, remains mandatory when resuming on different rank counts.
10Future Development Directions
The team is exploring GIL-free Python execution using thread-based runners, which showed 3.7x faster vision-language sample fetching on Python 3.14t. Additionally, they are developing a dynamic curriculum component. This will allow users to adjust mixture weights during a run while maintaining the same elastic determinism guarantees currently provided by the system.
11Ensuring Reliable Ablations
By providing elastic determinism, Zephon enables reliable ablation experiments in distributed environments. It removes the hardware-dependent variables that often plague large-scale training, allowing researchers to trust their results. This architecture is essential for modern workflows where agents frequently launch and interpret experiments in parallel across varying compute resources.