AI Infra Interviews logo
🕸️ Distributed Training
Foundational

Data Parallelism and DDP

Data parallelism gives every GPU a full copy of the model, feeds each a different slice of the batch, and averages the gradients with an all-reduce so every replica takes the same optimizer step. It is the first parallelism every training job uses, and the tokens-per-GPU arithmetic behind it decides whether the communication hides behind the backward pass or dominates the step.

TL;DR: Every rank holds the whole model and optimizer, computes gradients on its own micro-batch, and an all-reduce averages those gradients before an identical optimizer step on every rank. PyTorch DDP makes this cheap by grouping gradients into buckets and launching each bucket's all-reduce while the backward pass is still working on earlier layers, so the network runs in the shadow of the compute.

The idea in one picture

Replicate the model N times. Give each replica 1/N of the global batch. After backward, every replica holds a gradient computed on different data, and the only thing they need to agree on is the average. One all-reduce produces that average on every rank at once, so every optimizer step sees the same gradient and the replicas never drift apart. No parameter ever moves between GPUs; only gradients do.

rendering diagram…

The average rather than the sum is a real detail. The loss is a mean over the global batch, so its gradient is the mean of the per-rank gradients, and DDP divides by the world size inside the reduction. Sum instead and your effective learning rate scales with N; the run diverges at large world sizes.

The mechanism: buckets and overlap

A naive implementation waits for the whole backward pass, then all-reduces the full gradient. That serializes compute and communication and leaves the interconnect idle while the GPU is busy. DDP does three things instead.

It registers an autograd hook on every parameter, so it learns the moment each gradient is ready. It groups parameters into buckets (bucket_cap_mb, 25 MB by default) in reverse registration order, because backward produces the last layer's gradients first. And when every gradient in a bucket has landed, it launches that bucket's all-reduce on a separate CUDA stream while autograd keeps working on earlier layers. Only the final bucket's reduction is exposed, since nothing is left to compute behind it.

TIME backward layers 32-25 layers 24-17 layers 16-9 layers 8-1 all-reduce bucket 1 bucket 2 bucket 3 bucket 4 EXPOSED BUCKETS FIRE IN REVERSE LAYER ORDER, EACH UNDER THE NEXT CHUNK OF BACKWARD

Two knobs matter in practice. find_unused_parameters=True makes DDP walk the whole autograd graph every iteration to find parameters that got no gradient, which costs time and defeats the static bucket layout; turn it on only for models that skip parameters, and prefer static_graph=True when the graph never changes. With gradient accumulation you must wrap the non-final micro-steps in model.no_sync(), or DDP all-reduces after every micro-batch and you pay the full communication cost k times for k accumulation steps.

The numbers

Take Llama 3.1 8B, 8.03B parameters, gradients in BF16: 16.1 GB to reduce per step. A ring all-reduce moves 2(N-1)/N times the buffer through each rank, so on 8 ranks each GPU sends and receives about 28 GB, and on 64 ranks about 31.6 GB. The per-rank volume barely moves with N.

Time that against the fabric. Inside an H100 node, NVLink is 900 GB/s bidirectional per GPU as of 2026, so 28 GB takes at least 62 ms in each direction at the spec ceiling. Across nodes on a 400 Gb/s NIC per GPU (roughly 50 GB/s), 31.6 GB takes about 0.63 s. The compute for that step, at 16,384 tokens per GPU, is 6 × 8.03e9 × 16,384, about 7.9e14 FLOP, which an H100 at 40% of its 989 TFLOPS dense BF16 peak finishes in about 2 s. Intra-node the reduction is 3% of the step. Inter-node it is roughly 30%, and DDP's overlap is what makes most of that 30% disappear. Shrink the per-GPU batch to 4,096 tokens and the compute drops to 0.5 s while the reduction stays at 0.63 s: now you are communication-bound and no bucket schedule saves you.

Memory is the other limit, and it arrives earlier than most candidates expect. Mixed-precision Adam holds BF16 weights, BF16 gradients, FP32 master weights and two FP32 moments: 16 bytes per parameter. The 8B model is 128 GB of state per GPU before a single activation, which does not fit an 80 GB H100 and leaves almost nothing on an H200's 141 GB. On 80 GB parts, plain DDP with Adam tops out around a 4B-parameter model.

ModelParamsBF16 gradient reduced per stepTraining state per GPU at 16 B/paramFits an 80 GB H100 under DDP?
Llama 3.1 8B8.03B16.1 GB128 GBNo
Llama 3.1 70B70.6B141 GB1,130 GBNo, by 14x
A 4B-class model4B8 GB64 GBYes, with 16 GB left for activations

Above that, ZeRO or FSDP has to shard the state.

Where it breaks in production

The classic symptom is a step time that grows with world size while per-rank compute stays flat. In a torch.profiler trace the ncclDevKernel_AllReduce kernels should sit under the backward kernels on a second stream, not after them. If they trail the backward, something is blocking overlap, usually find_unused_parameters=True, an autograd graph that changes between iterations, or a bucket so large that it only fills at the end of backward.

What interviewers are listening for

The follow-up they hold back is "why reverse order?" The answer that shows you have read the code: backward computes gradients from the output toward the input, so the last-registered parameters are ready first, and buckets ordered that way start reducing earliest. The answer that sounds right and fails is "DDP overlaps communication with the forward pass." It cannot, because there is no gradient to send during forward.

They also want the memory sentence. Say "16 bytes per parameter in mixed precision with Adam" unprompted and you have told them you have sized a real job.

Common misconceptions

  • DDP splits the model. It does not; every rank holds all of it. Sharding the model is ZeRO, FSDP or tensor parallelism.
  • More GPUs mean more communication per GPU. The ring all-reduce moves 2(N-1)/N of the buffer per rank, 1.75x at N=8 and 1.97x at N=64. Latency grows with N; volume does not.
  • Gradient accumulation is free. Without no_sync() it multiplies the communication by the accumulation factor.

Key takeaways

  • Replicate the model, shard the batch, all-reduce the mean gradient, take identical optimizer steps.
  • Buckets fire in reverse registration order and overlap with backward; only the final bucket is exposed.
  • Per-rank traffic is 2(N-1)/N of the gradient size, so tokens per GPU per step decides whether communication hides behind compute.
  • 16 bytes per parameter with mixed-precision Adam means DDP alone stops around 4B parameters on an 80 GB GPU.
RELATED CONCEPTS
LESSONS THAT TEACH THIS
PRACTICE THIS IN REAL QUESTIONS