One Matrix, Two Cuts: Tensor Parallelism From Scratch

TP cuts the computation of a single layer, and the Column→Row order is dictated by the nonlinearity. Unlike most tutorials, this post derives the backward pass in full, which reveals a beautiful fact rarely spelled out. In TP, every weight gradient is communication-free, and communication only ever touches activations. The experiments verify a hand-written Column+Row MLP numerically exact against a single GPU, then measure the TP=2/4/8 communication share on pure PCIe.

1. Until now, computation was never cut

The previous four posts sharded storage. ZeRO/FSDP split parameters, gradients and optimizer state, but every forward pass all-gathers the parameters back, so the matrix multiply itself still runs whole on every GPU. When a single layer’s compute or activations are themselves too big or too slow, you need tensor parallelism (TP), which splits \(Y = XW^\top\), one matmul, across \(N\) GPUs, each doing \(1/N\) of it.

There are two natural directions to cut. Notation (PyTorch layout, series-wide): weight \(W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}\), activations \(X \in \mathbb{R}^{bs \times d_{\text{in}}}\), the layer computes \(Y = XW^\top\):

• Column cut (Megatron’s ColumnParallelLinear): split along \(d_{\text{out}}\), so rank \(k\) holds \(W_k = W[k\frac{d_{\text{out}}}{N}:(k{+}1)\frac{d_{\text{out}}}{N},\,:]\), which is complete rows. \(X\) is replicated, and the output comes out naturally column-blocked as \(Y_k = XW_k^\top\) with zero communication, but each rank has only \(1/N\) of the output.
• Row cut (RowParallelLinear): split along \(d_{\text{in}}\). This requires the input to arrive blocked the same way (\(X_k\)), and the output is a partial sum \(Y = \sum_k X_k W_k^\top\), one all-reduce away from complete.

A naming trap, resolved once. Megatron’s Column/Row refer to the math convention \(Y=XW\) with \(W\in\mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}\) (Column = the output dim), but in PyTorch’s [out, in] tensor layout that is precisely a dim-0 (row) cut. The one question that never misleads is “is each row (fan-in vector) still whole?” The Column cut says yes and the Row cut says no. This distinction becomes life-or-death in post #9.

2. The golden pair: two defects that cancel

Each cut alone has a defect. Column leaves the output blocked while Row demands a blocked input. Megatron’s insight is that the two defects cancel exactly, because Column’s blocked output is the blocked input Row wants:

The full MLP dataflow (\(Z = \mathrm{GeLU}(XW_1^\top)W_2^\top\)):

The GeLU lands exactly on the “each rank owns its complete block” state, so apply it element-wise and be done. One all-reduce per MLP forward, at the exit. Why can’t the order flip? Because GeLU is nonlinear:

Row-cut first, and fc1’s output is a partial sum that must be all-reduced before the GeLU, which forces communication into the middle of the layer, two collectives instead of one. The position of the nonlinearity dictates the cut order. That is the entire design of Megatron TP. Attention is the same story, more natural still. Each head is already an independent unit, so QKV projections are Column-cut (whole heads per rank), the output projection \(W_O\) is Row-cut, and softmax (per-head) sits in the middle communication-free.

3. The backward pass: derive every gradient, see where communication really is

Most tutorials skip this part, yet it answers why TP’s communication is as cheap as it is. Backward through the dataflow above (write \(\bar{A} \equiv \frac{\partial L}{\partial A}\), given \(\bar{Z}\), and note that \(g\)’s backward is identity, so \(\bar{Z}\) is replicated):

(1) The Row layer’s weight gradient costs no communication:

\(\bar{Z}\) is on every rank (replicated) and \(Y_k\) lives on this rank already, so both multiplicands are local and no communication is needed. Note also that \(\bar{W}_{2,k}\) is exactly the gradient of this rank’s own shard, so storage and update stay fully local (naturally orthogonal to ZeRO/FSDP).

(2) The intermediate activation’s gradient costs no communication:

This uses only this rank’s \(W_{2,k}\) and this rank’s activations, so it is local.

(3) The Column layer’s weight gradient costs no communication:

\(X\) is replicated (\(f\)’s forward is a copy) and \(\bar{Y}^{\text{pre}}_k\) is local, so this one is free again.

(4) The input’s gradient is the one and only communication:

Each rank can compute only its own term of a partial sum, hence an all-reduce. That is \(f\)’s backward.

Put the four together and a clean structural fact emerges:

All weight gradients are communication-free. Communication happens only at activation boundaries (\(f\)’s backward, \(g\)’s forward), and \(f/g\) are conjugates, so one copies forward and sums backward while the other does the reverse. The per-layer bill is 1 forward AR + 1 backward AR for the MLP, plus the same pair for attention, which makes 4 activation all-reduces per layer per step.

