Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotation

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. We trace this pathology to a specific property of the optimization landscape: the misalignment between the Hessian eigenbasis and the standard coordinate basis, which triggers oscillations in the update trajectories of coordinate-wise adaptive optimizers. We identify that these oscillations cause delayed updates to diverge from their true counterparts, invalidating their use for current iterations. This insight is formalized through theoretical analysis, including a convergence bound showing that basis misalignment amplifies the delay penalty, and substantiated with empirical evaluation. To address this, we propose basis rotation, a framework that rotates the optimizer's coordinate system to align with the Hessian eigenbasis, keeping delayed updates useful. We theoretically demonstrate that basis rotation minimizes basis misalignment, thereby counteracting the conditions that amplify delay penalties. Empirically, in training up to a 3B-parameter LLM, basis rotation reduces the required iterations by 81.7% compared to the best-performing asynchronous baseline.

Training large language models (LLMs) requires spreading the model across many GPUs. Running these GPUs simultaneously without waiting for each other maximizes hardware utilization, but causes a problem: each GPU ends up updating the model with outdated information ("gradient staleness"). This problem grows with the number of GPUs used, undermining the very scalability the approach was meant to provide. We identify the root cause: the direction of the loss landscape's curvature is misaligned with the coordinate system that the Adam optimizer relies on. This misalignment destabilizes the optimization path, and when stale gradients arrive, their damage is greatly amplified. To fix this, we propose basis rotation — a method that rotates the optimization space to realign these directions. In experiments on a 3B-parameter LLM, our method reaches the same training loss in 81.7% fewer iterations than the best existing baseline.