Tensor Parallelism

"They cut my weight matrix into eight strips and handed one to each of us. I only ever see an eighth of the answer; the all-reduce at the end is where we agree on what we computed."

A Shard That Believes It Is the Whole Model

In Section 16.1 the binding ceiling changed. Data parallelism, the subject of Chapter 15, assumes the whole model fits on one accelerator and replicates it; once a single layer's weight matrix no longer fits, replication is impossible and the model itself must be cut apart. Tensor parallelism is the most fine-grained way to make that cut. It does not move whole layers to different devices (that is pipeline parallelism, the subject of Section 16.3); it reaches inside one layer and partitions the matrix multiply that defines it, so that the layer's parameters, its activations, and the work of computing them are divided $T$ ways at once.

The idea rests on a single algebraic fact about matrix multiplication: a large matmul can be expressed as a sum or concatenation of smaller matmuls on slices of the operands. Tensor parallelism is the disciplined exploitation of that fact across $T$ devices, and the discipline matters, because a careless partition forces a communication after every linear layer, while the right one (the Megatron-LM pattern) forces only one all-reduce per pair of layers. We build the right pattern from the algebra, then run it.

1. Splitting a Single Matmul Two Ways Intermediate

Consider one linear layer, $Y = XW$, with input activations $X$ of shape $B \times d_{\text{in}}$ and a weight matrix $W$ of shape $d_{\text{in}} \times d_{\text{out}}$. There are two natural ways to slice $W$ across $T$ devices, and they behave very differently at combining time. The first splits $W$ by columns, the second by rows, and the whole art of tensor parallelism is choosing which to use where.

Column-parallel. Partition $W$ into $T$ column blocks, $W = [\,W_1 \;\; W_2 \;\; \cdots \;\; W_T\,]$, where block $W_t$ has shape $d_{\text{in}} \times (d_{\text{out}}/T)$ and lives on device $t$. Give every device the full input $X$. Each device computes a slice of the output columns,

No device holds the whole output, but no communication is needed to produce the slices: the result is the concatenation of what each device already has. The columns are simply distributed, $\frac{1}{T}$ of them per device.

Row-parallel. Now partition $W$ into $T$ row blocks, $W = [\,W_1^\top \;\; \cdots \;\; W_T^\top\,]^\top$, where block $W_t$ has shape $(d_{\text{in}}/T) \times d_{\text{out}}$. This time the input must also be split along its feature dimension, $X = [\,X_1 \;\; \cdots \;\; X_T\,]$ with $X_t$ of shape $B \times (d_{\text{in}}/T)$, and the product becomes a sum of partial products over the shared contraction dimension,

Here every device produces a full-width output, but each is only a partial result; the true $Y$ is their sum. Summing one same-shape tensor held on each device and returning the total to all of them is exactly an all-reduce, the collective introduced in Section 4.3. So a column-parallel layer needs no communication to compute its output, and a row-parallel layer needs one all-reduce. The Megatron insight is to chain them so the column-parallel concatenation feeds the row-parallel sum without ever gathering anything in between.

2. The Megatron Transformer Block Advanced

A transformer block is built from exactly the structure the previous insight rewards: two consecutive linear maps with a nonlinearity wedged between them. In the MLP sub-block the map is $Y = \mathrm{GeLU}(XA)\,B$, and in the attention sub-block it is the projection of the concatenated heads followed by an output projection. Megatron-LM (Shoeybi et al., 2019) makes the first matrix column-parallel and the second row-parallel, and arranges the partition so the nonlinearity in between acts independently on each device's slice.

