Title: Mixtures of Subspaces for Bandwidth Efficient Context Parallel Training
Venue: NeurIPS 2025
Accepted: yes
Keywords: Decentralized training

LLMs

Compressio

Distributed training
Abstract: Pretraining language models with extended context windows enhances their ability to leverage rich information during generation. Existing methods split input sequences into chunks, broadcast them across multiple devices, and compute attention block by block which incurs significant communication overhead. While feasible in high-speed clusters, these methods are impractical for decentralized training over low-bandwidth connections. We propose a compression method for communication-efficient context parallelism in decentralized settings, achieving a remarkable compression rate of over 95% with negligible overhead and no loss in convergence. Our key insight is to exploit the intrinsic low-rank structure of activation outputs by dynamically constraining them to learned mixtures of subspaces via efficient reparameterizations. We demonstrate scaling billion-parameter decentralized models to context lengths exceeding 100K tokens on networks as slow as 300Mbps, matching the wall-clock convergence speed of centralized models on 100Gbps interconnects.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the bandwidth bottleneck of context-parallel (CP) training for very-long-context transformers. CP divides a single long sequence across GPUs; every layer therefore requires an all-gather of full key/value (K,V) activations, which is trivial on datacentre interconnects but prohibitive on commodity networks. Similar to previous work, the authors show empirically that K,V lie on an extremely low-rank manifold (0.1-0.5% stable rank). They propose Mixtures-of-Subspaces (MoS), where each attention weight is factorised so that its outputs are constrained to an r-dimensional subspace. \nA short warm-up run provides an initial orthonormal basis for the low rank subspace $\\bar U$, and thereafter the model learns only a rotation $R(\\theta)$ of that basis together with a low-rank matrix $B$. This re-parameterisation keeps the columns of $U=R(\\theta)\\bar U $ orthonormal without per-step QR/SVD. Projecting $K,V$ onto this subspace reduces communication from $O(nd)$ to $O(nr)$.\n\nOn eight A100 GPUs linked by 300 Mbs Ethernet, MoS reduces CP traffic by ≈20x and matches the wall-clock convergence of an uncompressed baseline on 100 Gb s InfiniBand using an 800 M-parameter, 132 K-token model. Because weight-decay shrinks directions outside the learned subspace, the projection and rotation heads can be “unplugged” after training. Hereby, the final model is a standard transformer.\n\nSTRENGTHS AND WEAKNESSES:\n**Quality**\nSolid emprical results: Three data sets, 132 K-token context, and a ~20 × wall-clock speed-up on commodity 300 Mbs links. Interesting ablations: Layer-wise rank study, rotation-size sweep, comparisons to 4-bit and Top-10 % sparsity. Theory: Proof that the re-parameterised optimisation retains convergence guarantees.\nNo variance/error bars. Table 2 and Fig. 5 would be more convincing with ± 1 σ over at least three seeds. (I understand this is harder (costly) for the full runs in Table 1.)\n\n**Clarity**\nClear high-level motivation; algorithms and all hyper-parameters are spelled out; code is included in supplementary.\nFigures 2 and 3 are hard to read: the overlapping “cen-CP 100 Gbps” vs “dec-CP 300 Mbps (comp)” curves look identical but the axes are tiny. A zoom-in would be useful in the main paper. More detailed on explored hparams could also make it easier to reproduce experiments.\n\n**Significance**\nFirst method that makes Context paralism practical on commodity networks, potentially opening > 100 K-token training to a far wider pool of researchers. Orthogonal to existing attention kernels since it does not require architectural changes at inference.\nImpact hinges on the low-rank property holding for other modalities or tasks; scalability beyond the 800 M-parameter (and 3B) range is not demonstrated.\n\n**Originality**\nThe paper is the first to use low-rank structure specifically to minimise all-gather traffic in context-parallel training. This communication-first application of low-rank projections plus the “unplug-after-training” trick that removes the projector at inference is, to my knowledge, new and distinguishes the work from prior memory- or weight-centric compression schemes.\n\nQUESTIONS:\n1. Warm-start sensitivity:  How many uncompressed steps are needed to build \\bar U? Does the answer change with dataset choice?\n2. Unplugging: How sensitive is the final PPL to the number of fine-tune steps and weight-decay value used when removing the projector? I see Figure 4 shows an example of this, but how sensitive is it between the late and early “unplugging” steps in this figure?\n3. Training-loss vs validation gap: In Figs 2–3 the “cen-CP 100 Gbps” curve appears to reach lower training loss than MoS, yet Table 1 reports better validation perplexity for MoS. Can you comment on this discrepancy? A zoom-in plot (similar to App.\\ Fig 14) in the main paper would help.\n4. Variance and robustness: All main curves and ablations are single runs. Could you provide some error bars for some of the smaller ablation experiments?\n5. Relation to DeepSeek MLA: DeepSeek’s Mixed-Linear Attention (MLA) also projects K and V to a fixed low dimension, but as I understand it, it keeps the keys/values low-rank all the way through the attention kernel, giving a lasting memory/compute win, whereas MoS re-expands the compressed blocks after the all-gather and therefore targets communication rather than peak RAM.\nCould you clarify this distinction and comment on whether there is any practical benefit (or incompatibility) in combining MoS with MLA\n\nLIMITATIONS:\nYes, Section 7 discusses theory gaps and the unexplained success of very low-dimensional rotations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel method for bandwidth-efficient context-parallel (CP) training in decentralized environments. The authors propose dynamically compressing key and value activations using learned low-rank projections into mixtures of subspaces. By jointly optimizing both the subspaces and attention weights on a product manifold, and using per-chunk rotations with compact parameterization, the method achieves over 95% communication compression. The compressed activations can be transmitted efficiently even over low-bandwidth (e.g., 300Mbps) links, and decompressed without degrading model quality. Extensive experiments on long-context training (132K tokens) across datasets show that the proposed method achieves convergence comparable to centralized setups, with 20× speedup over uncompressed CP on commodity links.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n1. The use of dynamic mixtures of low-rank subspaces and efficient rotation-based reparameterization is both novel and theoretically grounded.\n2.  The paper includes careful breakdowns of design choices, including the role of warm starts, second-order approximations, and reparameterization.\n\nWeakness:\n1. The motivation is not persuasive. Why do we need decentralized long-context training over low-bandwidth connections? Why a long context is distributed across \"internet-grade\" linked devices if it has inherent long-range dependencies?\n2. Context parallel introduces frequent communication operations. Pipeline parallel might seem to be a better choice for low-bandwidth training with a much smaller number of communication volumes and frequencies.\n3. The surprising effectiveness of extremely low-dimensional rotations is acknowledged but not fully understood. The paper lacks comparison with other attention sparsification methods.\n\nQUESTIONS:\nPlease the the weakness part.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses a critical and previously unexplored bottleneck in training large language models (LLMs): the immense communication overhead of Context Parallelism (CP) in decentralized, low-bandwidth environments. Standard CP requires each device to broadcast its local key (K) and value (V) activation matrices to all other devices, which is prohibitively expensive over commodity internet connections. The authors' core contribution is a novel and highly effective compression method for these activations. The method is motivated by the empirical observation that the K and V activation matrices in large transformers have an intrinsically low-rank structure. The proposed solution leverages this by: Projecting the high-dimensional K and V activations onto a low-dimensional subspace, drastically reducing the amount of data to be communicated. Dynamically adapting this subspace for each input chunk via a lightweight, learnable rotation. This provides flexibility without sacrificing compression. Using an efficient reparameterization of these rotations based on Lie groups, which avoids costly matrix operations (like SVD/QR) and makes the method computationally cheap. Introducing a mechanism to \"unplug\" the auxiliary compression components after training, resulting in a standard transformer architecture at inference time with no overhead.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The paper tackles a highly relevant and timely problem. As context windows in LLMs expand beyond 100K tokens, CP is becoming essential. However, its application has been confined to expensive, centralized data centers. This paper is the first to propose a viable solution for CP in decentralized settings, potentially democratizing the training of long-context LLMs.\n2. The method is built upon a simple yet powerful insight: the low-rank nature of attention activations. This is convincingly demonstrated in Figure 1 with an analysis of LLaMA-70B. This strong empirical grounding makes the subsequent methodological choices feel well-justified and intuitive.\n3. The technical approach is very clever. The authors build a logical progression from a simple low-rank factorization to a more robust solution. The reparameterization of the projection subspace U via a fixed basis Ū and a learnable rotation R(θ) is an elegant way to maintain expressive power while making optimization tractable (avoiding expensive Riemannian gradient steps). The use of a low-dimensional Lie subgroup to parameterize R(θ) is a principled way to control the communication-flexibility trade-off.\n4. The results presented are outstanding and strongly support the paper's claims. Performance Parity (Figure 2): Matching the wall-clock performance of a 100Gbps system on a 300Mbps link is a remarkable achievement and the most compelling result in the paper. Scalability (Figure 3): The method is shown to scale effectively to a 3B parameter model and integrates seamlessly with other parallelism techniques like Pipeline Parallelism. Ablations and Baselines (Table 2, Figure 5): The authors conduct thorough ablation studies that validate their design choices (e.g., the importance of the warm-start, the adequacy of the 2nd-order approximation). The comparisons against sensible baselines like quantization and sparsification clearly demonstrate the superiority of the proposed approach.\n\nWeaknesses\n1. The method relies on a \"short, uncompressed warm-up phase\" to establish the initial subspace basis Ū. This step is crucial, as the ablations show that performance degrades significantly without it (\"No warm start\" in Table 2). However, the details are sparse. The sensitivity to the duration and context length of this phase is not explored, making it seem somewhat heuristic. A more detailed analysis or guideline on setting up this phase would strengthen the paper.\n2. The paper finds that a rotation parameterized by a single scalar (k=1) is sufficient. This is a surprising and powerful result for communication efficiency, but the paper offers little intuition as to why such a constrained, planar rotation is sufficient to adapt the subspace effectively across diverse data chunks. Is this a fundamental property, or could it be a limitation on more complex tasks or architectures?\n3. The experiments are large-scale and computationally expensive, so this is understandable. However, the results are based on single training runs. As acknowledged in the paper's checklist, there are no error bars or confidence intervals to assess the variance of the results (e.g., due to different random seeds for initialization).\n\nQUESTIONS:\n1. Could you provide more details on the warm-up stage? How was the duration (<500 iterations) and reduced context length determined? How sensitive is the final performance to these choices?\n2. The success with a single rotation parameter (k=1) is fascinating. Do you have a hypothesis for why this highly constrained adaptation works so well? Have you identified any failure cases or do you foresee scenarios (e.g., highly multi-modal data) where a larger k would be necessary?\n3. Figure 4 shows that late unplugging is effective. Is there a clear, automatable heuristic for determining the optimal point in training to remove the projection components? Does the subsequent \"polishing\" fine-tuning pass require a specific learning rate schedule to avoid disrupting the learned weights?\n4. The paper uses 98% and 95% compression for K and V respectively. Have you explored the trade-off curve between the compression rate (i.e., the rank r) and the final model perplexity? It would be interesting to see the point at which performance begins to degrade meaningfully.\n5. The proposed method compresses activations. Could it be productively combined with other techniques, such as weight quantization or gradient compression (for the infrequent weight syncs), to achieve even greater efficiency in a fully decentralized setting?\n\nLIMITATIONS:\nsee questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The paper attempts to solve the efficiency problem of long-context pretraining under low-bandwidth connections situation. Concretely it studies compression to reduce the communication problem, incurred by ring-style sequence parallelism. It achieves scaling billion-parameter models with more than 100k on slow networks.

