Model parallelism (tensor, pipeline): Worked examples

This page is the "hit every case" workout for the parent topic. Before we compute anything, let us list every kind of situation model parallelism can throw at you, so no scenario surprises you later.

We use only ideas already built in the parent note. A quick reminder of the two symbols that carry this whole page:

The scenario matrix

Every problem in this topic is one of the cells below. The last column names the example that covers it.

Case class	What makes it tricky	Covered by
Tensor split — column-parallel forward	need full $X$ , concat outputs	Ex 1
Tensor split — row-parallel forward	need all-reduce (a sum)	Ex 2
Tensor split — backward `all-reduce` on input grad	chain rule sums over slices	Ex 3
Degenerate input: $N = 1$	"parallelism" that does nothing	Ex 4
Degenerate input: $d_{o u t}$ not divisible by $N$	uneven slices / padding	Ex 4
Pipeline bubble — naive ( $M = 1$ )	worst-case idle time	Ex 5
Pipeline bubble — limiting $M \to \infty$	efficiency ceiling of 100%	Ex 5
Pipeline memory: GPipe vs 1F1B	activation storage $\propto M$ vs $\propto D$	Ex 6
Real-world word problem: does GPT fit?	mix TP + PP + memory budget	Ex 7
Exam twist: combined TP×PP×DP world size	which numbers multiply?	Ex 8

Now we walk each cell.

Forecast: guess the two numbers first — how many columns of $W$ per device, and how many columns of $Y$ each device produces?

Look at the figure: the tall weight block $W$ is sliced by a vertical chalk line into $W_{1}$ and $W_{2}$ .

Figure — Model parallelism (tensor, pipeline)

Slice the columns. $W = [W_{1} ∣ W_{2}]$ , each $W_{i} \in R^{6 \times 4}$ (since $d_{o u t} / N = 8/2 = 4$ ). Why this step? Column splitting is the "clean" split: each output column of $Y$ depends only on the matching column of $W$ , so no device needs another device's weights.
Broadcast the input. Both devices receive the full $X \in R^{4 \times 6}$ . Why this step? Every output column is a combination of all input features, so $X$ can't be sliced here — only $W$ can.
Local matmul. Device $i$ computes $Y_{i} = X W_{i} \in R^{4 \times 4}$ . Why this step? This is exactly the associativity fact $X [W_{1} ∣ W_{2}] = [X W_{1} ∣ X W_{2}]$ from the parent note.
Concatenate. $Y = [Y_{1} ∣ Y_{2}] \in R^{4 \times 8}$ .

Verify: shapes: $Y_{i}$ has $b = 4$ rows and $d_{o u t} / N = 4$ cols; $2 \times 4 = 8 = d_{o u t}$ . ✓ Weight memory per device $= 6 \times 4 = 24$ numbers vs $6 \times 8 = 48$ single-device — exactly $1/ N$ . ✓

Forecast: in Ex 1 we concatenated. Guess: do we concatenate or sum here?

Look at the figure — the horizontal slice of $W$ pairs with a matching vertical slice of $X$ .

Slice rows of $W$ and columns of $X$ together. $W = [W_{1} W_{2}]$ with $W_{i} \in R^{3 \times 8}$ , and $X = [X_{1} ∣ X_{2}]$ with $X_{i} \in R^{4 \times 3}$ . Why this step? A row of $W$ multiplies a column of $X$ ; if we cut $W$ 's rows we must cut $X$ 's columns to keep the dot products aligned.
Local matmul. Device $i$ computes $Y_{i} = X_{i} W_{i} \in R^{4 \times 8}$ — note $Y_{i}$ is already full size, not a slice. Why this step? Each device produces a partial contribution to every entry of $Y$ , not a separate piece.
All-reduce (a sum). $Y = Y_{1} + Y_{2}$ . Why this step? The matrix product is a sum over the shared inner dimension $d_{in}$ ; splitting that dimension splits the sum, so we must add the partial sums back — this is the "all-reduce" of the parent note.

Verify: $X W = [X_{1} ∣ X_{2}] [W_{1} W_{2}] = X_{1} W_{1} + X_{2} W_{2}$ . That is precisely $Y_{1} + Y_{2}$ . ✓ This is why the Transformer MLP uses column-parallel on $W_{1}$ then row-parallel on $W_{2}$ : the row-parallel step's all-reduce is the single needed communication — see Transformer Architecture.

Forecast: which direction (forward or backward) needed communication in Ex 1? Guess whether the other direction needs it now.

Recall the rule. $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^{⊤}$ . Why this step? Chain rule: to push the error one layer earlier we multiply by the transposed weights.
Substitute the split. $\frac{\partial L}{\partial X} = [G_{1} ∣ G_{2}] [W_{1}^{⊤} W_{2}^{⊤}] = G_{1} W_{1}^{⊤} + G_{2} W_{2}^{⊤}$ . Why this step? A column split of $Y$ makes $W^{⊤}$ a row split — so the product becomes a sum over devices.
All-reduce. Device $i$ computes $G_{i} W_{i}^{⊤} \in R^{4 \times 6}$ locally, then sum across devices. Why this step? Forward concat ⇒ backward all-reduce. The communication just moved to the other pass.

Verify: shapes: $G_{i} W_{i}^{⊤}$ is $4 \times 6 = b \times d_{in}$ , correct for $\frac{\partial L}{\partial X}$ . ✓ Also $\frac{\partial L}{\partial W _{i}} = X^{⊤} G_{i}$ is $6 \times 4$ , matching $W_{i}$ 's shape, and needs no communication. ✓

Forecast: guess (a) the memory saving when $N = 1$ , and (b) whether $8/3$ can give equal slices.

