Training on many GPUs.

Training on many GPUs means cutting one training job into pieces so that hundreds or thousands of chips can work on it together, passing numbers back and forth all the time so they stay in agreement.

It matters because the big models you use every day can’t be made any other way. A model the size of the large open ones needs more memory than ten of the biggest GPUs hold, and would take a single GPU centuries to train. Your cross-encoder meets a small version of the same problem the day you want to train it on two machines instead of one. And “how would you lay out a 70-billion-weight model on a thousand GPUs?” is a standard interview question at the labs that train these models.

It builds on How much memory does a model need? for counting bytes, Back of the envelope for estimating with powers of ten, and Processes, threads and the GIL, because every GPU here is driven by its own separate program. It sets up Train a small LM, end to end, where you split real training across processes, and Writing a design doc, where a plan for surviving broken machines becomes something a team can approve.

Today

• Work out, with two lines of arithmetic, why a 70-billion-weight model can’t be trained on one GPU: not enough memory, and not enough years.
• Draw the five group moves GPUs use to share data (the key one: everyone adds up their numbers and everyone gets the total), and explain why doing that sum by passing pieces around a ring sends about 2x the data per GPU, no matter how many GPUs there are.
• Explain the three ways to split training: give every GPU a copy of the model and different examples, cut each weight matrix across GPUs, or give each GPU a slice of the layers like an assembly line (data, tensor and pipeline parallelism). Say what each costs in talking between GPUs, and why cutting the matrices only works inside one machine.
• Explain how to stop every GPU keeping its own full copy of the weights, gradients and optimizer bookkeeping (ZeRO and FSDP): what gets split up at each stage, and what it costs.
• Pick a layout for a 70B model on 1,024 GPUs and defend it with numbers, including what happens when a GPU dies at 3 a.m.

Read How much memory does a model need? for bytes per weight, Back of the envelope for estimating with powers of ten, and Processes, threads and the GIL, since this whole lesson runs as separate processes. How much memory does training take? is the training version of the memory math; if you skip it, the one number you need is about 16 bytes per weight. Budget two evenings. Everything runs on your Mac or the Spark, and the cluster parts are done with numbers.

The Problem

Your cross-encoder trains fine on the Spark. It’s a few hundred million parameters, it fits, an epoch takes an hour. Now imagine someone hands you the real job: pretrain a 70-billion-parameter language model on 15 trillion tokens. You have a cluster of 1,024 H100s and a deadline.

Try the obvious thing first, in your head: put the model on one GPU and run the training loop you already know.

It fails before the first step. Training with the Adam optimizer in mixed precision costs about 16 bytes per parameter (weights, gradients, and the optimizer’s own memory, which you’ll take apart in a minute). That’s 70 × 10⁹ × 16 ≈ 1.1 TB. An H100 has 80 GB. The Spark has 128 GB of unified memory (spec sheet). You are off by more than 10x, and that’s before counting activations.

Suppose memory were free. The total compute for training a transformer is about 6ND6NDThe standard estimate of training FLOPs: 6 × (number of parameters N) × (number of training tokens D). Each token costs about 2N FLOPs in the forward pass and 4N in the backward pass.: 6 × 70 × 10⁹ × 15 × 10¹² ≈ 6.3 × 10²⁴ FLOPsFLOPsFloating-point operations: one multiply or one add. A matrix multiply of an (m × k) by a (k × n) matrix costs about 2mkn FLOPs.. An H100 peaks at about 989 trillion bf16 FLOPs per second (spec sheet, dense), and real training gets maybe 40% of that. Divide: about 1.6 × 10¹⁰ seconds, which is about 500 years. On the Spark, which is several times slower at bf16, it’s thousands of years.

On all 1,024 GPUs, perfectly used, it’s about 180 days. On 16,000 GPUs, about 11 days. So the job is not “train a model”. The job is “make 1,024 or 16,000 GPUs behave like one enormous GPU”, where every GPU holds only a slice of the problem and they talk to each other through cables that are much slower than their own memory.

That last assumption is where all the engineering lives. GPUs that wait for each other are GPUs you’re paying for and not using.

The Idea

The analogy: one textbook, many students

Picture a class that has to work through a very long textbook problem set, together, and agree on one answer key at the end of each page.

