Distributed training overview: Exercises — AI-ML deep dive

This page is a self-testing ladder. Each problem states what you know and asks one clean question. Solutions are hidden inside collapsible callouts — try first, then reveal. Levels climb from "can you name it?" up to "can you design a real system?".

Everything here builds on the parent overview. If a symbol feels unfamiliar, that is a signal to re-read the relevant callout there before pushing on.

Figure s03 — Communication cost $T_{comm}$ plotted against device count $k$ . The teal line is linear ring all-reduce $O (k)$ ; the plum line is tree-based $O (lo g k)$ ; the dotted orange line is the (unrealistic) constant assumption we sometimes use for clean arithmetic. Reading it: both real curves climb as GPUs are added, so the "constant $T_{comm}$ " is an optimistic floor, not the truth.

Level 1 — Recognition

Can you name the mechanism and read the vocabulary?

L1.1 — Match the strategy

Each device holds a full copy of the model but sees a different slice of the mini-batch. Name this strategy, and name the collective operation that averages the gradients.

Recall Solution

Data parallelism. The averaging step is the all-reduce operation: every device ends up holding the same summed-and-divided gradient $g_{global} = \frac{1}{k} \sum_{i = 1}^{k} g_{i}$ . Why not model parallelism? In model parallelism no device has a full copy — the giveaway phrase here is "full copy of the model."

L1.2 — Read the picture

Look at the figure below. One coloured band shows layers 0–11 living on GPU 0, the next band shows layers 12–23 on GPU 1, and so on. Is this tensor parallelism or pipeline parallelism? What is the idle-time problem it creates called?

Figure s01 — Pipeline parallelism. Four coloured horizontal bands each represent one GPU holding a contiguous group of layers (GPU 0 holds layers 0–11, GPU 1 holds 12–23, and so on). The downward arrows on the right show activations being passed device-to-device as data flows through the layer groups in order.

Recall Solution

This is pipeline parallelism (the model is split vertically by layers — inter-layer splitting). Tensor parallelism would split a single layer's matrix across devices (intra-layer). The idle-time problem is the pipeline bubble: while GPU 0 works on batch 1, GPUs 1–3 have nothing yet to do, so they sit idle until data reaches them.

L1.3 — Sync or async?

A worker reads parameters $θ_{t}$ , computes a gradient, but by the time it pushes the update the global model is already at $θ_{t + 5}$ . Which training mode is this, and what is the name for that gradient?

Recall Solution

This is asynchronous SGD. The gradient is a stale gradient — it was computed against an out-of-date copy of the weights. In synchronous SGD every device waits, so staleness cannot happen.

Level 2 — Application

Plug into the formulas.

L2.1 — Effective batch size

You run data parallelism on $k = 8$ GPUs, each processing a micro-batch of $b = 16$ images (no accumulation, so $a = 1$ ). What effective batch size $B$ does the optimizer "see"?

Recall Solution

$B = b \times k \times a = 16 \times 8 \times 1 = 128.$ All 8 gradients get averaged into one update, so the math is identical to training on a single batch of 128 images on one (imaginary huge) GPU.

L2.2 — Gradient accumulation steps

Your GPU can only fit a batch of $b = 2$ , but you want the training behaviour of $B = 32$ on a single GPU ( $k = 1$ ). How many micro-batches $a$ must you accumulate before each optimizer step?

Recall Solution

$a = \frac{B}{b} = \frac{32}{2} = 16 accumulation steps .$ You run 16 forward+backward passes, averaging their gradients $g_{accum} = \frac{1}{16} \sum_{i = 1}^{16} g_{i}$ , then do one weight update. Why divide by 16? Each micro-batch gradient $g_{i}$ is already the average loss-gradient over its own 2 samples. If you just summed the 16 of them, the total would be 16× too large — as if the learning rate were secretly multiplied by 16. Dividing by $a = 16$ restores the correct per-sample average, so the accumulated gradient has exactly the same scale as a real batch-of-32 gradient computed in one shot. That scale-matching is the whole point: same $η$ works unchanged. It is equivalent to batch 32, just 16× slower per update.

L2.3 — Speedup from the formula

Using $S = \frac{T _{compute}}{T _{compute} / k + T _{comm}}$ with $T_{compute} = 400$ ms, $k = 4$ , and $T_{comm} = 50$ ms, compute the actual speedup $S$ .

Recall Solution

