ML Systems Internals · Part 2

Distributed training

Q: Your 7B model trains fine on one GPU but you have 8 and want it 8× faster. Someone proposes ZeRO-3. Is that the right first move, and what does it change vs DDP?

No — plain DDP is the right first move. If a full copy already fits on one GPU, DDP replicates it across all 8 and all-reduces gradients each step, giving you the throughput win with the least complexity and communication. ZeRO-3/FSDP shards params, grads, and optimizer state across GPUs — that saves memory (letting you train models that don't fit) but adds all-gather/reduce-scatter comms you don't need when the model already fits. Reach for it only when you hit the memory wall, not the speed wall.

Q: You move a DDP job from a single 8-GPU node (NVLink) to 8 nodes over Ethernet and throughput collapses even though total GPU count is unchanged. Why?

You became comms-bound. DDP's per-step all-reduce moves ~2× the gradient bytes over the interconnect every step. NVLink (~hundreds of GB/s) hides that behind compute; commodity Ethernet (tens of Gb/s) is ~50× slower, so the all-reduce time now dwarfs the compute time — exactly the flip shown by comms_cost.py. The fix is a faster interconnect (InfiniBand), keeping chatty parallelism in-node, gradient accumulation to reduce sync frequency, and overlapping comms with compute.

One GPU can't hold a big model, and even if it could, training would take months. Distributed training splits the work across many GPUs — but the moment you do, communication, not compute, becomes the thing you engineer. This lesson maps the parallelism strategies and when each is actually required.

⏱️ ~2 hours🧪 5 labs🎯 Advanced→Tech-lead

Learning objectives

Explain why one GPU isn't enough — the memory wall and the time wall.
Distinguish data, tensor, and pipeline parallelism and what each splits.
Describe DDP (all-reduce gradients) and the memory problem it does not solve.
Explain ZeRO / FSDP sharding of optimizer state, gradients, and parameters.
Use gradient accumulation to simulate a large batch on small memory.
Reason about communication cost and why it dominates at scale.
Choose a parallelism strategy by model size and hardware.

⚙️ To run this for realThe Python labs below are pure stdlib and run anywhere with python3 — they model the arithmetic so you can reason about it offline. Only the PyTorch FSDP/DDP snippet needs real hardware:

2+ GPUs (single node) or a multi-node cluster with a fast interconnect (NVLink / InfiniBand) + pip install torch, launched with torchrun

Reading and the offline models teach the ideas without any GPUs — run the real snippet when you have a cluster.

1 · Why one GPU isn't enough essential

Two independent walls push training onto many GPUs. The memory wall: a model's weights, its gradients, and the optimizer's state (Adam keeps two extra tensors per weight) must all live in GPU RAM at once — a 7B model in mixed precision needs far more than the 80 GB on a single H100. The time wall: even if it fit, one GPU doing every FLOP for trillions of tokens would take months. Distributed training attacks both — but every split adds a communication cost that becomes the real design constraint.

The optimizer state is the hidden memory hogPeople size a model by its parameter count and forget that Adam stores momentum and variance per parameter, plus a copy of the gradients. In fp32 that's ~16 bytes per parameter (4 params + 4 grad + 8 optimizer). A 7B model is ~112 GB of training state before you count activations — which is exactly why DDP alone (which replicates all of it) hits a wall.

2 · Three ways to split the work essential

There are three orthogonal axes of parallelism. You can — and at scale you do — combine them ("3D parallelism").

🗺️ How to read this diagram

This diagram lays out the three orthogonal ways to spread training across GPUs. They answer different questions — "how do I go faster?" vs "how do I fit a model that doesn't fit?" — and at large scale you combine all three.

Data parallel (left) keeps a full copy of the model on every GPU and splits the batch — each GPU trains on different examples, then they average gradients. Great for speed; does nothing for a model that's too big.
Tensor parallel (middle) splits a single layer's matrices across GPUs so one huge matmul is computed in slices. Very chatty, so it's kept inside one node.
Pipeline parallel (right) splits the stack of layers across GPUs — GPU 0 runs the first layers, GPU 1 the next, and so on, like an assembly line.

In short: Ask of any strategy: what got split? Data parallel splits the batch; tensor parallel splits one layer; pipeline parallel splits the layers. That one question tells you which box you're in.

Strategy	What is split	What each GPU holds	Main cost
Data parallel	the batch	a full copy of the model	all-reduce gradients each step
Tensor parallel	a single layer's matrices	a slice of every layer	all-reduce inside every layer
Pipeline parallel	the stack of layers	a contiguous set of layers	activations across stage boundaries + bubble

Rule of thumb: data parallel when the model fits on one GPU and you just want to go faster; tensor and pipeline when a single copy no longer fits and you must carve the model itself across devices.

3 · Data parallelism & DDP advanced

DistributedDataParallel (DDP) is the workhorse. Each GPU holds a full replica of the model and processes a different shard of the batch. After the backward pass, every replica has different gradients (it saw different data), so they must be averaged. That averaging is an all-reduce: a collective operation (via NVIDIA's NCCL library) that sums each gradient across all GPUs and hands the identical result back to every one, so all replicas step in lockstep and stay identical.

Ring all-reduce moves ~2× the gradient bytesNCCL's ring all-reduce is bandwidth-optimal: each GPU sends and receives roughly 2×(N-1)/N of the full gradient size over the run — approaching 2× the model's gradient bytes regardless of GPU count. That per-step cost is fixed by the interconnect, which is why NVLink/InfiniBand matter so much.

Example code for learning — review, test, and adapt it before running against real or production systems. Commands can create, change, or delete resources. See the Terms & Disclaimer.

Python · all-reduce bytes moved per step (runs)

allreduce.pydef allreduce_bytes(param_count, bytes_per_grad=2, world_size=8):
    """Ring all-reduce moves ~2*(N-1)/N * gradient_size per GPU."""
    grad_bytes = param_count * bytes_per_grad
    factor = 2 * (world_size - 1) / world_size
    per_gpu = grad_bytes * factor
    return {
        "grad_size_GB": round(grad_bytes / 1e9, 2),
        "moved_per_gpu_GB": round(per_gpu / 1e9, 2),
        "note": "approaches 2x grad size as world_size grows",
    }

print(allreduce_bytes(7_000_000_000, world_size=8))
print(allreduce_bytes(7_000_000_000, world_size=64))

{'grad_size_GB': 14.0, 'moved_per_gpu_GB': 24.5, 'note': 'approaches 2x grad size as world_size grows'}
{'grad_size_GB': 14.0, 'moved_per_gpu_GB': 27.56, 'note': 'approaches 2x grad size as world_size grows'}

▶ How this works

This models the traffic DDP generates every single step. After each backward pass the GPUs must agree on one averaged gradient, and this is how many bytes cross the wire to do it.

grad_bytes = param_count * bytes_per_grad is the size of the full gradient (2 bytes each in mixed precision) — for 7B params that's 14 GB.
factor = 2 * (world_size - 1) / world_size is the ring all-reduce constant: each GPU sends and receives about this fraction of the gradient. As the GPU count grows it climbs toward 2×, never past it.
So the bytes moved per GPU are almost independent of how many GPUs you add — the cost is set by the model's gradient size and your interconnect speed.

What the output means: At 8 GPUs each moves ~24.5 GB; at 64 GPUs ~27.6 GB — barely more despite 8× the GPUs. That flatness is why all-reduce scales, and why the wire speed (not the GPU count) governs the per-step comms time.

Try this: Call it with bytes_per_grad=4 (fp32 gradients) and watch every number double — precision directly sets your communication bill.

4 · What DDP doesn't solve — ZeRO & FSDP advanced

DDP makes you faster but not bigger: every GPU still stores a full copy of the parameters, gradients, and optimizer state. If one copy doesn't fit, DDP can't help. ZeRO (Zero Redundancy Optimizer, from Microsoft's DeepSpeed) and PyTorch's FSDP (Fully Sharded Data Parallel) fix this by sharding that state across the data-parallel GPUs instead of replicating it. Each GPU owns 1/N of the state and gathers the rest only when needed, then frees it.

Stage	Shards across GPUs	Memory win
ZeRO-1	optimizer state	~4× on a typical Adam setup
ZeRO-2	+ gradients	~8×
ZeRO-3 / FSDP	+ parameters	~N× (near-linear in GPU count)

ZeRO-3 and FSDP are the same idea: fully shard params, grads, and optimizer state. During the forward/backward pass a layer's parameters are all-gathered just in time, used, then released — trading extra communication for a huge memory saving. This is what lets you train models far larger than any single GPU.

Python · memory per GPU: DDP vs ZeRO-3 (runs)

mem_shard.pydef memory_per_gpu(strategy, params, world_size,
                   bytes_param=2, bytes_grad=2, bytes_optim=12):
    """Per-GPU bytes for model state (params+grads+optimizer), excluding activations.
    Mixed precision: 2B param + 2B grad + 12B optimizer (fp32 master + Adam m,v).
    DDP replicates everything; ZeRO-3/FSDP shard all three across world_size GPUs."""
    full = params * (bytes_param + bytes_grad + bytes_optim)
    if strategy == "ddp":
        per_gpu = full                 # every GPU holds a full copy
    elif strategy == "zero3":
        per_gpu = full / world_size    # sharded across all GPUs
    else:
        raise ValueError(strategy)
    return round(per_gpu / 1e9, 1)     # GB

for n in (8, 64):
    ddp = memory_per_gpu("ddp", 7_000_000_000, n)
    z3 = memory_per_gpu("zero3", 7_000_000_000, n)
    print(f"world_size={n:>3}: DDP={ddp} GB/gpu  ZeRO-3={z3} GB/gpu")

world_size=  8: DDP=112.0 GB/gpu  ZeRO-3=14.0 GB/gpu
world_size= 64: DDP=112.0 GB/gpu  ZeRO-3=1.8 GB/gpu

▶ How this works

This is the heart of the lesson: it computes how much training state each GPU must hold under plain DDP versus ZeRO-3/FSDP. Training state = parameters + gradients + the optimizer's extra tensors (Adam keeps a master copy and two moments).

full = params * (2 + 2 + 12) — 16 bytes per parameter of state. For 7B params that's 112 GB, more than any single GPU has.
strategy == "ddp" returns the full amount on every GPU: DDP replicates all of it, so per-GPU memory never drops no matter how many GPUs you add.
strategy == "zero3" divides by world_size: ZeRO-3/FSDP shards the state so each GPU owns only 1/N of it.

What the output means: DDP stays pinned at 112 GB/GPU at both 8 and 64 GPUs — it will simply OOM. ZeRO-3 drops from 14 GB to 1.8 GB as you shard wider. This single contrast is why large-model training uses sharding, not DDP.

Try this: Bump params to 70B and re-run: DDP asks for over a terabyte per GPU (impossible), while ZeRO-3 at 64 GPUs still fits. That's the memory wall, and how sharding beats it.

DDP is constant; sharding scales with your GPU countNotice DDP's per-GPU state never drops — adding GPUs makes it faster, not smaller. ZeRO-3/FSDP falls off roughly as 1/N: the more GPUs you shard across, the less state each holds. That's the whole reason large-model training uses sharding, not plain DDP.

5 · Tensor & pipeline parallelism expert

Sharding state (ZeRO/FSDP) still runs a full copy of each layer's compute on one GPU. When a single layer is too big or too slow for that, you split the compute itself:

Tensor parallel splits an individual layer's matrix multiply across GPUs — e.g. each GPU computes a column-slice of a big [hidden × 4·hidden] MLP weight — then an all-reduce stitches the partial results back together within every layer. It's extremely communication-heavy, so it's kept inside one node where NVLink bandwidth is highest.

Pipeline parallel puts different layers on different GPUs (stage 1 = layers 0–7, stage 2 = layers 8–15, …). A microbatch flows through the stages like an assembly line. The catch is the pipeline bubble: while stage 1 works on the first microbatch, later stages sit idle. Splitting the batch into many microbatches shrinks the bubble but never fully removes it.

Python · pipeline bubble fraction (runs)

bubble.pydef bubble_fraction(stages, microbatches):
    """Idealized GPipe bubble: (stages - 1) / (microbatches + stages - 1).
    More microbatches -> smaller idle fraction; more stages -> bigger bubble."""
    return round((stages - 1) / (microbatches + stages - 1), 3)

for mb in (1, 4, 16, 64):
    frac = bubble_fraction(stages=4, microbatches=mb)
    print(f"4 stages, {mb:>2} microbatches -> {frac:.1%} of time idle in the bubble")

4 stages,  1 microbatches -> 75.0% of time idle in the bubble
4 stages,  4 microbatches -> 42.9% of time idle in the bubble
4 stages, 16 microbatches -> 15.8% of time idle in the bubble
4 stages, 64 microbatches -> 4.5% of time idle in the bubble

▶ How this works

Pipeline parallelism has a built-in inefficiency called the bubble: while the first stage works on the very first microbatch, every later stage is sitting idle waiting for work to reach it. This computes how much of the time is wasted that way.

(stages - 1) is the number of steps it takes to fill (and drain) the pipeline — the idle stretch at the start and end.
Dividing by (microbatches + stages - 1) spreads that fixed idle time over the whole run. More microbatches keep the pipeline full longer, so the idle fraction shrinks.
The loop sweeps 1 → 64 microbatches for a fixed 4 stages to show the effect.

What the output means: With 1 microbatch you waste 75% of the time; with 64 microbatches only 4.5%. The lesson: pipeline parallelism only pays off when you feed it many microbatches so the bubble becomes negligible.

Try this: Set stages=16 and re-run the sweep — a deeper pipeline needs even more microbatches to hide its (now larger) bubble.

6 · Gradient accumulation — a big batch on small memory expert

Large batches stabilize training, but a big batch means big activations, which means out-of-memory. Gradient accumulation fakes a large batch on small memory: run several small micro-batches, sum their gradients without stepping the optimizer, then step once after K micro-batches. The effective batch size is micro_batch × accumulation_steps × data_parallel_world_size — you trade wall-clock time for memory headroom.

Python · effective batch size (runs)

grad_accum.pydef effective_batch(micro_batch, accum_steps, world_size, seq_len=None):
    """Effective batch = per-GPU micro-batch * accumulation steps * # data-parallel GPUs.
    accum_steps lets you hit a target batch without more memory (just more steps)."""
    eff = micro_batch * accum_steps * world_size
    result = {"effective_batch": eff, "optimizer_steps_per_1M_samples": round(1_000_000 / eff)}
    if seq_len:
        result["effective_tokens_per_step"] = eff * seq_len
    return result

print(effective_batch(micro_batch=2, accum_steps=16, world_size=8))
print(effective_batch(micro_batch=1, accum_steps=32, world_size=64, seq_len=4096))

{'effective_batch': 256, 'optimizer_steps_per_1M_samples': 3906}
{'effective_batch': 2048, 'optimizer_steps_per_1M_samples': 488, 'effective_tokens_per_step': 8388608}

▶ How this works

Big batches train more stably, but a big batch means big activations and out-of-memory. Gradient accumulation fakes a large batch on small memory by summing gradients over several small micro-batches before stepping the optimizer once. This computes the batch you effectively trained on.

eff = micro_batch * accum_steps * world_size — the true batch is the per-GPU micro-batch, times how many you accumulate, times how many data-parallel GPUs are all contributing.
optimizer_steps_per_1M_samples shows the flip side: a bigger effective batch means fewer optimizer steps to cover the same data.
The optional seq_len converts the sample batch into tokens per step — the number that actually matters for LLM training recipes.

What the output means: A tiny per-GPU micro-batch of 2, accumulated 16× over 8 GPUs, trains as if the batch were 256 — with the peak memory of just 2. The second call reaches a 2048-sample / ~8.4M-token effective batch on 64 GPUs.

Try this: Halve micro_batch to survive an OOM, then double accum_steps to keep the same effective batch — same recipe, less peak memory, slightly slower.

Accumulation is free memory, not free timeAccumulating K steps gives you a K× larger effective batch with the same peak memory — but it also runs K forward/backward passes per optimizer step, so each step takes longer. It buys you a big-batch training recipe on hardware that couldn't otherwise hold one.

7 · Communication is the bottleneck expert

Compute per GPU stays roughly fixed as you add GPUs, but communication grows and eventually dominates. Whether your training is compute-bound or comms-bound depends on the ratio of FLOPs done to bytes moved, divided by your interconnect bandwidth. This is why the interconnect (NVLink at ~900 GB/s intra-node vs InfiniBand or, worse, Ethernet across nodes) often decides your real throughput.

Python · compute-bound or comms-bound? (runs)

comms_cost.pydef comms_cost(params, world_size, tflops_per_gpu, interconnect_GBps,
               tokens_per_step=8192, bytes_per_grad=2, flops_per_token_per_param=6):
    """Compare per-step compute time vs all-reduce comms time.
    compute ~ 6 * params FLOPs per token (fwd+bwd) over the step's tokens;
    comms ~ 2x grad bytes over the wire, once per optimizer step."""
    compute_flops = flops_per_token_per_param * params * tokens_per_step
    compute_s = compute_flops / (tflops_per_gpu * 1e12)
    grad_bytes = params * bytes_per_grad
    moved = grad_bytes * 2 * (world_size - 1) / world_size
    comms_s = moved / (interconnect_GBps * 1e9)
    bound = "comms-bound" if comms_s > compute_s else "compute-bound"
    return {"compute_s": round(compute_s, 3), "comms_s": round(comms_s, 3),
            "verdict": bound}

# Same job, fast NVLink vs slow Ethernet:
print(comms_cost(7_000_000_000, 64, tflops_per_gpu=300, interconnect_GBps=600))
print(comms_cost(7_000_000_000, 64, tflops_per_gpu=300, interconnect_GBps=12.5))

{'compute_s': 1.147, 'comms_s': 0.046, 'verdict': 'compute-bound'}
{'compute_s': 1.147, 'comms_s': 2.205, 'verdict': 'comms-bound'}

▶ How this works

This is the punchline of the whole lesson: it puts per-step compute time next to per-step communication time and declares which one you're bound by. Change only the interconnect and watch the verdict flip.

compute_flops = 6 * params * tokens_per_step uses the standard estimate of ~6 FLOPs per parameter per token (forward + backward); dividing by the GPU's TFLOP/s gives the compute seconds.
moved = grad_bytes * 2 * (world_size-1)/world_size is the all-reduce traffic from the earlier lab; dividing by the interconnect GB/s gives the comms seconds.
verdict is simply whichever is larger. The two calls use an identical job — same GPUs, same model, same 64-way DDP — and change only the wire.

What the output means: On NVLink (600 GB/s) comms is 0.046 s against 1.147 s of compute — compute-bound, the healthy case. On 10 GbE (12.5 GB/s) comms balloons to 2.2 s and dominates — comms-bound. Same GPUs; only the wire changed.

Try this: Sweep interconnect_GBps downward and find the value where compute and comms are equal — that's your crossover. Below it, faster GPUs won't help; you need a faster network or fewer syncs.

Comms is the bottleneck — engineer the interconnect, not just the GPUsThe same job is compute-bound on NVLink and badly comms-bound on Ethernet — the GPUs are identical; only the wire changed. At scale you overlap comms with compute (NCCL does this), keep the chatty parallelism (tensor) inside a node, and reach across nodes only for the cheaper collectives. If you remember one thing from this lesson: past a few GPUs, you are optimizing communication.

8 · Choosing a strategy by scale tech-lead

A lead picks the cheapest strategy that fits, and escalates only when a wall forces it. The decision is driven by whether one copy of the model fits on one GPU, and then by how many GPUs/nodes you have.

Python · pick a parallelism strategy (runs)

choose.pydef choose_strategy(model_fits_one_gpu, state_fits_one_gpu,
                    layer_fits_one_gpu, multi_node):
    """Escalate only as far as the memory walls force you."""
    if model_fits_one_gpu and state_fits_one_gpu:
        return "DDP (replicate; all-reduce grads) — simplest, just go faster"
    if not state_fits_one_gpu and layer_fits_one_gpu:
        return "ZeRO-3 / FSDP (shard params+grads+optim) — bigger than one GPU's memory"
    if not layer_fits_one_gpu and not multi_node:
        return "Tensor parallel within the node (+ FSDP) — a single layer is too big"
    return "3D parallelism: tensor (in-node) + pipeline (across nodes) + data/FSDP"

print(choose_strategy(True,  True,  True,  False))   # small model -> DDP
print(choose_strategy(True,  False, True,  False))   # state too big -> FSDP/ZeRO-3
print(choose_strategy(False, False, False, True))    # giant -> 3D parallelism

DDP (replicate; all-reduce grads) — simplest, just go faster
ZeRO-3 / FSDP (shard params+grads+optim) — bigger than one GPU's memory
3D parallelism: tensor (in-node) + pipeline (across nodes) + data/FSDP

▶ How this works

This encodes the tech-lead decision as a ladder: start with the simplest strategy that fits, and climb a rung only when a concrete memory wall forces you to. Each argument is an honest yes/no about what fits where.

If both the model and its full training state fit on one GPU, you only need speed → DDP. Simplest, least communication.
If the state no longer fits but a single layer still does → ZeRO-3 / FSDP shards params, grads, and optimizer state across GPUs.
If a single layer is too big → tensor parallel inside a node; and once you span multiple nodes at giant scale, you combine all three into 3D parallelism.

What the output means: The three calls walk up the ladder: small model → DDP; oversized state → FSDP/ZeRO-3; giant multi-node model → 3D parallelism. The rule is to escalate only as far as a wall actually pushes you.

Try this: Flip state_fits_one_gpu to False for your own model and confirm the recommendation jumps to FSDP — then justify it with the exact GB/GPU number from mem_shard.py.

The real PyTorch below shows the two most common rungs — DDP and FSDP — in the shape you'd actually launch with torchrun. It needs real hardware to run; read it for the structure, not to execute it here.

PyTorch · DDP + FSDP wrap (▶ needs multi-GPU + PyTorch)

train.py# needs multi-GPU + PyTorch: launch with `torchrun --nproc_per_node=8 train.py`
# NOT runnable on a laptop; shown for structure. Requires: pip install torch
import os, torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def setup():
    dist.init_process_group("nccl")              # NCCL = GPU collectives (all-reduce)
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def wrap(model, local_rank, shard=False):
    model = model.to(local_rank)
    if shard:
        # FSDP == ZeRO-3: shard params+grads+optimizer across the data-parallel group
        return FSDP(model, device_id=local_rank)
    # DDP: full replica per GPU; gradients all-reduced automatically in backward()
    return DDP(model, device_ids=[local_rank])

# training step is identical either way — the wrapper handles the comms:
#   loss = model(batch).loss; loss.backward()   # <- all-reduce / reduce-scatter here
#   optimizer.step(); optimizer.zero_grad()
# Use accumulation to grow the effective batch without more memory:
#   if (step + 1) % accum_steps == 0: optimizer.step(); optimizer.zero_grad()

Escalate up the ladder, don't start at the topDDP → FSDP/ZeRO → tensor+pipeline is an increasing-complexity ladder. Every rung adds communication and debugging cost. A lead starts at the lowest rung that fits, measures throughput and GPU utilization, and climbs only when a concrete memory or time wall demands it — never because a bigger setup sounds impressive.

✓ Knowledge check

Your 7B model trains fine on one GPU but you have 8 and want it 8× faster. Someone proposes ZeRO-3. Is that the right first move, and what does it change vs DDP?

Show answer

No — plain DDP is the right first move. If a full copy already fits on one GPU, DDP replicates it across all 8 and all-reduces gradients each step, giving you the throughput win with the least complexity and communication. ZeRO-3/FSDP shards params, grads, and optimizer state across GPUs — that saves memory (letting you train models that don't fit) but adds all-gather/reduce-scatter comms you don't need when the model already fits. Reach for it only when you hit the memory wall, not the speed wall.

✓ Knowledge check

You move a DDP job from a single 8-GPU node (NVLink) to 8 nodes over Ethernet and throughput collapses even though total GPU count is unchanged. Why?

Show answer

You became comms-bound. DDP's per-step all-reduce moves ~2× the gradient bytes over the interconnect every step. NVLink (~hundreds of GB/s) hides that behind compute; commodity Ethernet (tens of Gb/s) is ~50× slower, so the all-reduce time now dwarfs the compute time — exactly the flip shown by comms_cost.py. The fix is a faster interconnect (InfiniBand), keeping chatty parallelism in-node, gradient accumulation to reduce sync frequency, and overlapping comms with compute.

Exercise MS2.1 — Size the state and pick a strategy

Context: Choosing a parallelism strategy is a memory calculation first: you size the per-GPU state under each option, then let the number that fits pick the rung.

Your task: For a 13B model, compute per-GPU training state under DDP vs ZeRO-3 for world sizes 8 and 64, run the strategy chooser, and justify a starting rung with a specific memory number.

Requirements:

Compute per-GPU state under both DDP (replicated) and ZeRO-3 (sharded 1/G)
Do it for both world sizes 8 and 64
Run the chooser with honest answers about what fits
Name the rung you would start on and cite the memory number that justifies it

💡 Hint: The DDP-vs-ZeRO gap is roughly a factor of the world size on the replicated state; that gap is your argument.

Exercise MS2.2 — Find your comms wall

Context: Every distributed job has a bandwidth below which it flips from compute-bound to comms-bound, and finding that wall is what tells you whether the interconnect, not the GPU, is your limit.

Your task: Holding the GPU and a 64-way DDP job fixed, sweep the interconnect from NVLink (600 GB/s) down to 10 GbE (12.5 GB/s) to find where the job becomes comms-bound, then argue how accumulation pushes that crossover back.

Requirements:

Sweep interconnect bandwidth with GPU throughput and world size held fixed
Identify the bandwidth where comms time overtakes compute time
Explain that syncing less often (accumulation) reduces all-reduce frequency
State how that moves the compute-bound/comms-bound crossover

💡 Hint: Accumulation trades sync frequency for step count; fewer all-reduces per optimizer step buys back headroom on a slow link.

🪜 Practice ladder beginner → industry

Six graded exercises, easy to real-world. Try each before opening its solution.

Exercise 1 · Three ways to split the workBeginner

Context: Distributed training offers three orthogonal ways to split the work, and the first skill is matching a bottleneck to the axis that actually relieves it.

Your task: Write a classifier that maps a training bottleneck to the right split: data parallel (replicas), pipeline parallel (split by layer), or tensor parallel (split a layer's tensors).

Requirements:

Batch too slow but the model fits one GPU → data parallel (DDP)
Model too big by layers → pipeline parallel
A single layer too big → tensor parallel
Note that the two model splits exist only because something no longer fits one card

💡 Hint: Data parallelism is a throughput fix; the model splits are fit fixes — the constraint tells you which.

Show solution

Match the constraint to the split. Runnable:

def which_split(problem):
    table = {
        "batch too slow, model fits one GPU": "data parallel (DDP)",
        "model too big for one GPU by layers": "pipeline parallel",
        "a single layer too big for one GPU":  "tensor parallel",
    }
    return table.get(problem, "profile first -- unknown bottleneck")

for p in ["batch too slow, model fits one GPU",
          "a single layer too big for one GPU"]:
    print(p, "->", which_split(p))

Data parallelism speeds up throughput when the model fits; the two model splits exist only because a model (or a layer) is too big for one card.

Exercise 2 · DDP all-reduce volumeIntermediate

Context: DDP averages gradients with an all-reduce every step, and the volume moved — not the GPU count — is what interconnect bandwidth has to carry.

Your task: Model the per-GPU ring all-reduce volume as ≈ 2·(G−1)/G · grad_bytes and compute it for a 7B bf16 model across 8 GPUs.

Requirements:

Gradient buffer is one value per parameter (× bytes-per-param)
Per-GPU ring all-reduce moves ~2× the buffer, nearly independent of G
Show the volume approaching 2× as G grows
Conclude that interconnect bandwidth, not GPU count, sets step time

💡 Hint: The (G−1)/G factor saturates toward 1, which is exactly why ring all-reduce scales.

Show solution

The gradient buffer is one value per parameter; ring all-reduce is ~2× the buffer regardless of G (the (G-1)/G factor). Runnable:

def allreduce_gib(params_b, G, bytes_per=2):
    grad_bytes = params_b * 1e9 * bytes_per
    moved = 2 * (G - 1) / G * grad_bytes      # per-GPU ring all-reduce volume
    return moved / (1024 ** 3)

for G in (2, 8, 64):
    print(f"G={G:>2}: per-GPU all-reduce ~= {allreduce_gib(7, G):6.2f} GiB / step")
# approaches 2x the gradient buffer (~13 GiB) as G grows -- ~26 GiB moved per step

The volume is nearly independent of G — which is exactly why ring all-reduce scales, and why interconnect bandwidth (not GPU count) sets the step time.

Exercise 3 · ZeRO/FSDP memory shardingAdvanced

Context: DDP replicates the full optimizer state on every GPU; ZeRO/FSDP shards it, and that ~1/G cut is the reason big models fit at all.

Your task: Model the per-GPU memory saving of FSDP over DDP for a 7B Adam job on 8 GPUs, where replicated state becomes ~1/G under sharding.

Requirements:

Per-param training bytes = weight + master + grad + two Adam moments
DDP holds the full state on each GPU; FSDP holds ~1/G of it
Report DDP per-GPU, FSDP per-GPU, and the saving
Note FSDP pays extra comms (gather shards before each layer) for the memory cut

💡 Hint: Divide the replicated total by G for the sharded case; the trade is memory for communication.

Show solution

Adam state (18 B/param here) is replicated under DDP but sharded 1/G under FSDP. Runnable:

def per_gpu_gib(params_b, G, sharded):
    bytes_per = 2 + 4 + 4 + 4 + 4   # weight + master + grad + m + v (fp32 heavy)
    total = params_b * 1e9 * bytes_per / (1024 ** 3)
    return total / G if sharded else total

p, G = 7, 8
ddp  = per_gpu_gib(p, G, sharded=False)
fsdp = per_gpu_gib(p, G, sharded=True)
print(f"DDP  per-GPU: {ddp:6.1f} GiB (replicated)")
print(f"FSDP per-GPU: {fsdp:6.1f} GiB (sharded 1/{G})")
print(f"saving: {ddp - fsdp:6.1f} GiB per GPU")

FSDP trades extra communication (gather shards before each layer) for a ~1/G cut in per-GPU state — the reason big models train at all.

Exercise 4 · Gradient accumulation = big batch on small memoryExpert

Context: When your target batch does not fit in memory, gradient accumulation buys the big batch with time instead of memory by stepping the optimizer only every K micro-batches.

Your task: Compute the accumulation steps K needed to reach an effective batch of 512 from a micro-batch of 16, and show the effective batch is preserved across the GPUs.

Requirements:

effective = micro_batch × accum_steps × num_gpus
Solve for accum_steps, rounding up when it does not divide evenly
Verify the achieved effective batch matches the target
Note each step now costs K forward/backward passes

💡 Hint: It is one division; the insight is that the loss curve matches a true large batch at a fraction of the peak memory.

Show solution

effective = micro_batch × accum_steps × num_gpus. Solve for accum_steps. Runnable:

import math

def accum_steps(effective_batch, micro_batch, num_gpus):
    per_step = micro_batch * num_gpus
    if effective_batch % per_step:
        # round up; effective batch will be the next multiple
        return math.ceil(effective_batch / per_step)
    return effective_batch // per_step

K = accum_steps(effective_batch=512, micro_batch=16, num_gpus=4)
achieved = 16 * K * 4
print(f"accum_steps K = {K}, achieved effective batch = {achieved}")
# 16 * 8 * 4 = 512 -- same optimizer step as a 512 batch, at 1/32 the peak memory

Accumulation buys batch size with time instead of memory: the loss curve matches a true large batch, but each step takes K forward/backward passes.

Exercise 5 · Communication is the bottleneckProfessional

Context: Step time is compute plus communication, and as you add GPUs the compute per GPU falls while comms does not — so there is a crossover past which more GPUs make each step slower.

Your task: Model step time as compute (work/G) plus bandwidth-bound all-reduce time, and sweep G to find the point where comms overtakes compute.

Requirements:

Compute per GPU scales as work / G
Comms is the ring all-reduce volume divided by interconnect bandwidth
Sweep G and flag the first G where comms > compute
Explain that beyond the knee you pay bandwidth for compute you no longer need

💡 Hint: The crossover G is where you switch strategy; the all-reduce volume barely shrinks while compute keeps dropping.

Show solution

Model compute/GPU as work/G and comms as a bandwidth-bound constant that grows slowly; find where they cross. Runnable:

def step_ms(G, work_ms=800.0, grad_gib=13.0, bw_gib_s=300.0):
    compute = work_ms / G                      # perfect compute scaling
    moved = 2 * (G - 1) / G * grad_gib         # ring all-reduce volume (GiB)
    comms = moved / bw_gib_s * 1000            # ms at bw_gib_s GiB/s
    return compute, comms, compute + comms

print(f"{'G':>3} {'compute':>9} {'comms':>8} {'step(ms)':>9}")
for G in (1, 2, 4, 8, 16, 32, 64):
    c, m, tot = step_ms(G)
    flag = "  <- comms-bound" if m > c else ""
    print(f"{G:>3} {c:9.1f} {m:8.1f} {tot:9.1f}{flag}")

Past the crossover, adding GPUs makes each step slower: you are paying bandwidth for compute you no longer need. That knee is where you switch strategy.

Exercise 6 · Choose a strategy by scaleIndustry scenario

Context: Leading training across scales means encoding a decision table: models that fit one card go DDP, larger ones shard with FSDP, and the truly enormous need 3D parallelism — with a comms-risk flag as the cluster grows.

Your task: Build a planner that, given model size and GPU count, picks DDP / FSDP / 3D parallel and flags the communication risk, for cases from 1B on 8 GPUs to 400B on 1024 GPUs.

Requirements:

If the full Adam state fits one card → DDP (replicate)
Else if bf16 weights fit one card → FSDP/ZeRO (shard state)
Else → 3D parallel (tensor + pipeline + data)
Raise comms risk as GPU count grows and note a real profiler pass is still required (needs GPU)

💡 Hint: The thresholds are the serve/train footprints from ms1; the planner picks the axis but topology still decides the real knee.

Show solution

The lesson's scale ladder, as a planner. Runnable:

def strategy(model_b, gpus, card_gib=80):
    serve = model_b * 1e9 * 2 / (1024 ** 3)          # bf16 weights
    train = model_b * 1e9 * 18 / (1024 ** 3)         # Adam, per-GPU if replicated
    if train <= card_gib * 0.9:
        s = "DDP (replicate; simplest, fastest per step)"
    elif serve <= card_gib * 0.9:
        s = "FSDP/ZeRO (shard state; model weights still fit one card)"
    else:
        s = "3D parallel: tensor + pipeline + data (weights exceed one card)"
    risk = "HIGH" if gpus >= 64 else "moderate"
    return f"{model_b}B on {gpus} GPUs -> {s} | comms risk: {risk}"

for m, g in [(1, 8), (13, 32), (70, 128), (400, 1024)]:
    print(strategy(m, g))
# real runs verify with a profiler (torch profiler / nsys) -- needs GPU

The planner picks the axis; the real cluster still needs a profiler pass to confirm the comms knee (needs GPU), because interconnect topology, not the model, often decides.

✓ Checkpoint — you can move on when you can…

Explain the memory wall and the time wall that force multi-GPU training.
Distinguish data / tensor / pipeline parallelism and what each splits.
Explain DDP's all-reduce and the memory problem it does not solve.
Explain how ZeRO-3 / FSDP shard state and why it scales with GPU count.
Compute an effective batch with gradient accumulation.
Reason about when a job is compute-bound vs comms-bound and pick a strategy.

Educational content, provided as-is and without warranty. Code samples are examples — review, test, and adapt them before using in production. See the Terms of Use & Disclaimer. Use at your own risk.

Learning objectives

1 · Why one GPU isn't enough essential

2 · Three ways to split the work essential

3 · Data parallelism & DDP advanced

4 · What DDP doesn't solve — ZeRO & FSDP advanced

5 · Tensor & pipeline parallelism expert

6 · Gradient accumulation — a big batch on small memory expert

7 · Communication is the bottleneck expert

8 · Choosing a strategy by scale tech-lead

Exercise MS2.1 — Size the state and pick a strategy

Exercise MS2.2 — Find your comms wall

🪜 Practice ladder beginner → industry

✓ Checkpoint — you can move on when you can…

More in ML systems internals