Data loader method reduces idle time in distributed training by balancing workload
Coordination-Free Workload Regularization for Variable-Cost Distributed Training
Distributed, Parallel, and Cluster ComputingMachine Learning
Summary
Training AI models across many computers can slow down when some computers get harder tasks than others. The authors found that by letting each computer look ahead at its upcoming work and reorganize it to balance how hard the tasks are, the overall training speeds up without needing extra coordination between computers. Their method keeps the amount of data each computer processes the same but evens out how much time each takes, cutting down waiting times. This improvement helps make large-scale AI training more efficient on modern hardware.
What this means in practice
- •For distributed machine learning engineers: Speed up multi-GPU synchronous training by balancing variable-cost workloads without cross-node communication.
- •For data center operations teams: Improve utilization and reduce idle GPU time during large-scale AI training by optimizing data assignment at each compute node.
Authors
Tao Huang, Chen Zhu, Antonio Alonso, Daming Li, Adarsh Goura, Jaideep Ray, Yitong Zhou
Abstract
Synchronous data-parallel training usually assigns every rank the same number of records, even when record costs vary by orders of magnitude. For variable-length models, this converts input variation into idle time because each step waits for its most expensive rank-local batch. We present Batch Cost Lookahead (BCL), a coordination-free data-loading policy that buffers a bounded rank-local window and repartitions it into fixed-cardinality packs with balanced estimated cost. BCL preserves the trainer's record-count contract and requires neither cross-rank metadata exchange nor runtime modification. In simulation over public workload distributions, enabling BCL reduces per-rank estimated batch-cost variability by up to 82.8%. Across 13 NeMo-RL SFT corpora on eight H200 GPUs, rank-local BCL runs 1.050 to 1.167 times faster than pooled MFFD. Compiled H200 measurements on an internal sequence workload show speedups of 1.197 to 1.343 times over packed execution from one to four nodes. We also identify a memory-feasibility boundary: a p99 cap for high-variance or OOM-prone data bounds outliers before BCL balances the rest. These results show that substantial variable-cost imbalance can be removed at the data-loader boundary, while clarifying when data preparation or broader coordination is still necessary.