🎓 A1Learner
🔍

PyTorch The DDP Pattern

DistributedDataParallel is a wrapper that runs one model copy per worker and averages the gradients inside backward. The wrapper is often shortened to DDP in code and prose.

The wrapper, shown in full

This code is the pattern, shown and not run here. The wrap happens once, before the training loop.

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

dist.init_process_group("gloo", rank=rank, world_size=world_size)
model = DDP(model)

Where the average hides

The training loop keeps its usual shape under DDP. The averaging rides inside the backward call.

for xb, yb in loader:
    optimizer.zero_grad()
    loss = loss_fn(model(xb), yb)
    loss.backward()
    optimizer.step()

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])
loss_fn = nn.MSELoss()
shards = [(x[:2], y[:2]), (x[2:], y[2:])]


def fresh_model():
    model = nn.Linear(1, 1)
    with torch.no_grad():
        model.weight.fill_(0.5)
        model.bias.fill_(0.0)
    return model


grads = []
for xs, ys in shards:
    model = fresh_model()
    loss = loss_fn(model(xs), ys)
    loss.backward()
    grads.append((model.weight.grad.item(), model.bias.grad.item()))
    print(f"Shard gradient: weight {grads[-1][0]:.1f}, bias {grads[-1][1]:.1f}")

avg_w = sum(g[0] for g in grads) / len(grads)
avg_b = sum(g[1] for g in grads) / len(grads)
model = fresh_model()
loss = loss_fn(model(x), y)
loss.backward()
print(f"Manual average: weight {avg_w:.1f}, bias {avg_b:.1f}")
print(f"Full-batch gradient: weight {model.weight.grad.item():.1f}, bias {model.bias.grad.item():.1f}")
print("Manual average matches full batch:", avg_w == model.weight.grad.item() and avg_b == model.bias.grad.item())
print("Single process only: no DDP run is claimed here")

Key points

  • Use DistributedDataParallel to wrap a model before training starts
  • The wrapper hides the all-reduce inside the backward pass
  • Each worker still trains on its own shard of the batch
  • The optimizer steps once per worker, with the same averaged gradients
  • The verified run in this entry is the 2-process all-reduce from the last lesson
  • No DDP training run is claimed anywhere in this lesson

In this example, the all-reduce idea is rebuilt in DDP shape inside one process. The data is the same four rows split into two shards. A fresh model starts at weight 0.5 and bias 0.0 for each shard. Each shard runs its own backward pass and prints its gradient. The two shard gradients are then averaged by hand, line by line. That manual average is the step DDP performs inside backward. The averaged weight gradient prints -27.5, and the averaged bias gradient prints -9.5. A full-batch backward pass prints the same pair. The last line states the limit plainly, since no DDP run happened in this program.

Note

DDP needs one process per worker, started the same way mp.spawn started the all-reduce program

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

In the DDP pattern, when does the gradient averaging happen?