DDP's global batch is a promise, and set_epoch() is the one line that keeps it

Autograd’s backward hooks are what actually fire DDP’s gradient synchronization: as .grad gets populated during the backward pass, DDP has already registered a hook that triggers an AllReduce the instant a bucket of gradients is ready. That’s the mechanism. What it’s in service of is a much simpler idea, worth stating precisely before any of the machinery: every rank starts a step holding identical parameters, each rank processes a different slice of the batch, and AllReduce makes the resulting gradient update mathematically equivalent to having trained on the whole combined batch on one giant GPU. DDP is a replicated-state system, not a sharded one, every rank holds a full copy of the model, and the entire discipline of running it correctly reduces to one requirement: getting different data to different ranks, exactly once each, every epoch. Getting that requirement wrong doesn’t crash. It trains a real model that quietly never sees most of its own dataset.

Rendezvous and rank identity, the part that isn’t boilerplate#

Three integers define where a process sits: RANK is its unique index across the entire job, WORLD_SIZE is the total process count, LOCAL_RANK is its index on the current machine only. The convention is one process per GPU, and torch.cuda.set_device(local_rank) has to run before any collective touches that process, because a collective issued against the wrong device is exactly the kind of mismatch that produces a hang with zero symptoms rather than a clean error. torchrun injects all three as environment variables automatically; on a Slurm cluster, the same three numbers come from SLURM_PROCID, SLURM_NTASKS, and SLURM_LOCALID, and a launcher script that maps these incorrectly, two ranks landing on one GPU, or LOCAL_RANK computed wrong on a multi-node job, produces GPU utilization and memory numbers that look almost plausible while the job is subtly, silently wrong. Rendezvous, every rank finding every other rank and agreeing to form the process group, is all-or-nothing: it cannot complete partially, so treating the launcher as shell boilerplate rather than the actual first correctness-critical step in the job is a common, expensive misread.

The actual code rendezvous produces, the part every rank runs before training starts:

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend="nccl")   # rendezvous happens inside this call, using the
                                            # RANK / WORLD_SIZE / MASTER_ADDR env vars torchrun injected
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])   # the wrapper that registers the backward hooks

init_process_group is the literal function call rendezvous happens inside: it blocks until every rank in WORLD_SIZE has shown up and agreed to form the group, which is exactly why a single rank that fails to launch, a stale container image, a version mismatch, hangs every other rank here with nothing yet printed to any log. DDP(model, device_ids=[local_rank]) is the wrapper that does the hook registration described above; nothing about gradient synchronization exists until this line runs.

torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
  --master_addr=10.0.0.1 --master_port=29500 \
  train.py

This is the launcher itself, run once per node, --nproc_per_node spawning one process per GPU on that node and injecting RANK, WORLD_SIZE, and LOCAL_RANK into each process’s environment before train.py ever executes. A second node in the same job runs the identical command with --node_rank=1; torchrun computes each process’s global RANK from node_rank and nproc_per_node automatically, which is precisely the arithmetic a hand-rolled Slurm mapping has to get right on its own instead.

The sampler’s job: partition the dataset once, the same way, on every rank#

For a map-style dataset of length $D$ across $P$ ranks, each rank should get roughly $D/P$ samples, and critically, every rank has to derive its shard from the same global ordering, or the union of what every rank sees stops matching the dataset and ranks silently start seeing overlapping or missing examples. DistributedSampler solves this by building one global permutation of all $D$ indices, seeded deterministically from (seed, epoch), then handing rank $i$ every $P$-th slice of that single permutation. Two ranks never derive their own independent random shuffles; there is exactly one shuffle, computed identically everywhere, then sliced.

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=42)
loader = torch.utils.data.DataLoader(dataset, batch_size=micro_batch, sampler=sampler, shuffle=False)

shuffle=False on the DataLoader itself here is not a contradiction, it’s a delegation: the sampler owns ordering, and the DataLoader must not also shuffle on top of it, or the deterministic global permutation the sampler just built gets silently re-shuffled per rank, breaking the invariant that every rank’s shard is a clean, disjoint slice of one shared ordering.

