ZeRO and FSDP split the storage of a model across GPUs, but every GPU still does all of the math for its own tokens. Before each layer runs, the full weights are gathered onto every GPU, and every GPU stores the full activations of its own sequences. For a 70B model at 8,192 tokens that means 182 GB of activations per sequence, which no amount of sharding removes.

Tensor parallelism takes the opposite approach. Split each weight matrix itself across $t$ GPUs, have each GPU do $1/t$ of every matrix multiply on the same tokens, and combine the partial results with a collective. Each GPU then stores $1/t$ of the weights and computes and stores only part of the activations. Shoeybi et al. (2019) worked out, in the Megatron-LM paper, how to do this for a Transformer with only two all-reduces per layer in the forward pass, using nothing but the algebra of block matrices. This section derives their splits from scratch, checks them in NumPy and PyTorch, counts the communication, and explains why tensor parallelism almost never leaves a single node.

One matrix multiply, split two ways

Take a linear layer $\mathbf{Y} = \mathbf{X}\mathbf{A}$ with input $\mathbf{X} \in \R^{n\times d}$ ($n$ tokens, which is $bT$ for a micro-batch of $b$ sequences of $T$ tokens) and weight $\mathbf{A} \in \R^{d\times k}$, split over $t$ GPUs. There are two natural ways to cut $\mathbf{A}$.

Column parallel. Cut $\mathbf{A}$ into column blocks, $\mathbf{A} = [\mathbf{A}_1 \mid \cdots \mid \mathbf{A}_t]$ with $\mathbf{A}_i \in \R^{d\times k/t}$. Then

GPU $i$ holds the full input $\mathbf{X}$ and its block $\mathbf{A}_i$, and computes $\mathbf{Y}_i = \mathbf{X}\mathbf{A}_i$, the $i$-th block of output columns. No communication is needed in the forward pass. The output is left split along its last dimension.

Row parallel. Cut $\mathbf{A}$ into row blocks $\mathbf{A}_i \in \R^{d/t\times k}$, and the input correspondingly into column blocks $\mathbf{X} = [\mathbf{X}_1 \mid \cdots \mid \mathbf{X}_t]$ with $\mathbf{X}_i \in \R^{n\times d/t}$. Block multiplication gives

GPU $i$ holds $\mathbf{X}_i$ and $\mathbf{A}_i$ and computes a partial sum $\mathbf{X}_i\mathbf{A}_i$ of the full shape $n\times k$. The real output is the sum of the partial sums, so the GPUs must all-reduce them.

In plain words, splitting the weight's columns splits the output and needs no communication, while splitting its rows splits the input and leaves partial sums that must be added.

The backward pass mirrors this. With upstream gradient $\mathbf{G} = \partial\mathcal{L}/\partial\mathbf{Y}$ (§0.2), the column-parallel layer has $\partial\mathcal{L}/\partial\mathbf{A}_i = \mathbf{X}^\top\mathbf{G}_i$, local to GPU $i$, but $\partial\mathcal{L}/\partial\mathbf{X} = \sum_i \mathbf{G}_i\mathbf{A}_i^\top$, a sum over GPUs, so the backward pass needs an all-reduce. The row-parallel layer needs the full $\mathbf{G}$ on every GPU, which it has (the forward all-reduce gave every GPU the full $\mathbf{Y}$), and then $\partial\mathcal{L}/\partial\mathbf{X}_i = \mathbf{G}\mathbf{A}_i^\top$ and $\partial\mathcal{L}/\partial\mathbf{A}_i = \mathbf{X}_i^\top\mathbf{G}$ are both local. Each split needs exactly one all-reduce, one in the forward pass and the other in the backward pass.

import numpy as np

def column_parallel(X, A, t):
    return [X @ Ai for Ai in np.split(A, t, axis=0)]

def row_parallel(X, A, t):
    raise NotImplementedError

The Megatron MLP, one all-reduce

