Model parallelism (tensor, pipeline): Visual walkthrough

The parent note told you that tensor parallelism "partitions individual weight matrices... across multiple devices" and that the Transformer MLP uses a column-parallel then row-parallel pattern that needs only one all-reduce per block. That is the central result. This page earns every piece of it with a picture, starting from what a matrix multiply even looks like.

We will build up to this claim:

Prerequisites we lean on: Transformer Architecture (what an MLP block is), GPU Memory Hierarchy (why one GPU runs out of room), and the sibling idea Data Parallelism (which copies the whole model instead of cutting it). Parent: Model parallelism.

Step 1 — What a matrix multiply actually is

WHAT. Before splitting anything, we must agree on what $Y = X W$ means. A matrix is just a grid of numbers. Multiplying $X$ (the input, one row per example in the batch) by $W$ (the weights) produces $Y$ (the output).

Read the symbols slowly: $output b \times d_{o u t} Y = input b \times d_{in} X weights d_{in} \times d_{o u t} W$

$b$ ::: how many examples we push through at once (the batch).
$d_{in}$ ::: how many numbers describe each input example.
$d_{o u t}$ ::: how many numbers come out per example.

WHY. The whole reason we care about splitting matrices is that $W$ is enormous — in GPT-3 a single weight grid can hold hundreds of millions of numbers, more than one GPU can hold (see GPU Memory Hierarchy). So we must understand its shape before we cut it.

PICTURE. One output number is a row of $X$ dotted against a column of $W$ — slide across, multiply pairwise, add up. That "column of $W$ decides one column of $Y$ " fact is the seed of everything.

Step 2 — The one fact that lets us cut $W$ into columns

WHAT. We chop $W$ vertically into $N$ side-by-side slabs, each a bunch of columns: $W = [cols 1.. k W_{1} ∣ cols k + 1..2 k W_{2} ∣ \dots ∣ W_{N}], k = \frac{d _{o u t}}{N}$

The bars $∣$ ::: mean "placed next to each other, side by side" (concatenation of columns).
$W_{i}$ ::: device $i$ 's slab — shape $d_{in} \times \frac{d _{o u t}}{N}$ .
$k = d_{o u t} / N$ ::: how many output columns each device owns.

Then, from Step 1's fact that each column of $W$ independently makes one column of $Y$ : $Y = X [W_{1} ∣ W_{2} ∣ \dots ∣ W_{N}] = [Y_{1} X W_{1} ∣ Y_{2} X W_{2} ∣ \dots ∣ X W_{N}]$

WHY this tool — concatenation, not addition. We could have split $W$ into rows instead, but a column split has a magic property: each slab $W_{i}$ produces a whole, finished slice $Y_{i}$ of the output, needing no help from the other devices. No sum, no talking. That independence is exactly what we exploit next.

PICTURE. Colour each column-slab of $W$ a different colour; the output $Y$ inherits the same coloured stripes. Device 1 owns the cyan stripe, device 2 the amber stripe — they never look at each other's stripe.

Step 3 — Column-parallel costs a copy of $X$ , not a conversation

WHAT. Each device $i$ computes locally: $Y_{i} = X W_{i}$ Device $i$ ends up holding only its stripe $Y_{i}$ (shape $b \times \frac{d _{o u t}}{N}$ ).

$X$ ::: every device needs the full input — so we broadcast one copy to all.
$Y_{i}$ ::: only $1/ N$ of the output lives on each device.

WHY. This is the memory win the parent promised: each device stores $\frac{d _{o u t}}{N}$ output columns and $\frac{1}{N}$ of $W$ 's numbers instead of all of them. The price is duplicating $X$ — but $X$ is tiny compared to $W$ , so this is a great trade.

PICTURE. Two GPU boxes. The same $X$ arrow feeds both; each box lights up only its own coloured output stripe. Notice: no arrow runs between the boxes. That empty gap is the point.

Step 4 — GELU is applied element-by-element, so the split survives

WHAT. In the real MLP the first layer is $h = GELU (x W_{1})$ . $GELU$ is a nonlinearity: it takes each single number and bends it (roughly, it softly zeroes out negatives).

The crucial property: $GELU ([a ∣ b]) = [GELU (a) ∣ GELU (b)]$

$GELU$ ::: acts on one number at a time — it never mixes columns together.
$[a ∣ b]$ ::: our already-split hidden stripes.

WHY this matters. If $GELU$ mixed columns (like a softmax across the whole row would), we'd be forced to gather $h$ back together before applying it — a full communication step. Because $GELU$ is element-wise, device $i$ can apply it to its own stripe $h_{i}$ and stay split. This is precisely why the parent could say "No communication yet!" after $W_{1}$ .

PICTURE. Each hidden stripe passes through its own little GELU curve; the stripes never touch. The colour boundaries are preserved.

Step 5 — Now switch to a row-split for $W_{2}$ so the pieces recombine