Why are weight gradients free? Because the golden pair guarantees that each weight shard’s two multiplicands, its own activation block and the replicated boundary activation, are both local. That is not luck. It is what “aligning the cut geometry with the dataflow” buys. This view runs deeper than memorizing “Column pairs with Row,” and it generalizes, because to audit any new parallelism scheme’s communication you can stare at each gradient formula and ask “are the multiplicands local?”

4. The bill: TP moves activations, and that is its essential difference from DP

DP synchronizes gradients each step (\(\sim 2\Psi\) bytes, independent of batch), but TP moves activations \([bs \times h]\) every layer (proportional to batch, independent of parameter count). The ratio (\(L\) layers, \(\Psi \approx 12Lh^2\), 4 ARs per layer):

(the common \(2\frac{N-1}{N}\) factor cancels, same dtype assumed.) Past \(3h\) tokens per rank (GPT-2 Large: 3840), TP out-communicates DP, which in training is essentially always. And TP’s communication comes as serial, per-layer, latency-critical packets wedged between matmuls (hard to overlap), while DP’s is one large per-step transfer that bucketing hides (post #2). This is the quantitative version of post #0’s iron rule that tp stays inside the node.

5. Experiment: exactness verified, and the PCIe reality check

A hand-written Column+Row MLP (~40 lines of core, ships with the post), GPT-2 Large geometry (h=1280, ff=5120, 4096 tokens), checked against a single GPU:

The 1e-6 forward error is fp32 summation-order rounding (all-reduce adds in a different order than a fused matmul), and backward agrees to 1e-12. TP, like DP, is an exact algorithm, not an approximation, so “zero communication in the middle” costs no precision.

Speed (one MLP forward, compute/communication decomposed):

• Compute divides perfectly by N: 1.32 → 0.64 → 0.34 ms, so mathematically TP is impeccable.
• Communication grows with N: the all-reduce message is constant (\([4096\times1280]\) fp32 = 21 MB) but the ring gets longer and crosses more NUMA boundaries, so timing goes 1.07 → 1.66 → 2.09 ms. Reconciling with post #1, 21 MB in 2.09 ms → algbw ≈ 10 GB/s, precisely the measured 8-GPU all-reduce plateau ✓.
• Total time doesn’t move (2.39 → 2.30 → 2.43 ms), while the communication share climbs 45% → 72% → 86%. On a machine without NVLink, TP hands every saved FLOP to the wires, which is not TP’s failure but the measured proof of “TP must live inside a fast-interconnect domain.” On NVLink (~25× the bandwidth) the same experiment’s comm term divides by ~25 and TP=8’s share falls back to ~20%.

6. Reading along in real source

Megatron-LM keeps ColumnParallelLinear / RowParallelLinear in megatron/core/tensor_parallel/layers.py, with the \(f/g\) operators in mappings.py: copy_to_tensor_model_parallel_region = \(f\) and reduce_from_tensor_model_parallel_region = \(g\), names that literally say copy-in / reduce-out.

nanotron has TensorParallelColumnLinear (SplitConfig(split_dim=0), the tensor-row cut) and TensorParallelRowLinear (split_dim=1) in src/nanotron/parallel/tensor_parallel/nn.py. The differentiable primitives in distributed_differentiable_primitives.py implement \(f/g\) as autograd.Functions, which is §3’s derivation as code.

Our bench_tp.py strips those classes to the bone: one shard per rank, one all_reduce, all of §3 reproducible in 40 lines.

7. Summary

1. TP is the first scheme to cut computation. Column (rows whole, output blocked) + Row (rows severed, output partial-summed) form the golden pair, the nonlinearity’s position dictates the order, and the middle is communication-free.
2. Deriving the backward term by term shows that weight gradients are all communication-free and only activation boundaries communicate (\(f\)-backward + \(g\)-forward), 4 activation ARs per layer per step. The general audit method is to stare at each gradient formula and ask are the multiplicands local?.
3. TP comm ∝ activations (\(bsh\)) while DP comm ∝ parameters (\(\Psi\)), with ratio \(bs/3h\). TP is almost always the bigger bill, serial and hard to overlap, which is the quantitative basis for keeping tp inside nodes.
4. Measured, TP is an exact algorithm (error = rounding order). On PCIe, compute scales perfectly ÷N while total time stays flat (comm share 86%), which is the iron rule as an experiment.

Next comes Sequence & Context Parallelism. TP cut the weights, but each rank’s activations are still the full \([b \times s \times h]\). The parts TP cannot reach (LayerNorm, dropout) and the long-context attention problem both call for cutting along the sequence dimension, and they get two different schemes for two different problems.

Environment: 8× RTX PRO 6000 Blackwell, PyTorch 2.9.1, NCCL 2.27.5. Reproduce: torchrun --standalone --nproc_per_node={2,4,8} bench_tp.py. Plotting and schematic code accompanies the series.

All benchmark scripts, schematic generators, plotting code and raw result CSVs for this post live in assets/blog/code/05-tensor-parallel.

Appendix: The Code That Ran

Every number in this post comes from the scripts below, embedded verbatim. Plotting and schematic code plus the raw result CSVs live in the folder linked above.

"""
TP experiment (part 05 of the Illustrated Distributed Training series): a hand-written
Megatron-style Column+Row parallel MLP.

Verifies two things:
  1. Correctness: an MLP with Column-cut fc1 + Row-cut fc2 matches the single-GPU
     forward output and backward gradients exactly (bit-for-bit in fp64, up to
     rounding in bf16). "Zero communication in the middle" is not an approximation.
  2. Cost: 1 all-reduce per layer in forward (the g operator), 1 in backward
     (the f operator). Measures the per-layer time breakdown (compute vs comm)
     for TP=2/4/8.

Usage:
  torchrun --standalone --nproc_per_node={2,4,8} bench_tp.py --out ../results/tp.csv
"""

import argparse
import csv
import os
import time

import torch
import torch.distributed as dist
import torch.nn.functional as F

H = 1280               # hidden size (GPT-2 Large scale)
FF = 4 * H             # 5120
MBS, SEQ = 4, 1024
WARMUP, STEPS = 20, 50


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="../results/tp.csv")
    args = ap.parse_args()

    rank = int(os.environ["RANK"])
    world = int(os.environ["WORLD_SIZE"])
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group("nccl", device_id=device)

    torch.manual_seed(1337)  # same seed on every rank so weights match
    # full weights (fp32 so results can be reconciled exactly)
    W1 = torch.randn(FF, H, device=device) / H**0.5      # fc1 [out=FF, in=H]
    W2 = torch.randn(H, FF, device=device) / FF**0.5     # fc2 [out=H, in=FF]
    X = torch.randn(MBS * SEQ, H, device=device, requires_grad=True)

    # ---- single-GPU reference ----
    ref = F.gelu(X @ W1.t()) @ W2.t()
    ref_loss = ref.square().mean()
    ref_loss.backward()
    ref_grad = X.grad.clone()
    X.grad = None

    # ---- TP: Column-cut W1 (along out dim), Row-cut W2 (along in dim) ----
    shard = FF // world
    W1_k = W1[rank * shard:(rank + 1) * shard]            # [FF/N, H] whole rows
    W2_k = W2[:, rank * shard:(rank + 1) * shard]         # [H, FF/N] column slice

    Xtp = X.detach().clone().requires_grad_(True)
    # forward: f operator = identity (X already replicated), intermediate Y_k local, g operator = all-reduce
    Y_k = F.gelu(Xtp @ W1_k.t())                          # [B, FF/N] intermediate activation: zero comm
    Z_k = Y_k @ W2_k.t()                                  # [B, H] partial sum
    Z = Z_k.clone()
    dist.all_reduce(Z)                                    # g: the only forward comm
    loss = Z.square().mean()
    loss.backward()                                       # backward: dX is a partial sum
    gX = Xtp.grad.clone()
    dist.all_reduce(gX)                                   # f's backward: all-reduce dX

    fwd_err = (Z - ref).abs().max().item()
    grad_err = (gX - ref_grad).abs().max().item()

    # ---- timing: break down compute vs communication ----
    def timed(fn):
        s, e = torch.cuda.Event(True), torch.cuda.Event(True)
        dist.barrier(); torch.cuda.synchronize()
        s.record()
        for _ in range(STEPS):
            fn()
        e.record(); torch.cuda.synchronize()
        return s.elapsed_time(e) / STEPS

    def fwd_only():
        z = F.gelu(Xtp @ W1_k.t()) @ W2_k.t()
        return z

    def fwd_with_ar():
        z = F.gelu(Xtp @ W1_k.t()) @ W2_k.t()
        dist.all_reduce(z)
        return z

    for _ in range(WARMUP):
        fwd_with_ar()
    t_compute = timed(fwd_only)
    t_total = timed(fwd_with_ar)
    t_comm = t_total - t_compute

    if rank == 0:
        row = [world, fwd_err, grad_err, round(t_compute, 3), round(t_comm, 3),
               round(t_total, 3), round(t_comm / t_total * 100, 1)]
        newfile = not os.path.exists(args.out)
        os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
        with open(args.out, "a", newline="") as f:
            w = csv.writer(f)
            if newfile:
                w.writerow(["tp", "fwd_max_err", "grad_max_err", "compute_ms", "comm_ms", "total_ms", "comm_pct"])
            w.writerow(row)
        print("ROW:", row, flush=True)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()