HeLoCo: Efficient asynchronous low-communication training under data and device heterogeneity

2026-05-29

AI Summary

This work addresses the convergence degradation in asynchronous low-communication training caused by the dual heterogeneity of non-IID data distributions and device asynchrony, which leads stale pseudo-gradients to deviate from the current optimization trajectory. To mitigate this, we propose HeLoCo, the first method to systematically tackle both sources of heterogeneity within an asynchronous low-communication framework. HeLoCo employs an outer-loop momentum to track the optimization path, dynamically evaluates the directional consistency of incoming pseudo-gradients, and selectively corrects conflicting components while preserving aligned ones. This direction-aware correction mechanism substantially enhances convergence stability and efficiency. Experiments on multilingual language model training show that HeLoCo outperforms the asynchronous DiLoCo baseline by 7.5% under a fixed token budget, surpasses asynchronous momentum lookahead by 3.3% under a fixed time budget, and achieves a 22.1% improvement over synchronous baselines under severe system heterogeneity.

Abstract

Distributed Low-Communication (DiLoCo) training reduces communication overhead by allowing workers to perform multiple local optimization steps before sending pseudo-gradients to a global outer update. Its asynchronous variant further improves hardware utilization by removing synchronization barriers, but at the cost of stale pseudo-gradients computed from outdated model states. As a result, these updates can become misaligned with the current global optimization direction, particularly in heterogeneous systems. This issue becomes even more pronounced when data are non-IID, a setting that has not been well studied in asynchronous low-communication training. To address this limitation, we propose \textbf{HeLoCo}, a direction-aware correction method for asynchronous low-communication training that uses outer momentum as a reference for the current optimization trajectory and selectively adjusts incoming pseudo-gradients before the outer update. Updates that remain aligned are preserved, while directionally conflicting components are corrected. On multilingual language-model training with heterogeneous workers and non-IID data, HeLoCo consistently improves validation loss. It outperforms existing asynchronous DiLoCo-based baselines by up to 7.5\% at a fixed token budget, exceeds asynchronous momentum look-ahead by up to 3.3\% at a fixed wall-clock budget, and surpasses the synchronous baseline by up to 22.1\% under severe system heterogeneity. Our analysis further shows how staleness, worker speed, and data heterogeneity shape update quality and convergence in highly decentralized and heterogeneous training setups.

Problem

• asynchronous training
• low-communication
• data heterogeneity
• device heterogeneity
• non-IID

Innovation

• asynchronous low-communication training
• direction-aware correction
• outer momentum
• data heterogeneity
• non-IID