WHAT. The hidden state $h$ is split into column-stripes $h_{i}$ . The second weight $W_{2}$ maps $4 d_{m o d e l} \to d_{m o d e l}$ . We split $W_{2}$ the other way — into rows: $W_{2} = W_{2}^{(1)} W_{2}^{(2)} ⋮ W_{2}^{(N)}, W_{2}^{(i)} \in R^{\frac{4 d _{m o d e l}}{N} \times d_{m o d e l}}$

Because $h$ is split by columns and $W_{2}$ by matching rows, the multiply becomes a sum of partial products: $y = h W_{2} = [h_{1} ∣ h_{2} ∣ \dots] W_{2}^{(1)} W_{2}^{(2)} ⋮ = \sum_{i = 1}^{N} y_{i} h_{i} W_{2}^{(i)}$

$h_{i}$ ::: device $i$ 's hidden stripe (from Step 4).
$W_{2}^{(i)}$ ::: device $i$ 's matching row-slab of $W_{2}$ .
$y_{i} = h_{i} W_{2}^{(i)}$ ::: a full-width, partial output — right shape, but only part of the true value.
$\sum$ ::: add all the partials to get the real $y$ .

WHY row-split now. A row-split of $W_{2}$ lines up perfectly with the column-split of $h$ : dimensions match so each device can multiply its own $h_{i}$ by its own $W_{2}^{(i)}$ with no reshuffling of $h$ . The only debt we incur is the final $\sum$ .

PICTURE. Each device produces a faint, full-size output tile $y_{i}$ (correct shape, partial value). Stacked and summed, the faint tiles add up to the solid true output.

Step 6 — The single all-reduce (and why it's exactly one)

WHAT. The devices each hold a partial $y_{i}$ . Getting the true $y = \sum_{i} y_{i}$ onto every device is one collective operation called an all-reduce.

$y_{i}$ ::: each device's partial contribution.
The all-reduce ::: performs $\sum_{i} y_{i}$ and broadcasts it — one round of talking.

WHY only one. Trace the whole block: broadcast $X$ (input, cheap) → column-parallel $W_{1}$ (silent) → GELU (silent) → row-parallel $W_{2}$ (silent, produces partials) → one all-reduce. The clever column-then-row pairing pushed all the communication into a single point. A naive design that gathered after $W_{1}$ and after $W_{2}$ would need two rounds and move far more data.

PICTURE. Two GPUs holding faint partial tiles; amber arrows converge into a central " $Σ$ " node and fan back out carrying the solid full result to both.

Step 7 — Edge and degenerate cases (never leave the reader stranded)

WHAT / WHY / PICTURE, three quick corners:

$N = 1$ (one device). The sum $\sum_{i = 1}^{1} y_{i} = y_{1}$ collapses to plain $y = h W_{2}$ . No split, no all-reduce — tensor parallelism gracefully becomes ordinary computation. Good sanity check.
$d_{o u t}$ not divisible by $N$ . Our slab width $k = d_{o u t} / N$ must be a whole number of columns. If $d_{o u t} = 10, N = 3$ , you cannot split evenly. Real systems pad $d_{o u t}$ up to a multiple of $N$ (or choose $N$ dividing it) — e.g. attention heads, which are already an integer count, split cleanly.
The backward pass mirrors the forward. The parent note showed $\frac{\partial L}{\partial X} = \sum_{i} \frac{\partial L}{\partial Y _{i}} W_{i}^{T}$ — a sum over devices, i.e. another all-reduce, but now on the input-gradient. Forward's all-reduce sits at the block output; backward's sits at the block input. Symmetric bookends.

The one-picture summary

The full journey on one blueprint: broadcast $X$ → column-split $W_{1}$ (silent) → element-wise GELU (silent) → row-split $W_{2}$ (silent, partials) → one all-reduce → $y$ .

Recall Feynman retelling — say it back in plain words

Imagine a big multiplication job written on a giant sheet too heavy for one worker to hold. We cut the first weight-sheet into vertical strips and give each worker one strip plus a photocopy of the input. Each worker finishes their own vertical stripe of the answer without talking to anyone — that's the free lunch of a column split. Then a bending step (GELU) that only ever looks at one number at a time keeps those stripes separate, so still nobody talks. For the second weight-sheet we cut it into horizontal strips that line up with everyone's stripe, so each worker again multiplies just their own pieces — producing a full-size but incomplete answer. Finally everyone tosses their incomplete answer into one pile, we add them up, and hand the finished total back to all of them. That single "throw it in the pile and share the total" step is the one all-reduce. Because we cleverly went column-first then row, we only had to talk once per MLP block — and the backward pass does the exact same trick, just at the other end.

Quick self-check

Bubble-free understanding check

In a column-parallel layer, why does each device need the whole input

X

but only part of

W

Answer

Because a column of the output depends on all of

X

but only one column of

W

; the columns of

W

are independent, so each device owns some output columns and needs full

X

to compute them.

Why one all-reduce, not two?

Column-parallel

W_{1}

produces split output, element-wise GELU preserves the split, and row-parallel

W_{2}

turns the split back into a single summable set of partials — so communication is deferred to exactly one point.