Tensor-Parallelism with Partially Synchronized Activations

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. * This work was done while the author was at Intel. 39th Conference on Neural Information Processing Systems (NeurIPS 2025).

devices, the compute workload per device decreases, while communication payload per device remains relatively constant [7]. This means that the relative cost of communication grows as the compute-to-communication ratio decreases. In extreme cases, communication time can overcome computation time, and thus dominate the training process. For these reasons, even the largest language models are typically trained with a tensor-parallelism dimension of 8 [4], utilizing fast intra-node communication for the heavy all-reduce operations.

Improving tensor-parallelism efficiency is even more important given that the growth in compute power exceeds the growth in communication bandwidth [7]. This trend should further expose communication time in large-scale training. Minimizing tensor-parallelism communication can enable better hardware compute utilization and reduce overall training costs, especially when extending tensor-parallelism across nodes. This is in line with the recent trend of building multi-node systems with high bandwidth communication.

In traditional tensor-parallelism, the activation tensors after communication are identical on all devices, i.e., fully synchronized. In this work, we show that LLM training can converge without fully synchronizing the activation tensors in the all-reduce operation. This means that we allow activations to vary on different devices after communication. We show that without full synchronization, the current training practice needs to be slightly adjusted. Failing to do so leads to critical issues such as a mismatch between forward and backward passes and numerical issues, which often result in training divergence. Relying on this insight, we suggest the partial channel-reduce operation, in which only a subset of the channels in the hidden dimension of the activation tensors is reduced. Unlike regular all-reduce, activations are not identical on all devices after the partial channel-reduce operation. In the extreme case where no channels are synchronized in partial channel-reduce, the model resembles an ensemble, communicating only to compute the loss function and embeddings. In the case where all channels are reduced, the model is a vanilla transformer model. We introduce Communication-Aware Architecture for Tensor-parallelism (CAAT-Net) -a new model architecture that is tailored for tensor-parallelism by utilizing partial channel-reduce to decrease communication overhead. While CAAT-Net has a smaller communication overhead compared to an identical model with full all-reduce, the number of parameters and total compute stay the same.

We train a Llama2-7B model [8] with partial channel-reduce over 160B tokens and show that there is no significant degradation in nearly all evaluation benchmarks we tested, while reducing the communication payload by 50%. Furthermore, we train multiple variants of the 1.1B parameter TinyLlama model [9] and a smaller 130M parameter model. We study the effects of the number of synchronized channels and tensor-parallel dimension on accuracy. We find that a gradually reducing communication from full synchronization first yields a slight improvement in validation loss, but performance worsens when communication becomes too limited. Reducing the communication by 50% achieves either similar or slightly better validation loss for all models we tested. Finally, we show the training and inference speedup of our proposed method in various settings.

In summary, our contributions in this paper are as follows:

• We show that when using tensor-parallelism, LLMs can be trained without fully synchronizing activation tensors.

• We propose CAAT-Net, a novel architecture that significantly decreases communication traffic in training with tensor-parallelism by synchronizing only part of the activation tensors.

• We show that in various settings, CAAT-Net accelerates both training and inference, and achieves accuracy largely on par wit