Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotation

TL;DR

This paper introduces basis rotation to mitigate gradient staleness in asynchronous pipeline parallelism, significantly improving convergence speed and scalability in large-scale distributed training.

Contribution

The paper proposes a novel basis rotation method to correct gradient delays, enhancing the scalability and efficiency of asynchronous pipeline parallel training.

Findings

Achieves 76.8% fewer iterations in training a 1B-parameter LLM.

Restores effective curvature-aware optimization with basis rotation.

Theoretically and empirically demonstrates mitigation of gradient staleness effects.

Abstract

Asynchronous pipeline parallelism maximizes hardware utilization by eliminating the pipeline bubbles inherent in synchronous execution, offering a path toward efficient large-scale distributed training. However, this efficiency gain can be compromised by gradient staleness, where the immediate model updates with delayed gradients introduce noise into the optimization process. Crucially, we identify a critical, yet often overlooked, pathology: this delay scales linearly with pipeline depth, fundamentally undermining the very scalability that the method originally intends to provide. In this work, we investigate this inconsistency and bridge the gap by rectifying delayed gradients through basis rotation, restoring scalable asynchronous training while maintaining performance. Specifically, we observe that the deleterious effects of delayed gradients are exacerbated when the Hessian eigenbasis is misaligned with the standard coordinate basis. We demonstrate that this misalignment prevents coordinate-wise adaptive schemes, such as Adam, from effectively leveraging curvature-aware adaptivity. This failure leads to significant oscillations in the optimization trajectory and, consequently, slower convergence. We substantiate these findings through both rigorous theoretical analysis and empirical evaluation. To address this challenge, we propose the use of basis rotation, demonstrating that it effectively mitigates the alignment issue and significantly accelerates convergence in asynchronous settings. For example, our training of a 1B-parameter LLM with basis rotation achieves the same training loss in 76.8% fewer iterations compared to the best-performing asynchronous pipeline parallel training baseline.