PyTorch Distributed Training Tutorial

Distributed training lets you use multiple GPUs — on one machine or across many — to train models faster or to fit models that do not fit on a single GPU. PyTorch’s DistributedDataParallel (DDP) is the standard approach: each GPU runs a full copy of the model on a shard of the data, and gradients are averaged across GPUs after each backward pass. This guide covers the complete DDP setup, the torchrun launcher, and the patterns that make multi-GPU training reliable.

Why DDP Instead of DataParallel

PyTorch has two multi-GPU APIs: torch.nn.DataParallel (DP) and torch.nn.parallel.DistributedDataParallel (DDP). Always use DDP. DataParallel runs on a single process, uses Python’s GIL, and suffers from load imbalance across GPUs because one GPU acts as the parameter server. DDP runs one process per GPU, bypasses the GIL entirely, uses NCCL for efficient collective communication, and scales linearly with GPU count. The performance difference is significant — DDP is typically 2–3x faster than DataParallel at 4 GPUs, and DP does not scale beyond a single node.

DDP: The Minimal Working Example

# train.py — run with: torchrun --nproc_per_node=4 train.py
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def main():
    # Initialise the process group — reads LOCAL_RANK from environment
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    device = torch.device(f'cuda:{local_rank}')

    # Model — wrap in DDP after moving to device
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])

    # Data — each rank sees a non-overlapping slice
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    loader = torch.utils.data.DataLoader(
        dataset, batch_size=256, sampler=sampler,
        num_workers=4, pin_memory=True
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()

    for epoch in range(50):
        sampler.set_epoch(epoch)   # required for correct shuffling
        model.train()
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = criterion(model(x), y)
            loss.backward()        # gradient all-reduce happens here
            optimizer.step()

        # Only rank 0 saves checkpoints and prints metrics
        if local_rank == 0:
            print(f"Epoch {epoch+1} done")
            torch.save(model.module.state_dict(), f'ckpt_{epoch}.pt')

    dist.destroy_process_group()

if __name__ == '__main__':
    import os
    main()

Launching with torchrun

# Single node, 4 GPUs
torchrun --nproc_per_node=4 train.py

# Multi-node: 2 nodes, 4 GPUs each (8 GPUs total)
# On node 0 (master):
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0          --master_addr=192.168.1.100 --master_port=29500 train.py
# On node 1:
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1          --master_addr=192.168.1.100 --master_port=29500 train.py

torchrun (introduced in PyTorch 1.10, replacing torch.distributed.launch) sets the environment variables LOCAL_RANK, RANK, and WORLD_SIZE automatically. LOCAL_RANK is the GPU index within the node (0–3 for 4 GPUs); RANK is the global process rank; WORLD_SIZE is the total number of processes across all nodes.

Gradient Synchronisation and Communication

DDP’s core mechanism: after each backward pass, NCCL performs an all-reduce operation that sums gradients across all processes and divides by world_size, so every GPU ends up with identical averaged gradients before the optimizer step. This happens automatically — you call loss.backward() as usual and DDP handles the communication. The all-reduce is overlapped with the backward pass (gradients are communicated as they are computed, layer by layer from back to front), minimising the added latency.

# Temporarily disable gradient sync for gradient accumulation
accumulation_steps = 4
for i, (x, y) in enumerate(loader):
    # Disable sync for all but the last accumulation step
    with model.no_sync() if (i+1) % accumulation_steps != 0 else contextlib.nullcontext():
        out = model(x.to(device))
        loss = criterion(out, y.to(device)) / accumulation_steps
        loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Logging and Checkpointing in DDP

Only one process should write to disk or log — typically rank 0. Always guard file writes and print statements:

rank = dist.get_rank()
world_size = dist.get_world_size()

# Only rank 0 logs
if rank == 0:
    print(f"Loss: {loss.item():.4f}")
    wandb.log({'loss': loss.item()})

# Save the underlying model (not the DDP wrapper)
if rank == 0:
    torch.save(model.module.state_dict(), 'model.pt')

# Loading — load on CPU then move to device
state_dict = torch.load('model.pt', map_location='cpu')
model.module.load_state_dict(state_dict)

# Barrier — wait for rank 0 to finish saving before others proceed
dist.barrier()

Scaling the Learning Rate

With DDP, the effective batch size is batch_size_per_gpu × num_gpus. A larger effective batch size requires a higher learning rate to maintain equivalent training dynamics. The linear scaling rule: multiply the base LR by the number of GPUs. Combine with a linear warmup to avoid instability at the start of training:

base_lr = 1e-3
num_gpus = dist.get_world_size()
scaled_lr = base_lr * num_gpus   # scale linearly

optimizer = torch.optim.AdamW(model.parameters(), lr=scaled_lr)

# Warmup for first 5 epochs, then cosine decay
from torch.optim.lr_scheduler import OneCycleLR
scheduler = OneCycleLR(
    optimizer,
    max_lr=scaled_lr,
    epochs=50,
    steps_per_epoch=len(loader),
    pct_start=0.1    # 10% warmup
)

FSDP: Sharding the Model Itself

DDP replicates the full model on every GPU — this fails when the model is too large to fit on a single GPU. Fully Sharded Data Parallel (FSDP) shards model parameters, gradients, and optimiser state across GPUs, allowing models larger than a single GPU’s memory to be trained:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

# Wrap each transformer block independently (not the whole model at once)
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={LlamaDecoderLayer}
)
model = FSDP(model, auto_wrap_policy=auto_wrap_policy, device_id=local_rank)

FSDP is the standard approach for training large language models in PyTorch. It requires more careful configuration than DDP but enables training 7B+ parameter models on a cluster of A100 GPUs that could not hold the full model on any single card.