There are three ways to split the work. Give every student a full copy of the book and a different set of problems, then have them compare and average their answers (that’s data parallelism). Tear each page down the middle, so each student does half of every calculation and they shout partial results across the table (that’s tensor parallelism). Or give the first chapter to one student, the second to the next, and pass the work down the row like an assembly line (that’s pipeline parallelism).

Each method costs a different kind of talking. The whole art is matching the kind of talking to the kind of cable between the students.

First: what “talking” means. Collectives

Every GPU in a job runs the same Python script as its own process. Each process gets a number, its rankrankThe ID of one process in a distributed job, from 0 to world size minus 1. Usually one rank per GPU., and the total number of processes is the world sizeworld sizeHow many processes (usually GPUs) take part in the job.. They don’t share memory. They communicate by calling collective operationscollective operationsCommunication calls that every process in a group makes at the same time, together, like “everyone add up your numbers and everyone get the total”. Also called collectives.: functions that every rank calls at the same moment, and that only return once the group has finished the exchange. The library that does this on NVIDIA GPUs is NCCLNCCLNVIDIA Collective Communications Library, pronounced “nickel”. Implements collectives over NVLink, PCIe and network cards, choosing algorithms (rings, trees) for the hardware it finds.; on CPUs, PyTorch uses glooglooA collectives library from Meta that works on CPUs over ordinary TCP. Slower than NCCL, but runs anywhere, which is why you’ll use it on your Mac..

There are only five you need.

In plain words:

• Broadcast: one GPU has something, everyone gets a copy. Used once at the start so all GPUs begin with identical weights.
• Reduce: everyone has a vector, one GPU ends up with the sum.
• All-reduce: everyone has a vector, everyone ends up with the sum. This is the workhorse of data parallelism: sum (then average) the gradients so every GPU takes the same step.
• Reduce-scatter: everyone has a vector, and each GPU ends up with the sum of one slice. Like all-reduce, but each GPU keeps only its quarter of the answer.
• All-gather: each GPU has one slice, everyone ends up with all the slices glued together.

The last two look like curiosities until you notice the caption: a reduce-scatter followed by an all-gather is an all-reduce. That decomposition is both how the fast all-reduce works and the whole trick behind ZeRO. Hold on to it.

How to sum a vector across N GPUs without a traffic jam

Each GPU has a gradient of S bytes. Everyone needs the sum. How much does each GPU have to send?

The naive answer: everyone sends their whole gradient to GPU 0, GPU 0 adds them up and sends the result back. GPU 0 receives (N − 1) × S bytes and sends (N − 1) × S bytes. With N = 1,024 and a 140 GB gradient, GPU 0 is moving 143 TB per step while every other GPU waits. The cost grows with N. That’s the parameter-server traffic jam.

Arrange the GPUs in a circle: each one sends only to its right neighbour and receives only from its left neighbour. Cut every GPU’s gradient into N equal chunks. Then:

1. Reduce-scatter phase, N − 1 steps. At each step, every GPU sends one chunk to its right neighbour, who adds it to its own copy of that chunk. The chunks are staggered so that no two GPUs work on the same chunk at the same step. After N − 1 steps, each GPU holds one chunk that has passed through every GPU: the fully summed chunk.
2. All-gather phase, N − 1 steps. Now every GPU forwards its finished chunk around the ring, and the receiver just overwrites its stale copy. After N − 1 more steps, everyone has every finished chunk.

Count the bytes. 2(N − 1) steps, each sending a chunk of S/N bytes, so each GPU sends 2(N − 1)/N × S. Every GPU sends and receives at the same time, so all links are busy in parallel. If each GPU’s link moves B bytes per second, the time is about

T ≈ 2 × (N − 1)/N × S / B + 2(N − 1) × α

where α is the fixed latencylatencyThe fixed cost of one message, no matter how small: software overhead, setting up the transfer, the signal crossing the wire. Microseconds on good hardware. per step. The first term is the bandwidthbandwidthHow many bytes per second a link can move once data is flowing. The pipe’s width, not its length. term and is flat in N. The second grows with N. For big gradients the bandwidth term wins, which is why the ring is bandwidth-optimal; for small messages on thousands of GPUs, 2,046 latency hops start to hurt, which is why NCCL also has tree algorithms and builds rings hierarchically (inside a machine first, then across). You’ll see both terms on your Mac in Measure.

