AI Engineering from Scratch（英文原版）

When to use

The model fits on one device but you need more throughput (DDP). The model does not fit on one device (FSDP). Either case: a multi-rank training setup with the same code path.

Bring up the process group

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

gloo is the CPU backend; nccl is the GPU backend. Both implement the same collective surface.

Wrap the model

1. On rank 0, build the model from your seed.
2. Wrap it with the DDP shell.
3. The shell's __init__ calls dist.broadcast(p.data, src=0) for every parameter and buffer.
4. After every loss.backward(), the trainer calls sync_grads().
5. sync_grads() calls dist.all_reduce(p.grad, op=SUM) and p.grad.div_(world_size).
6. Optimizer step on every rank with the same averaged gradient.

Shard parameters (FSDP sketch)

1. Flatten each parameter, pad to a multiple of world_size.
2. Keep your shard locally; release the rest.
3. Before forward, dist.all_gather(...) to rebuild the full tensor on every rank.
4. After forward, drop the full tensor.

Failure modes

• Skipping the broadcast: ranks start from different inits, diverge silently.
• Forgetting to divide after sum: gradients scaled by world_size, optimizer steps too big.
• Using cross-device rename for checkpoints: not atomic; same lesson 47 trap.
• Mixing CPU and CUDA tensors on the same collective: backend mismatch, run hangs.