Now the feed-forward network of a Transformer block (§5.3), $\mathbf{Z} = \operatorname{GELU}(\mathbf{X}\mathbf{A})\mathbf{B}$ with $\mathbf{A} \in \R^{d\times 4d}$ and $\mathbf{B} \in \R^{4d\times d}$. Which way should each matrix be split?

Try splitting $\mathbf{A}$ by rows first. Then $\mathbf{X}\mathbf{A} = \sum_i \mathbf{X}_i\mathbf{A}_i$ is a sum of partial products, and we would have to all-reduce it before the GELU, because $\operatorname{GELU}(\mathbf{U}_1 + \mathbf{U}_2) \ne \operatorname{GELU}(\mathbf{U}_1) + \operatorname{GELU}(\mathbf{U}_2)$. That costs a synchronization in the middle of the block.

Split $\mathbf{A}$ by columns instead. GPU $i$ gets $\mathbf{X}\mathbf{A}_i$, a block of columns of the hidden activation, and because GELU is applied elementwise, $\operatorname{GELU}(\mathbf{X}\mathbf{A}) = [\operatorname{GELU}(\mathbf{X}\mathbf{A}_1) \mid \cdots \mid \operatorname{GELU}(\mathbf{X}\mathbf{A}_t)]$. Each GPU can apply it to its own block with no communication. The hidden activation is now split along its columns, which is exactly the input a row-parallel layer wants. So split $\mathbf{B}$ by rows, and

In plain words, each GPU runs a narrow MLP with $4d/t$ hidden units on all the tokens, and the narrow MLPs' outputs are added with one all-reduce.

Shoeybi et al. name the two conjugate communication points $f$ and $g$. Operator $f$ sits where the replicated input enters the parallel region. It is the identity in the forward pass and an all-reduce in the backward pass (the column-parallel $\partial\mathcal{L}/\partial\mathbf{X}$ sum). Operator $g$ sits where the partial sums leave. It is an all-reduce in the forward pass and the identity in the backward pass. So the MLP costs one all-reduce forward and one backward.

Biases need one piece of care. The first layer's bias $\mathbf{b}_1 \in \R^{4d}$ splits with $\mathbf{A}$'s columns, each GPU adding its own slice before the GELU. The second layer's bias $\mathbf{b}_2 \in \R^d$ is not split, and it must be added once, after the all-reduce, not on every GPU before it, or it would be counted $t$ times. Gated feed-forward networks such as SwiGLU split the same way, with the gate and up projections column-parallel (with matching column blocks, so that each GPU multiplies the right gate with the right value) and the down projection row-parallel.

import numpy as np

def tp_mlp(X, A, b1, B, b2, t):
    partials = [gelu(X @ Ai + b1) @ Bi + b2 for Ai, Bi in zip(np.split(A, t, axis=1), np.split(B, t, axis=0))]
    return sum(partials), partials

Attention, split by heads

Attention splits even more naturally, because its heads are already independent (§4.5). Give each GPU $h/t$ of the layer's $h$ heads. The query, key, and value projections are column-parallel, each GPU holding the columns of $\mathbf{W}_Q$, $\mathbf{W}_K$, and $\mathbf{W}_V$ that belong to its heads. Each GPU then computes full attention for its heads, on all tokens, with no communication at all. The output projection $\mathbf{W}_O$ is row-parallel, because the identity of §4.5, $[\mathbf{H}^{(1)}\mid\cdots\mid\mathbf{H}^{(h)}]\mathbf{W}_O = \sum_i \mathbf{H}^{(i)}\mathbf{W}_O^{(i)}$, says that the head outputs times their row blocks of $\mathbf{W}_O$ add up. One all-reduce, $g$, finishes the attention block, and $f$ sits at its input.

Two constraints follow. The number of heads must be divisible by $t$. With grouped-query attention and $n_{kv}$ key-value heads, each GPU also needs the key-value heads its query heads use, so $t$ should divide $n_{kv}$, or the key-value heads are replicated. Llama 3 70B, with 64 query heads and 8 key-value heads, splits cleanly up to $t = 8$, one key-value head per GPU.