People quote collective speed as bus bandwidthbus bandwidthThe number NCCL’s benchmark reports for all-reduce: message size × 2(N−1)/N ÷ time. It normalises away N so you can compare it directly to the link’s rated speed.: S × 2(N − 1)/N ÷ T. If bus bandwidth is close to the link speed, the collective is doing as well as physics allows.

Data parallelism: copy the model, split the batch

Now the simplest way to use N GPUs. Every GPU holds a full copy of the model. Each training step:

1. Each GPU takes a different slice of the batch (GPU 0 gets examples 0-31, GPU 1 gets 32-63, and so on).
2. Each GPU runs forward and backward on its slice and gets its own gradient.
3. All-reduce the gradients and divide by N. Now every GPU has the gradient of the whole batch, exactly what one giant GPU would have computed.
4. Each GPU runs the same optimizer step on the same gradient, so the copies stay identical without ever sending the weights.

This is DDPDDPDistributedDataParallel, PyTorch’s data-parallel wrapper. Broadcasts the initial weights from rank 0, then all-reduces gradients during backward. in PyTorch. The math is exact, not approximate: because the loss is an average over examples, the average of the per-slice gradients is the gradient of the full batch. You’ll check that to seven decimal places in Build It.

The clever part is timing. A naive version waits for backward to finish, then all-reduces 140 GB of gradients while the GPU sits idle. But backward produces gradients layer by layer, from the last layer to the first. The last layer’s gradient is final long before the first layer’s is computed. So DDP groups gradients into bucketsbucketsGroups of gradient tensors (25 MB each by default in DDP) that are all-reduced together as soon as all of them are ready. Bigger buckets waste less on latency; smaller ones start communicating sooner. and starts all-reducing each bucket the moment it’s full, while backward keeps computing the earlier layers. Communication hides behind compute. If the network is fast enough, the all-reduce costs almost nothing.

ZeRO and FSDP: stop keeping N copies of everything

Look at what data parallelism stores. Per parameter, in standard mixed-precision Adam:

With N data-parallel GPUs, every one of them stores all 16 bytes for every parameter. That’s N identical copies. ZeROZeROZero Redundancy Optimizer, from Microsoft’s DeepSpeed (Rajbhandari et al., 2019). Removes the duplicate copies in data parallelism by giving each GPU only a 1/N slice of the training state. asks: why? Each GPU only needs to update its own 1/N of the parameters, if someone else updates the rest. So shardshardSplit a tensor into N pieces and give each GPU one piece. Each piece is a shard. the training state, in three stages:

• Stage 1, shard the optimizer state. Each GPU keeps the fp32 master weights and Adam moments only for its 1/N of the parameters. Communication: instead of an all-reduce, do a reduce-scatter of gradients (each GPU gets the summed gradient for its slice), update your slice, then all-gather the new bf16 weights. Remember: reduce-scatter + all-gather is an all-reduce, so this costs the same bytes as plain DDP. Free memory, no extra communication.
• Stage 2, also shard the gradients. After the reduce-scatter, a GPU only ever needs the gradient for its own slice, so it can throw the rest away as backward goes. Still the same communication.
• Stage 3, also shard the weights. Now no GPU holds the full model at all. Before computing a layer, the GPUs all-gather that layer’s weights, use them, and immediately free them again. It’s “just in time” weights. This costs an extra all-gather in the backward pass (weights are needed again there), so about 1.5x the communication of DDP.

FSDPFSDPFully Sharded Data Parallel, PyTorch’s built-in version of ZeRO stage 3. The current API is fully_shard (often called FSDP2), which stores each parameter as a DTensor sharded along its first dimension. is PyTorch’s version of stage 3. You wrap each transformer block as a unit; before a block runs, FSDP all-gathers that block’s weights, and it prefetches the next block’s weights while the current one computes, so (again) communication hides behind compute.

FSDP fixes memory. It does not fix the other problem from the Predict: every step still all-gathers and reduce-scatters the whole model across all N GPUs over the slow network. And it can’t help with something we’ve ignored: activationsactivationsThe intermediate outputs of every layer saved during the forward pass, because backward needs them. They grow with batch size and sequence length, not with how you shard the weights.. One 4,096-token sequence through a 70B model leaves something like 90 GB of activations behind for backward by the standard estimate (the training memory lesson works this out). Sharding weights doesn’t shrink that. To shrink it, you have to split the layers themselves.

