Tensor-Parallelism with Partially Synchronized Activations
Authors: Itay Lamprecht
Asaf Karnieli
Yair Hanani
Niv Giladi
Daniel Soudry
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Models
Distributed Training
Tensor-Parallelism
Primary area: deep_learning
TL;DR: Training LLM’s with tensor-parallelism without completely synchronizing activations to accelerate training and inference.

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.

Tensor-Parallelism with Partially Synchronized Activations
Itay Lamprecht Asaf Karnielit Yair Hananit Niv Giladi* Daniel Soudry IItel, Israel
Department of Electrical and Computer Engineering - Technion, Haifa, Israel AWS Al Labs
{itay lamprecht,, yair www..haaanneinel.ccmm (ngiladi)@amazon com (asafkarnieli, daniel.. soudry)@gmaill.......
Abstract
Training and inference of Large Language Models (LLMs) with tensor- panal--paralleism requires substantial communication synchronize activations. Our findings sug- gest that with a few minor adjustments to current practices, LLMs can be trained without fully synchronizing activations, reducing bandwidth demands. We name this "Communication-Awar Architecture for Tensor censor-parallelisme (CAAT-Net). We train a 7B parameter CAAT-Net model and show that tensor- massor-araalll communi- cation 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. B 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.
Introduction
As Large Language Model (LLM) training continues to scale, often requiring the use of hundreds even thousands of devices, the need for efficient distributed training and inference techniques is constantly growing. -eaalo--arallellim T is a critical component the training of many well- known LLMs, such GPT-3 2]. Llama 3 4 and BLOOM (5), enabling the efficient utilization of
hardware resources to support models with billions of parameters. By partitioning weight tensors multiple devices, tensor-parallelism allows each device store only a fraction of the model's weights during computation. This significantly reduces the memory footprint per device, making feasible train extremely large models hardware with limited memory capacity. The efficiency of tensor- paraooraalllllismmmmmm in reducing training memory has made it a cornerstone technique for distributed training. During inference. where the LLM response time critical, distributing the computation over multiple devices using tensor- paralleiaal can significantly reduce latency.
Much of the total time spent training using LLMs is consumed by the communication overhead inherent in tensor-paralled parsllpprallll,sm, which arises from the need to synchronize activation or gradient tensors. Specifically, synchronization is done by all-reduce operations, where tensors from all devices are summed into a single tensor, which is copied to all devices. This operation occurs multiple times in each forward and backward pass and constitutes substantial fraction of the total workload time.
recent approach to addressing this issue is pipelining communication and computation (6) in such a way that the communication is overlapped with the computation and the footprint of communication is reduced. However, this approach is limited: when tensor- paralllsaalll is expanded to use more
"This work was done while the author was Intel.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

devices, the compute workload per device decreases, while communication payload per device remains relatively constant 77. This means that the relative cost of communication grows as the compute-to-commune co-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 tensor- cansor-paralllisin dimension of 8 4 utilizing fast intra-node communication for the heavy all-reduce operations.
Improving passor-parallelism efficiency is even more important given that the growth in compute power exceeds the growth in communication bandwidth 77. This trend should further expose communication time in large-scale training. Minimizing censor-parallelin better hardware compute utilization and reduce overall training costs, especially when extending tensor-pan parall-salllllismm across nodes. This is in line with the recent trend of building multi-node systems with high bandwidth communication.
communication can enable
In traditional censor-paralleiim. 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 a vanilla transformer model. We introduce Architecture for panso--paralllismm (CAAT-Net) that is tailored for tensor-paralle eenoo--arallelism by utilizing partial channel-reduce to decrease communication overhead. While CAAT-Net has a smaller communication overhead compared to an identical model with full relure, the number of parameters and total compute stay the same.
Communication-AAwenn
new model architecture
We train a Llama2-7B model S with partial channel-reduce over 160B tokens and show that there 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 smaller 130M parameter model. We study the effects of the number of synchronized channels and tensor-parae parallel dimension 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 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-parallelinn LLMs can be trained without fully synchroniz- ing activation tensors.
We propose CAAT-Net, a novel architecture that significantly decreases communication traffic in training with tensor- tensor-parallelism 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 with fully synchronized training.
2 Related Work
The challenge of efficient training and inference on a large scale has attracted much attention, both in the engineering and research fronts, in close correspondence with each other. While most of the literature focuses on reducing data-parallel communication, there have been a few works that focus on tensor-parallelism. Our approach is orthogonal to all methods covered in this section.
Pipelining computation and communication. Some methods accelerate training with tensor- parallelism by overlapping communication and computation, 66 10-12. In Domino DeepSpeed, the training batch split into smaller pieces, and data dependency is broken such that communication is not on the critical path and can be overlapped with computation. Alternatively, Ladder Residual

---

