AsyncMesh: Fully Asynchronous Optimization for Data and Pipeline Parallelism

TL;DR

AsyncMesh introduces fully asynchronous optimization techniques for data and pipeline parallelism in neural network training, reducing communication costs and maintaining performance on large-scale models.

Contribution

It proposes a novel asynchronous update framework with staleness mitigation strategies and provides convergence guarantees, enabling scalable training without co-location constraints.

Findings

Matches synchronous baseline performance on large models

Significantly reduces communication overhead

Provides convergence guarantees for asynchronous methods

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.