Per-device compute: $T_{compute} / k = 400/4 = 100$ ms. $S = \frac{400}{100 + 50} = \frac{400}{150} \approx 2.667.$ So instead of the ideal $4 \times$ , communication drags you down to about $2.67 \times$ .

Level 3 — Analysis

Compare, diagnose, reason about behaviour.

L3.1 — Efficiency drop

From L2.3 you found $S \approx 2.667$ with $k = 4$ . Compute the scaling efficiency. Then explain in one sentence why doubling to $k = 8$ (same $T_{comm} = 50$ ms) will make efficiency worse, not better.

Recall Solution

$Efficiency = \frac{S}{k} = \frac{2.667}{4} \approx 0.667 (66.7%) .$ At $k = 8$ : compute per device $= 400/8 = 50$ ms, so $S = 400/ (50 + 50) = 4.0$ , and efficiency $= 4/8 = 0.5$ (50%). Why worse: as $k$ grows, the shared compute term $T_{compute} / k$ shrinks toward zero, but the "constant" $T_{comm} = 50$ ms does not (and in reality it grows with $k$ ) — communication becomes a larger fraction of each iteration, so efficiency falls.

L3.2 — Bandwidth swap

BERT-Large has 340M float32 parameters. Each gradient is one float32 (4 bytes). Compute the raw gradient volume in gigabytes (use $1 GB = 1 0^{9}$ bytes). Then, if a 10 Gbit/s link gives $T_{comm} = 50$ ms and you upgrade to a 100 Gbit/s link, estimate the new $T_{comm}$ and the new speedup with $T_{compute} = 400$ ms, $k = 4$ .

Recall Solution

Volume $= 340 \times 1 0^{6} \times 4 = 1.36 \times 1 0^{9}$ bytes $= 1.36$ GB. Communication time is inversely proportional to bandwidth. A $10 \times$ faster link gives roughly $10 \times$ less time: $T_{comm}^{new} = \frac{50}{10} = 5 ms .$ New speedup: $S = \frac{400}{100 + 5} = \frac{400}{105} \approx 3.81$ , efficiency $\approx 3.81/4 = 0.95$ (95%). Lesson: the algorithm never changed — only the interconnect. Communication is a hardware bottleneck as much as a software one.

L3.3 — Why the pipeline stalls

A pipeline splits a model over 4 devices. The first micro-batch takes 4 "stage-times" to reach the last device. If you push only 1 micro-batch through, how many of the $4 \times 4 = 16$ stage-time slots are actually doing useful work, and what fraction is bubble?

Figure s02 — Pipeline-bubble grid. Rows are the 4 GPUs, columns are 4 time slots. Only the diagonal cells (teal, labelled "work") do useful computation — one busy device per time slot. All 12 off-diagonal cells (hatched, labelled "idle") are the bubble: devices waiting for data to arrive or having already finished.

Recall Solution

With a single micro-batch, exactly 4 slots are busy (the diagonal in the figure — each device works exactly once). The other $16 - 4 = 12$ slots are idle. Bubble fraction $= 12/16 = 0.75$ (75% idle!). This is why real pipeline training pushes many micro-batches in flight — the fixed startup/drain bubble gets amortised over more useful work, shrinking the idle fraction.

Level 4 — Synthesis

Design a scheme by combining pieces.

L4.1 — Hybrid layout

You must train a 100B-parameter model on 64 GPUs, arranged as 8 nodes × 8 GPUs. The model needs 8-way splitting just to fit in memory. Describe the hybrid layout: how many data-parallel replicas do you get, and what runs inside each replica?

Recall Solution

Use 8-way model parallelism inside each node (the model is spread over the 8 GPUs of one node so it fits), then replicate that whole node 8 times across the 8 nodes as data parallelism.

Data-parallel replicas: $64/8 = 8$ replicas.
Inside each replica: 8-way model parallelism. Gradients are all-reduced across the 8 replicas; within a replica the GPUs cooperate on one forward/backward of the split model.

L4.2 — Memory-and-batch plan

GPT-2 (1.5B params) OOMs at batch 2 on your 16 GB GPU, but batch 1 fits. You have 4 such GPUs and want an effective batch of 32. Combine data parallelism and gradient accumulation to hit exactly 32. Give the numbers.

Recall Solution