Tensor parallelism: tear each matrix down the middle

Most of a transformer’s compute is two matrix multiplies per block in the MLP and a few in attention. Megatron-LM (Shoeybi et al., 2019) found a way to split them across GPUs that needs very little talking.

The MLP is Y = GeLU(X · A), then Z = Y · B, where X is the activations coming in (tokens × hidden), A expands hidden to 4 × hidden, and B shrinks it back.

1. Split A by columns into A₁ and A₂, one per GPU. Each GPU computes X · Aᵢ, which is half the columns of the result. GeLU works element by element, so each GPU can apply it to its own half without seeing the other half. No communication yet.
2. Split B by rows into B₁ and B₂. Each GPU multiplies its half of Y by its half of B. A matrix multiply is a sum over the inner dimension, so Y₁B₁ + Y₂B₂ = Y · B. Each GPU has a partial sum of the full answer.
3. All-reduce the partial sums. Now both GPUs have Z.

Attention splits the same way: each GPU takes a subset of the heads (column-split the Q, K, V projections), computes attention for its heads entirely locally, and row-splits the output projection, then one all-reduce. So each transformer block costs two all-reduces in the forward pass and two in the backward.

What makes tensor parallelism special is what gets all-reduced and how often. Not the weights, the activations: tokens × hidden × 2 bytes. For one 4,096-token sequence with hidden size 8,192, that’s 64 MB, four times per block, eighty blocks, every micro-batch. And these all-reduces sit right in the middle of the forward pass: the next layer can’t start until the sum arrives, so nothing hides them. Tensor parallelism is a constant stream of medium-sized, blocking messages. It only works on the fastest links you have.

The cables: inside a box vs between boxes

Here’s the hardware fact that decides everything else. Approximate numbers, per GPU, per direction; check the spec sheets for exact values for your generation.

Gigabits divide by 8 to get gigabytes: 400 Gb/s is 50 GB/s. Inside a server, GPUs talk about ten times faster than they talk to GPUs in the next server. That one ratio is why the parallelism types get placed the way they do.

Your Spark has a NICNICNetwork interface card. On the Spark it’s an NVIDIA ConnectX-7, the same family used in big clusters, with two QSFP ports. for exactly this: NVIDIA’s DGX Spark user guide has a page on connecting two Sparks over their ConnectX-7 ports into a two-node cluster (check the current doc for supported cables and setup). It’s the smallest real multi-node setup you can own, and its link is about 10x slower than the Spark’s own ~273 GB/s memory (spec sheet). Same shape as a real cluster, at a smaller scale.

Pipeline parallelism: an assembly line of layers

The third way: cut the model by depth. GPU 0 gets layers 1-20, GPU 1 gets 21-40, and so on. Each GPU sends its output activations to the next one. That’s one point-to-point message per stage boundary per micro-batch, tiny compared to tensor parallelism’s four all-reduces per layer, so it tolerates the slower links between servers.

The catch is obvious once you draw it. With one batch, GPU 1 can’t start until GPU 0 finishes, GPU 2 waits for GPU 1, and on the way back the order reverses. At any moment, one GPU works and the rest wait. The idle time is the pipeline bubblepipeline bubbleThe idle time in a pipeline when some stages have no work because the micro-batches haven’t reached them yet (at the start) or have already left (at the end)..

The fix, from GPipe (Huang et al., 2019): split the batch into m micro-batchesmicro-batchesSmall slices of the batch that move through the pipeline one after another, so different stages can work on different slices at the same time. Gradients are added up over all of them before the optimizer step.. GPU 0 does micro-batch 1, hands it on, and immediately starts micro-batch 2 while GPU 1 works on micro-batch 1. The line fills up, runs full, and drains. With p stages and m micro-batches, the fraction of time wasted is

bubble = (p − 1) / (m + p − 1)

Play with it and notice two things. First, the measured idle fraction matches the formula exactly; the bubble is the fill and drain triangles at the two ends. Second, look at what GPipe does: all m forwards, then all m backwards. Stage 0 must store activations for all m micro-batches at once, and we just said m should be large. 1F1B (“one forward, one backward”, from PipeDream, Narayanan et al.) starts each micro-batch’s backward as early as possible and alternates. Same bubble, but stage 0 never holds more than p micro-batches’ activations. That’s why everyone uses 1F1B or its interleaved variant, which gives each GPU several small, non-adjacent chunks of layers to shrink the bubble further at the cost of more messages.