(a) $N = 1$ . Memory saving $= 1/ N = 1/1 = 1$ : no saving. One device holds all $6 \times 8 = 48$ weights and does the whole matmul. Why this step? $N = 1$ is the "identity" of tensor parallelism — the framework runs but there is literally one slice, so there is no concat, no all-reduce, no benefit. Always sanity-check that your split count exceeds 1 before blaming communication overhead.
(b) $N = 3$ , $d_{o u t} = 8$ . $8/3 = 2.67$ is not an integer. Why this step? Slices must be whole columns. Real frameworks (Megatron-LM) require $d_{o u t}$ divisible by $N$ , or they pad $d_{o u t}$ up to the next multiple.
The fix. Pad to $d_{o u t} = 9$ so each device gets $9/3 = 3$ columns; the extra output column is masked/ignored.

Verify: (a) saving ratio $= 1$ . ✓ (b) $9/3 = 3$ integer; padding overhead $= 9/8 - 1 = 12.5%$ wasted columns. ✓

Forecast: guess the idle fraction at $M = 1$ , and what number the fraction approaches for huge $M$ .

Look at the figure: red gaps are idle "bubble" time; feeding more micro-batches shrinks them.

Naive, $M = 1$ . $\frac{4 - 1}{1 + 4 - 1} = \frac{3}{4} = 0.75$ . Why this step? With one big batch, three of four stages sit idle waiting in line — 75% waste, the worst case named in the parent note.
$M = 16$ . $\frac{3}{16 + 3} = \frac{3}{19} \approx 0.158$ . Why this step? Chopping into 16 parcels keeps late stages busy while early stages still work; idle drops to ~16%.
Limit $M \to \infty$ . $\frac{3}{M + 3} \to 0$ . Why this step? As $M$ grows the fill/drain cost is amortised over more work, so utilization approaches 100% — the ceiling you can never quite reach but can approach.

Verify: $3/4 = 0.75$ ; $3/19 = 0.1579 \dots$ ; and $lim_{M \to \infty} 3/ (M + 3) = 0$ . ✓ Efficiency (= $1 -$ bubble) rises $0.25 \to 0.842 \to 1$ . ✓

Forecast: GPipe does all forwards then all backwards; 1F1B interleaves. Guess which one stores $32$ units and which stores $4$ .

GPipe. All $M$ forwards run before any backward, so stage-0 must keep activations for all $M = 32$ micro-batches → $32$ units. Why this step? Backward needs the forward activations; if backward is postponed to the very end, nothing can be freed early. Memory $\propto M$ .
1F1B (PipeDream). In steady state each stage holds at most $D$ micro-batches' activations, so $\approx 4$ units. Why this step? Interleaving "1 forward, 1 backward" frees a micro-batch's activations as soon as its backward completes; only the in-flight ones ( $\leq D$ ) remain. Memory $\propto D$ .
Ratio. GPipe / 1F1B $= 32/4 = 8 \times$ more memory for the same schedule. Why this step? This $M / D$ factor is the entire reason 1F1B exists — see Activation Checkpointing for the orthogonal compute-for-memory trade.

Verify: GPipe $= M = 32$ ; 1F1B $= D = 4$ ; ratio $= 8$ . ✓ Uses Gradient Accumulation-style micro-batch splitting, but memory behaviour differs by schedule.

Forecast: guess the total training memory in GB first, then how many ways you must slice it.

Total memory. $16 B \times 20 \times 1 0^{9} = 320 \times 1 0^{9}$ B $= 320$ GB. Why this step? Adam's optimizer states dominate — this is the "2.1 TB" arithmetic of the parent note at smaller scale. See GPU Memory Hierarchy.
Per-device with split $N$ . Need $320/ N \leq 40$ , i.e. $N \geq 8$ . Why this step? Tensor parallelism slices all of weights, grads and optimizer states by $N$ , so the requirement is a simple division.
Smallest integer. $N = 8$ . Why this step? $N$ must be a whole device count; $8$ is the least integer $\geq 8$ . With Mixed Precision Training (2 B weights) the budget would shrink and $N$ could drop.

Verify: $320/8 = 40 \leq 40$ . ✓ And $320/7 = 45.7 > 40$ , so $N = 7$ fails — $8$ is indeed minimal. ✓

Forecast: guess whether the three numbers add or multiply.

Total world size. $N \times D \times R = 8 \times 4 \times 2 = 64$ GPUs. Why this step? The three parallelisms are orthogonal axes: every (tensor-slice, pipeline-stage, replica) triple is a distinct GPU, so counts multiply, not add.
How many hold "stage 0, slice 0". Fix the tensor slice ( $1$ choice) and the pipeline stage ( $1$ choice); only the data-replica axis is free → $R = 2$ GPUs. Why this step? Data parallelism duplicates the exact same slice across replicas; those are the copies that must sync gradients via all-reduce.
Sanity of the split. Each GPU holds $1/ (N \cdot D) = 1/32$ of the model (TP splits within a stage, PP splits across stages; DP does not shrink per-GPU model size). Why this step? Only TP and PP reduce the model footprint; DP trades memory for throughput.

Verify: $8 \cdot 4 \cdot 2 = 64$ . ✓ Replicas of one specific slice $= R = 2$ . ✓ Model fraction per GPU $= 1/ (8 \cdot 4) = 1/32$ . ✓

Recall Quick self-test

Column-parallel forward reassembles with ::: concatenation (no communication in forward) Row-parallel forward reassembles with ::: an all-reduce (a sum across devices) Naive pipeline bubble fraction with $D = 4, M = 1$ ::: $3/4 = 75%$ GPipe activation memory scales with ::: $M$ (number of micro-batches) 1F1B activation memory scales with ::: $D$ (number of pipeline stages) Combining TP, PP, DP multiplies world size to ::: $N \times D \times R$