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.

Mixed-Precision Training

Mixed precision is the single highest-leverage change you can make to a training

job: on modern GPUs it cuts memory roughly in half and speeds up matmul-heavy

models 1.5–3× with little or no accuracy loss. This guide explains what is

actually happening and how gpu_training.amp keeps you on the correct path.

Two ideas people conflate

"Mixed precision" bundles two independent mechanisms. Understanding that they are

separate is the whole game:

1. autocast runs selected operations (mostly matmuls and convolutions) in a

low-precision dtype — float16 or bfloat16 — while keeping a master copy of

the weights in float32. This is where the speed and memory savings come

from.

2. loss scaling (GradScaler) multiplies the loss by a large factor before

backward() so that small gradients do not *underflow to zero* in float16,

then divides them back out before the optimizer step. This is a numerical

workaround for one specific dtype.

The critical fact: loss scaling is only needed for float16. bfloat16 has

the same 8-bit exponent as float32 (just fewer mantissa bits), so its dynamic

range is identical and gradients do not underflow. A bfloat16 run needs **no

scaler at all**.

fp16 vs bf16: which to use

Rule of thumb: if your GPU supports bf16 (Ampere or newer), prefer it — it is

strictly simpler because there is no scaler to manage, no inf/NaN from

overflow, and no warmup of the scale factor. Use fp16 on older hardware

(e.g. V100, T4) where bf16 is unavailable.

The toolkit encodes this decision so you cannot get it wrong:

from gpu_training.amp import needs_grad_scaler
needs_grad_scaler("fp16")  # True
needs_grad_scaler("bf16")  # False

The correct training step

MixedPrecisionTrainer wraps the three precision modes behind one interface. You

never construct a scaler or pick a dtype by hand:

from gpu_training.amp import MixedPrecisionTrainer

trainer = MixedPrecisionTrainer(optimizer, precision="bf16",
                                grad_accum_steps=4, max_grad_norm=1.0)

for micro_step, (x, y) in enumerate(loader):
    with trainer.autocast():          # autocast in bf16 (or nullcontext for fp32)
        loss = loss_fn(model(x), y)
    trainer.backward(loss)            # scales loss, divides by accum, backward
    grad_norm = trainer.step(model.parameters(), micro_step)  # steps on boundary

What the trainer does internally, in order:

1. autocast() returns torch.autocast(device_type, dtype=...) for fp16/bf16

or a nullcontext() for fp32, so the same loop body works for every precision.

2. backward() divides the loss by grad_accum_steps (so accumulated

gradients average rather than sum) and calls scaler.scale(loss).backward().

When the scaler is disabled (bf16/fp32), scale() is the identity — same code

path, no special-casing.

3. step() only acts on accumulation boundaries (should_step). On a

boundary it unscales, then clips, then steps, in that exact order.

Why unscale-before-clip matters

Gradient clipping bounds the gradient *norm*. If you clip while gradients are

still multiplied by the loss-scale factor (often 2¹⁶), you are clipping against a

threshold that is 65,536× too large — i.e. you are not clipping at all. The

trainer calls scaler.unscale_(optimizer) first so that

clip_grad_norm_ sees true gradient magnitudes:

self._scaler.unscale_(self.optimizer)                       # undo loss scaling
grad_norm = torch.nn.utils.clip_grad_norm_(params, max_norm) # now correct
self._scaler.step(self.optimizer)                            # skips step on inf/NaN
self._scaler.update()                                        # adapt the scale factor

scaler.step() is also where fp16 safety lives: if any gradient overflowed to

inf/NaN, the scaler skips the optimizer step for that iteration and lowers

the scale factor, so a single overflow does not poison your weights.

How the scale factor adapts (fp16)

GradScaler starts at a high scale (default 2¹⁶) and:

• Halves it whenever an overflow is detected (skipping that step).
• Doubles it after a configurable run of overflow-free steps (default 2000).

This auto-tuning means you rarely touch it. Two practical consequences:

• The first few steps of an fp16 run may be skipped while the scaler finds a

workable factor — this is normal, not a bug. Do not be alarmed if the first

reported grad_norm is None/inf once or twice.

• The scaler's state is part of your checkpoint. MixedPrecisionTrainer.state_dict()

/ load_state_dict() round-trip it so a resumed fp16 run continues with the

right scale factor instead of re-warming from 2¹⁶.

Memory savings and the batch finder

Halving activation precision roughly halves activation memory, which is usually

the dominant term during training. The freed memory lets you raise the batch

size — and the toolkit finds the new ceiling for you instead of making you guess:

from gpu_training.profiling import find_max_batch_size

# `fits(batch)` runs a forward+backward and reports False on CUDA OOM.
best = find_max_batch_size(fits, min_batch=1, max_batch=1024, safety_margin=0.85)

The safety_margin (default 0.85–0.9) is deliberate: a batch that *barely*

passes a single probe often OOMs once real training churns the allocator

(fragmentation, cuDNN workspace growth, autograd's variable peak). Leave the

headroom. See examples/train_amp.py for the finder wired into a real loop.

If even min_batch does not fit, the finder raises immediately with an

actionable message — at that point you need activation checkpointing,

gradient accumulation for a larger effective batch, or a smaller model /

shorter sequence length.

Accuracy considerations

Mixed precision is usually accuracy-neutral, but keep a few things in float32:

• The loss computation and softmax/cross-entropy are autocast to fp32 by

PyTorch automatically — do not force them to half precision.

• Normalization statistics (BatchNorm/LayerNorm) accumulate in fp32 under

autocast; this is handled for you.

• Reductions over long sequences (e.g. attention logits) benefit from bf16's

range; this is one more reason to prefer bf16 when available.

If you see NaN losses specifically with fp16, the usual culprits are: missing

GradScaler (gradients underflowed), an exploding learning rate (add/lengthen

warmup via linear_warmup_cosine_lr), or an operation that should have stayed in

fp32. Switching to bf16 makes most of these disappear.

Quick decision tree

Does your GPU support bf16 (Ampere/Hopper, recent AMD/TPU)?
├─ Yes → use precision="bf16"  (no scaler, simplest, robust)
└─ No  → use precision="fp16"  (GradScaler required; unscale before clip)

Out of memory?
├─ Lower per-device batch and raise gradient_accumulation_steps
├─ Run find_max_batch_size() to size the batch to memory
└─ Still OOM at batch 1 → activation checkpointing / smaller model

Pair this with the Distributed Training Guide

to combine mixed precision with multi-GPU DDP — the two compose cleanly, and

MixedPrecisionTrainer works identically whether the model is DDP-wrapped or not.