Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotation

TL;DR

This work identifies gradient delay as a depth-correlated bottleneck in asynchronous pipeline parallelism and pinpoints basis misalignment between the Hessian eigenbasis and the coordinate basis as the key cause of deteriorated convergence under delay. It introduces basis rotation to realign the optimization space so that curvature-aware adaptivity (e.g., Adam-type methods) remains effective despite stale gradients, leveraging block-diagonal and Kronecker-factorized Hessian assumptions for tractable implementation. The authors provide both theoretical convergence guarantees and extensive empirical validation on Transformer-scale models up to 1B parameters, showing that basis rotation dramatically reduces iteration counts to reach target loss (e.g., 71.6%–76.8% fewer iterations) and restores scalability in asynchronous pipeline training. The approach, supported by efficient eigenbasis-estimation strategies and robustness to practical constraints (e.g., without weight stashing), promises substantially more efficient large-scale training in real-world distributed environments.

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.

Figures (16)

• Figure 1: (a--b) Schematic diagrams showing how micro-batches (blocks 1--7) are processed over time: a micro-batch travels from stage 1 to stage 4 in a forward pass (blue), and then goes back to stage 1 through a backward pass (green). Once the gradient becomes available after the backward pass, the model is updated (red dots indicate the time points). Asynchronous pipelining removes idle periods by processing subsequent micro-batches immediately after completing a backward pass without waiting for the completion of the pipeline cycle. (c) An illustration of model update with delayed gradient at stage 1: $w_3$ is updated to $w_4$ with $\nabla f(w_0; B_4)$.
• Figure 2: Summary of this work. (a) Impact of pipeline depth (i.e., number of stages) on convergence of asynchronous pipeline parallel LLM pre-training. In all cases, the model itself is kept the same while the number of stages is divided to be different. Increased delay leads to significant degradation on convergence speed. (b) Basis rotation substantially accelerates convergence in the presence of a large delay (here, for the case of 32 stages).
• Figure 3: Impact of basis alignment on the effect of delay. (Top) Optimization trajectories of SGD and Adam with and without delay. (Bottom) Update of Adam along the dominant direction (red arrows in the top panel). (a) When the Hessian eigenbasis is aligned with the standard coordinate basis, Adam yields a stable trajectory and shows robustness to delay. (b) When two bases are misaligned, Adam exhibits strong oscillations along the dominant direction and suffers significantly from delay. See \ref{['app:sub:toy-exp-details']} for experimental details.
• Figure 4: (a) Optimization trajectory of Adam on the spiral loss. The optimizer maintains a stable trajectory in basis-aligned regions (green boxes) but exhibits severe oscillations in misaligned regions (red boxes). (b) Slowdown ratio $T_{\text{delay}}/T_{\text{no-delay}}$ in different regions. The ratio is minimized near basis-aligned regions (green dotted lines), whereas it is maximized in misaligned regions (red dotted lines). This demonstrates that basis misalignment significantly amplifies the deleterious effects of delayed gradients.
• Figure 5: Performance of different methods when increasing the number of stages $P$ for the same model (with $95$M parameters). (a) Basis rotation maintains stable performance while baselines suffer significantly from delay. (b) Basis rotation shows much faster convergence under large delay compared to baselines (for $P=32$). (c-e) Basis rotation reduces slowdown, i.e., the iteration ratio required to reach target loss for $P=32$ relative to $P=1$, by a large margin compared to baselines. More results including other baselines are presented in \ref{['fig:result-appendix-optmain', 'fig:result-appendix-slowdown', 'fig:result-appendix-val', 'fig:result-appendix-dc']} of \ref{['app:additional-exp-results']}.

Theorems & Definitions (18)

• Theorem 2.3
• Theorem 3.1
• Lemma 3.1
• Lemma 3.2
• Lemma 3.3
• Lemma 3.4
• Lemma 3.5
• proof
• Lemma 4.1
• Lemma 4.2