Distributed training

1 · Why one GPU isn't enough essential

You've probably felt the small version of this already. You fine-tune a model on your laptop or a single cloud GPU, and one day torch greets you with CUDA out of memory — not because your data got bigger, but because the model did. You shrink the batch size, it fits, you exhale. Now scale that frustration up by a thousand: a modern LLM doesn't just refuse to train on one card, it can't even be loaded onto one. The question this whole chapter answers is: when a single GPU can't hold or can't finish the job, how do you spread it across many without the coordination cost eating all your gains?

Here's the mental model to carry through every section. Training a model is a giant, repeated loop — read a batch, compute a forward pass, compute gradients backward, update the weights — and it runs on two finite resources: memory (does everything fit?) and time (will it finish this decade?). Distributed training exists because for large models you run out of both, and there is no single trick that fixes both cleanly. Every strategy you'll meet is really a different answer to "which resource am I short on, and what am I willing to pay to get more of it?" The price you pay is always the same currency: communication — GPUs talking to each other over a wire — and that bill is the villain of this entire chapter.

The common mistake beginners make is to reach for more GPUs the moment training feels slow. But "slow" and "too big" are different diseases with different cures. If the model fits and you just want it done sooner, you replicate it and split the work (that's data parallelism, coming in §3). If the model doesn't fit, adding replicas does nothing — you need to carve the model itself apart (sharding and model parallelism, §4–5). Diagnosing which wall you hit — memory or time — before you pick a tool is the single most valuable habit in this whole chapter, and it's why we open by naming both walls precisely.

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.

2 · Three ways to split the work essential

Imagine a restaurant kitchen swamped on a Friday night. There are three genuinely different ways to cope, and every distributed-training strategy is one of them in disguise. You can clone the whole kitchen — build a second identical kitchen next door and send half the orders there (that's data parallelism: more throughput, same menu). Or a single dish is so enormous — a wedding cake — that no one cook can carry it, so several cooks each take one tier and assemble at the end (that's tensor parallelism: one item split across hands). Or you keep a single line of cooks but give each a station — one chops, one fries, one plates — and dishes flow down the line (that's pipeline parallelism: split by stage). The kitchen never changes its recipes; it only changes how the work is carved up.

Now the precise version. There are exactly three orthogonal axes along which you can split training, and "orthogonal" is the load-bearing word: they're independent choices, so you can stack them. The one question that instantly tells you which axis you're looking at is "what got split?" Data parallel splits the batch (the data), keeping a full model copy per GPU. Tensor parallel splits a single layer's matrices, so every GPU holds a thin slice of the same layer. Pipeline parallel splits the stack of layers, so each GPU owns a contiguous block of the model's depth. Ask "what got split?" of any real system and you can immediately name its strategy — and because the three are orthogonal, a frontier training run typically uses all three at once, which the industry calls 3D parallelism.

The mistake to avoid: treating these as competitors where you pick "the best one." They solve different problems. Data parallelism is a speed tool and does nothing for a model that's too big; tensor and pipeline parallelism are fit tools that let an oversized model exist at all, but they add communication and complexity you don't want unless you're forced into them. Read the diagram and table below not as a ranking but as a menu keyed to your bottleneck — the right choice is dictated by what's broken, not by what sounds most powerful.

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

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

Picture eight students grading the same enormous exam. To go faster, you split the stack of papers into eight piles, one per student, and everyone grades in parallel. But there's a catch that has nothing to do with grading: at the end, every student has only seen their pile, so before anyone can update the shared answer key they must pool what they learned and agree on one combined result. That pooling step — not the grading — is where the whole scheme lives or dies, and it's exactly the situation data parallelism creates. This is the most common way to train, and understanding why the syncing step dominates is the difference between an engineer who "uses 8 GPUs" and one who knows what those 8 GPUs are actually spending their time on.

Concretely: every GPU holds a complete copy of the model and processes a different shard of the batch. Because each copy saw different examples, after the backward pass they all hold different gradients — and if you let them each update independently, the eight copies would immediately drift apart and stop being the same model. So the gradients must be averaged into one shared result that every GPU applies identically, keeping all replicas in lockstep. That averaging is a collective operation called an all-reduce, and the phrase to burn in is "sum-then-share": it sums each gradient value across all GPUs and hands the identical total back to every one.

Why does all-reduce get its own callout and its own runnable lab below? Because it is the recurring cost of data parallelism — it happens every single step, forever, and its price is set by how fast your GPUs can talk, not by how fast they can compute. The counter-intuitive fact that trips people up: adding more GPUs barely changes how many bytes each one has to move for the all-reduce (the clever "ring" algorithm sees to that), so the communication cost is governed almost entirely by your model's gradient size and your interconnect bandwidth. That's the seed of this chapter's thesis — past a handful of GPUs, you stop optimizing compute and start optimizing the wire — and the lab makes the numbers concrete.

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.

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'}

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

Go back to the eight students, but change the problem. DDP made them grade faster — but what if the answer key itself is a thousand-page binder that no single student's desk can hold? Handing all eight students their own full copy doesn't help; the binder was the problem, and now you have eight copies of an unfittable binder. That is exactly DDP's blind spot: it hands every GPU a complete copy of the model, the gradients, and the optimizer's bookkeeping. If one copy of that pile doesn't fit on one GPU, DDP is useless no matter how many GPUs you add — it only ever multiplies replicas, it never makes one replica smaller.

The fix has an elegant name and a simple idea: ZeRO — the Zero Redundancy Optimizer, from Microsoft's DeepSpeed — and its PyTorch-native twin, FSDP (Fully Sharded Data Parallel). Instead of every GPU redundantly storing the whole binder, each GPU keeps only 1/N of it and borrows the missing pages from its neighbours only at the exact moment it needs them, then throws them away again. Back to the students: rather than eight thousand-page binders, each student keeps a hundred-and-twenty-five pages, and when the class hits a section someone else owns, that owner briefly reads their pages aloud. Nobody hoards the whole thing, so the memory each desk needs falls roughly in proportion to the number of desks.

The precise term for "keep only your slice" is sharding, and the crucial insight — the one the table and lab below drive home — is that ZeRO comes in stages that shard progressively more. ZeRO-1 shards only the optimizer state (the biggest hog, as §1 warned); ZeRO-2 adds the gradients; ZeRO-3 (identical in spirit to FSDP) shards the parameters too, so nothing is fully replicated. The trade you're making is explicit and worth saying out loud: you spend extra communication — gathering each layer's shards just in time, using them, releasing them — to buy an enormous memory saving. The common misconception is that sharding is "just a faster DDP"; it is not — it is a memory tool that costs you comms, and reaching for it when you only needed speed adds gather/scatter traffic you didn't have to pay. Read the table as a memory-vs-communication ladder, not a free lunch.

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.

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.

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

5 · Tensor & pipeline parallelism expert

Here's the wall sharding alone can't climb. ZeRO/FSDP cleverly spreads the storage of the model across GPUs — but when a layer actually runs, its shards are gathered back onto one GPU and computed there. That's fine until a single layer is so large that even one layer's worth of matrix multiply won't fit in one GPU's memory or takes agonizingly long on one card. At the frontier — think the giant feed-forward and attention matrices inside a 100B+ model — you hit exactly this. When the unit of work itself is indivisible-by-storage-tricks, you have to split the computation, and that's what tensor and pipeline parallelism do. They're the top rung of the ladder precisely because they're the most invasive and the most communication-hungry.

Think of tensor parallelism as several accountants sharing one colossal spreadsheet formula. The formula is a single matrix multiply — say a [hidden × 4·hidden] weight in an MLP — too big for one accountant, so each takes a vertical band of columns, computes their partial answer, and then they must add up the partials to get the true result. That "add up the partials" is an all-reduce, and here's the sting: it happens inside every single layer, on every forward and backward pass. That is a torrent of communication, which is why tensor parallelism is kept within one node, where the ultra-fast NVLink fabric can absorb it — reach across slower node-to-node links with this and throughput collapses.

Pipeline parallelism is the assembly line from §2 made real: GPU 0 owns layers 0–7, GPU 1 owns layers 8–15, and a batch flows through the stages in sequence. It's far less chatty than tensor parallelism — you only pass activations at the seams between stages — but it has its own signature tax, the pipeline bubble. When work starts, stage 2 sits idle waiting for stage 1 to finish the first chunk; at the end, stage 1 is idle while the last chunk drains out the far end. That idle time is pure waste. The cure, and the thing the lab below quantifies, is to chop the batch into many small microbatches so the line stays full most of the time — but the bubble shrinks toward zero without ever quite reaching it. The mistake to avoid is running a deep pipeline with too few microbatches, which leaves half your expensive GPUs idle; the bubble formula tells you exactly how many you need.

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.

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

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

A recipe you found says "use a batch of 512," and your GPU OOMs at 32. This is one of the most common walls in practice, and it feels like a dead end: the training recipe that everyone reports good results with simply won't fit on your hardware. Beginners respond by shrinking the batch and quietly hoping it's fine — but batch size isn't a free knob. Small batches give noisier gradient estimates, and the loss curve you get can differ meaningfully from the recipe's. What you actually want is the statistics of a 512-batch with the memory footprint of a 32-batch. Gradient accumulation is the trick that gives you exactly that, and it's worth its own section because it's the everyday tool that lets modest hardware reproduce big-lab recipes.

The idea is almost embarrassingly simple once you see it. Remember the training loop: forward, backward (which adds gradients into each weight's .grad), then optimizer.step() to apply them and zero_grad() to wipe the slate. Normally you step after every batch. In gradient accumulation you run several small micro-batches and let their gradients pile up — you deliberately skip the step and the zeroing — then step just once after K micro-batches. Because gradients sum, the accumulated gradient over K micro-batches of size m is mathematically the same as one gradient computed on a batch of K×m. Think of it as filling a bucket one cup at a time instead of needing a bucket-sized tap: the bucket ends up equally full, you just took more trips.

So the effective batch size — the number the loss curve actually "sees" — is micro_batch × accumulation_steps × data_parallel_world_size, and the lab below makes that arithmetic concrete, including the LLM-relevant conversion to tokens per step. The trade, and the misconception to guard against, is spelled out in the callout after: accumulation buys memory with time, not for free. K micro-batches mean K forward/backward passes per optimizer update, so each step takes K× longer — you haven't dodged the work, you've just fit it through a narrower memory pipe. That's usually a trade worth making when the alternative is not being able to run the recipe at all.

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.

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}

7 · Communication is the bottleneck expert

This is the war story every scaling team eventually tells. A team buys twice the GPUs expecting roughly twice the training speed, and instead throughput barely moves — sometimes it gets worse. The GPUs are pristine and idle-hot, sitting there waiting. What happened is the theme this whole chapter has been building toward: past a certain point, your training is no longer limited by how fast your GPUs can compute, but by how fast they can talk to each other. If you internalize one idea from this page, make it this one — it reframes every decision, from which cloud instance to rent to how often you sync.

Here's the tension, made mechanical. Split a fixed amount of compute across more GPUs and each GPU's compute time falls — that's the win you paid for. But the communication per step (the all-reduce from §3) does not fall correspondingly; it's set by your gradient size and your wire speed, and it stays stubbornly flat as you add GPUs. So step time is a race between two clocks: compute, which shrinks, and comms, which doesn't. Whichever is larger is your bound. Early on compute dominates and everything scales beautifully; eventually the two clocks cross, and after that every extra GPU is paying bandwidth for compute savings you no longer need — you're comms-bound.

The lever that decides where that crossover lands is the interconnect — the physical fabric the GPUs communicate over — and the gap between options is enormous, not marginal. NVLink inside a single node runs at hundreds of GB/s and can hide the all-reduce entirely behind compute; commodity Ethernet across nodes can be tens of times slower and turns the very same job comms-bound. The lab below runs one identical job — same model, same GPUs, same 64-way DDP — and changes only the wire, and watches the verdict flip from healthy to hopeless. That's the honest lesson for a systems engineer: at scale you engineer the network as deliberately as the GPUs, you keep the chattiest parallelism (tensor) inside the fast fabric, and you overlap comms with compute wherever the framework lets you. Buying faster GPUs for a comms-bound job is the classic expensive mistake.

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.

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'}

8 · Choosing a strategy by scale tech-lead

Here's the failure mode that separates a junior from a lead. Handed a training project, the eager engineer reaches straight for the most impressive-sounding stack — "we'll do 3D parallelism with tensor and pipeline and FSDP" — for a 3B model that would have trained perfectly happily under plain DDP on eight cards. Now they've signed up for a mountain of communication overhead, brittle launch scripts, and days of debugging pipeline bubbles, all to solve a problem they never had. The skill this section teaches is the opposite instinct: pick the cheapest strategy that fits, and climb the ladder only when a concrete wall physically forces you up a rung.

Everything from this chapter collapses into a short decision procedure driven by two questions, asked in order. First: does one full copy of the model plus its training state fit on one GPU? If yes, you only have a speed problem — replicate and split the batch with DDP, the simplest and least chatty option. If the state doesn't fit but the model's weights still do, shard the state with ZeRO-3 / FSDP: you've traded some communication for the memory to exist. Second question, only if you're still stuck: does a single layer fit on one GPU? If not, you must split the compute — tensor parallelism inside a node, and once you're spanning many nodes at giant scale, you combine all three axes into 3D parallelism. Each rung up adds communication and operational cost, so each rung must be earned by a wall you actually hit.

The reason to encode this as an explicit ladder — which the lab below does in code — is that it turns an intimidating "what parallelism should we use?" into a sequence of honest yes/no questions about what fits where, each answerable with the memory arithmetic from §1 and §4. That's also why a real lead never picks a rung by reputation: they cite the specific GB-per-GPU number that justifies it. "We need FSDP because DDP asks for 112 GB/GPU and our cards have 80" is an argument; "FSDP is what the big labs use" is not. Start low, measure GPU utilization and throughput, and escalate only when the numbers demand it.

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.

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

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.

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()

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))

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

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")

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

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}")

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