Title: Tensor-Parallelism with Partially Synchronized Activations
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models

Distributed Training

Tensor-Parallelism
Abstract: Training and inference of Large Language Models (LLMs) with tensor-parallelism requires substantial communication to synchronize activations. Our findings suggest that with a few minor adjustments to current practices, LLMs can be trained without fully synchronizing activations, reducing bandwidth demands. We name this “Communication-Aware Architecture for Tensor-parallelism” (CAAT-Net). We train a 7B parameter CAAT-Net model and show that tensor-parallel communication can be reduced by up to 50% with no significant drop in pretraining accuracy across nearly all evaluated benchmarks. We also experiment with smaller 130M and 1.1B models to show the robustness and scalability of our method. We find that, in some scenarios, validation loss can even improve when reducing communication. Finally, we demonstrate how CAAT-Net accelerates both training and inference workloads across various settings and model sizes.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces CAAT-Net, a method for LLM training and inference ini a distributed setup (tensor parallelism) in a more efficient way that reduces communication overhead. Tensor-parallelism in LLMs requires extensive communication to synchronize activations across devices, which becomes a bottleneck as models scale. CAAT-Net addresses this by implementing partial channel-reduce where only a subset of activation channels are synchronized thereby decreasing bandwidth demands without less impacting model accuracy. To enable this optimization the paper proposes the architectural changes on synchronization point, including specific modifications to MLP and attention layers and gives implementation considerations like forward-backward mismatch and the necessity of fp32 accumulation for numerical stability. Authors state that CAAT-Net can achieve speedups in both training and inference for 1B and 7B parameter models.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n\n1. While tensor parallelism provides more computability and larger scalability to model training, the distributed setup is often bottlenecked by communication between devices, e.g., all-reduce. This problem becomes significant in larger scale training working with a number of devices since the number of communication linearly increases. This paper addresses an important issue in the efficiency of training large models.\n\n2. Clear background explanation about tensor parallelism with and without all-reduce operation\n\n\nWeaknesses\n\n\nAlthough the paper tries to address the important problem, I have a lot of concerns on evaluations.\n1. The paper only shows accuracy results on p=0.5 setup on tensor parallel 8. While this setup well preserves the results compared to the baseline, it is not sufficient to say that all the speedups in Figure3 are meaningful results.\n2. The evaluation on zero-shot tasks typically generate a single token which aren’t sufficient to validate the quality of trained model. Showing the impact on generation accuracy when the model is required to produce a longer sequence could firmly validate the effect of partial channel-reduce.\n3. The evaluated models are limited to LLaMA2 and TinyLlama models. Testing on other recent models (e.g., Llama 3.3, Qwen 2.5, …) could validate its generalizability.\n\nQUESTIONS:\n1. What if the normalization method applied on the transformer block is not RMSNorm? Is CAAT-net still applicable to that cases?\n\n2. I’m confused about the statement that even though the implementations are mathematically identical, there is a large gap in training convergence. Why does this problem come up with 16bit precision although the two implementations are identical? If this problem is from the range of fp16, can bf16 also solve this problem?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces CAAT-Net, a novel approach to reduce the communication overhead in tensor-parallel training and inference of LLMs. Traditional tensor-parallelism requires full synchronization of activation tensors across devices using expensive all-reduce operations. CAAT-Net relaxes this constraint by synchronizing only a subset of activation channels. Thereby it reduces communication costs while maintaining model accuracy. The authors validate their method through extensive experiments on LLaMA2-7B, TinyLLaMA, and a smaller 130M-parameter model, and analyze how the synchronization factor and tensor-parallel degree affect performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. This paper proposes an effective modification to tensor-parallelism that halves communication cost without requiring model compression or accuracy trade-offs.\n2. Experiments on models of different scales and tasks showing comparable or improved zero-shot accuracy.\n3. Benefits observed in both training and inference, which is crucial for deployment.\n\nWeaknesses:\n1. Requires non-trivial changes to standard tensor-parallel frameworks such as Megatron.\n2. While results on 7B models are encouraging, it remains to be seen if the same gains hold for 100B+ parameter models, where communication bottlenecks are most severe.\n3. Lack of experiments on Nvidia GPUs, which limited the generation of this paper.\n\nQUESTIONS:\n1. How well does CAAT-Net interact with other techniques such as quantization, gradient checkpointing, or activation compression?\n2. You mention that inference must use the same tensor-parallel configuration and p as training. Is there a principled way to support variable p at inference time or to convert a CAAT-Net model back to a fully synchronized model (p = 1) without retraining?\n\nLIMITATIONS:\nPlease refer to weakness and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes CAAT-Net, a Communication-Aware Architecture for Tensor-parallelism, designed to reduce the communication overhead inherent in training and inference of Large Language Models (LLMs). Traditional tensor-parallelism fully synchronizes activation tensors across devices, requiring substantial communication. CAAT-Net introduces partial channel–reduce operations that synchronize only a subset of channels, significantly reducing communication demands without compromising model accuracy. This design achieves up to a 50 % reduction in communication payload (as shown on Llama2-7B and TinyLlama) while preserving model accuracy. The paper’s channel-level communication analysis is innovative, but the manuscript would be strengthened by a more thorough empirical evaluation: additional metrics (e.g., end-to-end throughput, latency breakdowns, scalability across more model sizes) and ablation studies are needed to fully characterize CAAT-Net’s behavior under diverse workloads.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n### Quality\n- The paper rigorously formulates the **partial channel–reduce** operation for both MLP and attention layers. \n- Numerical stability concerns (e.g., fp32 accumulation) are clearly identified and addressed in the implementation (§4).\n\n### Clarity\n- Visualizations highlight when and how the all-reduce shifts between forward and backward passes under partial synchronization.  \n\n### Significance\n- By synchronizing only a fraction \\(p\\) of channels (e.g.\\ \\(p=0.5\\)), CAAT-Net achieves up to **50 %** communication reduction and delivers **14 %** throughput gains on Gaudi3 clusters.  \n\n### Originality\n- **Channel-level synchronization** is a novel axis.\n\n---\n\n## Weaknesses\n\n### Quality\n- **Figure redundancy**: Figures 1 and 2 convey very similar information. Merging or re-annotating them would improve flow.\n- **Numerical details**: The impact of fp32 accumulation on specific normalization equations—especially under fp8 training—remains underexplored.\n\n### Clarity\n- The discussion around fp32 accumulation flags its importance but omits which exact equations need 32-bit restoration. A detailed breakdown would help readers reproduce results.\n- Overall writing can be improved a lot.\n\n### Significance\n- Even though the channel level is a new axis, but the author lacks more comprehensive analysis (e.g., each channel's / head sensitivity to the all-reduce). There is no reason we conduct all-reduce for continuous dimension.\n- Evaluation is limited to zero-shot accuracy and validation loss on 1 B/7 B models. Extending experiments to larger models (e.g.\\ Llama 3.3, Llama 4, Qwen 3) and to NVIDIA GPUs with NVLink would strengthen real-world applicability.\n\n### Originality\n- The related-work comparison lacks a head-to-head ablation against techniques like Top-K activation compression, Domino fusion and Ladder Residual.  \n\n### Generalizability\n- Inference-time tensor-parallel dimension mismatches (§3.3) pose deployment challenges. Empirical results quantifying accuracy vs. synchronization fraction \\(p\\) during serving would clarify practical trade-offs.\n\nQUESTIONS:\n1. **Overlap Efficiency Claim**  \n   - **Issue:** Partial synchronization is claimed to reduce computation–communication overlap.  \n   - **Request:** In my understanding, even the model is larger, we still have enough computation time to overlap with communication. Overlapping is still effective here. Please correct this part or show some more robust evidence here.\n\n2. **Comparison to Desynchronized-Attention / Ladder Schemes**  \n   - **Issue:** No head-to-head comparison with existing ladder or desynchronized-attention approaches.  \n   - **Request:** Include a representative experiment (or cost model) contrasting bandwidth reduction vs. accuracy/throughput trade-offs between CAAT-Net and other methods (e.g., **Ladder-Residual**).  \n\n3. **Justification for Partial Synchronization**  \n   - **Issue:** It’s unclear why synchronizing only a subset of channels preserves correctness.  \n   - **Request:** Provide:\n     - A theoretical argument (e.g., variance or information-flow analysis), or  \n     - An ablation study showing that removing random channel subsets up to fraction \\(1 - p\\) incurs < 0.1 % accuracy loss. \n     - Author shall design more fine-grained algorithm.\n\n4. **Channel-Selection Strategy & Figure 2.C**  \n   - **Issue:** Figure 2.C picks contiguous channels for all-reduce without justification.  \n   - **Request:** Add an ablation comparing contiguous vs. interleaved channel selection (e.g., a hybrid split into B\\(_{11}\\), B\\(_{12}\\), B\\(_{13}\\) with communication on B\\(_{11}\\) and B\\(_{13}\\)).  \n\n5. **Scalability & Inference Flexibility**  \n   - **Issues:**  \n     1. How does the optimal sync fraction \\(p\\) change for much larger models (30 B+ parameters)?  \n     2. Can inference run on a different device count than training without full fine-tuning?  \n   - **Requests:**  \n     - (a) Provide cost projections or small-scale experiments on a 30 B-parameter model.  \n     - (b) Propose a strategy (e.g., dynamic \\(p\\) adjustment or lightweight channel-wise calibration) to avoid complete retraining when device counts change.  \n     - (c) I don't think it's appropriate to put this part in future study.\n\nLIMITATIONS:\nPartially yes. The authors acknowledge the limitations, particularly regarding inference flexibility. There is no so muchnegative societal impacts or broader generalizability issues would strengthen the paper. But The authors are encouraged to explicitly address whether partial synchronization could introduce biases or reduce robustness in specific use cases. And the authors shall compare related methods more directly to reveal this work's potential limitations.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces CAAT-Net (Communication-Aware Architecture for Tensor-Parallelism), a new method for reducing communication overhead in tensor-parallel training and inference of Large Language Models (LLMs). The key idea is to partially synchronize activation tensors across devices instead of performing full all-reduce operations. The authors propose a new primitive called partial channel–reduce, which synchronizes only a portion of the hidden dimension channels (controlled by parameter p), thereby reducing communication bandwidth by up to 50%.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\n\n1.The paper proposes CAAT-Net that applies partial channel–reduce to reduce synchronization overhead. CAAT-Net lowers communication volume during tensor-parallel training and inference.\n\n2.The paper describes modifications to the backward pass required by partial synchronization, including changing the placement of the all-reduce operation and using fp32 accumulation to mitigate numerical differences.\n\n3.The paper provides a design of the partial channel–reduce mechanism, including its integration into both MLP and attention layers. \n\nWeaknesses：\n\n1.In section 3.2, the variance scaling adjustment in partial synchronization is briefly mentioned, but lacks deeper theoretical justification or ablation to validate its necessity and effect.\n\n2.The experimental evaluation focuses on a limited set of model scales (130M, 1B, 7B) and primarily on the LLaMA architecture. Including results on other widely used open-source models and stronger baselines would further strengthen the empirical claims.\n\n3.In the partial channel–reduce mechanism, only a subset of channels is synchronized across devices. It is not clearly discussed how the private (unsynchronized) activations contribute to downstream layers.\n\n4.The paper frames CAAT-Net as an architectural design, but the actual modification primarily involves a replacement of all-reduce with partial-reduce and a slight adjustment to initialization. It is unclear to what extent this constitutes a new “architecture” versus a communication primitive. Clarifying this distinction and justifying the architectural framing would improve clarity.\n\nQUESTIONS:\nQuestions：\n\n1.Could the authors provide more theoretical or empirical justification for the √r scaling of private channels at initialization? Would removing this scaling lead to divergence or worse convergence? A controlled ablation would help clarify.\n\n2.Why are direct comparisons to pipelining approaches or Top-K compression baselines not included? The controlled experiment could quantify whether partial synchronization outperforms in realistic scenarios.\n\n3.Have the authors considered evaluating CAAT-Net on other open-source architectures? This would help demonstrate whether the proposed partial synchronization mechanism generalizes beyond the LLaMA-based models currently used in the experiments. Adding such results could strengthen the empirical validation.\n\n4.While Section 5.2 explores different p values on a 130M model, the main large-scale experiments (1B and 7B) fix p = 0.5 without further justification. It remains unclear whether p = 0.5 is optimal or just a safe default. A more systematic ablation of p on larger models would help clarify the tradeoff between communication cost and accuracy.\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}