DDP is the right tool for 95% of distributed training needs: multiple GPUs on one node or a small cluster, where the model fits on a single GPU and you want to scale throughput by data parallelism. Use the minimal example above as your template — it handles all the correct setup (init_process_group, DistributedSampler, set_epoch, rank-guarded logging and checkpointing) and has been the starting point for countless production training runs. Switch to FSDP when the model itself exceeds single-GPU memory. Use torchrun for launching — it handles the environment variable setup that manual scripts often get wrong, and supports fault-tolerant restarts for long multi-node runs.

Environment Setup for DDP

DDP requires NCCL for GPU communication. Verify your environment before writing any training code:

pip install torch torchvision  # NCCL is bundled with the CUDA build

python -c "import torch; print(torch.cuda.nccl.version())"  # e.g. (2, 18, 1)
python -c "import torch; print(torch.cuda.device_count())"  # number of GPUs

For multi-node training, all nodes must have network connectivity and the same PyTorch version. The master node’s IP and port must be reachable from all worker nodes — check with ping and verify the port is not blocked by a firewall. NCCL will hang silently if nodes cannot communicate, which is the most common multi-node setup failure.

Mixed Precision with DDP

Automatic Mixed Precision (AMP) works transparently with DDP — the GradScaler runs on each GPU independently, and DDP’s all-reduce happens after the backward pass on float32 gradients (not float16, which avoids precision issues in the communication step):

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()

    with autocast():
        out = model(x)
        loss = criterion(out, y)

    scaler.scale(loss).backward()   # DDP all-reduce on unscaled grads
    scaler.step(optimizer)
    scaler.update()

The combination of DDP and AMP is the standard training setup for large models — DDP provides horizontal scaling across GPUs, AMP reduces memory per GPU and enables larger batch sizes. Together they allow training runs that would be impossible with either technique alone.

Debugging DDP Hangs

The most frustrating DDP issue: the training script hangs silently with no error message. This almost always means a collective operation (all-reduce, barrier, broadcast) was called by some ranks but not others, or NCCL cannot establish connectivity. Debugging steps: set NCCL_DEBUG=INFO to see NCCL’s communication attempts; set TORCH_DISTRIBUTED_DEBUG=DETAIL to see which collectives are being called; add explicit print(f"rank {rank}: reached step X") calls to identify which rank is hanging. Most hangs are caused by Python exceptions on some ranks (the exception terminates the process but other ranks wait forever) or by different code paths on different ranks (one rank enters a conditional the others do not).

NCCL_DEBUG=INFO torchrun --nproc_per_node=4 train.py 2>&1 | grep -E "NCCL|error|hang"
TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=4 train.py

Common DDP Pitfalls

The issues that most commonly break DDP setups. Forgetting sampler.set_epoch(epoch) — without this, all epochs use the same shuffle permutation, reducing effective data diversity and slowing convergence. Saving the DDP wrapper instead of model.module — torch.save(model.state_dict()) saves DDP-prefixed keys; save model.module.state_dict() for a clean checkpoint compatible with non-DDP loading. Calling dist.barrier() incorrectly — a barrier that only some ranks reach causes a deadlock; every barrier must be called by all ranks. Not scaling the learning rate — the effective batch is per_gpu_batch × world_size; the same LR as single-GPU training will underfit. Non-deterministic results — set seeds per rank (torch.manual_seed(42 + rank)) so each GPU gets different data augmentation but reproducible results across runs.

Monitoring Multi-GPU Training

import torch.distributed as dist

def log_metrics(loss, rank, world_size):
    """Aggregate loss across all ranks and log from rank 0."""
    loss_tensor = torch.tensor(loss, device='cuda')
    dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
    if rank == 0:
        print(f"Global avg loss: {loss_tensor.item():.4f}")

# In training loop:
log_metrics(loss.item(), local_rank, dist.get_world_size())

Use nvidia-smi to monitor all GPUs simultaneously during training: watch -n 1 nvidia-smi shows utilisation, memory, and temperature for all GPUs. All GPUs should show similar utilisation — if one is consistently at 100% and others at 60%, the sampler or data loading is not distributing work evenly.

DDP vs Model Parallelism vs Pipeline Parallelism

DDP (data parallelism) replicates the full model on each GPU and splits the data — the model must fit on a single GPU. Model parallelism splits the model itself across GPUs — different layers on different GPUs, data passing sequentially. Pipeline parallelism is a variant where different GPUs handle different pipeline stages with micro-batches flowing through simultaneously. Each approach suits different scale constraints. For most practitioners working with models under ~10B parameters that fit on a single A100 (80GB), DDP is the right tool. FSDP is the practical bridge to larger models without needing to implement model or pipeline parallelism manually. True model and pipeline parallelism (using torch.distributed.pipeline or Megatron-LM) are needed only for frontier-scale training (70B+ parameters) where no single GPU can hold even one layer of the model.

DDP is the right tool for the vast majority of distributed training: multiple GPUs on one node or a small cluster, where the model fits on a single GPU. Use the minimal example above as your starting template — it handles process group initialisation, DistributedSampler, set_epoch, and rank-guarded checkpointing correctly. Switch to FSDP when the model exceeds single-GPU memory. Use torchrun to launch — it handles environment variable setup reliably and supports fault-tolerant restarts for long multi-node runs.

Starting with a clean single-GPU implementation and then wrapping it in DDP is always easier than trying to write a distributed training loop from scratch — the DDP wrapper is thin enough that the debugging workflow stays essentially the same, and you can validate correctness on one GPU before committing to a multi-GPU run.

The linear scaling rule for learning rate and the sampler.set_epoch() call are the two details most often missed by practitioners new to DDP — both have a meaningful impact on training quality and are easy to forget when adapting a single-GPU script.