Visual walkthrough — Distributed training overview

The parent note Distributed training overview handed you a formula: split the batch across $k$ devices, average their gradients, update everyone with the same average. It claimed this equals training on one giant batch. This page proves that claim from scratch — with a picture at every step — so you never have to take it on faith.

We build every symbol before using it. If you have never seen a gradient, we build the gradient. Let's go.

---

Step 1 — What is one training step, on ONE machine?

WHAT. Before we split anything, let's draw the thing we are about to split. A neural network (see 3.1-Neural-network-basics) is a machine with a big bag of adjustable numbers. We call that whole bag $\theta$ (the Greek letter "theta"). Training means nudging $\theta$ so the network's mistakes shrink.

WHY this matters. You cannot parallelize a process you cannot describe. So first we describe one honest step of learning on a single GPU.

PICTURE. Look at the figure. The horizontal axis is one knob inside $\theta$ (imagine all the others frozen). The curved black line is the loss — how wrong the model is. The red dot is where we are now. The red arrow points downhill: that is the direction we want to walk.

$$\underset{\text{new knobs}}{\underbrace{\theta }}\leftarrow \underset{\text{old knobs}}{\underbrace{\theta }}-\underset{\text{step size}}{\underbrace{\eta }}\underset{\text{uphill slope}}{\underbrace{{\nabla }_{\theta }\mathcal{L}}}$$

The minus sign is the whole idea: slope points uphill, we subtract it, so we move downhill. That is gradient descent in one line.

---

Step 2 — Where does the gradient come from? A batch of examples.

WHAT. We never compute the loss on the entire dataset at once — too slow. We grab a handful of examples, called a mini-batch $B$, of size $b$. The gradient is the average of the gradient over those $b$ examples. Crucially, that handful must be a random grab from the dataset.

WHY average, not sum? Because we want a number that doesn't change meaning when we change $b$. If we summed, a batch of 100 examples would give a gradient 100× larger than a batch of 1, and our step size $\eta$ would have to be re-tuned every time. Averaging keeps the scale fixed.

WHY random? The mini-batch gradient is only a trustworthy stand-in for the true full-dataset gradient if the batch is an unbiased sample — i.e. randomly drawn so it doesn't systematically over-represent, say, only cat images. A biased grab points the arrow the wrong way.

PICTURE. Each example is a small grey arrow (its own opinion about which way is downhill). The single red arrow is their average — the batch's combined vote.

The bold $\mathbf{g}$ reminds us it is not one number but a whole list (one slope per knob) — a vector.

---

Step 3 — The one fact everything rests on: the average of averages

WHAT. Here is the mathematical heart. Suppose we have one big batch — call it $\mathcal{B}$ (a fancy capital B, to keep it distinct from the mini-batch $B$ of Step 2) — containing $b\times k$ examples. Now chop $\mathcal{B}$ into $k$ equal groups ${\mathcal{B}}_1,\dots ,{\mathcal{B}}_k$, each of size $b$. Claim: the average over the whole big batch equals the average of the $k$ group averages.

WHY it's true — and why it's the only reason data parallelism works. If this fact were false, splitting the batch would give a different gradient than the big batch, and the two would train into different models. So we must check it, not assume it.

PICTURE. On the left, all $b\times k$ arrows averaged in one shot → one red arrow. On the right, average each group first (three faint arrows), then average those three → the same red arrow. Same tip, exactly.

First, a shorthand so the sums stay readable. For a single example numbered $j$ (with input ${\mathbf{x}}_j$ and label $y_j$), write its own gradient as:

$${\nabla }_j:={\nabla }_{\theta }\mathcal{L}(f({\mathbf{x}}_j;\theta ),y_j)$$

• ${\nabla }_j$ ::: the gradient contributed by one example $j$ — one grey arrow from Step 2.
• $:=$ ::: "is defined to be" — the left side is just a new name for the right side.

Now let ${\mathbf{g}}_i$ be the average of these per-example gradients over group ${\mathcal{B}}_i$ (each group has size $b$):

$${\mathbf{g}}_i=\frac{1}{b}\sum _{j\in {\mathcal{B}}_i}{\nabla }_j⟹\sum _{j\in {\mathcal{B}}_i}{\nabla }_j=b{\mathbf{g}}_i$$

Now the average over the whole big batch $\mathcal{B}$, term by term:

$${\mathbf{g}}_{\text{big}}=\underset{\begin{array}{c}\text{divide by}\\ \text{total count}\end{array}}{\underbrace{\frac{1}{bk}}}\underset{\begin{array}{c}\text{sum every example,}\\ \text{grouped}\end{array}}{\underbrace{\sum _{i=1}^k\sum _{j\in {\mathcal{B}}_i}{\nabla }_j}}=\frac{1}{bk}\sum _{i=1}^{k}b{\mathbf{g}}_i=\frac{1}{k}\sum _{i=1}^k{\mathbf{g}}_i$$

Read the cancellation: the $b$ inside the sum meets the $b$ in the denominator and dies. What survives is exactly the average of the $k$ group gradients — which is what each GPU will compute. This is the whole game.

---

Step 4 — Now hand each group to its own GPU

WHAT. We physically place group ${\mathcal{B}}_i$ on GPU $i$. Every GPU holds a full copy of the same $\theta$ and computes only its group's gradient ${\mathbf{g}}_i$. Nobody talks yet.