Putting it together: 3D parallelism

Now the layout falls out of the cable table:

• Tensor parallelism inside a server (TP = 8 on an 8-GPU box). It needs the fastest link, all the time. Going past 8 puts those blocking all-reduces on a 10x slower network.
• Pipeline parallelism across servers. Small point-to-point messages, tolerable on InfiniBand. Choose p to make the model (weights and activations) fit.
• Data parallelism outermost (usually with ZeRO-1 or FSDP on top). Its big gradient exchange happens once per step and overlaps with backward, so it can live on the slowest tier.

The number of GPUs is the product: G = TP × PP × DP. For 1,024 GPUs with TP = 8 and PP = 8, DP = 16. Each GPU is identified by three coordinates, and each collective runs within the group that shares two of them: a TP group is 8 GPUs in one box, a DP group is the 16 GPUs at the same position in each replica. PyTorch calls this grid a device meshdevice meshA logical N-dimensional grid of ranks, like 16 × 8 × 8, where each dimension is one kind of parallelism. Collectives along one dimension only involve the ranks in that row..

A real data point: Meta’s Llama 3 paper reports training the 405B model on 16,384 H100s with TP = 8, PP = 16 and DP = 128 for most of pretraining, reaching roughly 40% of peak throughput (roughly, from their table; read it yourself).

Long sequences: sequence and context parallelism

One more axis, briefly. Some parts of a block (layer norm, dropout) aren’t split by tensor parallelism, so every TP rank stores the full activations for them. Sequence parallelism (Korthikanti et al., 2022) splits those along the sequence dimension instead, and replaces each all-reduce with a reduce-scatter plus all-gather, which costs the same bytes. Free memory again, the same trick as ZeRO.

For very long contexts (128K tokens and up), attention’s activations dominate. Context parallelism splits the sequence itself across GPUs: each GPU holds a block of tokens, and keys and values are passed around a ring so every query block eventually sees every key block (ring attention, Liu et al., 2023). It’s the ring idea again, applied to attention instead of gradients. Llama 3 used it for its long-context phase.

When something breaks (and it will)

Here’s the part people skip in tutorials, and the part that dominates the real job.

A GPU is very reliable. A thousand of them are not. Meta’s Llama 3 paper is unusually honest about this: during a 54-day stretch on 16,384 GPUs they had 466 job interruptions, 419 of them unexpected, and about 78% of the unexpected ones were hardware: faulty GPUs, GPU memory, network links, hosts. That’s about 8 unplanned failures a day, one every three hours. Synchronous training has a nasty property here: the all-reduce needs every GPU, so one dead GPU stops all 16,383 others.

So the system has three jobs.

Checkpoint often, and cheaply. A checkpointcheckpointA saved copy of everything needed to resume training exactly: weights, optimizer state, data loader position, random seeds, step number. for a 70B model is the full 1.1 TB training state. Written in parallel, each GPU saves only its own shard (PyTorch’s distributed checkpointing does this), so the work is spread over all of them. Asynchronous checkpointing copies the state to CPU memory in a second or so and writes it to storage in the background while training continues. How often? A classic rule (Young, 1974; refined by Daly) says the best interval is about √(2 × checkpoint cost × mean time between failures). With a 1-minute checkpoint cost and a failure every 3 hours, √(2 × 1 × 180) ≈ 19 minutes. Each failure then costs half an interval of lost work on average, plus the time to restart.

Restart fast, ideally without a human. elastic trainingelastic trainingTraining that survives losing or adding machines: the job notices a worker died, regroups the survivors (or swaps in a spare), reloads the last checkpoint and continues, without someone restarting it by hand. in PyTorch comes from torchrun: workers meet at a rendezvousrendezvousThe meeting point where workers find each other at startup and agree on ranks. torchrun runs one; if a worker dies, the survivors meet again and restart from the checkpoint., and if one dies, torchrun restarts the group and your script resumes from the latest checkpoint. Big clusters keep spare nodes on standby so the job comes back at the same size. Restart time matters as much as checkpoint frequency: loading 1.1 TB, re-initializing NCCL across 1,024 GPUs and warming up can take many minutes, and at 8 failures a day those minutes add up. This is a design-doc problem, not a code problem; Writing design docs is where you’d write it down.

