Model parallelism (tensor, pipeline): Exercises

This page is a self-test. Each problem states everything it needs; the solution is hidden inside a collapsible callout so you can try first, then check. Work them in order — they climb from "can you name it" up to "can you design a system."

Parent: Model Parallelism (Tensor, Pipeline).

Before we start, we lock the symbols so nothing appears un-earned:

Definition Symbol dictionary for this page

$b$ ::= batch size (how many samples flow through together).
$d_{in}, d_{o u t}$ ::= number of input / output features of a linear layer.
$d_{m o d e l}$ ::= the "width" of a Transformer (feature count of the main stream).
$N$ ::= number of devices (GPUs) we split a single layer across (tensor parallelism).
$D$ ::= number of pipeline stages (GPUs stacked depth-wise).
$M$ ::= number of micro-batches (a mini-batch chopped into $M$ small pieces).
$[A ∣ B]$ ::= horizontal concatenation (glue matrices side by side, more columns).
all-reduce ::= every device sends its partial result, they get summed, everyone gets the sum.
all-gather ::= every device sends its slice, everyone ends up holding the full glued-together result.
bubble ::= wasted GPU time where a device sits idle waiting for data.

The core memory arithmetic we reuse everywhere:

Level 1 — Recognition

Exercise 1.1 (L1)

A single linear layer computes $Y = X W$ . One sentence: does tensor parallelism cut $W$ within this layer, or does it place this whole layer on one device and the next layer elsewhere?

Recall Solution

Tensor parallelism cuts $W$ within the layer — every device keeps a slice of the same weight matrix. Placing whole layers on different devices is pipeline parallelism (inter-layer). Mnemonic: Tensor = Thin slices of one matrix; Pipeline = Piles of layers.

Exercise 1.2 (L1)

A 7-billion-parameter model is stored in FP16. How many gigabytes do the parameters alone take? (Use $1 GB = 1 0^{9}$ bytes.)

Recall Solution

$7 \times 1 0^{9} params \times 2 bytes = 14 \times 1 0^{9} bytes = 14 GB$ . Why FP16 = 2 bytes: 16 bits = 2 bytes; that is the whole reason mixed precision (Mixed Precision Training) halves weight memory versus FP32.

Exercise 1.3 (L1)

In the GPipe timeline, symbols $F_{i}$ and $B_{i}$ appear. What does each stand for, and why must every $F_{i}$ finish before its matching $B_{i}$ starts?

Recall Solution

$F_{i}$ = forward pass of micro-batch $i$ ; $B_{i}$ = backward pass of micro-batch $i$ . The backward pass needs the loss, which only exists after the forward pass has produced an output. So $F_{i}$ always precedes $B_{i}$ — cause before effect.

Level 2 — Application

Exercise 2.1 (L2)

A layer has $W \in R^{1024 \times 4096}$ (so $d_{in} = 1024$ , $d_{o u t} = 4096$ ). We do a column-wise split across $N = 4$ devices. Give the shape of each device's slice $W_{i}$ and the shape of the local output $Y_{i}$ if the batch is $b = 8$ .

Recall Solution

Column split cuts the output dimension: each device holds $d_{o u t} / N = 4096/4 = 1024$ columns. $W_{i} \in R^{1024 \times 1024}, Y_{i} = X W_{i} \in R^{8 \times 1024} .$ Full output is $Y = [Y_{1} ∣ Y_{2} ∣ Y_{3} ∣ Y_{4}] \in R^{8 \times 4096}$ , glued along columns (all-gather). See figure below.

Exercise 2.2 (L2)

A Transformer MLP block uses $W_{1} : d_{m o d e l} \to 4 d_{m o d e l}$ and $W_{2} : 4 d_{m o d e l} \to d_{m o d e l}$ with $d_{m o d e l} = 2048$ . Split across $N = 8$ devices (column-parallel $W_{1}$ , row-parallel $W_{2}$ ). How many parameters does one device store for this block, and what fraction of the full block is that?

Recall Solution

