Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

2026-06-03

AI Summary

This work addresses the activation communication bottleneck that limits pipeline-parallel training of large language models under low-bandwidth network conditions. The authors propose MAPL, a novel method that enables each pipeline stage to independently learn and dynamically optimize an orthogonal compression subspace. Task-adaptive low-rank compression is achieved through orthogonal projection learning constrained on the Stiefel manifold, while reconstruction fidelity is enhanced via factorized anchor embeddings combined with residual vector quantization. A streaming codebook synchronization protocol is further introduced to reduce communication overhead. Evaluated on LLaMA models ranging from 150M to 1B parameters, MAPL achieves negligible performance degradation even at high compression ratios, significantly outperforming existing approaches such as Subspace Networks.

Abstract

Pipeline parallelism enables training of large language models that exceed single-device memory, yet inter-stage activation communication becomes the dominant bottleneck when trained on low-bandwidth networks. Recent work in this area has proposed using fixed orthogonal projections to compress activations. However, this still results in a significant performance degradation and requires a number of non-standard adaptations to constrain the optimization. A natural alternative is to learn a low rank projection for each pipeline stage, however maintaining the necessary orthogonality of these projectors during training remains a challenge. We present Manifold Aware Projection Learning (MAPL), a method that treats inter-stage compression as a learnable orthogonal projection under explicit Stiefel manifold (orthogonal matrices) constraints. Rather than prescribing a fixed global subspace, MAPL lets each pipeline stage discover and continuously adapt its own task-optimal compression subspace via manifold-constrained steepest descent. To recover token-specific signals at stage boundaries, we introduce per-stage factorized anchor embeddings that allow for full-rank activation reconstruction with negligible communication overhead. We further show that we can incorporate residual vector quantization after projection with a streaming codebook synchronization protocol that amortizes dictionary communication. Across LLaMA models from 150M to 1B parameters we show that MAPL can be easily applied to the existing pipeline and can achieve high compression with neglibile performance degradation with a drastically improved tradeoffs in performance vs. compression compared to Subspace Networks.

Problem

• pipeline parallelism
• activation compression
• orthogonal projection
• communication bottleneck
• low-bandwidth networks

Innovation

• Manifold Aware Projection Learning
• Stiefel manifold
• pipeline parallelism
• learnable orthogonal projection
• activation compression