Distributed Training Guide

A practical guide to multi-GPU training with PyTorch DistributedDataParallel

(DDP), written to pair with the code in src/gpu_training/. It assumes you can

already train on one GPU and want to scale out correctly — not just "make it

run", but make the math, the data sharding, and the checkpoints right.

Why DDP (and not DataParallel)

PyTorch ships two multi-GPU APIs. Reach for the right one:

• DataParallel (DP) — single process, one thread per GPU. The Python GIL

and the scatter/gather of every batch to one "master" GPU make it slow and

memory-lopsided. It is effectively deprecated. Do not use it for new work.

• DistributedDataParallel (DDP) — one process per GPU. Each process holds a

full model replica, computes gradients on its own data shard, and the

gradients are averaged across processes with an all-reduce that overlaps with

the backward pass. This is the standard, and it scales from 2 GPUs on one box

to thousands across a cluster.

For models too large to fit one GPU even at batch size 1, you graduate to

FSDP (Fully Sharded Data Parallel), which sharies parameters, gradients, and

optimizer state across ranks. FSDP reuses the same launch and sampler machinery

described here; only the model wrapper changes.

The process group mental model

DDP is built on collective communication. Every process ("rank") joins a

*process group* and participates in collectives like all-reduce. Three numbers

define a rank's identity, and torchrun injects them as environment variables:

gpu_training.ddp.DistributedContext.from_env() parses exactly these. It also

falls back to a single-process context when they are unset, so the same script

runs unchanged in a notebook and under torchrun:

from gpu_training.ddp import init_distributed, resolve_device

ctx = init_distributed(backend="nccl")   # nccl for GPUs, gloo for CPU
device = resolve_device(ctx)             # cuda:<local_rank>

init_distributed calls torch.cuda.set_device(local_rank) for you. Skipping

that step is the classic first bug: every rank lands on cuda:0, you run out of

memory, and ranks 1..N sit idle.

Launching with torchrun

torchrun spawns the per-GPU processes and wires up rendezvous. The toolkit's

launcher module generates the exact command so you never fight the flags:

from gpu_training.launcher import single_node_command, multi_node_commands

# Single node, 4 GPUs:
print(single_node_command("examples/train_ddp.py", 4, ["--epochs", "3"]))
# -> torchrun --standalone --nproc_per_node=4 examples/train_ddp.py --epochs 3

# Two nodes, 8 GPUs each (run the matching line on each node):
for cmd in multi_node_commands("examples/train_ddp.py", nnodes=2,
                               nproc_per_node=8, rdzv_endpoint="head-node:29500"):
    print(cmd)

Key points for multi-node:

• --rdzv_endpoint is the host:port of the rank-0 node. Every node must

point at the same endpoint and share the same --rdzv_id.

• --rdzv_backend=c10d is the modern built-in rendezvous; you do not need a

separate etcd server.

• --max_restarts turns on *elastic* behavior: if a worker dies, torchrun

restarts the whole group from the last checkpoint instead of hanging forever.

Sharding data correctly

Each rank must see a different slice of the dataset, or you are just training

the same batch N times. Use a DistributedSampler:

from gpu_training.ddp import make_distributed_sampler

sampler = make_distributed_sampler(dataset, ctx, shuffle=True, seed=42)
loader = DataLoader(dataset, batch_size=cfg.per_device_batch_size,
                    sampler=sampler, shuffle=(sampler is None))

Two non-obvious rules:

1. Call sampler.set_epoch(epoch) every epoch. Without it the sampler uses

the same shuffle seed each epoch and every rank sees the same order forever.

train_one_epoch() does this for you.

2. Do not also pass shuffle=True to the DataLoader when a sampler is set —

PyTorch forbids it. The helper returns None in single-process mode so the

shuffle=(sampler is None) idiom works in both cases.

Effective batch size: the number that actually matters

This is where most "my distributed run won't converge" issues live. The batch

size your optimizer sees is not the per-GPU batch size:

effective_batch = per_device_batch_size × gradient_accumulation_steps × world_size

If you tuned a learning rate at effective batch 256 on one GPU, then run the same