The bug that never raises an error. DistributedSampler builds its permutation from (seed, epoch), and epoch defaults to whatever it was last set to, 0, if nothing ever calls sampler.set_epoch(epoch) inside the training loop. Miss that one line and every epoch reuses the identical seed pair, which means every epoch produces the exact same global permutation, which means every rank sees the exact same subset of the dataset, in the exact same order, epoch after epoch. Loss still goes down. The run looks completely normal from the outside. What’s actually happening is a model training on a dataset $P\times$ smaller than intended, forever, because $D/P$ examples per rank repeated identically every epoch is not the same computation as $D/P$ different examples drawn from a fresh global shuffle each epoch. This is structurally the same failure shape as the missing zero_grad() bug: a one-line omission, no exception anywhere, a loss curve that looks fine, and a model that is quietly not doing what you think it’s doing.

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)          # this line is the entire fix
    for batch in loader:
        train_step(batch)

Verifying the invariant directly, rather than trusting it. The correctness property a distributed sampler must satisfy is simple enough to check by hand on a small synthetic dataset:

seen_by_rank = {r: set(DistributedSampler(range(23), num_replicas=4, rank=r, seed=0).__iter__()) for r in range(4)}
all_seen = set().union(*seen_by_rank.values())
for r1 in range(4):
    for r2 in range(r1 + 1, 4):
        assert seen_by_rank[r1].isdisjoint(seen_by_rank[r2])   # no rank duplicates another rank's work
print(len(all_seen), "of 23 indices covered")   # 24 -- padding rounds 23 up to a multiple of 4

$23$ isn’t divisible by $4$, so DistributedSampler pads by repeating the first few indices of the permutation until the total is divisible by world_size, by default, rather than silently giving one rank a shorter batch than the others. drop_last=True takes the opposite trade, dropping the remainder instead of padding, and a global drop_last at the sampler level is a different knob from a local drop_last inside a single rank’s own DataLoader, conflating the two is a common source of “why do my ranks disagree on step count” confusion. Iterable and streaming datasets get none of this for free, since there’s no fixed-length permutation to slice, and need explicit rank-aware sharding written by hand, with the same disjointness invariant enforced manually.

The global batch is one formula, and changing one input changes the optimization problem#

$B_{\text{global}} = B_{\text{micro per GPU}} \times P_{\text{data-parallel}} \times N_{\text{grad accumulation}}$

Llama 3 405B’s own published configuration makes this concrete rather than abstract: in its 8K-context pretraining stage, the parallelism mesh ran with a data-parallel width of $P=64$ (TP=8, CP=1, PP=16, DP=64, across 8,192 GPUs), holding a fixed 16M-token global batch. If someone resizes that job to 32 data-parallel ranks without touching microbatch size or accumulation steps to compensate, $B_{\text{global}}$ silently halves, and that is not a free performance change, it’s a different optimization problem, one that can call for a different learning rate, a different warmup schedule, or both. “Scale the cluster” and “keep training the same model” are not automatically the same action, and the global-batch formula is the exact place that distinction lives.

Gradient synchronization: buckets, overlap, and no_sync()#

DDP doesn’t wait for the entire backward pass to finish before communicating. It groups parameters into buckets, and the instant a bucket’s gradients are fully populated, it fires an AllReduce for that bucket while backward computation for earlier layers is still running, overlapping communication with compute rather than serializing them. This overlap, not merely “running on multiple GPUs,” is a large part of why DDP throughput doesn’t collapse as the model grows.

Synchronizing gradients on every microbatch during gradient accumulation is pure waste: only the final accumulated gradient, right before the optimizer step, needs to be correct across ranks. no_sync() suppresses the AllReduce hooks for every microbatch except the last:

for i, micro_batch in enumerate(micro_batches):
    context = model.no_sync() if i < len(micro_batches) - 1 else contextlib.nullcontext()
    with context:
        loss = model(micro_batch)
        loss.backward()          # accumulates into .grad locally, no AllReduce fired
optimizer.step()                   # runs against the one, fully-accumulated, now-synced gradient
optimizer.zero_grad()

Skip no_sync() and DDP still produces a correct result, it just pays for $N_{\text{accum}}$ AllReduce calls per optimizer step instead of one, bandwidth spent synchronizing intermediate gradient state nobody needed synchronized.

find_unused_parameters=True exists for models where not every parameter participates in every forward pass, conditional computation, certain multi-task heads, and its cost is real: DDP has to traverse the autograd graph on every forward pass to work out which parameters were actually used this time, extra work paid on every single step, not a one-time setup cost. Reach for it only when the graph is genuinely dynamic, never as a default defensive setting.