pipelines communication and computation by changing the transformer architecture, such that each layer uses an earlier version of the imput from two layers back. This allows to compute while the previous layer's results are still being communicated. While these approaches achieve speedup in inference and training, they are limited - there are many cases in which the compute time is not sufficient to fully hide tensor-parte rensor-parallel communication. For example, with growing tensor-parte parallll dimension, the computation per device is reduced while the communication is not, so overlapping efficiency is reduced as well. Furthermore, other parallelism types interfere with tensor-par tensor-parallel related communication. One of many examples is contextpparallelism where communication time alone exceed compute time [13] When adding censor-parallelie after attention blocks that use context- paralltllallllssmmm there is no compute left to pipeline both types of communication. For this reason view our method, which reduces the total censor-paralle communication, complementary to pipelining methods.
Additionally, in NVIDIA GPUs such as H100/H200, pipelining communication and computation comes with a cost - the Stream Multiprocessors (SM), which perform the computation, are also needed for communication. When pipelining communication and computation, some SM cycles are spent on communication-rlllted tasks instead of computation. Furthermore, while pipelining computation and communication can accelerate training and inference, our method reduces total communication, and can potentially improve other aspects such as power consumption.
Compressed communication. Another approach is to address the communication overhead by compressing the communication itself at the censor-parallel or appline-parallle stages [4 [S]. This can be done using various methods, such as auto-encoders or sending the Top-K elements of the tensors. While compression succeeded in other aspects of optimization, it has yet been demonstrated that compression works for vensor-parallede severe degradation in model accuracy with even minimal compression using Top-K or random masking to compress activations (our observations can be found in Section 5.3) Additionally, compression introduces additional computation overhead which can reduce the speedup achieved
to the best of our knowledge. In fact, we observed
by reducing communication. Our method differs from these methods because it does not compress
activations. Although not all activations are shared between devices, those that are not shared are still used.
Asynchronous optimization. A different approach, mainly focused on mitigating overhead intro- duced by data-parall parralel communication, introduces asynchronism to the optimization by not waiting for all communication or straggling devices. In this case, different devices store all of the model weights, but each device can hold slightly different copy. Asynchronous training is inherently more scalable than synchronous training by being robust to all kinds of worker and communication faults. However, this scalability comes with a price: convergence difficulties and generalization deterioration [6]. In contrast, when training with censor-parallelimm each device stores different subset of the model weights. Therefore, while we do not completely synchronize activations, do not encounter the problem of multiple workers storing a different copy of the same weight, and training remains synchronous.
Post-training techniques. There have been many attempts to reduce tensor-pane aasoor-arralll communication using post-training techniques, such as quantization [17 [8] and selectively removing complete communication points after the attention operation [19. While these techniques are promising, they are relevant only for decreasing the communication bottleneck during inference.
3 CAAT-Net
To reduce communication bandwidth in LLM training and inference, we replace the all-reduce operation, designed to synchronize activations between devices, with a partial channel-reduce operation, such that only a subset of the channels in the hidden dimension is synchronized between devices. In partial channel-reduce, channels that are sent to other devices are referred to as shared channels, while those that are unique per device are private channels. The shared channels are chosen at initialization be the first h p channels, where h is the hidden dimension size and p is a synchronization factor controlling how much of the hidden dimension is synchronized. For p = 1 the partial channel-reduce turns into a full all-reduce operation (no private channels). The operation is visualized in Figure (c). In this section we present a novel communication-awanl model architecture which utilizes partial channel-reduce to accelerate LLM training and inference, A theoretical speedup analysis of our method is available in Appendix C
3

---

Tensor Parallet
CAAT-Net
Original Model
Device
X1
Device
Z1
X
Z
X
P
Values sent to device
Device
X2
Device
Z2
Values calculated on device
Figure CAAT-Net model architecture. Left. We exemplify our approach on two-layer fully- connected neural network single device. Middle. When using ansor-parallelism with two devices, the input activation is identical on both devices. Each device uses its own set of weights to multiply the inputs, yielding intermediate activations, which are then reduced into identical copies of Z on both devices. Right. CAAT-Net receives different input activations. X1 and X2, for Device and Device 2, respectively. These yield intermediate activations that partially synced between the devices, producing 21 and Z2 on Device and Device 2, respectively. Private channels are marked P. and shared channels marked S.
3.1 Architecture
While changing the communication primitive from all-reduce to partial channel-reduce seems like minor change, it effectively changes the model architecture. Partial channel-reduce fundamentally alters the model's computation graph and information flow. This distinguishes CAAT-Net from communication optimizations like pipelining or compression, which preserve or approximate the underlying mathematical operations. Consequentially. the Multi Layer Perceptron (MLP) and attention layers must be defined differently when taking partial channel-reduction into account. In this section define the MLP using partial channel-reduce. The definition of the attention operation described in Appendix A.
Standard all-reduce in MIP. Before mathematically describing the MLP using partial chan- nel-reduce, we describe the MLP using regular all-reduce a tensor-paralle setting. For simplicity, we examine case where the lensor-paralles dimension is 2. For the general case, refer to Appendix Given consecutive MLP weight matrices and B, input tensor X and an activation function a, the output of the MLP layer Z is given by
A.2
(X
YB
(1)
where Y is the output of the first MLP layer. The basic architecture of a transformer MLP is visualized Figure (left. When using tensor-parallelism. we partition A along columns, i.e.,
A = [Ai, A2]
(2)
so the computation of Y.
Y1 = a (XA1)
Yg2 = a (XA2)
Y= [1, Y2)
(3)
We partition B along its rows, i.e.
B
B
(4)
B2
and, after an all-reduce, the output of the MLP is
7 = Y, B1 + V2
(5)

---

where Y, is computed on the first device and YzB₂ is computed on the second. The implementation of all-reduce with tensor- tensor-parallee parallelism is visualized in Figure [ (middle).
Partial channel-reduce in MLP. We first note that with partial channel-reduce, the input to the MLP will be different per device (as the previous attention layer will also use partial channel-reduce in its output). Therefore, the inputs to each device are denoted X1 and X2, as is shown in Figure |
(right). Thus, the output of the first linear layer when using partial channel-reduce is:
Y1 = a (X1A1)
Y/2 = a (X2A2)
Y = [Y,, Y2)
(6)
Next, we need to partition B along both its columns and rows
Bu B21
B
B12 B22
Then, the outputs of the MLP with partial thannel-reduse denoted Z1 and Z2, are different per device and are calculated by
Z1 V Bn + Y2112, Y B21] Zz M. + Y2 Y₁ B22)))
(8)
where Bu and B21 are on the first device, and B and B22 are on the second device. The values in Bu + B22 are in the shared channels, and B21 and B2 are both in the private channels, which are visualized in Figures andd]
3.2 CAAT-Net Inference
Models trained with CAAT-Net use partial channcl-reduce in inference as well. This leads to communication reduction and speedup when serving models, as we show in Section [5.2
If the model is served with the same tensor-paree censor-parallel dimension that it is trained with, serving the model is straightforward. This is different when using the model in inference with censor-paralle dimension different than in training. As an example, we examine a scenario where the model is trained with a tensor-parallel dimension of 2, and inference is performed with single device. In training, the output of a transformer layer is different on each device. In inference, one device needs handle both of these copies. This can be done using 'logical devices'. In this case, a single physical device can simulate multiple tensor-parallel ranks, and sequentially calculate the output of each layer for every 'logical' tensor-parallel device. Furthermore, the partial channel-reduce is replaced with local summation inside the single device. It is also possible to return to a full all-reduce operation during inference through fine tuning, increase the value of p to Then, a regular transformer architecture is achieved, and inference can be run on any tensor tensor-parallel dimension.
4 Partial Synchronization: Implementation
To train LLMs without full synchronization of activation tensors, there are two main adjustments that must be done to current training frameworks (I. The first is adjustment to the backward implementation of tensor-pan parnll-paaalllissm, which necessary to avoid arward-backwaad mismatch in training. The second is accumulating tensors in 32 bit (fp32) in the all all-reduce of the backward pass. We find empirically that this is necessary to avoid numerical issues training.
Forward serraddbache bbakwardd mismatch. To train L.MSs with partial activation synchronization, we need to further examine the backward pass of a transformer model. In traditional L.MM training, the all-reduce operation in the backward pass is done on the neural gradients produced by an MLP or attention layer, as shown in Figure 2.. Explicitly calculating the gradients propagated through the network, we find that the all-reduce operation can be applied in multiple locations in the backward pass without altering the backpropagation algorithm. We also find that in the case partially synchronized activations, the reduction operation can only be applied in one place. To show this analytically we examine the input to an MLP layer on device m, denoted X,, as a function of the output of previous attention layer on device m before the all-reduce operation, denoted Zm.. with R being the residual connection, have
Xm, = Norm
Σ Zm + R m70
(9)
5

---

(
Z
Norm,
X1
(c)
Residual
p.h p)-
Token
Z1
Reduce
7, +72
7
Z
Norm,
Residual
Batch
Partial Channel
Attention
MLP
Hidden
Shared Channels
Private Channels
(b)
Z1
Norm,
X1
Residual
Zz
21 Z2
Zz
h
Norm, Residual
7.2
X2
Figure 2: Partial synchronization and partial channel-reduce. (a) Vanilla transformers in current training frameworks. The operation f is an all-reduce in the forward pass and identity in the backward pass. The operation an all-reduce in the backward pass and identity in the forward pass. (b) With partial synchronization, h denotes the reduction operation in both the forward and the backward pass, since both must be done at the same location. Synchronization of the normalization function parameters is necessary. (c) Partial channel-reduce with parameter p over 2 devices.
where Norm can be any normalization function (typically LayerNorm 20 or RMSNorm 21 We
examine the loss as a function of X, for m, and Xm as function Z1, and calculate derivative of the loss function w.r.t Z1. For simplicity, we ignore below the residual connections (in Appendix B we show the derivation including the residual). Using the chain rule:
OC (Xi (Zi) ,.., XM (Z1))
(Xm) OXm (21)
OL (X
0x
Σ
Σ
(10)
0Z1
Ox,
0Z,
Ox,
az
ax because in In the second equation, recalled that 02, are equal for all m (denoted now 82, censor-parallelin with full all-reduce operations X1 (71) X₂ (Z1)
axmm
XM (71). and thus
their derivative w.r.t. Z are also equal. Here, the summation over m is the all-reduce operation the backward pass. Our key insight here is that, since the matrix multiplication distributes over addition, i.e.
Σ m OX
OX
Σ
@0(Xm))
ax
(1I)
OZ
OX
3Z
the all-reduce operation can be used before back- back-propagating through the normalization function (right side of the equation), after the normalization function (left side of the equation). In standard training frameworks, such as Megatron-LM, the all-reduce operation is used before propapproag through the normalization function (i.e., the 'g' operation in Figure 2p). When training with partial synchronization, the assumption that 2x aZ.. equal for all longer holds. For this reason, the all- reduce operation in the backward pass must occur after the calculation of the normalization function derivative (i.e., the operation in Figure 2b) so it would match the location of the summation the forward pass (i.e., the 'f" operation in Figure 2). to get the correct gradients.
Similarly, examine the normalization function parameter update, denoted B, with a full all-reduce operation:
ac (X1 (3)
Xas (B))
@C (Xm)
axmm
Σ
(12)
08
8Xm
0
m
In the gradient descent step, after each iteration, before updating the weights, the updates 03 need be synchronized with all-reduce operation. In popular training frameworks, the fact that ax.,. is
"3

---

identical for all m devices is used once again. There is no need for gradient synchronization after each step, if the all-reduce happens before backpropagating through the normalization function:
Σ {
OL ( 0X
(0xmm))
(13)
05
m where, again, ae are equal for all (denoted as ag ). To summarize, when using full all-reduce, the following are equivalent:
Reduce the neural gradients before they backpropagate through the normalization function.
Reduce the gradients after they backpropagate through the normalization function, and reduce normalization function parameters before the weight update.
However, when training with partial synchronization, such as partial channel-reduce, only the second approach is possible. Our altered backward implementation detailed in Figure 2b. It is important to note that the number of normalization function parameters per layer is typically ((h), where h is the hidden size, and the synchronization of normalization function parameters happens after every step. For these reasons, the additional communication of normalization function parameters is negligible in comparison to the tensor-paralle all-reduce operation.
Numerical stability via 32 bit gradient accumulation. In the previous section we presented two different implementations of the backward pass. If we fully synchronize activations in the forward pass, these implementations are mathematically identical. Despite this equivalence, we found in our experiments that there can still be large gap in training convergence when using our alternative implementation, even when activations are fully synchronized. This gap arises due to numerical differences between the implementations, and it can be completely mitigated by accumulating the reduced gradients in 32 bit precision instead of 16 bit precision. While LLMs are conventionally trained in 16 bit precision, there are some operations where 32 bit precision is critical, such as in the accumulation in matrix multiplication operations, or normalization layers. While the tensor-parallel gradient all-reduce typically not one of these operations, we find that when moving the all-reduce in the backward pass, 32 bit accumulation is crucial. It is important to note that while accumulation needs to be done in 32 bit precision, the communication itself can be done in 16 bit precision, and values are simply upcast after communication and before accumulation. Loss curves comparing training experiments with all-reduce accumulation in fp32 and bfloat are available Appendix
Private channel scaling. Partial channel reduce leads to differences in the statistics of shared and private channels, which affects signal propagation in the network. Consider the case of partial channel-reduce with 2 tensor-paralle reduction (Yn B1, B1,, Yi B21) has zero mean and variance o² and are independent between devices. The activation variance for shared channels
devices. Assuming each element in the MI.P outputs before
Var(Y, + Y2B12) Var(Y;Bii) + Var(Y2B12) 2003.
(14)
and for private channels is:
Var(Y,B21) = =
(15)
This variance mismatch can lead to uneven signals across channels in both the activation and gradient caleulations. To correct this mismatch, we multiply the activations in the private channels by corrective factor of /2 (or /r we train with tensor-parallel dimension of r). This way, after the partial
the activations have identical variances over all channels. Ablation studies for private channel scaling available at Appendix [E.]] While we find slight improvement in validation loss across values of p. private channel scaling is not mandatory to achieve sufficient accuracy results.
5 Experiments
Experiments* were conducted with Intel Gaudi3 HPU accelerators. Gaudi3 has 128GB on- bo-bo memory. Each device has 525 GB/s intra-node connection and 75 GB/s inter-node connection.
* Code is available at http://gihnalcoommiim//Mgaaarrr
7

---

Table CAAT-Net vs baseline: Zero-shot aceuracy after pretraining, 7B parameter models, with
= 0.5 and tensor- tensor-parallel 8.
Model
LAMBADA (acc)
Hellaswag (acc)
WinoGrande (acc)
PIQA (acc)
Baseline
61.34 + 0.68
45.85 + 0.50
61.48 + 1.37
72.91 + 1.06
CAAT-Net
61.05 0.68
46.10 + 0.50
62.19 + 1.36
72.86 + 1.04
OpenBookQA (acc)
BOOL-Q (acc)
WikiText (ppl)
Validation Loss
Baseline
26.60 = 1.98
64.89 = 0.83
12.51
1.01 1.00
CAAT-Net
24.00 = 1.87
62.51 = 0.85
12.46
5.1 Large Scale Training
We trained a variation of Llama2-7b 8 with partial channel-reduce. Training was conducted from scratch over 160B tokens from the RedPajama dataset, spanning 8 nodes, each containing accelerators. We chose partial channel reduce hyperparameter of p 0.5 and tensor- massooraaaalle dimension 8, with all other hyperparameters identical to those used in training of the original model. We choose these values of p and tensor-parre censor-araallll dimension an example which expected, based on our ablation studies, to yield good accuracy results, while significantly reducing network bandwidth.
We evaluate the model's performance on a diverse set of common sense tasks selected from the Language Model Evaluation Harness framework [22] The chosen benchmarks include tasks such as HellaSwag [23], WikiText103 (24, LAMBADA 255, WinoGrande (26). PIQA [27], OpenBookQA
28 and BoolQ [29]. The zero-shot results are available in Table D All accuracy results are not
statistically significant, except for Bool-Q, which is marginally significant (p-value = 0.046 using Welch's 's t-test).
To further study the scalability and robustness our method, we train 130M and 1.1B models in multiple settings. The 1.IB model used is TinyL.aam [. and is trained over 100B tokens. The 130M parameter model a small custom model based on the L.LaMA architecture. It has 16 attention heads and hidden dimension size of 768. It was trained with initial learning rate of 6 10" with the Adam optimizer. The training was performed with a global batch size of 256 and a sequence length of 1024. The architecture consists of 12 transformer layers with multi-head attention. It is trained over 7.8B tokens. Both models are trained on the RedPajama dataset, using the GPTSentencePieee tokenizer.
1.10
1.94
Normalized Loss
130M
TP
1B
1.92
TP 8
1.05
7B
1.90
16
1.88
1.00
1.86
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
p
p
Figure 3: Training accuracy in multiple scenarios. Left. Validation loss of 130M and 1.1B models for different values of p. and of the 7B model with 0.5, normalized to the loss at p 1. Right. Validation loss for the 130M model with varying values of p and cenoo-paraallll dimension (TP).
We train both models with varying values of p. with tensor-parallel dimension of 8. Results are available in Figure 3 (left). We find that for large values of (over 0.75), the 130M and 1.1B models behave similarly. For intermediate values p (between 0.75 and 0.25), the smaller 130M model has a slight improvement in validation loss compared to the baseline, while the 1.IB model is close to or slightly above the baseline. For smaller values of p, there is degradation for both models. Furthermore, we added the 7B model we trained with p = 0.5, which shows slight improvement w.r.t the baseline, showing potential in further scaling our method. Zero-shot accuracy for all I.1B experiments is available in Appendix E.2
8

---

We also study the effects of tensor parallel dimension on the 130M model. Results are available in tensor-parallel parallel dimensions, gradually decreasing p from initially leads to slight improvement in validation loss, but performance deteriorates once p becomes too low. As the tensor-parallel dimension increases, this degradation begins at higher value of p. For all tensor-parallel dimensions, p = 0.5 does not degrade or improves validation loss.
Figure 3 (right). For all tensor-
Finally, we train an additional model, other than the models from the Llama architecture, to show the
robustness of our method. Specifically, we trained GPT3-XL 2 with tensor-parallel 8 and p 0.5, on the RedPajama dataset using the GPTSentencePiece tokenizer, rotary positional embeddings and with global batch size of 512. We trained for a total of 50B tokens. Similarly to results on Llama, we find that CAAT-Net maintains accuracy results. The full results are available in Appendix E.2
5.2 Speedup
We examine Llama2-7b with partial channel-reduce measure inference and training speedup. In this section we show relative speedup (i.e., speedup compared to the baseline). For absolute measurements, see Appendix [E.3] Furthermore, for inference, we experiment with larger models (34B and 70B) and achieve similar results to those reported in this section. Full results for these models are available in Appendix E.4
Inference. To measure inference speedup, we report the Time-to-First-Tokee (TTFT) for varying batch sizes, using a prompt length of 2K tokens. We use partial thannel-reduce with tensor- tensor-parallel 8 and 16. We compare results against the baseline, which is Intel's Optimum-Har Habanna without any of adjustments. The results are shown in Figure (bottom)) For tensor-parallel 16 and p 0.25 we measure a maximum speedup of 26% with batch size 32. In the large scale training setting of p = 0.5 and censor-parallel 8, which we show does not degrade accuracy, we obtain a maximum speedup of 14% for batch size 32.
Furthermore, to show that accelerating workloads using CAAT-Net is not specific to Intel Gaudi hardware, we conduct experiments on NVIDIA hardware, Our experiments were done using the gpt-fast repository, and consisted of replacing all-reduce with partial channel-reduce during inference. We conducted experiments on 8 NVIDIA H100-80GB: 80GB-HBGBM and NVIDIA A A100-SXM4-800 80GBB both with NVLink. Results are shown in Figure 4 (top right). We find that our method speeds up inference TTFT by up to 13% on NVIDIA hardware with p = 0.25 and tensor- aansor-prralll dimension 8. Speedup Gaudi is more significant in our experiments, but believe further software performance optimizations on NVIDIA hardware can bridge this gap. Full measurements, including for additional batch sizes, are available in Appendix E.6
Alongside accuracy considerations, practitioners should consider hardware constraints when selecting p. For example, Gaudi3, speedup results were more significant when selecting p such that the communication volume a multiple of the communication buffer size. Specifically, when experimenting with Llama2-7B with batch size 16, we identified that the optimal performance speedup is achieved with a value of p which is a multiple of 2 (i.c., the number of channels reduced is a multiple of 16). As a result, selecting p 0.703125, which is equal to 90- 2-7, achieves better speedup than p = 0.7.
Training. To measure training throughput speedup, we report the tokens per second (TPS), for tensor- paralell dimensions 4, 8, and 16, with varying values of p. Tensor- -araalel experiments were conducted with micro-batch size of 4, while tensor-parallel 8 and 16 experiments were conducted with micro-batch size of The results are shown in Figure 4 (top left). We compare results against the baseline, which is Intel's Megatron-LM fork without any of our adjustments. For training in tensor-pun paasorrparaalll 16, our method can improve throughput by up to 14% when selecting p = 0.25 and by 9% when selecting p 0.5.
All experiments were conducted using private channel scaling, as described in Section 4 Disabling provides an additional 1-2% speedup, with only minor accuracy degradation (see Appendix [E3) making training without private channel scaling reasonable alternative. It is important to note that in training speedup experiments, for all values of and the baseline, accumulation in the backward pass was done in 16-bit precision due to technical implementation limitations. For this reason, the training throughput results in this section should be seen as performance projection. We do not expect accumulation in fp32, if communication is kept 16-bit precision, to significantly affect the speedup reported in this section.
9

---

Training Speedup
Inference Speedup
30
p-0
p-0.5
Baseline
30
H100
A100
Gaudi3
Speedup [%6]
p=0.25
+ p=0.75
Speedup [%)
20
20
10
10
0
4
8
16
0.00
0.25
0.50
0.75
1.00
Tensor-Parallelism Dimension
TP 8 Inference Speedup
TP 16 Inference Speedup
75
p-0
pm0.5
Baseline
p-0
p-0.5
Baseline
Speedup [%]
40
p-0.25
p-0.75
Speedup [%)
p=0.25
p=0.75
50
20
25
2
25
2
2
20
2
2
2 Batch Size
2
2
Batch Size
Figure 4: Speedup in training and inference for a Llama 7B model. Top Left Training speedup for varying values of and tensor- tensor-parallel dimension. Top Right. Inference Time-To To-First- Token (TTFT) speedup on different hardware, using batch-size 128. Bottom Left. Inference TTFT speedup using tensor-parallel 8 as function of batch size, for different values of p. Bottom Right. Inference TTFT speedup using tensor-parallel 16 as function of batch size, for different values of p.
5.3 Comparison to compression methods
To compare our method to compression methods, we use two alternative approaches for compress- ing activation communication. While CAAT-Net preserves low validation loss, we report severe degradation for both compression methods, even with minimal compression. The first method is reducing communication using random mask. The values that are multiplied by O in the random mask not communicated. This is effectively applying dropout before communication. The sec- ond method using Top-K mask, which selects the Top-K entries in each token of the activation tensor and discards the others. In both cases, we apply the mask in the forward pass before commu- nication. To avoid a heward-bbkkwwrr mismatch, the same mask is applied in the backward pass as well. Furthermore, the Top-K K and random masks are less efficient in reducing communication than CAAT-Net. See Appendix D for details.
3.0
CAAT-Net
2.8
Top-K Random Masking
2.6
2.4
2.2
2.0
0.0
0.2
0.4
0.6
0.8
1.0
Figure 5: Comparison to compression meth- ods. Validation loss vs. p for 130M models using CAAT-Net, Top-k masking and random masking.
To be consistent with CAAT-Net experiments, we denote by p the probability that an activation is not zeroed (so if p 1, applying the mask is equivalent applying the identity function). We train the 130M parameter Llama model detailed in Section 5. for a total of 2.6B tokens. We experiment with different values of p. and report significant degradation for Top-K and random masks. Results available in Figure 5
Conclusion
In this paper, we show that with minor changes to current training frameworks, activations do not need to be fully synchronized in tensor-parallel communication for training to converge. We propose CAAT-Net, which significantly decreases tensor-parallel related network traffic LLM training with minor to no degradation in accuracy. We experiment with smaller networks to explore the effects of tensor-parallel dimension and the synchronization factor. Finally, we show how our method speeds up inference and training workloads.
10

---

Acknowledgments and Disclosure of Funding
The research DS was funded by the European Union (ERC, A-B-C-Deep, 101039436). Views and opinions expressed are however those of the author only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. DS also acknowledges the support of the Schmidt Career Advancement Chair Al.
References
[1] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron- Training multi-billion parameter language models using model parallelism, 2020. URL httpt::/aatati..oot/t org abs////////9///////// 10tp//9000000000333
[2] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners, 2020. URL httthhtts://aaat org/ abs 2005. 14165
[3] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL https::/taaattie org/ abs 2302 13971
[4] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models, 2024. URL https::/atxxi..oot wwwprxiiv.vorr/eeo 2b0/ .1783
[5] BigScience Workshop. :,
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b-parameter open-acces, multilingual language model, 2023. URL
https://arra org hbtp//212.0010.0 05100
[6] Guanhua Wang, Chengming Zhang, Zheyu Shen, Ang Li, and Olatunji Ruwase. Domino: Eliminating communication Ilm training via generic tensor slicing and overlapping, 2024.
httpss////// wwwp:ooog//aaoo abs///as///// 2409. 15241
[7] Suchita Pati, Shaizeen Aga, Mahzabeen Islam, Nuwan Jayasena, and Matthew D. Sinclair. Computation communication scaling for future transformers on future hardware, 2023. URL
sstpp://arxxwworr//bbb.s 22225
[8] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, et al. Ilama 2: Open foundation and fine-tuned chat models, 2023. https:ssss/// org abs/aass///0// 220p//9928877008883 092888
[9] Peiyuan Zhang, Guangtao Zeng. Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024. URL htt.:httts://arttt org/ / abs /0bs//2441 2401. 02385
[10] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Sarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads, 2022. URL
https://arxiv... /bt.///////00002000
[11] Shengwei Li, Zhiquan Lai, Yanqi Hao, Weijie Liu, Keshi Ge, Xiaoge Deng, Dongsheng Li, and Kai Lu. Automated tensor model parallelism with overlapped communication for efficient foundation model training, 2023. URL httptttp///aaaa //tp/v///iv//v.////////a/////a org/abs//////ooooooooo.oo... 2305. 16121
[12] Muru Zhang, Mayank Mishra, Zhongzhu Zhou, William Brandon, Jue Wang, Yoon Kim, Jonathan Bagan-Kelley Shuaiwen Leon Song, Ben Athiwaratkun, and Tri Dao. Ladder-residual: parallelism-aware architecture for accelerating large model inference with communication overlapping, 2025, URL https: http//aaxvvv...//e/v/eeee/eaaav/ org/ Sarggabs//20 abtp//1../222....// 06589

---

[13] Amy Yang. Jingyi Yang, Aya Ibrahim, Xinfeng Xie, Bangsheng Tang, Grigory Sizov. Jeremy Reizenstein, Jongsoo Park, and Jianyu Huang. Context parallelism for scalable million-token inference, 2025. httpthttts///rtttt org/ abs/ 2411. 01783
[14] Song Bian, Dacheng Li, Hongyi Wang, Eric Xing, and Shivaram Venkataraman. Does compressing activations help model parallel training?, 2023. UtL/ttpp://arr/v..rr/// 2301.02654
[15] M. I. Rudakov, A. N. Beznosikov, Ya. A. Kholodov, and A. V. Gasnikov. Activations and gradients compression for model-parallel training. Doklady Mathematics, 108($2):5222$281 December 2023, ISSN 1-36-. doi: 10.1134/110645562227711111 org/10. 010101113///0004422227111111
URL http::/dxxxoo
[16] Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry. At stability's edge: How to adjust hyperparameters to preserve minima selection in asynchronous training of neural networks?, 2020. URL htttps://arxiy ort/ 1909 12340
[17] Jan compression for tensor parallel IIm inference, 2024. URL httph:tttttiiiii wwtp:/iww.rgg//////////
Michael Truong Le, Oliver Hausdörfer, and Alok Verma. Communication
09510
[18] Qingyuan Li, Bo Zhang, Liang Ye, Yifan Zhang, Wei Wu, Yerui Sun, Lin Ma, and Yuchen Xie. Flash communication: Reducing tensor parallelization bottleneck for fast large language model inference, 2024. URL htt.::tttt/s/rtttleee org/ abs/ abs/2412.... 2412. 04964
[19] Han-Byul Kim, Duc Hoang, Arnv Kundu, Mohammad Samragh, and Minsik Cho. Spd: Sync-point drop for efficient tensor parallelism of large language models, 2025. URL https
//arxiv. org/ /2022 20727
[20] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL attps://arxxx org/ abs 16077.6650000
[21] Biao Zhang and Rico Sennrich. Root mean square layer normalization, 2019. URL https
//arxiv.org baarg/b///// 19100 07467
[22] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoif, et A framework for few-shot language model evaluation, 12 2023. URL https://zenodo. wogg/records/10225633
[23] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence?, 2019. httpss////aax org. hatgg//b///// 0tt///190000077
[24] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL httpp::/t//// wwtp:////////bb/aaooooooooss wttp://699..0700
[25] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context, 2016. /ttpssts:///arx org/ abs/ 1606.06031
[26] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale, 2019. URL httpt:/tt:/:arrt http:://aaavvoaraaabbbbb org/abs/1//0
10641
[27] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. URL /ttts:://arxi... abs/
1911. 1961.1164
[28] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering, 2018. URL https
//arxiv.org/ass/1800 02789
[29] Christopher Clark, Kenton Lee, Ming-Wei Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019.
URL
ht https:://arxiti www.aaiwvv.../raaea. abs/ 19te/1995.1111 10044
12

---

A Mathematical Description of MLP and Attention Using Partial Channel-Reduce
A.1 Attention
In this section, we describe attention layer using partial channel- Like in section [.1| examine a case where the tensor caasor-arralle dimension is 2. Furthermore, for simplicity, we assume the model consists of only 2 attention heads.
Standard all-reduce in Attention. Below, describe the output of the attention layer with a full all-reduce. The attention head outputs, Hi and H2, for the first and second heads, respectively, are
H1 Attention (XWF,, xwY)
H2 Attention (XW XW2 x
(16)
In tensor-parallen parallelli,,,,, the output projection of the attention layer O, is sharded between devices. The output Z after all-reduce is:
01 O2
H1, H2) Z- O
7 H10 + HzO2,
(17)
where H1O, is calculated on the first device and H₂O₂ calculated on the second one.
Partial channel-reduce Attention. As in the MLP case, when introducing partial channel-reduce, the input to the attention is not necessarily identical on each device. The inputs on each device are denoted X₁ and X2: H1 Attention (X,W X,W x(W)
(18)
H₂ Attention (2W X2W2 X2WY)
The outputs of the attention layer with partial channel-reduce are different per device, and are denoted Zi and Zz:
O = [O Oz1
Y = [H1, H2).
012 O22
(19)
Zi = HH + H2O12, H1Om]
Z2 = H1On + H2012, H2Oml
A.2 Partial channel-reduce with many devices
In this section, relax the assumptions made in Section 3.1 and Appendix A.I to allow more than 2 tensor-paut cassorlarralle devices.
MIP. An MIP operation in transformers with full all-reduce operation can be described as follows. Given weight matrices E RAx and B E R where I' the Feed Forward Network (FFN) hidden dimension, and input activations X ggxexh the MLP can be written in index notation
Za = Σo
Σ ExjAA
Bui
(20)
where is the sequence length, b is the batch size, and is the nonlinearity. To introduce the partial channel-reduce operation, we first rewrite Eq. 20 in a notation that incorporates tensor-parallelism.
Zel =o
Σ Xiş Asstom
Batm
(21)
Zu
Zam
m
13

---

where A is A written in column- Salummn-prrallll notation such that the last dimension corresponds to tensor- parallel rank, and similarly B is B written in row-parallel notation. The summation over m is the all-reduce operation. Finally, we write the MLP with the partial channel-reduce operation as follows:
Zitm = Eo Xijm Xijm Askm
Bxim
(22)
E
Zitm
if < floor (hp)
Zilm
Zilm
if Z floor (hp)
Here, the summation over m is the partial channel-reduce operation. The main differences between Eq 21 and Eq 22 are:
The sum over Z only over the first p h channels of the hidden dimension.
The imput to the MLP is different per device when using partial channel-reduce, because the previous attention layer also used partial channel-reduce in its output. Thus, the input to MLP, X, has additional index, corresponding to tensor-parallel device rank. Similarly, has an additional index, m, because each device contains different parameters in the private channels after the partial channel-reduce operation.
Attention. We present the attention operation with partial channel-reduce. For simplicity, assume that the tensor-pane tensor-parallel degree is the number of heads, though this assumption can easily be relaxed. The attention operation with full operation is all-reduce is:
Hijm Attention (xWi XW m.
XWW),
Zilm
Σ
Hijm Ojim
(23)
Zu Σ Zsim
where Hejm is item in the mh attention head output, and Öztm is an item in the attention output projection, after reshaping the attention output projection O to be written in row-parallel notation. Hence, the output of the attention operation with the partial channel-reduce operation is:
Hijm = Attention (XmWK m. X,m WQ m.
Xm W.).
Zulm
HegmO
(24)
Σ Zetm
if < floor (hp) if floor (hp)
Zam
ilm
The main differences between Eq 23 and Eq
24 are:
The sum over Ž is only over the first p h channels of the hidden dimension.
Apart from the first attention layer, the input to the attention different per device when using partial channel-reduce, because the previous MLP layer also used partial channel-reduce in its output. Thus, the input to the attention, X, has an additional index, m, corresponding to the paraalll rank. Similarly. Z has additional index because each device contains different parameters after the partial channel-reduce operation.
Analyzing Farward-Backward Mismatch Considering Residual Connections
For transformer with all-reduce, we examine the input to an MLP layer on device m, denoted Xm, function of the output of previous attention layer on device m before the all-reduce operation,
14

---

denoted Zm with R being the residual connection from before the attention block, we have
Σ Zm + R
Xm = norm
(25)
R is the residual connection from before the attention block. We look the loss as a function of X for all m, and of the residual connection (that skips the next layer) on the m device, Rm. We also look Xm and Rm as function of Z1, and calculate the derivative of the loss function w.r.t Z using the chain rule:
(x (Z1) ... XM (Z1),
Ru (Zz) ... RM (Zi))
(X
8Xm (Zi)
OC (Rm).
2Rm (21)
Σ
0Xm
07,
2Rm
071
(26)
(Xm)
0Xm (Z1)
+ OC (Bm)
)
0Xm
0Z1
dRm
OC (Xm)
ax
0L (Rm)
+
0Xm
dz1
Rmm
where in the second equation used the fact that
RR (Z1)
0
Σ 2m 1,
(27)
d2i
OZi
and in the last equation, we equated @Xm for all m (denoted now as O2, 722
8X because in
with full all-reduce operations X, (21) = X2 (2,) = X (Z1). and thus their derivative w.r.t
Zi are also equal. Here, the summation over m the all-reduce operation in the backward pass. Our key insight here is that, due to the distributivity of matrix multiplication:
OC (Xm)
OL (Rmm)
OC (Xm) -E VIL
0x
OL (Ri))
Σ m
+
+
Σ m
(28)
ax
07
2R..
OX m
37
ŌRm
In other words, the all-reduce operation can happen before back-propagating through the normal- ization function (right side of the equation), or after the normalization function (left side of the equation). On the right hand side of the equation, the summation of the residual gradient happens in the all-reduce of the next layer, i.e the layer after the residual is accumulated back into the activation tensors. In standard training frameworks, such as Megatron-LM, the all-reduce operation happens before back- propagating through the normalization function. When training with CAAT-Net, the assumption that 8Xm are equal for all m no longer holds. For this reason, the all-reduce operation in OZ.. the backward pass must be after the calculation the normalization function derivative.
Similarly, examine the normalization function parameter update, denoted B, with a full all-reduce operation. Because the residual splits from the activations before the normalization function is applied, the normalization function parameter update is not a function of the residual states:
AC (X1 (3) XM (8))
OC (Xm))
axmm
Σ
(29)
03
OX
03
In the gradient descent step, this means that after each iteration, before updating the weights, the updates 03 need be synchronized themselves with an all-reduce operation. In popular training frameworks, the fact that 0x are equal for all m is utilized once again. When the all-reduce happens before backpropagating through the normalization function, there is need for gradient synchronization after each step:
03
OC (Xm)
0X
Σ
(30)
axm,
0g
ax... Once again, equated 03 for all m (denoted now
ax
03
15

---

C
Speedup Analysis
We suggest simple model further understand the relation between compute and communication times, and analyze the potential speed-up of our method. We make the following assumptions:
Compute is modeled only as GEMM operations - specifically in the MLP and attention. Communication is modeled only as the tensor-parallel all-reduce. We neglect all other computations (embeddings, language model head, etc.) and communications (pipeline parallel, data parallel, etc.)
We assume constant ratio between compute and communication capacity, C, given in units of FLOP/s GB//s
We assume that the communication time roughly does not change when we adjust the tensor- parallel dimension. This assumption is true in Gaudi processors when using multi-node connections, due to inter- and intra-node communications overlapping during all-reduce operations. Additionally, we assume that computation time is inversely proportional to the tensor tensor-parallel dimension. This assumption only approximately true since computation time depends many factors, such as tensor shapes.
We examine the forward pass of a transformer with simple MLP, where the first weight matrix in the MLP is of dimension h x 4h and the second weight matrix is of dimension 4h h. This type of MLP is common in LLM's such as GPT and BERT. We consider batch size 1, as batch size does not affect the results in this section. When multiplying number of operations (multiplications and additions) is m p (2n 1), which we approximate to 2m p n. Based on this, the total number of computation operations transformer layer with sequence length s, n heads and hidden size h is:
x p matrix by a p matrix, the total
MLP ops = 8sh2 + 8sh² = 16sh²
Attention ops = 6sh² + proj
h + 2sh
(31)
attention
proj
4s + 8sh²
To calculate the total communication payload, we recall that in each full all-reduce operation, s activations sent and received by each device, and this happens twice every transformer block. Thus the total number of compute operations G, and communication payload as a function of p. per device, are:
G 24sh 2 + h
(32)
P(p) 2shp,
where is the paraalel dimension and p is the CAAT-Net synchronization factor. Consequently, the total time for computation and communication of a single transformer layer in the forward pass is
T(p) G + P(p)
(33)
We find that the speed-up introduced by our method is:
T(1) -(1)-T(p) (1
(34)
T(1)
12h+2x Cr
Our method is more significant for large C and r, but diminishes for very large h and 8. This is because computation is quadratic in h and s while the communication is linear. If communication is hidden behind all of the computation, then there and computation times are equal, and there is less motivation to decrease to improve speed-up. If G > P()). then the optimal value of p is p* Otherwise we can set G P(p). Considering both cases, we find that the value of p" is:
maximal value of p at which communication
12h 2s
p² min
1).
(35)
Cr
16

---

D Top-K and Random Masks Communication Reduction
Analyzing the communication volume reduction in these experiments, note that only the commu- nication in the forward pass is compressed. This is because we use vanilla tensor-paralle maasoo-parallelimm and not alternate implementation detailed in Section [ so communication doesn't happen in the same location in the forward and backward passes. For this reason, unlike CAAT-Net, communication compression in the forward pass does not imply communication compression in the backward pass. Additionally, all-reduce happens in steps reduce-scatter and all-gather. Applying a mask before communication equivalent to compressing the reduce-scatie operation, which accounts for only half of the total communication in the all-reduce. For these reasons, compressing communication using Top-K or random masking with parameter p leads to (100-(1-P) 5% reduction in tensor- parallel communication volume. In CAAT-Net with parameter p communication compressed in the reduce-sccee scatter and all-gather of both the forward and backward passes. The total tensor- massorl communication CAAT-Net is reduced by (100 -)))%
E
Additional Results
E.1
Effects of All-Reduce Bit Precision in the Backward Pass
In this section we show that 16 bit accumulation in the backward pass all-reduce leads to loss divergence. To show this, we train the variants of the TinyLlama model, as is described in Section The first is trained using the Intel's Megatron-LM fork repository, without any additional changes. The second variant trained after applying the changes detailed in Section pass all-reduce accumulation in bfloat 6 precision. The third is trained after applying all changes in the Section including 32 bit accumulation the backward pass. Results in Figure show severe degradation with 16 bit accumulation after 15K training steps. Despite the mathematical equivalence between the two implementations, numerical differences lead instability.
with the backward
10
bfloat16 accumulation
float32 accumulation
Baseline
8
Loss
4
2
0
5000
10000 Step
15000
20000
Figure 6: Effeets of All-Reduce Accumulation Precision in the Backward Pass on TinyLlama
Train Loss.
E.2 Tiny Llama and GPT-XL Accuracy Evaluation for varying values of
report the full zero-shot accuracy evaluation of TinyLlama using varying values of p. The evaluation metrics are those described in Section 5.1 Overall p 0.75 achieves the best accuracy, with the highest score out of 8 metrics. Results are available Table Additionally, we report full zero-shot accuracy evaluation of GPT-XL for 0.5. Results available in Table
17

---

Table 2: Zero-shot accuracy results on GPT-XL
LAMBADA (acc)
Hellaswag (acc)
WinoGrande (acc)
PIQA (acc)
50.16 + 0.70
36.07 + 0.48
52.41 + 1.40
68.06
1.11
0.5
48. 17 0.69
36.26 + 0.48
53.83 + 1.40
67.79 + 1.09
OpenBookQA (acc)
BOOL-Q (acc)
WikiText (ppl)
Validation Loss
19.20 + 1.76
60.85 + 0.85
19.25
1.28
0.5
18.60 = 1.74
61.01 = 0.85
19.00
1.27
Table 3: Zero-shot accuracy results TinyLlama for varying values of
LAMBADA (acc)
Hellaswag (acc)
WinoGrande (acc)
PIQA (acc)
48.26 = 0.70
35.55 + 0.48
55.25 = 1.40
67.03 + 1.10
0.875
46.43 0.69
35.89 + 0.48
53.04 = 1.40
64.96 = 1.11
0.75
4 0.70
35.88 + 0.48
52.80 1.40
68.99 + 1.08
0.625
48.92 0.70
35.61 0.48
53.20 1..0
67.41 1.09
0.5
48.01 + 0.70
35.65 + 0.48
51.26 = 1.40
67.68 + 1.09
0.375
45.74 0.69
35.48 + 0.48
52.88 1.40
68.28 + 1.09
0,25
48,83 = 0.70
35.63 0.48
50.83 1.40
67.74 # 1.09
0.125
44.25 + 0.69
34.63 + 0.47
50.83 = 1.38
66.27 + 1.10
0
33.96 + 0.66
30.49 + 0.46
50.28 = 1.41
64.09 + 1.12
OpenBookQA (acc)
BOOL-Q (acc)
WikiText (ppl)
Validation Loss
21.20 1.83
59.57 + 0.86
19.10
1.28
0.875
20,00 = 1.79
60.24 0,86
19.05
1.28
0.75
20.60 1.81
61.04 0.85
18.84
1.27
0.625
21.40 + 1.84
60.12 + 0.86
19.08
1.28
0.5
21.00 1.82
59.24 + 0.86
20.12
1.28
0.375
19.40 + 1.77
57.13 + 0.85
19.30
1.29
0.25
20.60 + 1.81
51.53 0.87
19.60
1.29
0.125
20.60 = 1.81
54.07 0.87
20.36
1.31
17.80 + 1.71
56.97 + 0.87
28.37
1.49
E.3 Private Channel Scaling Ablations
In this section present ablation study for private channel scaling. We train the 130M model presented in Section over 7.8B tokens, with a tensor- tensor-parallel dimension of and p = 0.5. We record the validation loss at the end of training, with and without private channel scaling, and see slight improvement when using private-channel scaling over most values of p. Results are presented in Figure E.3
E.4 Additional Speedup Results
In this section, we report inference speedup 34B and 70B Llama models. These models use grouped query attention (GQA) of 8, so it possible to evaluate our speedup only on tensor-parallel dimension of 8 (our lower multiple of 2). Absolute measurements are shown in Tables 4 and 5
18

---

2.10
Private Channel Scaling
2.05
No Private Channel Scaling
2.00
Loss
1.95
1.90
1.85
0.0
0.2
0.4
0.6
0.8
1.0
p
Figure Partial Channel Scaling, Validation loss vs. p for with and without private channel scaling, on a 130M model.
Table 4: Absolute measurements of Llama 70B inference Time-to-First-Tokes (TTFT) for tensor-parallel Measured milliseconds, with varying batch sizes and values of p.
Batch size
Bascline
p 0.75
0.5
0.25
No communication
16
1402
1301
1224
1153
1093
32
2769
2595
2436
2281
2092
64
5513
5290
4982
4654
4259
128
10962
10293
9817
9457
8731
256
20606
19242
18055
16854
15107
Table 5: Absolute measurements of Llama 34B inference Time-to-First-Tokee (TTFT) for tensor-parallel 8. Measured in milliseconds, with varying batch sizes and values of p.
Batch size
Baseline p 0.75
0.5
0.25
No communication
16
730
682
637
590
532
32
1519
1429
1328
1226
1074
64
2977
2861
2675
2453
2178
128
6081
5738
5424
5081
4544
256
11104
10270
9575
8870
7942
E.5
Absolute Performance Metrics
In Section 5.2 of the paper we reported speedup in training and inference on Llama-7b model, relative to the baseline. Here we report the absolute measurements, from which the speedup was calculated.
Table 6: Absolute measurements of training throughput. Measured in tokens per second (tps) for different tensor-parallel (TP) dimensions, and values of p.
TP dimension
Baseline 0.75
p 0.5
p 0.25
No communication
4
34904
36883
38586
39728
40944
8
64395
67495
70403
72917
76646
16
96642
98625
106050
112456
126227
19

---

Table Absolute measurements of inference Time-to-First-Toker (TTFT) for lensor-paralle
16. Measured in milliseconds, with varying batch sizes and values of p.
Batch size
Baseline
0.75
p
0.5
p
0.25
No communication
16
214
209
187
165
100
32
443
386
358
326
193
64
746
696
637
565
391
128
1388
1291
1174
1060
792
256
2625
2393
2178
2059
1630
Table Absolute measurements inference sured in milliseconds, with varying batch sizes and values of
for tensor-paralle Mea-
Batch size
Baseline
p = 0.75
0.5
0.25
No communication
16
204
189
181
160
139
32
411
374
352
320
279
64
853
786
727
641
571
128
1645
1571
1475
1295
1133
256
3320
3071
2916
2713
2432
E.6
Measurements on NVIDIA Hardware
In this section present the full inference speedup measurements H100 and A100, as detailed in Section 5.2 of the paper.
Table 9: TTFT on H100 and A100. Measured in milliseconds, for different batch sizes and p values.
Device
Batch Size
=
p = 0.75
p == 0.5
p= 0.25
p
0.0
173
166
157
149
136
32
341
338
312
296
270
H100
64
679
663
623
589
545
128
1360
1295
1280
1179
1088
16
406
415
373
353
328
32
795
764
727
694
643
A100
1568
1508
1436
1363
1281
128
3117
2999
2927
2722
2553
20

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA] is not grounds for rejection. While the questions are phrased acknowledge that the true answer is often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in Appendix. If you answer [Yes] question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Contributions claimed in the abstract and introduction are shown in Section 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitation detialed in Section 3.3, 5 and 6
21

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The only theoretical claim, which appears in Section 4, is accompanied by full and thorough derivation.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in Appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
All of the steps to reproduce our claims are detailed in sections 4 and 5.
Guidelines:
The answer NA means that the paper does not include experiments.
22

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification:
Code will be published after acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ /ttppp.cc/pblic/guideeee http://oeeuuiiisssoooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
23

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer:
[Yes]
Justification: Relevant details are in section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in Appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The paper contains the standard error for accuracy experiments in section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or Figures symmetric error bars that would yield results that are of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding Figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources are detailed in section 5.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The paper conforms to the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: The paper does not address societal impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
The paper does not release data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
25

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All of the code, data and models used in the paper are open-source and properly cited.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pataasess.paperswtt paperswi thcode. cow/coo////ooooo /atasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer[[NA]
Justification: Code will be published after acceptance.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
26

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer:[NA
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the L.M used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy /ettps://neurr cc/ //oofffrennees22255555 for what should or should be described.
27


Reviews:
[{"review_id": "Ek5UOtReP6", "reviewer": "Reviewer_YSCj", "summary": "This 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%.", "questions": "Questions：\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.", "limitations": "None.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths：\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.", "was_revised": true, "final_justification": "I have read the authors' response. Thanks for the rebuttal. The response mainly addressed my concerns. Considering other reviewers' comments and rating scores, I will raise my rating to 4 accordingly.", "rebuttal": "We thank the reviewer for the constructive and thoughtful feedback. We appreciate the recognition of our contributions to reducing communication overhead and the required changes to the backward pass. We respond to the reviewers comments and questions below. Furthermore, we have expanded our experiments to include more diverse models and broader ablations over $p$. \n\n--- \n\n**(1)**\n\n**Questions**\n\nThe 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\nIn 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\n**Answer**\nWe acknowledge that the core innovation involves modifying the communication primitive. While this seems like a slight adjustment, in practice it constitutes an architectural change. Partial channel–reduce fundamentally alters the model's computation graph and information flow. This distinguishes CAAT-Net from communication optimizations like pipelining or compression, which preserve or approximate the underlying mathematical operations. For example, the MLP in CAAT-Net can no longer be expressed as two standard linear transformations with a non-linearity between them, as shown in Section 3. Instead, we need to redefine both the MLP and attention mechanisms to handle device-specific inputs and outputs.\n\nMoreover, supporting this design requires non-trivial modifications beyond the communication operation and initialization. These modifications are detailed in Section 4.\n\nRegarding the private (unsynchronized) activations: they are not discarded or approximated. They flow through the same downstream layers as shared activations, contributing fully to the model’s expressivity.\n\n---\n\n\n**(2)**\n\n**Question**\nIn 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\n**Answer**\nWe thank the reviewer for this question. The variance scaling adjustment addresses a statistical property of partial channel-reduce that we can clarify with a more detailed theoretical analysis.\n\n**Theoretical Justification:**\n\nConsider the activations in $Z_1$ after partial channel-reduce. Assuming the MLP outputs before reduction ($Y_1B_{11}$, $Y_2B_{12}$, $Y_1B_{21}$) have zero mean and variance $\\sigma^2_A$ and are independent:\n\nFor shared channels:\n\n\\begin{equation}\n    \\text{Var}(Y_1B_{11} + Y_2B_{12}) = \\text{Var}(Y_1B_{11}) + \\text{Var}(Y_2B_{12}) = 2\\sigma^2_A\n\\end{equation}\n\nFor private channels:\n \n\\begin{equation}\n    \\text{Var}(Y_1B_{21}) = \\sigma^2_A\n\\end{equation}\n\nThis variance mismatch can lead to uneven gradient signals across channels. By scaling private channel weights by $\\sqrt{2}$ at initialization, we achieve uniform variance of $2\\sigma^2_A$ across all channels.\n\n**Empirical Validation:**\n\nWe validated this approach on 130M parameter models ($p=0.5$, 5B tokens), observing the following improvements:\n\n|                        | TP=16 | TP=8  |\n| ---------------------- | ----- | ----- |\n| Regular Initialization | 1.94  | 1.927 |\n| Altered Initialization | 1.926 | 1.924 |\n\nThe improvements are modest but support the theoretical motivation.\n\n**Response to Reviewer Concern:**\n\nWhile this initialization provides theoretical and empirical benefits, it represents a secondary contribution compared to the core architectural changes in Section 3 and 4. We propose moving this discussion to the appendix with expanded ablation studies for the camera-ready version, allowing the main text to focus on the essential modifications.\n\n---\n\n**(3)** \n\n**Question**\n\nThe 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\n**Answer**\n\nDifferent architecture:\n\nTo address this concern, we extended our experimentation to GPT3-XL [A], a 1.3B parameter model. We find that CAAT-Net maintains and even slightly improves model accuracy. Due to character limitation, we refer the reviewer to the section 3 of our rebuttal to reviewer 7RQG to see the full experiments and data. \n\nModel scale:\n\nIf the concern pertains to speedup measurements on larger models, we refer the reviewer to Appendix D.1, where we report inference speedups on larger LLaMA models, including LLaMA-34B and LLaMA-70B. On LLaMA-70B, our method achieves a 13\\% speedup at $p=0.5$ and 23\\% at $p=0.25$, demonstrating the effectiveness of our approach at scale. \n\nIf the concern instead relates to training-time accuracy at such scales, we note that most prior work is evaluated on smaller models (130M–7B parameters), given the immense cost of training. For instance, training our 7B model already requires substantial compute resources and access to a large cluster, which is beyond the capabilities of many academic groups. Notably, recent prominent works like GaLore [B] and Rho1 [C] also focus on models up to 7B parameters.\n\nOur results demonstrate that the method scales effectively from 130M to 7B parameters, consistently preserving accuracy while improving performance.\n\n---\n\n**(4)**\n\n**Question**\n\nWhile 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.\n\n**Answer**\n\n$p=0.5$ was indeed chosen as a safe default for the large experiments. To further study the effects of tensor-parallel dimension and $p$, we added many data points to Table 2. We show that our method performs well on small models when $p$ is sufficiently large, sometimes even yielding slight improvements in validation loss. Furthermore,  we include additional experiments on TinyLlama with varying values of $p$, demonstrating that these trends scale consistently. Due to character limitation, we refer the reviewer to the section 3 of our rebuttal to reviewer 7RQG to see the full experiment setup and data. \n\n---\n\n**(5)**\n\n**Question**\n\nWhy 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\n**Answer**\n\nComparison to pipelining approaches:\n\nCAAT-Net differs fundamentally from pipelining methods such as Domino. While pipelining methods aim to hide communication latency by overlapping it with computation, CAAT-Net directly reduces the amount of data communicated.\n\nPipelining can be effective on low tensor-parallel dimensions with fast interconnects, but this approach does not scale well. As tensor-parallelism increases, the computation per device decreases, while the communication doesn't. Consequently, communication can dominate total runtime. For example, in our Llama2-7B experiment with tensor-parallelism of 16, setting $p=0$ achieves over 50\\% inference speedup. This implies that more than half of the workload time is spent on communication. This imbalance becomes even more pronounced at larger tensor-parallel scales. For this reason, we view CAAT-Net as complementary—not competing—to pipelining methods.\n\nIn terms of accuracy, CAAT-Net preserves and in some cases improves accuracy, particularly at larger values of $p$, as shown in the paper. Domino preserves accuracy because it does not alter the transformer architecture. \n\nComparison to Compression Methods:\n\nAccuracy: Unlike pipelining approaches, Top-K compression does reduce communication bandwidth. However, as shown in Appendix E, CAAT-Net achieves significantly better accuracy than Top-K compression.\n\nSpeedup: In general, computing Top-K is very slow and its use overshadows communication benefits. As an example, we present throughput results in TFLOP/s for the experiments on the 130M model described in Appendix E. To be consistent with CAAT-Net notation, $p$ is the percentage of the original tensor communicated (so $p=0$ is no communication and $p=1$ is full communication). To calculate the Top-K we use torch.topk(). Using the Top-K operation completely degrades throughput. Although this experiment is on a small model, this severe performance degradation is evident and holds for larger models as well. We can add more experiments on larger models for the camera-ready version if the reviewer thinks it is necessary. \n\n| \\$p\\$ | Throughput (TFLOP/s) |\n| ----- | -------------------- |\n| 1     | 101.1                |\n| 0.9   | 2.9                  |\n| 0.7   | 3.8                  |\n| 0.5   | 5.6                  |\n| 0.3   | 9.8                  |\n| 0.1   | 24.7                 |\n\n---\n\n[A] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,\nArvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are\nfew-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.\n\n[B] J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian,\n“Galore: Memory-efficient llm training by gradient low-rank projection,” arXiv preprint arXiv:2403.03507, 2024.\n\n[C] Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian\nJiao, Nan Duan, et al. 2024b. Rho-1: Not all tokens\nare what you need. arXiv preprint arXiv:2404.07965."}, {"review_id": "d2PszUgKvn", "reviewer": "Reviewer_f4Cy", "summary": "This 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.", "questions": "1. **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.", "limitations": "Partially 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.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## 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.", "was_revised": true, "final_justification": "The technical concerns I had with this paper are mostly resolved. I will raise my point.\n\nBut I hope to see the final manuscript include our discussion in the full paper/appendix and include the points below.\n\n* Ladder-residual is also trained from scratch. I hope to see at least a discussion in the related work or a comprehensive analysis in the appendix.\n* I am still doubting ‘we show no accuracy degradation’, please include more accuracy training results for larger models.", "rebuttal": "We thank the reviewer for the thoughtful and detailed feedback. We are pleased that the reviewer recognizes the novelty and innovation of channel-level synchronization, as well as the importance of the numerical stability considerations we addressed. In response to the reviewer's valuable suggestions, we have expanded our empirical evaluation. These new experiments provide a more comprehensive characterization of CAAT-Net's performance and robustness across diverse workloads. Additionally, we address the reviewer's concerns, particularly those related to partial reduction on contiguous channels, further clarifications on numerical stability, and limitations at inference time.\n\n---\n\n**(1)** \n\n**Contiguous channel selection**\n\nWe thank the reviewer for these important questions about channel selection strategy. There are two compelling reasons for contiguous channel selection:\n\n**Theoretical equivalence**:\nThe reduced channels are selected at initialization. Since channels are randomly initialized, any subset of channels is equivalent.\nSpecifically, we could permute all weight matrices to make any arbitrary channel subset contiguous, yielding a statistically equivalent initialization. Therefore, contiguous selection imposes no fundamental limitation.\n\n**Implementation efficiency**:\nIn our work we aim to use existing communication collectives efficiently. Collectives operate most efficiently on contiguous memory\nregions. Interleaved selection would require either additional operations before communication or custom communication kernels.\n\nGiven this theoretical equivalence and practical efficiency advantage, we focused on contiguous selection.\n\n--- \n**(2)** \n\n**Numerical stability** \n\n We thank the reviewer for seeking clarification on the fp32 accumulation requirements.\n\nThe specific operation requiring fp32 accumulation is the gradient summation across devices in the backward pass, corresponding to the left-hand side of Equation 15:\n$$\\sum_m \\left( \\frac{\\partial \\mathcal{L}(X_m)}{\\partial X_m} \\cdot \\frac{\\partial X}{\\partial Z_1} \\right)$$\n\nThis summation represents the all-reduce operation that aggregates gradients from all tensor-parallel devices after they have\nback-propagated through the normalization layer. Importantly, the communication itself can be in 16bit precision, but the all-reduce summation after the communication must be in 32bit.\n\nThe normalization function itself remains unchanged. We mentioned normalization operations only as an analogy for other computations that typically require higher precision during training.\n\nRegarding fp8 training: while we have not extensively tested CAAT-Net under fp8, the same principle would likely apply. The gradient\nreduction step would require higher precision to maintain numerical stability.\n\n--- \n\n**(3)** \n\n**Figure redundancy**\n\nWe thank the reviewer for the valuable feedback. Figure 1 visualizes the modified two layer MLP architecture at a high level, while Figure 2 provides a detailed explanation of the partial channel-reduce implementation and the corresponding backward pass adjustments. We appreciate the reviewer's suggestion and will consider merging these figures to improve clarity and flow in the\ncamera-ready version.\n\n--- \n**(4)** \n\n**Why does synchronizing only a subset of channels preserves correctness.**\n\nWe thank the reviewer for this fundamental question about the correctness of partial synchronization.\n\nCAAT-Net does not preserve correctness in the sense of approximating the original fully-synchronized model. Instead, it defines a new architecture specifically tailored for training on multiple devices using partial synchronization. As in all papers which introduce novel architectures, we validate our claims empirically. Our empirical results demonstrate that this specific architectural change preserves learning effectiveness. \n\n---\n\n**(5)**\n\n**Evaluation of larger models and using Nvidia GPU's** \n\n**Speedup**\n\nPlease find speedup for 34B and 70B models in Appendix D (in the supplementary materials). On Llama-70B, our method delivers a 13% speedup at $p = 0.5$ and a 23% speedup at $p = 0.25$, highlighting its continued effectiveness at larger model scales. \n\nRegarding Nvidia GPU's, we conducted additional experiments to show real-world applicability. Due to character limits in our response, we refer the reviewer to section 3 of our rebuttal to 8fqL for full experiment setup and results. We observe up to a 9% speedup for $p = 0.5$ and a 14% speedup for $p = 0.25$ on NVIDIA hardware. Importantly, these results are preliminary and not fully optimized, and we believe they can be further improved.\n\n**Accuracy** \n\nPrecisely evaluating accuracy/optimal value of $p$ for larger models requires full training. Most of the LLM training literature experiments with smaller scale models, because of the high compute cost. Even training a 7B model requires a large cluster and is very costly. We show in our paper that our method scales well from 130M to 7B parameters, and expect these trends to continue. For example, for $p=0.5$, we found that accuracy was slightly improved with respect to the baseline for both the\n130M model and the 7B model. In the regime that we investigated, we did not find evidence that CAAT-Net accuracy deteriorates with increasing model size for a given $p$.\n\n---\n\n**(6)**\n\n**Scaling cost projections**\n\nIn general, the utility of our method depends on many factors, such as hardware, architecture, sequence length, tensor-parallel dimension, etc. Please see Appendix C in the supplementary materials for a detailed speedup analysis, where we model our methods utility as a function of architecture, hardware and tensor-parallel rank.\n\n---\n\n**(7)**\n\n**Comparison to Top-K, Domino and Ladder Residual** \n\nWe thank the reviewer for this question, and use this opportunity to highlight CAAT-Net's advantages compared to other\nmethods. \n\n**Comparison to pipelining methods**\n\nThere is an important distinction between Ladder-residual and Domino to our method. Ladder-residual and Domino both aim to move communication off the critical path. In practice, both of these approaches do not reduce bandwidth at all. They just enable pipelining between communication and computation. Our approach reduces communication bandwidth.\n\nFor specific hardware (with fast interconnect) and for low tensor-parallel dimensions, pipelining can be sufficient. However, this\napproach is not scalable. When increasing the tensor-parallel dimension, computation per device decreases, while communication does not. Total communication time can quickly overcome total computation time. For example, in our paper we show inference speedup for Llama2-7B with tensor-parallel 16. There, setting $p=0$ gives over 50% speedup, so more than half of the time is spent on communication. For larger tensor-parallel dimensions this problem is even more pronounced. For this reason, we do not view CAAT-Net as a competing approach to pipelining optimization, but a complementary one. \n\nRegarding accuracy, for large enough $p$ our method preserves and can even improve accuracy in some cases, as we show in the paper. \nDomino does not affect accuracy as it does not change the underlying transformer architecture.\nThe Ladder-Residual paper reports similar accuracy for pretraining a 1.2B parameter model, and measures degradation when training a 3.5B parameter model. These results are based on Table 3 of their paper.\n\n**Comparison to compression methods** \n\nAccuracy: Unlike Ladder-residual and Domino, Top-K reduces communication bandwidth. Results for Top-K compression are available in Appendix E (supplementary materials). CAAT-Net significantly outperform Top-K compression.\n\nSpeedup: In general, computing Top-K is very slow and its use overshadows communication benefits. As an example, we present throughput results for the experiments described in Appendix E. Due to character limits in our response, we include the experiment setup and results to section 5 of our response to reviewer YSCj. \n\n---\n\n**(8)**\n\n**Inference Flexibility**\n\nInference can run on different device counts than training without finetuning, but with important limitations and trade-offs.\n\nThe most straightforward approach is using logical devices. A single physical device can simulate multiple tensor-parallel ranks. For\nexample, if trained with TP=2, inference on TP=1 requires the single device to perform both sets of computations sequentially, replacing the partial channel-reduce with local summation. This can increase computation time but preserves model correctness.\n\nChanging the synchronization factor $p$ at inference time is not supported, as this effectively changes the model architecture. A model trained with $p=0.5$ expects specific private and shared channel configurations. Altering $p$ is equivalent to serving a completely different model.\n\nWhile we have not conducted extensive empirical studies on inference flexibility, our current results suggest the primary deployment path is maintaining the training tensor-parallel configuration for optimal performance. CAAT-Net represents a model tailored to a specific parallelization scheme, whereas traditional approaches separate model architecture from parallelization strategy. While this compromises flexibility, we argue that in an age of increasing model sizes, this is a relevant approach to achieve effective scaling. This is similar to Quantization Aware Training (QAT), in which a model is trained to be served at a specific precision, and is not meant to be used in different precisions [B]. However, these deployment limitations can be addressed through logical devices or fine-tuning approaches to adjust $p$. \n\n---\n\n[A] PyTorch Labs. gpt-fast, 2024. URL https://github.\ncom/pytorch-labs/gpt-fast. Accessed: 2024- 09-29.\n\n[B] K. Du, Y. Zhang, and H. Guan, \"From quantized dnns to quantizable\ndnns,\" CoRR, vol. abs/2004.05284,2020. [Online]. Available:\nhttps://arxiv.org/abs/2004.05284"}, {"review_id": "t3GzTB63K1", "reviewer": "Reviewer_8fqL", "summary": "This 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.", "questions": "1. 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?", "limitations": "Please refer to weakness and questions.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "I would give this paper a borderline accept because it is a complete work with easy-use design, sufficient experiments, and well-written manuscripts.", "rebuttal": "We thank the reviewer for taking the time to read our paper carefully and provide thoughtful feedback. We appreciate the recognition of the originality and significance of our work. Specifically, we value the acknowledgment of its effectiveness in reducing communication overhead without sacrificing accuracy, across both training and inference. Below, we address the reviewer’s concerns in detail. As requested, we have included experiments on Nvidia GPUs, which further strengthen our claims and demonstrate the broader applicability of our method. Furthermore, we refer the reviewer to large-scale speedup experiments that appear in the appendix.\n\n----------\n\n**(1)**\n\n**Weakness**  \nRequires non-trivial changes to standard tensor-parallel frameworks such as Megatron.\n\n**Response**  \nOur method requires changing only a few lines of code (less than 50 for training, and less than 10 for inference). Additionally, our code will be open source upon acceptance. For these reasons, we do not view this to be a weakness.\n\n----------\n\n **(2)**\n\n**Weakness**  \nWhile 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.\n\n**Response**  \nIf the concern is about speedup measurements on 100B+ models, we would like to refer the reviewer to Appendix D.1, where we report inference speedups for larger Llama models, including Llama-34B and Llama-70B. On Llama-70B, our method achieves a 13% speedup for $p=0.5$ and a 23% speedup for $p=0.25$, demonstrating that our approach remains effective even at larger scales. While we have not yet run evaluations on 100B+ models, we can run these experiments in the near future, before the camera-ready deadline.\n\nIf the concern is about training-time accuracy on 100B+ models, we note that most prior works in the literature evaluate methods at smaller scales (e.g., 125M to 7B parameters), due to the prohibitive computational cost of training 100B+ models. For example, prominent works such as GaLore [A] and Rho1 [B] demonstrate their results primarily on models up to 7B parameters. Even training our 7B model requires substantial compute resources and a large cluster, which is beyond the reach of many academic groups. Training 100B+ models can cost hundreds of thousands of dollars. \n\nNevertheless, we show in our main results that our method scales well from 130M to 7B parameters, consistently preserving accuracy while improving performance.\n\n----------\n\n**(3)**\n\n**Weakness**  \nLack of experiments on Nvidia GPUs, which limited the generation of this paper.\n\n**Response**  \nTo address the reviewer’s concern and show the applicability of our method to other hardware, we conducted additional experiments with Nvidia GPUs. Our experiments were done using the gpt-fast [C] repository, and consisted of replacing all-reduce with partial channel-reduce during inference. We conducted experiments on both:\n\n-   8 × NVIDIA H100-80GB-HBM3 with NVLink\n    \n-   8 × NVIDIA A100-SXM4-80GB with NVLink\n    \n\nExperiments were conducted on Llama2-7B using tensor-parallel 8 and a sequence length of 2048.\n\n**H100: TTFT (s) and Relative Speedup vs. p=1.0:**\n\n| Batch Size | $p=1.0$           | $p=0.75$          | $p=0.5$           | $p=0.25$          | $p=0.0$           |\n| ---------- | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- |\n| 16         | 0.173s, **1.00×** | 0.166s, **1.04×** | 0.157s, **1.09×** | 0.149s, **1.14×** | 0.136s, **1.22×** |\n| 32         | 0.341s, **1.00×** | 0.338s, **1.01×** | 0.312s, **1.09×** | 0.296s, **1.13×** | 0.270s, **1.21×** |\n| 64         | 0.679s, **1.00×** | 0.663s, **1.02×** | 0.623s, **1.08×** | 0.589s, **1.13×** | 0.545s, **1.20×** |\n| 128        | 1.360s, **1.00×** | 1.295s, **1.05×** | 1.280s, **1.06×** | 1.179s, **1.13×** | 1.088s, **1.20×** |\n\n**A100: TTFT (s) and Relative Speedup vs. p=1.0**\n| Batch Size | $p=1.0$           | $p=0.75$          | $p=0.5$           | $p=0.25$          | $p=0.0$           |\n| ---------- | ----------------- | ----------------- | ----------------- | ----------------- | ----------------- |\n| 16         | 0.406s, **1.00×** | 0.415s, **0.98×** | 0.373s, **1.08×** | 0.353s, **1.13×** | 0.328s, **1.19×** |\n| 32         | 0.795s, **1.00×** | 0.764s, **1.04×** | 0.727s, **1.09×** | 0.694s, **1.13×** | 0.643s, **1.19×** |\n| 64         | 1.568s, **1.00×** | 1.508s, **1.04×** | 1.436s, **1.08×** | 1.363s, **1.13×** | 1.281s, **1.18×** |\n| 128        | 3.117s, **1.00×** | 2.999s, **1.04×** | 2.927s, **1.06×** | 2.722s, **1.13×** | 2.553s, **1.18×** |\n\n\nWe can see up to **9% speedup for $p=0.5$** and **14% speedup for $p=0.25$** on Nvidia hardware. Importantly, these results are preliminary and not fully optimized, and we believe they can be further improved.\n\nIf the reviewer is asking about training accuracy, we point out that there should not be any significant discrepancy in accuracy between Nvidia and Intel hardware.\n\n----------\n\n**(4)**\n\n**Question**  \nHow well does CAAT-Net interact with other techniques such as quantization, gradient checkpointing, or activation compression?\n\n**Answer**  \nWe appreciate the reviewer’s interest in how our method interacts with techniques such as quantization, gradient checkpointing, and activation compression. While these combinations are indeed important and promising directions for future work, it is not feasible within the scope of this study to comprehensively evaluate our method alongside the large and continually growing set of existing techniques. Our current work focuses on thoroughly establishing the core effectiveness of our method, and we leave exploration of its integration with other optimization or compression methods to future research.\n\n----------\n\n**(5)**\n\n**Question**  \nYou 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\n**Answer**  \nWe appreciate the reviewer’s comment and the opportunity to clarify our wording.\n\nWe state in the paper that changing the value of $p$ at inference time effectively alters the model architecture and requires fine-tuning. However, we do not claim that inference must strictly use the same tensor-parallel configuration as training. We only note that using the same configuration ensures the inference benefits from reduced communication compared to the $p=1$ baseline.\n\nIf the tensor-parallel configuration differs at inference, there are two primary options:\n\n-   **Post-training adjustment of $p$** via fine-tuning, which allows transitioning back to $p=1$ (fully synchronized) or other values.\n    \n-   **Logical TP mapping at inference**: An alternative is to simulate the training-time TP configuration during inference, without altering $p$. For instance, if the model was trained with TP=2 but needs to be served with TP=1, the single inference device can emulate multiple logical TP devices. In this case, partial channel-reduce operations at attention and MLP layers can be replaced with intra-device summation. While this approach avoids retraining, it may incur extra computation and memory usage on the device, and its efficiency depends on the specific deployment environment.\n    \n\nIn practice, we recommend maintaining the training-time tensor-parallel configuration for inference, though it is not strictly required. If the reviewer deems it necessary, we can update the paper to make this point clearer.\n\n---\n\n[A] J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian, “Galore: Memory-efficient llm training by gradient low-rank projection,” arXiv preprint arXiv:2403.03507, 2024.\n\n[B] Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, et al. 2024b. Rho-1: Not all tokens are what you need. arXiv preprint arXiv:2404.07965.\n\n[C] PyTorch Labs. gpt-fast, 2024. URL https://github. com/pytorch-labs/gpt-fast. Accessed: 2024- 09-29."}, {"review_id": "f1MSSVDnnD", "reviewer": "Reviewer_7RQG", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors clearly resolved technical concerns regarding numeric precision, quality-performance relationship.\nHowever, I still have a concern on evaluations of this paper (Lack of results: showing the results with validation loss, presenting only on zero-shot tasks).\nTherefore, I raised score from 3 to 4 (still borderline).", "rebuttal": "We thank the reviewer for their careful evaluation and constructive comments. We appreciate that the reviewer recognizes the importance of the problem we aim to solve, and the significance of our approach. To address the reviewers concerns regarding evaluation, we expanded the setups in which we evaluate our model, including experimentation with additional architectures and  varying values of $p$.\n\n---\n\n**(1)**\n\n**Question** - What if the normalization method applied on the transformer block is not RMSNorm? Is CAAT-net still applicable to that cases?\n\n**Answer** CAAT-Net is applicable to any normalization function, as the derivations in Section 4 are independent of the specific normalization method used. We used RMSNorm as a standard example, but the approach generalizes to any normalization function. The calculations in Section 4 should indeed be presented with a general normalization function for clarity. We will update Section 4 to use general notation in the camera-ready version. Futhermore, we have added training experiments with GPT3-XL [A] using LayerNorm to demonstrate this generalizability (see section 5 of this rebuttal).\n\n---\n\n**(2)**\n\n**Question** - 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\n**Answer** Although the two implementations are mathematically equivalent, they are not identical when implemented in finite precision. The changed order of floating-point operations introduces different rounding errors that compound during training. We find that when moving the all-reduce in the backward pass, the errors that accumulate throughout training are significant, and using fp32 accumulation solves this issue.\n\nRegarding bf16: our baseline experiments use bf16, and we observed the convergence issues in this setting.\n\n---\n\n**(3)**\n\n **Weakness** - 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 Figure 3 are meaningful results.\n\n**Response** - In figure 3 we ran multiple experiments which don't require full model training. Running all of the experiments in figure 3 on large models would require an unrealistic amount of time and computation resources. However, to address this concern, we have run a few additional experiments. We hope that this convinces the reviewer that CAAT-Net achieves good accuracy in multiple settings, which we see robustly in our experimentation. First, we replaced Table 2 in our paper with a graph to include a more detailed visualization of the trends in $p$ and TP for a 130M model. Unfortunately, we cannot share the graph in this rebuttal due to the guidelines, so it is presented here as a table. (The results are slightly different than those in the paper because a different learning rate scheduler was used):\n\n| $p$    | $1$    | $0.9$  | $0.8$  | $0.7$  | $0.6$  | $0.5$  | $0.4$  | $0.3$  | $0.2$  | $0.15$ | $0.1$  | $0.05$ | $0.03$ | $0.01$ | $0.001$ |\n|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|---------|\n| TP=2   | 1.887  | 1.890  | 1.890  | 1.883  | 1.879  | 1.876  | 1.872  | **1.870**  | 1.871  | 1.874  | 1.878  | 1.880  | 1.889  | 1.899  | 1.970   |\n| TP=4   | 1.884  | 1.884  | 1.882  | 1.879  | 1.873  | 1.867  | **1.863**  | 1.872  | 1.878  | 1.884  | 1.893  | 1.916  | 1.924  | 1.952  | 2.056   |\n| TP=8   | 1.889  | 1.892  | 1.879  | 1.875  | 1.871  | **1.876**  | 1.880  | 1.886  | 1.895  | 1.904  | 1.920  | 1.950  | 1.978  | 2.020  | 2.152   |\n| TP=16  | 1.887  | 1.880  | 1.879  | 1.877  | **1.872**  | 1.878  | 1.886  | 1.903  | 1.918  | 1.931  | 1.953  | 1.992  | 2.011  | 2.059  | 2.264   |\n\n\nThe best validation loss per tensor-parallel configuration is highlighted. The reviewer can see training accuracy remains intact, and can even improve, for large enough values of $p$.\n\nAdditionally, we conduct additional experiments with varying values of $p$ on TinyLlama to show that these trends scale relatively well. In these experiment TinyLlama was trained over 45B tokens. The best value per task is highlighted.\n\n| **Model** | **LAMBADA (acc)**     | **Hellaswag (acc)**   | **WinoGrande (acc)**   | **PIQA (acc)**        | **OpenBookQA (acc)**  | **BOOL-Q (acc)**      | **WikiText (ppl)** | **Validation Loss** |\n|-----------|------------------------|------------------------|-------------------------|------------------------|------------------------|------------------------|---------------------|----------------------|\n| $p=1$ | **43.59 ± 0.69**           | 33.41 ± 0.47           | 51.22 ± 1.40            | 66.65 ± 1.10           | 17.80 ± 1.71           | **59.05 ± 0.86**           | **21.44**              | **1.36**                 |\n| $p=0.7$ | 42.73 ± 0.69           | **33.59 ± 0.47**           | 51.70 ± 1.40            | 66.05 ± 1.10           | 18.20 ± 1.73           | 58.17 ± 0.86           | 21.50              | **1.36**                 |\n| $p=0.5$ | 42.64 ± 0.69           | 33.04 ± 0.47           | 51.70 ± 1.40            | **67.25 ± 1.09**           | 17.60 ± 1.70           | 56.57 ± 0.87           | 22.66              | 1.37                 |\n| $p=0.3$ | 40.33 ± 0.68           | 33.62 ± 0.47           | **52.80 ± 1.40**        | 66.76 ± 1.10           | **19.40 ± 1.77**       | 56.97 ± 0.87           | 21.80              | 1.37                 |\n| $p=0.1$ | 39.94 ± 0.68           | 33.14 ± 0.47           | 51.22 ± 1.40            | 66.32 ± 1.10           | 17.20 ± 1.69           | 56.97 ± 0.87           | 22.75          | 1.40             |\n\n---\n\n**(4)** \n\n**Weakness** 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\n**Response** To address this concern, we trained a version of the 1.3B parameter GPT-XL with partial channel-reduce. As per the specific requests for Qwen and Llama3 - both have an architecture which is very similar to Llama2, so training GPT3-XL is a better candidate to show that our method is robust. We trained GPT3-XL with using tensor-parallel 8, on the RedPajama dataset using the GPTSentencePiece tokenizer, rotary positional embeddings and with a global batch size of 512. We trained for a total of 45B tokens. Similarly to results on Llama, we find that CAAT-Net maintains and even slightly improves accuracy results. \n\n| **Model** | **LAMBADA (acc)**     | **Hellaswag (acc)**   | **WinoGrande (acc)**   | **PIQA (acc)**        | **OpenBookQA (acc)**  | **BOOL-Q (acc)**      | **WikiText (ppl)** | **Validation Loss** |\n|-----------|------------------------|------------------------|-------------------------|------------------------|------------------------|------------------------|---------------------|----------------------|\n| $p=1$     | **50.15 ± 0.70**       | 35.93 ± 0.48           | 52.49 ± 1.40            | 67.41 ± 1.09           | **21.4 ± 1.84**        | **61.74 ± 0.85**       | 19.58              | 1.30                 |\n| $p=0.5$   | 49.56 ± 0.70           | **36.27 ± 0.48**       | **53.20 ± 1.40**        | **68.06 ± 1.09**       | 20.80 ± 1.82           | 59.66 ± 0.86           | **19.35**          | **1.29**             |\n\n---\n\n**(5)**\n\n**Weakness** The evaluation on zero-shot tasks typically generate a single token which aren't sufficient to validate the quality\nof 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. \n\n**Response** We thank the reviewer for the valuable suggestion to evaluate longer sequence generation. While we acknowledge that multi-token generation can offer additional insights, we focused on single-token predictions to align with the standard zero-shot evaluation protocols widely used in the literature. Established methods such as SparseGPT [B] , SmoothQuant [C], and recently Ladder Residual [12] (to name a few) have all been evaluated using solely this approach. In our view, single-token evaluations provide a reliable and well-established proxy for assessing model quality.\n\n---\n\n[A] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.\n\n[B] E. Frantar and D. Alistarh. 2023. Massive language models can be accurately pruned in one-shot. arXiv preprint\narXiv:2301.00774.\n\n[C] Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. 2023. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pages 38087–38099. PMLR."}]

Author rebuttal:

Meta-review: This paper addresses the high communication cost of tensor-parallelism in training LLMs by introducing minor adjustments that reduce communication by 50% without sacrificing accuracy, validated on 1B and 7B models. The approach is simple yet effective, accelerating both training and inference. While the method lacks theoretical backing, the experimental setup is reasonable and the algorithm itself is sound.

Concerns about the scale and model size are not critical here: reducing tensor-parallel communication is an important challenge, and this work makes a valuable contribution toward addressing it. Although it is not systematically proven that all transformers would benefit, the paper’s ablations help clarify the method’s limitations. Overall, the contribution is novel, meaningful and timely, and the rebuttal adequately addressed the main weaknesses.
Decision: Accept (poster)