Put together, a Transformer layer under tensor parallelism communicates four times per training step, two all-reduces in the forward pass ($g$ after attention and after the MLP) and two in the backward pass ($f$ at their inputs), each of the full activation $\mathbf{X} \in \R^{bT\times d}$. The cell below checks the attention split numerically.

Embeddings and the loss

The embedding table and the output layer are the largest single matrices in many models, $V\times d$ each (for Llama 3 70B, $128{,}256 \times 8192$, a billion parameters), so they are split too.

Vocabulary-parallel embedding. Split the table by rows, so GPU $i$ holds the embeddings of token ids in its range of $V/t$. Each GPU looks up the tokens that fall in its range, writes zeros for the others, and an all-reduce sums the results, so each token's row comes from exactly one GPU.

Vocabulary-parallel cross-entropy. The output layer is column-parallel, so GPU $i$ ends with logits for its $V/t$ vocabulary entries. Gathering the full $bT\times V$ logits would be expensive (gigabytes, as §10.1 showed). Instead compute the loss from three small all-reduces. The cross-entropy for a token with logits $\mathbf{z}$ and target $y$ is $\log\sum_v e^{z_v} - z_y$. Each GPU computes its local maximum, and an all-reduce with the max operation gives the global maximum $\mu$ for numerical stability. Each GPU computes $\sum_{v\in\text{shard}} e^{z_v - \mu}$, and an all-reduce sums them. The GPU that owns $y$ contributes $z_y$ and the others contribute 0, and an all-reduce sums those. Every all-reduce is over $bT$ numbers instead of $bTV$.

import numpy as np

def vocab_parallel_ce(logit_shards, targets):
    z = np.concatenate(logit_shards, axis=1)                  # gathers everything, and is unstable
    return float(np.mean(np.log(np.exp(z).sum(1)) - z[np.arange(len(targets)), targets]))

What tensor parallelism costs

Memory. Each GPU holds $1/t$ of every layer's large matrices, so weights, gradients, and optimizer states all shrink by $t$ (norm gains and a few biases are replicated, which is negligible). Activations shrink too, but not completely. Korthikanti et al. (2022) redo the count of §10.1 under tensor parallelism. The tensors inside the parallel regions are split, $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, the attention matrices, the $\mathbf{W}_O$ input, and the two wide MLP tensors, $24Tbd$ bytes plus $5hT^2b$. The tensors outside them are replicated on every GPU, the two norm inputs, the inputs to the first linear of each block, and the two dropout masks, $10Tbd$ bytes. So

The replicated $10Tbd$ is a floor that does not fall with $t$, and we remove it below.

Communication. Each all-reduce is of a $bT\times d$ activation in 16 bits, $2bTd$ bytes, so each GPU sends $2(t-1)/t \cdot 2bTd$ bytes per all-reduce, and four of them per layer per step. Compare with the compute. A layer's forward pass costs about $24bTd^2$ FLOPs for its weight products (§4.7), split over $t$ GPUs, against two all-reduces. With per-GPU throughput $F$ and bandwidth $\mathcal{B}$, the ratio of communication time to compute time is

In plain words, the bigger the model's width and the faster the link, the cheaper tensor parallelism is, and every extra GPU in the group makes it more expensive.

The number of tokens has cancelled. Unlike data parallelism, no batch size can hide this communication, because it scales with the tokens exactly as the compute does. Worse, it sits on the critical path. The next layer cannot start until the all-reduce finishes, so it is hard to overlap.

For a width-8192 model at $t = 8$ over NVLink, communication takes about 0.6 times as long as compute at peak throughput (less at realistic throughput, but it adds directly to the step time). Over InfiniBand the same split takes 5.6 times as long as the compute. That is the reason for the rule every practitioner knows. Tensor parallelism stays within a node, usually at $t = 8$, the number of GPUs sharing an NVSwitch. Larger NVLink domains, such as 72-GPU racks, relax the rule. Splitting matrices finer also makes each GPU's matrix multiplies smaller and less efficient, another reason not to push $t$ past the point where memory is solved.