Find the slow ones. A stragglerstragglerA worker that’s slower than the rest: a GPU throttling because it’s hot, a flaky network link retransmitting, a noisy neighbour. In synchronous training everyone waits for it. doesn’t crash, it just makes every step as slow as itself. Each all-reduce waits for the slowest of N, which is the tail latency problem again: the step time is the max over 1,024 GPUs, not the average. Teams track per-rank step times and kick out the slow node.

Silent errors are worse still: a GPU that returns wrong numbers without crashing. They show up as loss spikes or NaNs. Large training runs watch the loss curve the way you’d watch a patient’s heart monitor.

Build It

You have one GPU, so you’ll fake the cluster with processes. On the Mac or the Spark, torchrun starts several Python processes, gives each a rank, and the gloo backend lets them talk over local TCP. It’s slow (local TCP, not NVLink), but the code is the same code you’d run on 1,024 GPUs with backend="nccl", and correctness is exactly the same.

mkdir -p ~/dist && cd ~/dist
python3 -m venv .venv && source .venv/bin/activate
pip install torch

# hello.py
import torch, torch.distributed as dist

dist.init_process_group("gloo")  # torchrun sets RANK, WORLD_SIZE, MASTER_ADDR for us
rank, world = dist.get_rank(), dist.get_world_size()
x = torch.tensor([float(rank)])
dist.all_reduce(x)  # default op is SUM
print(f"rank {rank}/{world}: sum of ranks = {x.item()}")  # 0+1+2+3 = 6
dist.destroy_process_group()

torchrun --standalone --nproc_per_node=4 hello.py

# ring.py
import time, torch, torch.distributed as dist

def ring_all_reduce(t: torch.Tensor) -> torch.Tensor:
    """Sum t across all ranks, in place, with a ring. Return t."""
    r, n = dist.get_rank(), dist.get_world_size()
    right, left = (r + 1) % n, (r - 1) % n
    chunks = list(t.chunk(n))  # views into t: editing a chunk edits t
    # Phase 1, reduce-scatter: n-1 steps.
    #   At step s, which chunk index do you send right? Which do you receive from the left?
    #   Add the received data into that chunk.
    # Phase 2, all-gather: n-1 steps.
    #   Now forward finished chunks; the receiver copies instead of adding.
    ...
    return t

if __name__ == "__main__":
    dist.init_process_group("gloo")
    torch.manual_seed(dist.get_rank())
    x = torch.randn(1_000_003)          # deliberately not divisible by 4
    mine = ring_all_reduce(x.clone())
    ref = x.clone(); dist.all_reduce(ref)
    # print the max absolute difference on rank 0

# ddp_check.py
import torch, torch.nn as nn, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def make_model():
    torch.manual_seed(0)  # identical init everywhere
    return nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 1))

# 1. init gloo; every rank builds the same full batch X (64 x 32), Y (64 x 1) from one seed
# 2. reference: a plain model, MSE loss on the FULL batch, .backward()
# 3. DDP(make_model()), MSE loss on this rank's 1/world slice of the batch, .backward()
# 4. print the max |grad difference| between the DDP model and the reference, per rank

# fsdp_peek.py
import torch, torch.nn as nn, torch.distributed as dist
from torch.distributed.fsdp import fully_shard
from torch.distributed.device_mesh import init_device_mesh

dist.init_process_group("gloo")
mesh = init_device_mesh("cpu", (dist.get_world_size(),))
torch.manual_seed(0)
model = nn.Sequential(*[nn.Linear(256, 256) for _ in range(4)])
for layer in model:
    fully_shard(layer, mesh=mesh)  # each layer is one all-gather unit
fully_shard(model, mesh=mesh)

w = model[0].weight
# print: type(w), the global shape, w.to_local().shape, w.placements
# then run one forward + backward + AdamW step and print "ok"

# tp_mlp.py
import torch, torch.distributed as dist, torch.nn.functional as F

# every rank: same seed -> same X (8 x 64), A (64 x 256), B (256 x 64)
# ref = gelu(X @ A) @ B                       (the unsplit answer)
# A_i = this rank's column slice of A        (torch.chunk along dim=1)
# B_i = this rank's row slice of B           (along dim=0)
# Z = gelu(X @ A_i) @ B_i, then one collective -> Z should equal ref