This last point is what makes the scheme exact rather than approximate. A nonlinearity does not distribute over a sum, so $\mathrm{GeLU}(\sum_t X_t W_t) \neq \sum_t \mathrm{GeLU}(X_t W_t)$; if the activation sat after a row-parallel layer it would force a gather, apply, re-scatter, three communications instead of one. By placing the column-parallel matrix first, each device holds a complete, correct slice of hidden units ($\mathrm{GeLU}(XA_t)$ is the true GeLU of those units, because column $t$ depends only on $A_t$), applies the nonlinearity locally, and feeds the result straight into its own row block of the second matrix. The communication is deferred to a single point: the all-reduce that sums the row-parallel output. Counting both sub-blocks, the forward pass of one transformer block costs exactly two all-reduces (one for attention's output projection, one for the MLP), and by symmetry the backward pass costs two more.

Figure 16.2.1 makes the dataflow concrete: blue column blocks feed green local activations feed orange row blocks, and only at the bottom does anything cross between devices. The memory accounting is the payoff. Each device stores $\frac{1}{T}$ of $A$, $\frac{1}{T}$ of $B$, and $\frac{1}{T}$ of the hidden activations, so a block that needed $M$ bytes of weights and activations on one accelerator now needs about $\frac{M}{T}$ per device. That is precisely the relief Section 16.1 argued for: the layer that did not fit now fits, $T$ times over.

3. Tensor Parallelism From Scratch Intermediate

The code below implements the full column-parallel then row-parallel block on $T$ simulated devices, with the combining all-reduce, and checks the result against the single-device layer. There is no framework and no network; the "devices" are loop iterations, so the only thing being tested is the algebra of Section 1 and the locality of the nonlinearity from Section 2. The GeLU is applied per device, exactly as Megatron requires.

import numpy as np

rng = np.random.default_rng(0)
B, d_in, d_h, d_out, T = 4, 16, 32, 16, 4   # batch, in, hidden, out, devices

X  = rng.standard_normal((B, d_in))
W1 = rng.standard_normal((d_in, d_h))    # column-parallel weight
W2 = rng.standard_normal((d_h, d_out))   # row-parallel weight

def gelu(z):
    return 0.5 * z * (1.0 + np.tanh(0.7978845608 * (z + 0.044715 * z**3)))

# Single-device reference: Y = gelu(X @ W1) @ W2
ref = gelu(X @ W1) @ W2

# Tensor-parallel across T devices.
# Column-parallel: split W1 by COLUMNS -> each device owns d_h/T hidden units,
# applies its own nonlinearity locally, NO communication needed yet.
# Row-parallel:    split W2 by ROWS to match -> each device produces a PARTIAL
# output over the full d_out; the combining all-reduce SUMS the partials.
W1_shards = np.split(W1, T, axis=1)   # each (d_in,  d_h/T)
W2_shards = np.split(W2, T, axis=0)   # each (d_h/T, d_out)

partials = []
for t in range(T):
    h_t = gelu(X @ W1_shards[t])      # local activation slice, device t
    y_t = h_t @ W2_shards[t]          # partial full-width output, device t
    partials.append(y_t)

allreduced = np.sum(partials, axis=0)  # the ONE forward all-reduce

print("devices T            :", T)
print("hidden units / device:", d_h // T)
print("max abs difference   :", f"{np.max(np.abs(allreduced - ref)):.2e}")
print("relative error       :", f"{np.linalg.norm(allreduced - ref) / np.linalg.norm(ref):.2e}")

devices T            : 4
hidden units / device: 8
max abs difference   : 1.07e-14
relative error       : 1.90e-16

The relative error is at the level of floating-point addition, the same signature of exactness seen for data parallelism in Section 1.1. Splitting the operator changed nothing about the answer; it only changed which device holds which piece and inserted one all-reduce to put the pieces back together. That all-reduce is the whole story of where tensor parallelism can run.

4. Why It Must Live on the Fast Interconnect Advanced

Count the communications. A model with $L$ transformer blocks fires $2L$ all-reduces in the forward pass and $2L$ in the backward, every single training step. For a large model $L$ is in the dozens or low hundreds, so tensor parallelism issues hundreds of all-reduces per step, each one blocking: the row-parallel layer cannot finish until the sum has returned to all devices. Data parallelism, by contrast, issues a handful of all-reduces per step (one per gradient bucket) and overlaps them with the backward computation. Tensor parallelism cannot hide its all-reduces nearly as well, because each one sits on the critical path between two matmuls that depend on it.

This volume of blocking, latency-sensitive communication has a hard consequence: tensor parallelism is viable only on the fastest interconnect available, which inside a modern accelerator node is NVLink (or its switched form, NVSwitch), offering an order of magnitude more bandwidth and far lower latency than the network between nodes. The topology argument of Section 4.9 applies in full force here: a collective is only as fast as the slowest link it crosses, so an all-reduce that must traverse the inter-node network would stall every layer on the slow wire. The practical rule that falls out is sharp and widely followed: keep the tensor-parallel degree $T$ within one node, at or below the number of GPUs sharing that node's NVLink fabric (commonly 8). Beyond one node, you switch to a coarser-grained parallelism, pipeline (Section 16.3) or sharded data parallel (Section 16.4), whose communication is rarer and easier to hide.

# pip install megatron-core ; assumes the process group spans one node's GPUs
from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear

# fc1 splits its OUTPUT (hidden) dim across the tensor-parallel group;
# fc2 splits its INPUT dim to match and all-reduces the partial outputs.
fc1 = ColumnParallelLinear(d_in,  d_hidden, gather_output=False)   # no gather: feed slices on
fc2 = RowParallelLinear(d_hidden, d_out,    input_is_parallel=True) # all-reduce happens here

def mlp(x):
    h, _ = fc1(x)          # each rank holds its hidden-unit slice
    h = gelu(h)            # nonlinearity applied locally, exactly as in Code 16.2.1
    y, _ = fc2(h)          # RowParallelLinear fires the combining all-reduce
    return y

5. Communication, Memory, and the Limits of the Technique Intermediate

It is worth stating plainly what tensor parallelism does and does not buy. It divides the weights, the optimizer state, and the activations of each parallelized layer by $T$, which is the relief that lets an oversized layer fit. It keeps the computation exact, with no approximation beyond floating-point rounding, as Output 16.2.1 showed. And it is perfectly load-balanced, since every device does an identical $\frac{1}{T}$ slice of the same matmul, unlike pipeline parallelism where some stages can idle. Against those gains stands one cost that dominates everything: communication volume. The all-reduce inside every layer makes tensor parallelism the most communication-intensive of the model-parallel strategies, which is why it never travels alone past one node and why it is almost always composed with the other axes rather than used by itself.

That composition is the subject of the rest of the chapter. Real large-model training stacks tensor parallelism within each node, pipeline parallelism across nodes, and sharded data parallelism across the whole cluster, an arrangement called 3D parallelism that Section 16.9 assembles. Each axis owns the scale at which its communication is affordable: tensor parallelism on NVLink inside the node, pipeline and data parallelism on the slower network between nodes. Getting that mapping right is the difference between a model that trains and one that idles its GPUs waiting on the wrong wire.

We now have the finest-grained cut of a model: split the operator itself, pay one all-reduce per layer-pair, and keep it on NVLink. The next section steps up one level of granularity to pipeline parallelism, which cuts the model between layers rather than inside them, trades the in-layer all-reduce for a much rarer point-to-point handoff, and so reaches across nodes where tensor parallelism cannot. That story begins in Section 16.3.