Uneven inputs are the same barrier problem in a new costume#

If dataset size isn’t evenly divisible by world size and padding isn’t used, one rank can run out of data a step before the others. That rank stops calling backward(); every other rank is still waiting at the AllReduce barrier those hooks would have fired. This is the identical silent-hang mechanism covered in the NCCL post, just triggered by a data-length mismatch instead of a dead process, every surviving rank blocks forever because NCCL has no default timeout. Join() is DDP’s purpose-built fix: ranks that exhaust their data early enter a context where they participate in the collectives late-finishing ranks still need, without a real gradient update, so the barrier resolves cleanly instead of hanging.

When DDP is the right answer, and when it stops being one#

DDP is the correct default when the full model and full Adam optimizer state fit on every GPU, and that “fits” is a concrete, computable number, not a feeling: a 7B-parameter model in Adam needs roughly 112GB of persistent state per GPU by the direct 16-bytes-per-parameter accounting, and DDP replicates that entire figure onto every rank. That’s precisely the wall DDP hits, and precisely the reason ZeRO and FSDP exist: not a more advanced default to reach for out of enthusiasm, but the direct engineering answer to persistent state that no longer fits when replicated whole.

Common mistakes#

Setting shuffle=True on the DataLoader in addition to DistributedSampler(shuffle=True): the sampler already owns global ordering, and a second shuffle on top of it breaks the disjoint-shard guarantee the sampler was built to provide.

Forgetting sampler.set_epoch(epoch): the run doesn’t crash, it trains on the identical dataset subset every epoch, indistinguishable from healthy training by loss curve alone.

Resizing the number of data-parallel ranks without revisiting the global-batch formula: the effective batch size changes, and with it the optimization problem, even though nothing about the model code changed.

Treating find_unused_parameters=True as a safe default: it’s a per-step graph traversal cost paid on every forward pass, worth enabling only for models whose computation graph genuinely varies between steps.

What this takes to be frontier-job-ready#

The technical axis is holding the full mental model at once: rank identity and rendezvous, the sampler’s disjointness guarantee, the global-batch formula, and gradient-sync overlap are one connected system, not four separate facts, and a bug in any one of them (a launcher rank mismatch, a missing set_epoch(), an unadjusted batch size after a resize) produces a run that trains without ever raising an exception.

The operational axis is stated close to verbatim in real hiring language for this kind of role: Mistral’s own engineering bar for its research-engineering team is explicit that “you don’t panic when you see OOM errors or when NCCL feels like not wanting to talk,” alongside a concrete stack requirement of “PyTorch or JAX” plus “DeepSpeed / FSDP / SLURM / K8s (distributed training frameworks)”. Google DeepMind’s research-engineer postings for large-model work state the same expectation more plainly still: “experience training large-scale models on accelerators (TPUs, GPUs) in a distributed environment” is listed as a hard requirement, not a nice-to-have, for roles working on production-scale models.

The autonomy axis is the sample-ID audit exercise below: nobody hands you a pre-verified sampler in production, the discipline is checking the disjointness and coverage invariant yourself, on your own dataset, before trusting a multi-day run to it.

Try it yourself#

Beginner. Using the synthetic sampler-audit code above, change num_replicas from 4 to 5 against the same 23-index dataset, and explain in one sentence why the padded index count changes the way it does.

Intermediate. Build the missing-set_epoch() bug directly: instantiate a DistributedSampler with shuffle=True, iterate it for three simulated “epochs” without ever calling set_epoch(), and confirm the yielded index order is identical each time. Then add the call and confirm it isn’t.

Advanced. Given Llama 3’s published 8K-context configuration (TP=8, CP=1, PP=16, DP=64, 16M tokens per global batch), compute the global batch size if the job were resized to DP=32 with everything else held fixed, and state precisely what has to change elsewhere (microbatch size, accumulation steps, or both) to keep $B_{\text{global}}$ constant across the resize.

---

The one-sentence version: DDP’s entire job is keeping one promise, that every rank trains on a distinct, disjoint slice of the same global shuffle, and the mechanisms that keep it, rendezvous, DistributedSampler, set_epoch(), the global-batch formula, gradient bucketing with no_sync(), all exist because breaking that promise doesn’t crash the job, it just quietly trains a different, worse model than the one you think you’re running. DDP replicates everything; the next question is what happens when the thing being replicated no longer fits.