DDP under the hood
DDP is replicated local training held together by a strict collective contract: same parameters, matched backward work, partitioned data, averaged gradients.
An eight-GPU job is slower than one GPU and loses progress on failure.
This chapter moves from prove one-step equivalence to keep rank behavior symmetric. We will make the idea concrete in a small code example before placing it back inside a production system with eight GPUs with restart under 15 minutes.
Ensure all ranks take the same branches around forward and backward. Restrict logging and artifact writes to rank zero without skipping required collectives elsewhere.
Verify the behavior illustrated in ddp_epoch.py, then measure decision utility, slice reliability, and system latency.
One rank skips a batch and misses an all-reduce, leaving every other rank hung. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
DistributedDataParallel feels local—each rank calls `backward()`—because hooks turn gradient readiness into coordinated all-reduces behind the scenes.
Every DDP rank owns a complete model replica and processes a different local batch. During backward, gradients are grouped into buckets; as a bucket becomes ready, an all-reduce sums it across ranks and DDP exposes the averaged result. Identical starting weights plus identical averaged gradients let every rank take the same optimizer step.
Correctness depends on symmetry. Ranks must execute compatible graphs and reach collectives in the same order. The distributed sampler must partition data without unintended overlap, and `set_epoch` must change its deterministic shuffle each epoch. Global batch is local batch multiplied by world size and accumulation steps, so scaling can change optimization even when the code looks identical.
- distributed samplers
- all-reduce
- global batch
- FSDP
Follow the reasoning before touching the code.
- 01
Prove one-step equivalence
Compare a fixed global batch on one GPU and several ranks. Check averaged loss, gradients, and updated parameters within the expected numerical tolerance.
- 02
Partition examples deliberately
Use a distributed sampler, call `set_epoch`, and decide how to handle a final uneven batch. Log example IDs on a small fixture to catch duplicates or omissions.
- 03
Keep rank behavior symmetric
Ensure all ranks take the same branches around forward and backward. Restrict logging and artifact writes to rank zero without skipping required collectives elsewhere.
import os, torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
dist.init_process_group('nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
dataset = TensorDataset(torch.randn(128, 8), torch.randn(128, 1))
sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=8, sampler=sampler)
model = DDP(torch.nn.Linear(8, 1).to(local_rank), device_ids=[local_rank])
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
sampler.set_epoch(0)
for x, y in loader:
optimizer.zero_grad(); loss = torch.nn.functional.mse_loss(
model(x.to(local_rank)), y.to(local_rank))
loss.backward()
optimizer.step()A minimal DDP epoch with rank-local data and synchronized gradients.
What to notice
- No explicit all-reduce appears in the loop; DDP’s autograd hooks synchronize parameter gradients during `loss.backward()`.
- Each rank sees only its sampler partition, while the optimizer step uses gradients averaged across all ranks.
YOUR TURN Pause and predict before revealing the answer.
One rank conditionally skips `loss.backward()` for an empty batch. What happens to the other ranks?
Need a hint?
Their backward passes are waiting for every participant in the same gradient collectives.
Check your reasoning
The other ranks can hang in all-reduce because the skipping rank never enters the matching collectives. Handle empty input symmetrically or coordinate a shared stop/skip decision before backward.
Lessons 1 in this course.
Connect the chapter to the course’s quantitative spine.
DDP is replicated local training held together by a strict collective contract: same parameters, matched backward work, partitioned data, averaged gradients. The scaling efficiency below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Scaling efficiency
| Symbol | Meaning / shape / unit |
|---|---|
W | worker/GPU count |
throughput_W | distributed samples per second |
throughput_1 | single-worker baseline |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: W, throughput_W, throughput_1.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c13(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
All-reduce scaling lab
Add workers and estimate delivered speedup after communication and stragglers.
Assumption: 120 ms local-compute step; ring communication and logarithmic straggler penalty.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| 1 workers | 120.3 ms | 1.0× |
| 2 workers | 125.7 ms | 1.9× |
| 4 workers | 131.4 ms | 3.7× |
| 8 workers | 137.8 ms | 7.0× |
| 16 workers | 145.6 ms | 13.2× |
Trace the complete operating path.
Carry the chapter’s central idea into the system: DDP is replicated local training held together by a strict collective contract: same parameters, matched backward work, partitioned data, averaged gradients. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Prove one-step equivalence” into a written invariant and a failing test.
- 02
Promote ddp_epoch.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Keep rank behavior symmetric” as the release checkpoint: Ensure all ranks take the same branches around forward and backward. Restrict logging and artifact writes to rank zero without skipping required collectives elsewhere.
- 04
Serve or materialize behind explicit eight GPUs with restart under 15 minutes.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
One rank skips a batch and misses an all-reduce, leaving every other rank hung. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Fault-tolerant distributed trainer
Specify DDP correctness, scaling experiments, worker failure, and complete checkpoint recovery.
- Gradient/data equivalence
- 1/2/4-GPU efficiency report
- Kill-and-resume proof
- Hang/straggler runbook
Read primary material with a purpose.
Return to the opening failure.
Log per-rank data/compute/collective time, preserve complete training state, and keep a one-GPU equivalence test.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce efficient ddp, parallelism choices, elastic sampling, and resumable checkpoints.