LLM Lab Notes 02: DDP Minimal Training — Multi-GPU in One Line

Code files: llm_lab/01_ddp/mini_train.py (plus the profiled variant mini_train_profile.py). After Note 01 got the process group up, this experiment runs a real model through a multi-GPU training loop.

1. Experiment Goal

1. Wire the process group initialization from Note 01 into a real training loop;
2. Wrap the model with DistributedDataParallel and run 5 steps of multi-GPU training;
3. Break down what DDP does for us behind the scenes of the "training loop."

2. Code Walkthrough

Full code (mini_train.py, 37 lines):

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist

dist.init_process_group("nccl")

rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])

torch.cuda.set_device(local_rank)

model = nn.Linear(1024, 1024).cuda()
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])   # ① one-line wrapper

optimizer = optim.Adam(model.parameters(), lr=1e-3)                           # ② optimizer

x = torch.randn(32, 1024).cuda()
y = torch.randn(32, 1024).cuda()

for step in range(5):                                                         # ③ training loop
    optimizer.zero_grad()

    out = model(x)

    loss = ((out - y) ** 2).mean()

    loss.backward()                                                           # ④ backward
    optimizer.step()                                                          # ⑤ update

    if rank == 0:
        print(f"Step {step}, Loss {loss.item()}")

dist.destroy_process_group()

① DistributedDataParallel(model, device_ids=[local_rank])

This is the core line of the whole experiment. At initialization, DDP does two things that are critical in distributed training:

• Parameter broadcast: rank 0 broadcasts its initial weights to all processes so everyone starts at the same point;
• ReducHook registration: a sync hook is attached to every parameter's gradient, automatically performing gradient AllReduce during backpropagation.

② Create the optimizer after wrapping with DDP

A very common pitfall: you must wrap with DDP first, then create the optimizer. Because model.parameters() remains usable after wrapping — but if you create the optimizer first and then wrap, the optimizer holds references to the pre-wrapping parameters, and DDP's parameter broadcast will never reach the optimizer state.

③④⑤ The "invisible communication" inside the training loop

The loop looks identical to single-GPU training — and that is exactly DDP's design goal: transparent to the caller, with communication hidden inside backward. When loss.backward() runs, each parameter's gradient triggers an AllReduce as soon as it is computed, so at optimizer.step() every GPU holds the globally consistent gradient, and the updated parameters naturally stay in sync.

3. How It Works: Two Timing Strategies for Gradient Sync

DDP uses asynchronous gradient AllReduce by default: backpropagation runs in reverse parameter order, and each parameter's gradient starts communicating as soon as it is ready, overlapping with the forward computation of earlier layers. This is gradient bucketing: small gradients are merged into buckets, and one communication handles a large bucket — reducing the number of messages and improving bandwidth utilization.

Compared to the naive approach (compute all gradients, then communicate once), bucket overlap makes communication overhead nearly "invisible" inside compute — a key reason DDP's throughput pulls ahead of the older DataParallel.

4. Results

Running 5 steps on 4 GPUs, expected output shape (printed by rank 0):

Step 0, Loss 1.0396
Step 1, Loss 0.9983
Step 2, Loss 0.9597
Step 3, Loss 0.9231
Step 4, Loss 0.8885

(Loss decreases steadily; exact values depend on random initialization, but the trend and magnitude should match.)

How to tell it succeeded:

• All GPUs finish training cleanly, with only rank 0 printing;
• Loss decreases step by step — the gradient sync path works (if gradients were inconsistent across GPUs, the loss would jump around or diverge);
• To verify sync correctness explicitly, compare model.module.weight across GPUs after a step — all 4 should be element-wise equal.

The profiled variant (mini_train_profile.py, Linear(4096, 4096)) times each step in three segments:

Step 0
 forward : 3.21 ms
 backward: 5.87 ms
 update  : 0.42 ms

5. Engineering Takeaways

1. "One-line wrapper" ≠ "zero cost": DDP's gradient AllReduce moves 2× the model parameter size per step. Negligible for a 1024×1024 toy model, but for a 70B model that is ~280 GB of cluster traffic per step — exactly what Chapter 11's NCCL tuning article addresses.

2. device_ids and CUDA_VISIBLE_DEVICES interplay: in Slurm environments it's common to set CUDA_VISIBLE_DEVICES=$SLURM_LOCALID and then set_device(0); both styles are equivalent, but never mix the two, or device indexing skews and you lose half an hour debugging.

3. Divergent training loss → suspect gradient sync first: when multi-GPU training diverges, verify per-GPU gradient consistency (e.g., with all_reduce comparison) before touching the learning rate — this is usually a sync problem, not an optimization problem.

4. From 5 steps to 10K GPUs: same framework, different engineering: the loop structure in mini_train.py maps line-by-line to production frameworks (Megatron-LM, DeepSpeed); the differences are only in engineering depth — parallelism strategy, communication overlap, checkpointing, and failure recovery.

Next: DDP's gradient sync relies on AllReduce — but how does AllReduce "participate in backpropagation"? Note 03 hand-writes an AllReduce autograd function to verify it: LLM Lab Notes 03: Writing Collective Communication Autograd.