Techniques that hide some of this communication do exist. Asynchronous tensor parallelism (or collective matmul, Wang et al. 2023) splits each matrix multiply into chunks and overlaps sending one chunk with computing the next, and fused kernels can overlap the all-reduce with the following computation. They reduce the exposed time but do not change the bytes.

Sequence parallelism removes the replicated activations

Look again at the $10Tbd$ bytes that tensor parallelism leaves replicated. They belong to the LayerNorms, the dropouts, and the residual adds, which work on each token independently. Korthikanti et al. observed that these regions can be split along the sequence instead, each GPU handling $T/t$ of the tokens. They called this sequence parallelism.

The transitions change. Entering a tensor-parallel region, the column-parallel layer needs every token, so the sequence-split activations must be all-gathered. Leaving it, the row-parallel layer's partial sums must be added, and the result is to be split by sequence again. An all-reduce followed by keeping only your own tokens is exactly a reduce-scatter. So the conjugate pair $f$, $g$ is replaced by a new conjugate pair. Korthikanti et al. call them $g$, an all-gather in the forward pass and a reduce-scatter in the backward pass, and $\bar g$, a reduce-scatter in the forward pass and an all-gather in the backward pass. (They rename Megatron's $f$ and $g$ as $f$ and $\bar f$, so watch the letters when you read the two papers side by side.)

Here is the beautiful part. Tensor parallelism did one all-reduce at each point. Sequence parallelism does one all-gather and one reduce-scatter instead, and an all-reduce is a reduce-scatter plus an all-gather (§10.2). So the communication volume is unchanged, and every activation is now split $t$ ways,

and with selective recomputation or FlashAttention, $34\,Tbd/t$. (To remove the whole $10Tbd$, Korthikanti et al. also store only the local sequence shard of the norm's output, the input of the first column-parallel layer, and all-gather it again during the backward pass when its weight gradient needs it. That one extra all-gather per block is the only added traffic, and they overlap it with computation.)

def tp_layer_cost(T, b, d, h, t, sp=False, recompute="none"):
    act = T * b * d * (34 + 5 * h * T / d) / t
    return int(act), 0.0

In PyTorch

The $f$ and $g$ operators are four lines each as custom autograd functions, which is how Megatron-LM implements them. The code below needs a tensor-parallel process group and is shown for reference.

PyTorch also has tensor parallelism built in, on top of distributed tensors. You describe each submodule's split, and parallelize_module shards the weights and inserts the collectives. With Shard(1) layouts (the sequence dimension) at the boundaries, the same plan gives sequence parallelism. The PrepareModuleInput entries gather the tokens once per block, at its input, rather than once for each of the three projections. This plan follows PyTorch's tensor-parallel tutorial for a Llama-style model, whose attention takes the hidden states and the rotary frequencies.

On a single process we can still check the core claim of this section, that the sharded computation gives exactly the dense gradients, by letting autograd play the role of the collectives. Summing the partial outputs is $g$, whose backward is the identity, and giving each shard its own copy of the input makes $f$'s backward a sum of their gradients.

Read more

• Shoeybi et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. The column and row splits, the f and g operators, and an 8.3B model trained with 8-way tensor parallelism.
• Korthikanti et al. (2022). Reducing Activation Recomputation in Large Transformer Models. Sequence parallelism, the activation formulas under tensor parallelism, and selective recomputation.
• Narayanan et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. Section 3 shows why tensor parallelism should stay inside a node.
• Wang et al. (2023). Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models. Collective matmuls that hide tensor-parallel communication.
• PyTorch. Large Scale Transformer model training with Tensor Parallel. The parallelize_module interface and sequence parallelism on distributed tensors.
• Tazi et al. (2025). The Ultra-Scale Playbook. Its tensor and sequence parallelism chapters measure throughput as the degree grows across the node boundary.

Watch next

▶

▶