Zephon Enables TorchTitan Resume Across Different GPU Counts
Deterministic resume across cluster sizes removes a real blocker for teams scaling training mid-run or migrating between GPU allocations without wasting checkpoint progress. Practitioners running large-scale pretraining can adopt Zephon's in-job data pipeline to verify batch reproducibility via SHA-256 hashes. The approach was published by DatologyAI engineers Josh Wills, Maximilian Böther, and Ties Robroek, and compares favourably with workflows that require regenerating preprocessed datasets whenever hardware changes.
Zephon's integration with TorchTitan allows training to resume on a different GPU count while preserving identical global batches. In a test with a 1B-parameter model, a checkpoint saved on 16 GPUs was resumed on 8, 32, and 64 GPUs, and the SHA-256 hash of every global batch matched the original run. Zephon performs tokenization, packing, shuffling, and data mixing inside the training job, removing the need for a preprocessed corpus copy. Constraints include unchanged data, tokenizer, seed, mixture, context length, and global batch size, with lane counts dividing evenly across data-parallel replicas. Training losses can still differ because floating-point reduction order depends on GPU count.
Zephon’s TorchTitan integration preserves identical global batches when training resumes on a different GPU count.
In our latest blog’s 1B-parameter model test, training resumed from a 16-GPU checkpoint on 8, 32, and 64 GPUs, and the SHA-256 hash of every global batch matched the original run.
Zephon tokenizes, packs, shuffles, and mixes data inside the training job, avoiding a preprocessed copy of the corpus.
The data, tokenizer, seed, mixture, context length, and global batch size must stay unchanged when resuming. The lane count must divide evenly across data-parallel replicas, and each optimizer step’s batch must split evenly across lanes. Training losses can still differ because floating-point reduction order depends on GPU count.
✍️ Josh Wills, Maximilian Böther, Ties Robroek (DatologyAI)
Read more: https://t.co/JKn45hQarm