Full block parameters (ignore bias): $W_{1} : 2048 \times 8192 = 16, 777, 216, W_{2} : 8192 \times 2048 = 16, 777, 216.$ Total $= 33, 554, 432$ params. Both are split by $N = 8$ , so one device stores $33, 554, 432/8 = 4, 194, 304$ params $= 1/8$ of the block.

Exercise 2.3 (L2)

GPipe with $D = 4$ stages and $M = 16$ micro-batches. Compute the bubble fraction $\frac{D - 1}{M + D - 1}$ .

Recall Solution

$\frac{D - 1}{M + D - 1} = \frac{4 - 1}{16 + 4 - 1} = \frac{3}{19} \approx 0.158 = 15.8% .$ Meaning ~16% of GPU time is idle bubble; ~84% is useful work.

Level 3 — Analysis

Exercise 3.1 (L3)

For a column-parallel layer, the input gradient is $\frac{\partial L}{\partial X} = \sum_{i = 1}^{N} \frac{\partial L}{\partial Y _{i}} W_{i}^{⊤}$ . Explain in words why a sum (all-reduce) is unavoidable here, referencing the figure — and state what communication the forward pass needed by contrast.

Recall Solution

Look at figure s01: each output column block $Y_{i}$ depends on the full input $X$ . So when gradient flows back, every column block pushes a correction onto every input feature. Each device only knows its own $W_{i}$ and $\frac{\partial L}{\partial Y _{i}}$ , so it can only compute one term $\frac{\partial L}{\partial Y _{i}} W_{i}^{⊤}$ — a partial answer. The true $\frac{\partial L}{\partial X}$ is the sum of all partials, so devices must add their contributions: all-reduce. Forward, by contrast, produced independent column blocks that only needed to be glued (all-gather), not summed. Asymmetry: forward = concat, backward-to-input = sum.

Exercise 3.2 (L3)

Two engineers debate. A: "GPipe with $M = 32$ is strictly better than $M = 8$ because the bubble is smaller." B: "Not free." Quantify both bubble fractions for $D = 4$ , and state the cost B is pointing at.

Recall Solution

$M = 8 : \frac{3}{8 + 3} = \frac{3}{11} \approx 27.3%, M = 32 : \frac{3}{32 + 3} = \frac{3}{35} \approx 8.6%.$ So $M = 32$ does slash bubble from ~27% to ~9%. But GPipe stores activations for all $M$ micro-batches until the backward sweep. Activation memory scales linearly in $M$ : $M = 32$ needs 4× the activation memory of $M = 8$ . B's point: efficiency bought with memory. Fix uses Activation Checkpointing to recompute activations instead of storing them.

Exercise 3.3 (L3)

The bubble fraction $\frac{D - 1}{M + D - 1}$ has a limit as $M \to \infty$ . Find it, and explain physically what it means to "keep adding micro-batches forever."

Recall Solution

$lim_{M \to \infty} \frac{D - 1}{M + D - 1} = \frac{D - 1}{\infty} = 0 .$ Physically: with infinitely many micro-batches, the pipeline stays permanently full — no stage ever waits, so idle time → 0% of a job that runs forever. In reality $M$ is capped by memory (Ex 3.2) and by the mini-batch size, so we approach 0 but never reach it. The figure shows the curve flattening toward zero.

Level 4 — Synthesis

Exercise 4.1 (L4)

A 20B-parameter model is trained with Adam in FP32. Compute total memory for (a) parameters, (b) Adam optimizer states, (c) grand total including gradients (FP32). Then state the minimum number of 80 GB GPUs needed if you split evenly by tensor parallelism (ignore activations). Use $1 GB = 1 0^{9}$ bytes.

Recall Solution

$P = 20 \times 1 0^{9}$ params.

(a) Parameters FP32: $20 \times 1 0^{9} \times 4 = 80 \times 1 0^{9}$ bytes $= 80 GB$ .
(b) Adam states = 8 bytes/param: $20 \times 1 0^{9} \times 8 = 160 \times 1 0^{9} = 160 GB$ .
Gradients FP32 = 4 bytes/param: $20 \times 1 0^{9} \times 4 = 80 GB$ .
(c) Grand total $= 80 + 160 + 80 = 320 GB$ . Split evenly across GPUs of 80 GB each: $320/80 = 4$ GPUs minimum (ignoring activations; real systems need more headroom).