# pipe.py  (run: torchrun --standalone --nproc_per_node=4 pipe.py 16)
import sys, time, torch, torch.distributed as dist
# m = int(sys.argv[1]); time the whole loop between two barriers on rank 0
# print measured time, the ideal (m + p - 1) * 0.02, and the bubble 1 - m*0.02/time

def step_time(tp: int, pp: int, dp: int, micro_batch: int = 1) -> dict:
    # compute = 6 * P * tokens / (G * peak * mfu)
    # m = (sequences per DP replica) / micro_batch; bubble = (pp-1)/(m+pp-1); compute /= (1 - bubble)
    # TP comm: 4 all-reduces per layer per micro-batch, each of micro_batch*seq*hidden*2 bytes,
    #          on (layers/pp) layers, over NVLink if tp <= 8 else the inter-node network
    # DP comm: ring all-reduce of this GPU's gradient shard, 2 * P / (tp*pp) bytes, inter-node
    # memory/GPU: 16 * P / (tp*pp) without ZeRO; (4 + 12/dp) * P / (tp*pp) with ZeRO-1
    ...

Measure

Fill in your own numbers. Time on an idle machine; take the faster of two runs.

For reference, one run on an M-series MacBook, 4 gloo processes over loopback, PyTorch 2.14:

What the numbers should tell you:

Small messages are all latency. 4 KB and 64 KB take about the same time, 16x apart in size. That’s the 2(n − 1)α term: six hops of a millisecond or so each through local TCP and Python. NCCL over NVLink has α of a few microseconds, a thousand times smaller, but the same shape holds: below some size, a collective costs a fixed amount no matter how little you send. That’s why DDP batches gradients into 25 MB buckets instead of sending each tensor separately.

Big messages approach a bandwidth ceiling. From 16 MB to 64 MB, time grows about 2-3x for 4x the data. Bus bandwidth climbs toward around 1 GB/s, which is what local TCP through the kernel with memory copies can do. On an H100 node, NCCL’s all_reduce_perf shows a bus bandwidth of a few hundred GB/s at large sizes (check yourself if you get cluster access). Same curve, several hundred times higher.

Your ring may beat gloo at large sizes. Mine did, with noisy timings. gloo picks its own algorithm and isn’t tuned for this case. Don’t read much into the ranking; read the shape.

DDP gradients. Expect differences around 3 × 10⁻⁸ (I got 2.98 × 10⁻⁸ on all four ranks). That’s float32 rounding from adding in a different order, not an approximation in the method.

FSDP. Each rank’s local weight is (64, 256) out of (256, 256): exactly a quarter, split along dim 0.

Pipeline. My fake-compute run gave a measured bubble of 0.80, 0.54 and 0.32 for m = 1, 4, 16, against the formula’s 0.75, 0.43 and 0.16. The gap is each hop’s send/recv overhead (about a millisecond on gloo), which the formula assumes is zero. The trend is the point: more micro-batches, smaller bubble. Real pipelines also pay this hop cost, which is one reason you don’t use more stages than you need.

The calculator. My version (same assumptions as in the task, no overlap, no activation memory) printed:

Read it row by row. Pure DP doesn’t fit and its gradient exchange is longer than the compute. TP = 16 is the tell: spreading tensor parallelism across two nodes makes its communication 20x worse, the whole lesson in one cell. Small TP with deep PP fits easily but loses a third of the cluster to the bubble.

The interesting rows are the TP = 8 ones, and there isn’t one right answer among them. TP = 8, PP = 1 with ZeRO-1 has no bubble at all, and on paper it wins, if its 0.69 s gradient exchange really hides behind backward and the activations fit. TP = 8 with PP = 4 or 8 pays a ~9-10% bubble but cuts each GPU’s gradient traffic 4-8x and leaves lots of memory headroom. Now add activations. A standard estimate (Korthikanti et al., 2022) is about 34 × hidden bytes per token per layer in bf16 without recomputation, so one 4,096-token sequence through one layer is about 1.1 GB, and TP = 8 with sequence parallelism divides that by 8. Multiply by layers per stage and by how many micro-batches a stage holds at once (up to p with 1F1B) and check each row against 80 GB. Which one would you pick, and what would you measure first to find out if you were right? That’s the plan in the next task.

