Data Parallelism and DDP

TL;DR: Data parallelism runs a model replica on each worker and combines gradients from different data. Ordinary DDP averages those gradients so replicas with matching weights and optimizer state can take the same update. It saves training time by sharing the batch; it does not divide the model's memory across workers.

Two workers, one update

Training repeats one loop: run a batch through the model, measure the error (the loss), compute a gradient for every weight, meaning how much and in which direction that weight should move to reduce the loss, and let the optimizer apply the update. Data parallelism splits the batch across workers and combines their gradients before the update.

Suppose two workers start with identical weights. Each computes a two-element gradient from its own equally sized batch. A rank is a worker's ID in the communication group; in this example each rank uses one GPU.

The input examples remain local. Each replica runs the complete model, so ordinary DDP does not exchange intermediate activations. The training program must distribute the input, for example with a distributed sampler; wrapping a model in DDP does not do that automatically. Accidentally feeding both ranks the same examples still produces synchronized gradients, but you have duplicated the work.

There is also communication outside backward. Default DDP initialization synchronizes parameters and buffers. Runtime buffer synchronization can broadcast values such as BatchNorm running statistics. Optimizer state is managed separately: after resuming, identical weights with different Adam moments are not identical training replicas.

When an average of averages is wrong

The two-rank example assumes equal contributions to the loss. If rank 0 averages two valid tokens and rank 1 averages six, averaging their local means gives each rank equal influence, not each token.

For a scalar illustration, let those local mean gradients be 1 and 3:

equal-rank mean = (1 + 3) / 2 = 2
token mean      = (2 × 1 + 6 × 3) / 8 = 2.5

For a global token-mean objective, normalize summed token losses by the total valid-token count across the group, accounting for DDP's division by group size. Padding tokens excluded from the loss must also be excluded from that count. Gradient accumulation extends the denominator across all micro-batches in the optimizer step.

This is a correctness check before a performance check. Compare a tiny distributed update with the corresponding combined batch, allowing floating-point rounding differences. A loss curve alone can hide the wrong weighting.

Buckets create opportunities for overlap

Backward produces gradients progressively. DDP groups them into buckets, then reduces ready buckets while other backward computation continues. Bucket scheduling must preserve a consistent collective order across ranks. Readiness, message size, available bandwidth and late ranks all affect completion; there is no guarantee that only one bucket remains after backward.

Consider this hypothetical serialized communication resource. Buckets become ready at 2, 4, 6 and 8 milliseconds, and each reduction takes 3 milliseconds:

Backward computation ends at 8 ms, but reduction ends at 14 ms. The exposed tail is 6 ms, even though one reduction takes only 3 ms. Earlier communication queued up. The DDP backward call must still wait for its required reductions; the 8 ms marker is compute completion, not the API returning. This schedule is a teaching example, not a prediction of a particular GPU.

In a torch.profiler trace, inspect bucket readiness and completion on multiple ranks. A smaller bucket may start earlier but create more launches. A larger bucket may use bandwidth better but arrive too late. Compare end-to-end step time after warmup, rather than choosing the shortest individual collective.

Count the actual gradient buffer

If a hypothetical one-billion-parameter model communicates FP32 gradients, its gradient payload is 4 GB in decimal units. A ring all-reduce across eight ranks sends 2 × 7/8 × 4 = 7 GB per rank and receives the same amount. BF16 gradients would halve that payload, but an autocast forward does not by itself establish the gradient storage or communication dtype.

Model Memory Footprint gives the separate capacity check. Plain DDP keeps full parameter, gradient and optimizer storage on each rank, plus activations and temporary buffers. One mixed-precision Adam layout costs 16 bytes per parameter; other layouts differ. Adding DDP ranks cannot rescue a replica that already exceeds one rank's memory budget. ZeRO and FSDP changes that ownership.

Accumulation and changing graphs

With several micro-batches per update, ordinary DDP synchronizes on every backward pass. Put both forward and backward for non-final micro-batches inside no_sync() to defer gradient synchronization until the final pass. This changes the number of reductions, not the need to normalize the loss correctly.

Unused parameters and changing control flow require care. find_unused_parameters can identify gradients that will not arrive. A static-graph setting is appropriate only when the used/unused parameter set and training graph behavior stay fixed. It is not a general switch for hiding network time.

Check your understanding

Would identical reduced gradients repair a rank that restored the wrong optimizer state? No. Its update can still differ. Would doubling ranks halve gradient bytes per rank? No: the eight-rank ring factor is 1.75; the sixteen-rank factor is 1.875. More workers mainly divide the data workload, while collective scheduling and topology decide the communication cost.

Key takeaways

• Equal rank means represent a global mean only with equal loss weights.
• DDP synchronizes gradient buckets; initialization and buffers add other communication.
• Four 3 ms reductions can leave a 6 ms tail when they queue behind an 8 ms backward pass.
• Check per-rank memory first, then measure whether the reduction schedule improves total step time.