Exercise 4.2 (L4)

Design question. You have $D = 4$ pipeline stages, mini-batch $32$ , micro-batch size $1$ (so $M = 32$ ). Steady-state 1F1B (PipeDream 1F1B) stores only $D$ micro-batches' worth of activations instead of $M$ . If one micro-batch's activations cost $1.5$ GB, how much activation memory does 1F1B save versus GPipe at peak?

Recall Solution

GPipe peak: stores all $M = 32$ → $32 \times 1.5 = 48 GB$ .
1F1B steady state: stores $D = 4$ → $4 \times 1.5 = 6 GB$ .
Saving $= 48 - 6 = 42 GB$ (a $8 \times$ reduction). Why: 1F1B runs a backward as soon as a forward completes, freeing that micro-batch's activations early instead of hoarding all forwards first.

Level 5 — Mastery

Exercise 5.1 (L5)

Combined regime. A cluster has $16$ GPUs arranged as tensor-parallel groups of $N = 4$ inside pipeline stages of $D = 4$ (so $N \times D = 16$ ). The MLP block per stage uses one all-reduce (tensor, cost $c_{T}$ ) and the pipeline sends activations $D - 1$ times (cost $c_{P}$ per hop). For $M$ micro-batches, write the total communication count in terms of $M$ , and explain which cost you would attack first if $c_{T} ≫ c_{P}$ .

Recall Solution

Tensor communication happens inside every micro-batch's forward and backward on every stage. Per micro-batch: forward all-reduce + backward all-reduce = $2$ all-reduces per MLP block. Over $M$ micro-batches this is $2 M$ all-reduce events per stage.
Pipeline communication: activations hop between adjacent stages $D - 1$ times per micro-batch forward, and again on backward → $2 (D - 1)$ hops per micro-batch, $2 M (D - 1)$ total.
If $c_{T} ≫ c_{P}$ : attack the tensor all-reduce first. It fires $2 M$ times and is the expensive one. Practical fixes: keep tensor-parallel groups inside a single node (NVLink is fast), raise pipeline $M$ (cheap $c_{P}$ ), and use Mixed Precision Training to halve bytes moved in each all-reduce.

Exercise 5.2 (L5)

Efficiency budget. GPipe with $D = 8$ , $M = 32$ . (a) Bubble fraction. (b) If activation checkpointing adds 33% recompute overhead but lets you double $M$ to $64$ within the same memory, compute the new bubble. (c) Net: does the trade help throughput? Give the useful-work fraction each way (useful = $1 - bubble$ , discounted by recompute overhead where present).

Recall Solution

(a) $M = 32$ : bubble $= \frac{8 - 1}{32 + 8 - 1} = \frac{7}{39} \approx 0.1795 = 17.95%$ . Useful work $\approx 82.05%$ . (b) $M = 64$ : bubble $= \frac{7}{64 + 8 - 1} = \frac{7}{71} \approx 0.0986 = 9.86%$ . Raw useful $\approx 90.14%$ . (c) Checkpointing recompute makes each unit of forward work cost $1.33 \times$ . Effective throughput scale $\approx$ useful-fraction $\div$ overhead: $no-ckpt: \frac{0.8205}{1.0} = 0.8205, ckpt: \frac{0.9014}{1.33} \approx 0.6777.$ Conclusion: despite the smaller bubble, the 33% recompute tax outweighs the bubble savings here — raw throughput drops (~0.82 → ~0.68). Checkpointing wins only when the alternative is not fitting at all (out-of-memory), not as a pure speed play. This is the deep point: Activation Checkpointing trades compute for memory, not for speed.

Recall Quick self-check clozes

Tensor parallelism cuts the weight matrix within one layer; pipeline parallelism cuts the model across layers (depth-wise). Column split of $W$ makes output blocks that are combined by all-gather (concatenate); the input gradient is combined by all-reduce (sum). GPipe bubble fraction ::: $\frac{D - 1}{M + D - 1}$ Adam optimizer state cost per parameter ::: 8 bytes (2 states × 4 bytes) 1F1B stores activations for ::: $D$ micro-batches, not $M$