At about 3.9 s per step and ~4.2M tokens per step, 15T tokens is 15 × 10¹² / 4.2 × 10⁶ ≈ 3.6 million steps, about 160 days at 50% utilization (Figure 001’s ~180 days assumed 40%). Make your calculator print that line too; a plan that doesn’t say how many days it takes isn’t a plan.

Break It

1. Deadlock the ring. Replace isend with a blocking send in your ring, so every rank sends before it receives. On small tensors it may still work, because gloo buffers small messages; on 64 MB it hangs forever. Real training jobs hang like this when ranks disagree about which collective comes next (one rank skips a step because of a data-dependent if, everyone else waits in an all-reduce). The symptom is a job using 0% GPU with no error. TORCH_DISTRIBUTED_DEBUG=DETAIL and NCCL’s timeout are how you find it.

2. Break gradient equivalence. In ddp_check.py, change the loss to reduction="sum". DDP still averages gradients across ranks, so you now get 1/world of the right gradient. Then go back to mean, but give rank 0 a slice of 10 examples and the others 18. The mean of per-rank means is no longer the global mean. Both are real bugs people ship: the training still runs, the loss still goes down, and the model is quietly different from what you think you trained.

3. Forget to seed the init. Remove torch.manual_seed(0) from make_model. Run plain all-reduce data parallelism by hand (not DDP) and the copies drift apart from the first step. Then try it with DDP: it still matches, because DDP broadcasts rank 0’s weights at construction. That broadcast in Figure 002 is doing real work.

4. Put tensor parallelism on the slow link. In plan.py, set TP = 16 (or 32). Watch TP communication jump from under a second to tens of seconds per step. Then change the inter-node bandwidth to the Spark’s ConnectX-7 (~25 GB/s) and to 10 GbE (~1.25 GB/s) and see which configs survive. This is the quantitative version of “TP inside the node”.

Explain It Back

Answer out loud, in plain words, as if in an interview. Then check.

Go Deeper

• Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism” (2019), arXiv:1909.08053. The column/row split of Figure 005, explained in two pages. Then the follow-up, Narayanan et al., “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM” (2021), arXiv:2104.04473, the paper that worked out how to combine TP, PP and DP and introduced the interleaved schedule.
• Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models” (2019), arXiv:1910.02054. The source of the 16 bytes/param accounting and the three stages. Very readable.
• Huang et al., “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism” (2019), arXiv:1811.06965. Micro-batches and the bubble formula.
• Llama Team, Meta, “The Llama 3 Herd of Models” (2024), arXiv:2407.21783. Read the infrastructure and parallelism sections: the real 4D layout, the failure table, and what goes wrong at 16K GPUs. The best public account of a large run.
• Hugging Face, “The Ultra-Scale Playbook: Training LLMs on GPU Clusters” (2025), on huggingface.co (nanotron). A long, practical, interactive walk through every parallelism type with real benchmark numbers. The closest thing to a textbook for this topic.
• PyTorch docs: Distributed communication package, the DDP design note, and the FSDP2 fully_shard tutorial. Check the FSDP pages against your version; the API moved from FullyShardedDataParallel to fully_shard.
• Patarasuk and Yuan, “Bandwidth Optimal All-reduce Algorithms for Clusters of Workstations” (2009), and Andrew Gibiansky, “Bringing HPC Techniques to Deep Learning” (Baidu Research blog, 2017). The proof that the ring is optimal, and the post that brought it to deep learning.
• NVIDIA, DGX Spark User Guide, “ConnectX-7 Networking” on docs.nvidia.com. How to cable two Sparks together if you ever want a real two-node run.

Connections

This builds directly on How much memory does a model need? (the bytes-per-parameter accounting that ZeRO shards), and on Processes, threads and the GIL: every rank here is a separate process with its own memory, which is exactly why they need collectives instead of shared variables. The two-term cost model, latency plus size over bandwidth, is the same one from How slow is slow?, now for cables between GPUs instead of RAM and disk. Stragglers are Why p99 matters at cluster scale.

Next: Train a small LM, end to end is where you’d use DDP for real (even two processes on the Spark teach you the launch, the seeding and the checkpointing). And the checkpoint-and-restart strategy belongs in a document: Writing design docs is the place to practise turning “something fails every three hours” into a plan someone can approve.