Here all three multipliers of $B = b \times k \times a$ are in play. Per GPU, max micro-batch that fits: $b = 1$ . Across $k = 4$ GPUs (data parallel), one synced step processes $b \times k = 1 \times 4 = 4$ samples. To reach $B = 32$ , accumulate over $a$ steps: $b \times k \times a = 32 \Rightarrow 1 \cdot 4 \cdot a = 32 \Rightarrow a = 8$ accumulation steps. Plan: each GPU runs 8 forward/backward passes of batch 1, averaging its own 8 gradients; then one all-reduce across the 4 GPUs; then one weight update. Effective batch $= b \cdot k \cdot a = 1 \cdot 4 \cdot 8 = 32$ . ✓

L4.3 — When does async win?

You have a cluster where one "straggler" GPU is 3× slower than the others. In synchronous SGD, every iteration waits for it. Argue whether asynchronous SGD helps, and name its cost.

Recall Solution

Sync SGD is limited by the slowest device: every iteration pays the straggler's time. With a 3× slower GPU, the whole cluster runs at (roughly) the straggler's pace — the fast GPUs idle. Async SGD helps here: workers push updates whenever ready, so fast GPUs are not blocked by the straggler → better hardware utilization. The cost: stale gradients — a fast worker may apply a gradient computed against outdated weights, which can slow or destabilise convergence. So async trades deterministic correctness for throughput.

Level 5 — Mastery

One integrated question, reasoned end to end.

L5.1 — Full system budget

A team trains a model where a single full iteration would take $T_{compute} = 1200$ ms on one device. They scale to $k = 6$ data-parallel GPUs. The gradient tensor is 2 GB, and their interconnect moves 40 GB/s during all-reduce (treat $T_{comm} = volume / bandwidth$ ). (a) Compute $T_{comm}$ . (b) Compute $S$ and efficiency. (c) They consider adding gradient accumulation of $a = 4$ micro-batches to grow the effective batch — does that change $S$ per update, and what happens to communication frequency? Reason it out.

Recall Solution

(a) $T_{comm} = \frac{2 GB}{40 GB/s} = 0.05 s = 50 ms .$ (b) Per-device compute $= 1200/6 = 200$ ms. $S = \frac{1200}{200 + 50} = \frac{1200}{250} = 4.8, Efficiency = \frac{4.8}{6} = 0.8 (80%) .$ (c) Gradient accumulation runs $a = 4$ forward/backward passes before a single all-reduce. Why does the single-device baseline become $4 \times$ larger? Because "one update" now means processing $4 \times$ as many samples ( $B = b \cdot k \cdot 4$ ). To compare fairly, the baseline must also do that same total work on one device — and a lone GPU doing 4× the samples takes 4× the time, so its per-update time is $4 \times 1200 = 4800$ ms. We scale the numerator precisely because the unit of work per update grew.

Compute per update grows 4× (four micro-batches), but communication happens once per update instead of once per micro-batch → the all-reduce is amortised over 4× more compute.
The ratio $T_{comm} / T_{compute}$ shrinks, so efficiency improves even though wall-clock per update goes up.
Numerically, per update: distributed time $= 4 \times 200 + 50 = 850$ ms versus one-device baseline $4800$ ms, giving ratio $\frac{4800}{850} \approx 5.65$ , efficiency $\approx 5.65/6 = 0.94$ (94%). Big-picture insight: accumulation is a lever that trades update frequency for communication efficiency — fewer, bigger updates make the fixed sync cost matter less.

Recall Quick self-check ledger (reveal after finishing)

L2.1 :: B = 128 L2.2 :: a = 16 accumulation steps L2.3 :: S ≈ 2.667 L3.1 :: efficiency ≈ 0.667 at k=4, drops to 0.5 at k=8 L3.2 :: 1.36 GB; new comm 5 ms; S ≈ 3.81, efficiency ≈ 0.95 L3.3 :: 4 useful slots, bubble fraction 0.75 L4.2 :: b=1, k=4, a=8 accumulation steps L5.1 :: T_comm = 50 ms; S = 4.8; efficiency 0.8; with accum ratio ≈ 5.65, eff ≈ 0.94

Connections

Parent: Distributed training overview
The gradient being averaged comes from 3.9-Gradient-descent-optimizers on top of 3.1-Neural-network-basics.
Big batches interact with 3.3.5-Batch-normalization statistics.
Model-parallel giants like Transformers motivate tensor/pipeline splits.
Hardware context: 4.2-GPU-acceleration and 5.1-Training-at-scale.