AsyncMesh: Fully Asynchronous Optimization for Data and Pipeline Parallelism

TL;DR

AsyncMesh enables fully asynchronous data- and pipeline-parallel training on a 2D mesh by combining AsyncPP-style pipeline updates with sparse, EMA-corrected data averaging. The method uses a look-ahead weight extrapolation for pipeline staleness and an EMA-based delay correction for sparse averaging, yielding consensus on expectation and convergence to a fixed point of the consensus objective $F(W) = \sum_{i=1}^m F(W_i;\mathcal{D}_i)$ with $W_i=W$. Theoretical results bound the consensus error and establish convergence under diminishing staleness, while extensive experiments on large language models (up to $1\times 10^9$ parameters) show parity with fully synchronous baselines and substantial communication reductions, including successful training of a 1B-parameter model in AsyncMesh. This work enables scalable, bandwidth-efficient distributed training across heterogeneous and bandwidth-constrained environments, expanding the practicality of large-scale model training beyond tightly coupled clusters.

Abstract

Data and pipeline parallelism are key strategies for scaling neural network training across distributed devices, but their high communication cost necessitates co-located computing clusters with fast interconnects, limiting their scalability. We address this communication bottleneck by introducing asynchronous updates across both parallelism axes, relaxing the co-location requirement at the expense of introducing staleness between pipeline stages and data parallel replicas. To mitigate staleness, for pipeline parallelism, we adopt a weight look-ahead approach, and for data parallelism, we introduce an asynchronous sparse averaging method equipped with an exponential moving average based correction mechanism. We provide convergence guarantees for both sparse averaging and asynchronous updates. Experiments on large-scale language models (up to \em 1B parameters) demonstrate that our approach matches the performance of the fully synchronous baseline, while significantly reducing communication overhead.

Figures (16)

• Figure 1: SyncMesh vs AsyncMesh, for a 3-stage, 2- replica setup (only 2 stages: S1 and S3, are shown for clarity). Notably, AsyncMesh eliminates idle time due to communication for both and . In synchronous , devices are idle while parameters are averaged, whereas asynchronous eliminates this idle time by using the "old" average. Moreover, in Async, each stage alternates between forward and backward passes without any communication delay.
• Figure 2: Results on WT with varying mesh configurations with AsyncPP for all methods except FullSync. In all scenarios, our method matches the performance of the fully synchronous method, while outperforming the fully asynchronous baseline AsyncSPARTA.
• Figure 3: Validation perplexity scores at 30k iterations for the $4 \times 2$ mesh. Our method outperforms AsyncSPARTA, and matches FullSync while eliminating the communication overhead. Except FullSync, all other methods use AsyncPP.
• Figure 4: Perplexity gap (lower the better) between our method and DPAvg on WT for different configurations. Our method is robust to a range of configurations and scales favourably with the number of replicas and stages. Here, 100* indicates, async-delay = 10 with avg-interval = 10, simulating 100-step effective delay, as async-delay = 100 did not converge.
• Figure 5: Effect of quantization for different methods, for the base model with $4\times 2$ mesh on WT. Quantization degrades the performance for DPAvg (DPAvg- FP8 did not converge) while sparse averaging (even with 10-step delay) is robust to it. While intriguing, it may be explained by the fact that since the quantization error is introduced only for a small subset (5% in this case) at each iteration, the effect of quantization on training is negligible.

Theorems & Definitions (12)

• Theorem 1: Sparse averaging ensures consensus
• proof
• Theorem 2: Delayed averaging with ensures consensus
• proof
• Lemma 1
• proof
• Theorem 3
• proof
• Lemma 2
• proof