47. DDP-Style Gradient All-Reduce Simulation

Problem

Data-parallel training (PyTorch's DistributedDataParallel) runs an identical copy of the model on every worker, each on its own shard of the batch. After backward(), every worker has a different local gradient for the same parameter, because each saw different data. DDP's all-reduce step fixes that: it sums each parameter's gradient across all workers and divides by world_size, so that afterward every worker holds the exact same averaged gradient and takes the exact same optimizer step -- keeping every replica's weights identical throughout training.

Implement the simulated all-reduce: given each worker's local gradient dict for the same model replica, return the single averaged-gradient dict every worker would hold locally after a real all_reduce(op=SUM) followed by dividing by world_size.

Source: src/47_ddp_gradient_allreduce_sim.py

def ddp_allreduce_average_grads(
    worker_grads: list[dict[str, torch.Tensor]],
) -> dict[str, torch.Tensor]: ...

>>> w0 = {"w": torch.tensor([1.0, 2.0])}
>>> w1 = {"w": torch.tensor([3.0, 4.0])}
>>> ddp_allreduce_average_grads([w0, w1])
{'w': tensor([2., 3.])}

Step-by-Step Approach

1. Note that every worker's gradient dict has the same parameter names and shapes (DDP requires every replica to start from identical weights), so the first worker's keys are the full set of parameters to reduce.
2. For each parameter name, gather that parameter's gradient tensor from every worker into a list, in worker order.
3. Stack that list along a new leading dimension with torch.stack, turning world_size separate tensors of shape S into one tensor of shape (world_size, *S).
4. Take the mean along the new dimension 0 -- this single call does both the SUM-across-workers and the divide-by-world_size at once.
5. Store the result under that parameter's name in the output dict.
6. Repeat for every parameter name and return the combined dict.

The key insight is that a real all-reduce's ring or tree communication pattern is purely a bandwidth-efficiency optimization -- the arithmetic result it leaves on every worker is exactly the elementwise mean across workers, so a single-process simulation only needs to reproduce that final numeric answer, not the communication topology, to correctly test DDP-style gradient synchronization logic.

Reference solution

def ddp_allreduce_average_grads(
    worker_grads: list[dict[str, torch.Tensor]],
) -> dict[str, torch.Tensor]:
    world_size = len(worker_grads)
    param_names = worker_grads[0].keys()

    averaged = {}
    for name in param_names:
        # torch.stack adds a new leading "worker" dimension so a single
        # mean(dim=0) does the SUM-then-divide-by-world_size in one call --
        # numerically identical to what a real ring/tree all_reduce(SUM)
        # followed by a /world_size would leave on every worker.
        stacked = torch.stack([worker[name] for worker in worker_grads], dim=0)
        averaged[name] = stacked.mean(dim=0)

    return averaged

Key Functions & Tricks

• torch.stack(tensors, dim=0) — combines a Python list of same-shaped tensors into one tensor with a new leading dimension, unlike torch.cat which needs an existing dimension to concatenate along.
• tensor.mean(dim=0) — reduces across the stacked "worker" dimension in one call, computing both the sum and the divide-by-world_size together.
• dict comprehension over parameter names — a clean way to apply the same per-parameter reduction across an arbitrary, dynamically-discovered set of parameter names rather than hardcoding them.
• worker_grads[0].keys() — reads the parameter name set from the first worker, relying on the guarantee that every worker's replica has identical parameter names.

How to Recognize This Pattern