per_device_batch_size on 8 GPUs, your effective batch silently becomes 2048 —

8× larger — and your learning rate is now far too small. Use the config helpers

to keep this explicit and to hit a target batch regardless of GPU count:

from gpu_training.config import TrainingConfig, gradient_accumulation_for_target

accum = gradient_accumulation_for_target(target_effective_batch=512,
                                         per_device_batch_size=8, world_size=4)
cfg = TrainingConfig(per_device_batch_size=8, gradient_accumulation_steps=accum,
                     world_size=4)
assert cfg.effective_batch_size >= 512
print(cfg.describe(num_samples=1_000_000))

A common scaling heuristic is linear LR scaling: when you multiply the

effective batch by *k*, multiply the base learning rate by *k* and lengthen

warmup. The toolkit's linear_warmup_cosine_lr gives you the warmup + cosine

curve to do this cleanly.

Gradient accumulation under DDP

Accumulation lets you reach a large effective batch on small GPUs by summing

gradients over several micro-batches before stepping the optimizer. Under DDP

there is a performance trap: DDP runs an all-reduce on every backward() by

default, so naive accumulation syncs gradients you are about to keep

accumulating. The fix is to suppress the sync on non-stepping micro-batches with

model.no_sync() and only sync on the final micro-batch:

from gpu_training.amp import should_step

for micro_step, batch in enumerate(loader):
    stepping = should_step(micro_step, cfg.gradient_accumulation_steps)
    sync_context = nullcontext() if stepping else model.no_sync()
    with sync_context:
        with trainer.autocast():
            loss = loss_fn(model(x), y)
        trainer.backward(loss)
    trainer.step(model.parameters(), micro_step)   # no-ops until `stepping`

should_step(micro_step, accum) is the single source of truth for "is this the

boundary"; it returns True after micro-steps accum-1, 2·accum-1, ….

Logging metrics across ranks

Each rank computes loss on its own shard, so rank 0's loss is not the global

loss. Reduce before you log:

from gpu_training.ddp import reduce_scalar_dict

reduced = reduce_scalar_dict({"loss": local_loss}, ctx, average=True)
if ctx.is_main_process:
    print(f"global loss = {reduced['loss']:.4f}")

Guard all printing, progress bars, and metric uploads with ctx.is_main_process

to avoid eight identical log streams.

Checkpointing without corrupting files

All ranks hold identical weights, so only rank 0 writes. Having every rank

write the same path to a shared filesystem at once corrupts the file. The

CheckpointManager enforces rank-0 writes, writes atomically (.tmp then

os.replace), keeps the last *N*, and turns save() into a barrier on the other

ranks so nobody races ahead of the write:

from gpu_training.ddp import unwrap_model
from gpu_training.checkpoint import CheckpointManager

ckpt = CheckpointManager("checkpoints", ctx=ctx, keep_last_n=3)
ckpt.save(state={"model_state": unwrap_model(model).state_dict(),
                 "optimizer_state": optimizer.state_dict(),
                 "epoch": epoch},
          step=global_step)

Note unwrap_model(model) — it strips the DDP wrapper so the saved keys are not

prefixed with module. and the checkpoint loads into a plain (non-DDP) model for

inference.

Tearing down

Always destroy the process group at the end (and in your exception handler) so

NCCL releases resources cleanly:

from gpu_training.ddp import cleanup_distributed
cleanup_distributed()   # idempotent; safe to call once on every rank

Troubleshooting checklist

NCCL environment variables worth knowing

• NCCL_DEBUG=INFO — print the topology NCCL discovered; first stop when collectives hang.
• NCCL_SOCKET_IFNAME=eth0 — pin the network interface on multi-NIC hosts.
• NCCL_P2P_DISABLE=1 — disable peer-to-peer if you hit driver/topology issues (slower, but a useful diagnostic).
• CUDA_VISIBLE_DEVICES — restrict which GPUs a node uses; combine with --nproc_per_node to match.

Start with NCCL_DEBUG=INFO whenever a multi-GPU job stalls — most "hangs" are a

network-interface or topology misconfiguration that NCCL will tell you about.