WHY. Step 3 proved we may compute the pieces separately. This step just spreads those pieces across hardware (see 4.2-GPU-acceleration) so they run at the same time instead of one after another.

PICTURE. Four GPUs, each a black box holding an identical copy of $\theta$ (same colour), each fed a different slice of data, each producing its own grey arrow ${\mathbf{g}}_i$. The arrows point in different directions — because each GPU saw different data.

---

Step 5 — The sync: all-reduce averages the arrows

WHAT. Every GPU sends its ${\mathbf{g}}_i$ around; they get summed and divided by $k$; the identical result ${\mathbf{g}}_{\text{global}}$ lands back on every GPU. This combined "sum-then-share" is called all-reduce.

WHY "all-reduce" and not "gather to one boss"? We could send every arrow to GPU 0, let it average, then broadcast back. But then GPU 0's link is a bottleneck and the others idle. All-reduce is a clever ring where each GPU sends and receives at once, so the total network traffic is spread evenly. Same math, far less waiting.

PICTURE. The four grey arrows funnel into one red arrow ${\mathbf{g}}_{\text{global}}$, copies of which fly back to all four GPUs. After this, every GPU holds the same red arrow.

Now every GPU applies the same update from Step 1, using this shared arrow:

$$\theta \leftarrow \theta -\eta {\mathbf{g}}_{\text{global}}$$

Same starting $\theta$, same ${\mathbf{g}}_{\text{global}}$, same $\eta$ → the copies stay bit-for-bit identical after every step. No drift, ever. The mistake in Step 4 is now cured.

---

Step 6 — The edge case: sync costs time (so it's not a free 4×)

WHAT. Steps 1–5 assumed talking is instant. It isn't. Sending arrows over the network takes time $T_{\text{comm}}$. This is the degenerate case where more GPUs stop helping.

WHY include it. A derivation that ignores its own failure mode is a trap. We must show when the beautiful equivalence stops paying off.

PICTURE. Two stacked timelines. Top: one GPU computes for a long solid black bar. Bottom: each of $k$ GPUs computes for $T_{\text{compute}}/k$ (short bar) but then a red bar $T_{\text{comm}}$ is glued on for the sync. When the red bar dominates, the short compute bar barely helps.

Two limits to feel the shape:

• No talking cost ($T_{\text{comm}}\to 0$): $S\to k$. Perfect linear scaling.
• Talking dwarfs compute ($T_{\text{comm}}\gg T_{\text{compute}}/k$): $S\to T_{\text{compute}}/T_{\text{comm}}$, a constant — adding more GPUs stops helping entirely.

This is exactly why fast interconnects (NVLink, InfiniBand) matter: they shrink the red bar. It's the same batch-statistics reasoning that makes 3.3.5-Batch-normalization and big-batch training in 3.4.1-Transformers-overview worth the coordination.

---

The one-picture summary

One figure, the whole story: identical model copies (same colour) each eat a different data slice, each emit a grey gradient arrow, all-reduce fuses them into one shared red arrow, and every copy steps by that same arrow — landing them all on the same new $\theta$. The little red sync bar in the corner is the price we pay.

Recall Feynman retelling — say it like you're teaching a friend

Imagine four students each grading a different quarter of the same exam pile. Each student computes the average correction their quarter suggests. If they just averaged those four averages, they'd get the exact same answer as one teacher grading the whole pile alone — as long as each student had an equal number of papers, drawn at random. That "average of equal-size random averages equals the big average" is the one fact (Step 3) that lets us split the work at all.

So: give every GPU a full copy of the model and an equal, randomly-chosen slice of the batch. Each computes its own gradient arrow. Then they do an all-reduce — pass arrows around a ring, sum them, divide by the count — so everyone ends up holding the same merged red arrow. Because they all started identical (weights and optimizer memory) and all step by the same arrow, they stay identical forever. No drift.

The catch: passing arrows around costs time. If that sync time is small next to the shrunken per-GPU compute, you get nearly the full $k\times$ speedup. If the network is slow, the sync bar swallows your gains and more GPUs stop helping. That's the entire theory of data-parallel training in one breath.

Recall Quick self-check

Why must every device get the same-size sub-batch? ::: So the $b$ cancels cleanly in Step 3 — unequal sizes would weight some devices' votes wrongly, and the split would no longer equal the big-batch gradient. Why must the data be partitioned randomly? ::: Averaging only recovers the true full-batch gradient if each slice is an unbiased sample; a biased shard (e.g. one class per GPU) makes even the perfect average point the wrong way. What operation shares the averaged gradient to all devices? ::: All-reduce (sum every local gradient, divide by $k$, deliver the result to every device). Why don't the model copies drift apart? ::: They start identical (weights and optimizer state) and every device applies the same ${\mathbf{g}}_{\text{global}}$ with the same $\eta$, so they update in lockstep. Does averaging gradients suffice for Adam/Momentum? ::: Yes only if every replica's optimizer state started identical — then feeding the same shared gradient keeps velocity and second-moment memories in lockstep too. When does adding GPUs stop speeding things up? ::: When $T_{\text{comm}}$ dominates $T_{\text{compute}}/k$; then $S$ flattens to $T_{\text{compute}}/T_{\text{comm}}$.