STRENGTHS AND WEAKNESSES:
Strengths:
-  The problem on ring-style sequence parallelism is significant.
- The paper is well written and easy to follow.

Weakness:
- The convergence is done on a small model, compared to the models mentioned in the introduction, e.g. Llama4.  The convergence can behave differently.
- It does not studies Ulysses style sequence parallelism, where the communication is less a problem, compared to ring-style.

QUESTIONS:
(1) Is it possible that the author provides experiment result on larger models?  The convergence can be different for larger model (the paper pretrains a 800M model, where the reviewer suggest the scale should be > 7B to draw significant applicable conclusion), especially when the paper method is based on compression.
(2) Is it possible to show effectiveness on larger cluster? The current experiments are done on 8xA100, where typically models are pretrained at the scale of > 100 H100 GPUs.
(3) How does the method applies to Ulysses style sequence parallelism, which typically induces less communication?
(4) (Minor) The figure 3 and figure 4 ratio are different. Could the author align them for a more compact view of the paper?

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a compression method enabling efficient context-parallel training of large language models across decentralized devices with limited bandwidth. By dynamically exploiting the low-rank structure of attention activations through learned mixtures of subspaces, the method achieves over 95% compression without sacrificing model convergence or performance. Experiments demonstrate successful training at scale, matching centralized high-bandwidth systems in convergence speed.\n\nSTRENGTHS AND WEAKNESSES:\nI really liked how effectively they tackle the major bottleneck in decentralized training, namely communication. Reducing bandwidth usage by over 95% while preserving accuracy is seriously impressive. The method is also flexible, integrating seamlessly into existing transformer architectures, so you can drop their special components after training and just deploy a normal model. Plus, they show practical scaling to billion-parameter models, which is exactly what decentralized setups need.\n\nOn the downside, though, the technique relies heavily on an initial \"warm-up\" phase, which might be tricky in unpredictable real-world setups. Another concern is that the theoretical guarantees lean on assumptions (like low-rank activations) that might not hold perfectly across all tasks or domains. Finally, while compression rates are great, their rotation-based subspace method introduces extra computational complexity, even though they've tried to minimize it.\n\nQUESTIONS:\nHow does the method proposed relate to PowerGossip (https://arxiv.org/abs/2008.01425) and Photon (https://arxiv.org/abs/2411.02908)?\n\nYou mention an initial warm-up phase is crucial to your method; how sensitive is your approach to variations in this phase (such as reduced length or interrupted warm-up), and have you explored strategies to mitigate dependency on it in dynamic decentralized environments?\n\nCould you clarify how robust your method is when the low-rank assumption of attention activations is violated or weaker than expected? What happens to performance and convergence in such scenarios?\n\nLIMITATIONS:\nNothing too strong to push on.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-08"}