Distributed Training with PyTorch DDP

Data parallelism lets several workers train copies of the same model on different examples. Each worker computes gradients locally; DDP combines those gradients before every worker updates its own copy. This article follows that process through data splitting, training, validation, and accumulation. The examples use one process per device and an ordinary, unsharded optimizer.

One process per GPU

For multi-GPU training, PyTorch recommends DistributedDataParallel (DDP) over DataParallel, including on one machine. DataParallel uses threads in one process, replicates the module during forward, and gathers outputs; DDP uses persistent replicas in separate processes. Their overhead depends on the workload. There is no universal two-GPU threshold beyond which scaling fails.

A rank identifies a process across the whole job; local rank identifies it on its machine; world size is the total process count. On two machines with two workers each, global ranks are 0–3, while each machine has local ranks 0 and 1. In the single-node launch below, local rank selects a GPU from the devices visible to that process. It need not be the physical GPU number if CUDA_VISIBLE_DEVICES remaps devices.

The single-device update order and memory measurements are covered in GPU Training: Mixed Precision, Gradient Accumulation, and Profiling. DDP adds communication between those training loops. It does not split the input automatically; the sampler below assigns examples to each worker.

What DDP synchronizes

With default initialization synchronization, DDP starts replicas from matching parameters and buffers. It installs gradient hooks when the wrapper is constructed. During a synchronized backward pass, ready gradients are grouped into buckets for all-reduce, which combines values across workers and makes the result available to each one. Communication can overlap with remaining backward computation. Under the default reduction, every worker receives the average gradient across ranks.

For equal local batch sizes and an example-wise mean loss, that average is the gradient of the combined batch mean, provided the forward computations match. If two workers each process two examples and obtain a gradient component of 2 and 6, the combined value is (2 + 6) / 2 = 4. Unequal counts change the weighting: with one example on the first worker and three on the second, the example-weighted mean is (1 × 2 + 3 × 6) / 4 = 5. Averaging the two local means still gives 4. BatchNorm, stochastic operations, and floating-point reduction order can also prevent exact agreement with a single large-batch run.

Each worker needs matching optimizer state, hyperparameters, and update decisions as well as matching gradients to keep parameters aligned. In ordinary synchronized steps, DDP reduces gradient buckets instead of broadcasting newly updated parameters. It also broadcasts buffers such as BatchNorm running statistics from rank 0 by default; “only gradients are communicated” would omit that work. For a fixed dense model and dtype, gradient payload scales with parameter count, but elapsed communication time also depends on bucket sizes, process count, interconnect, and overlap. It is not necessarily one collective per optimizer step.

A complete training script

Save the following block as train.py. It defines its own model and synthetic data, so no download is needed. Each process generates the same dataset; the training sampler chooses a different shard. The example uses float32 and two epochs to make the distributed control flow visible. Its scores are not a model-quality benchmark.

import os
from datetime import timedelta
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, TensorDataset, Subset
from torch.utils.data.distributed import DistributedSampler


def setup():
    local_rank = int(os.environ["LOCAL_RANK"])
    use_cuda = torch.cuda.is_available() and os.environ.get("USE_CPU") != "1"
    device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
    if use_cuda:
        torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl" if use_cuda else "gloo", timeout=timedelta(seconds=120))
    return device, dist.get_rank(), dist.get_world_size()


def build_model():
    return nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 3))


@torch.no_grad()
def evaluate_distributed(model, loader, device):
    model.eval()
    # Synchronize buffers once, then avoid DDP forward collectives on uneven shards.
    for buffer in model.module.buffers():
        dist.broadcast(buffer, src=0)
    totals = torch.zeros(3, dtype=torch.float64, device=device)
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        logits = model.module(xb)
        totals[0] += F.cross_entropy(logits, yb, reduction="sum").double()
        totals[1] += (logits.argmax(1) == yb).sum()
        totals[2] += yb.numel()
    dist.all_reduce(totals, op=dist.ReduceOp.SUM)
    loss_sum, correct, count = totals.tolist()
    if count == 0:
        raise ValueError("validation set is empty")
    return loss_sum / count, correct / count, int(count)


def main():
    device, rank, world_size = setup()
    try:
        torch.manual_seed(42)
        g = torch.Generator().manual_seed(7)
        x = torch.randn(229, 8, generator=g)
        teacher = torch.randn(8, 3, generator=g)
        y = (x @ teacher).argmax(1)
        train_set = TensorDataset(x[:128], y[:128])
        val_set = TensorDataset(x[128:], y[128:])
        if world_size > len(train_set):
            raise ValueError("too many processes for this training set")
        sampler = DistributedSampler(train_set, num_replicas=world_size,
                                     rank=rank, shuffle=True, seed=9, drop_last=True)
        loader = DataLoader(train_set, batch_size=16, sampler=sampler, num_workers=0)
        val_shard = Subset(val_set, list(range(rank, len(val_set), world_size)))
        val_loader = DataLoader(val_shard, batch_size=16, num_workers=0)
        model = DDP(build_model().to(device),
                    device_ids=[device.index] if device.type == "cuda" else None)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
        for epoch in range(2):
            sampler.set_epoch(epoch)
            model.train()
            for xb, yb in loader:
                xb, yb = xb.to(device), yb.to(device)
                optimizer.zero_grad(set_to_none=True)
                loss = F.cross_entropy(model(xb), yb)
                loss.backward()
                optimizer.step()
            val_loss, val_acc, count = evaluate_distributed(model, val_loader, device)
            if rank == 0:
                print(f"epoch {epoch} val_n {count} loss {val_loss:.4f} acc {val_acc:.4f}")
                torch.save(model.module.state_dict(), "ddp_weights.pt")
    finally:
        dist.destroy_process_group()


if __name__ == "__main__":
    if "RANK" in os.environ:
        main()
    else:
        print("Launch this file with torchrun to execute distributed training.")

Run this command in a terminal with two visible CUDA GPUs and a PyTorch build with NCCL support. torchrun starts two copies of the script and provides the rank and rendezvous environment variables. Setting the CUDA device before process-group initialization binds each worker to its local GPU. Gloo provides a CPU path for checking distributed logic when CUDA is unavailable.

torchrun --standalone --nproc_per_node=2 train.py

To deliberately test two CPU workers on a POSIX shell, use the command below. Launch only as many CUDA workers as there are visible GPUs. A multi-node job additionally needs shared rendezvous settings and a launch on each machine; --standalone here is for one machine.

USE_CPU=1 torchrun --standalone --nproc_per_node=2 train.py

The loader batch size, 16, is per process. With two workers, full steps use 32 examples. The learning rate is a starting choice for this example, not automatically multiplied by world size. Linear learning-rate scaling is a tuning heuristic; optimizer, batch size, warmup, and the objective affect whether it helps. Keeping the global batch fixed instead requires reducing the local batch as workers are added.

The script logs and writes weights only on rank 0, after every rank has participated in validation. Saving model.module.state_dict() avoids the wrapper’s module. prefix. A wrapped state dict is also usable if that prefix is removed when loading into an unwrapped model. This file contains weights only: resuming training also needs optimizer and scheduler state, progress, and relevant random states. Each rank must restore the appropriate state and device mapping.

Sample counts and validation

DistributedSampler(drop_last=True) drops enough indices to make the training dataset divisible by world size. With False, it pads with repeated indices instead. Both options give each rank the same sample count. The loader’s separate drop_last option decides whether to discard a short local batch. A short batch does not itself imply different numbers of backward calls. Calling sampler.set_epoch(epoch) changes the shuffle seed for that epoch; without it, the index order repeats. The effect on final quality depends on the training problem.

For validation, duplicated or omitted examples can alter the score. The script therefore assigns indices by range(rank, len(val_set), world_size), without padding or dropping. Its 101 validation examples split into 51 and 50 for two workers. Each worker sums loss, correct predictions, and sample count; all-reduce adds these totals, and division happens afterward. A mean of local accuracies would weight the two shards equally instead of weighting each example equally. Locally computed accuracy is valid for that shard, but does not by itself describe the full set.

Uneven validation shards may have different numbers of forwards. The evaluator synchronizes buffers once, then calls model.module to avoid DDP’s per-forward synchronization. This is suitable for the ordinary evaluation-mode model shown here, whose forward has no distributed collectives or changing state. A custom model that communicates during forward needs a matching evaluation protocol. Even a rank with zero validation examples must participate in the final reduction. Metrics such as AUROC cannot be obtained by adding local AUROCs; they need an appropriate distributed metric implementation or gathered predictions and targets.

Keeping workers in step

Collective operations must be reached by participating ranks in a compatible order. If a worker skips backward or exits early, another worker can wait for a collective that never arrives, raise a communication error, or fail after a timeout. The problem is not guaranteed to remain silent forever. Record each rank’s progress and inspect the first worker error; a later NCCL timeout can be a consequence of that earlier failure.

A shared DistributedSampler and matching loader settings avoid a common source of unequal training lengths. Rank-dependent filtering or branching can still break the protocol; loader drop_last=True cannot repair arbitrary unequal shards. Do not put the whole validation call inside if rank == 0 when it contains collectives. Likewise, unused-parameter handling must match the model: find_unused_parameters=True can help with parameters absent from a backward graph, but does not fix mismatched collective order or unequal loop lengths.

For intentionally uneven training inputs, DDP’s join() context can shadow DDP communication after a worker exhausts its data. It does not automatically shadow extra collectives in a custom loop or SyncBatchNorm. Such cases need a coordinated policy; throw_on_early_termination=True lets all ranks terminate together instead of attempting to continue those additional operations.

Gradient accumulation with DDP

By default, every backward pass synchronizes gradient buckets. During accumulation, wrap both forward and backward for intermediate micro-batches in no_sync(), then let the last backward synchronize the accumulated gradients. The last partial window must also synchronize and update. This function can replace the inner training loop above; call it once per epoch after sampler.set_epoch(epoch). Its returned list contains the global example count for each update.

from contextlib import nullcontext


def accumulated_epoch(model, loader, optimizer, device, accum_steps):
    if not isinstance(accum_steps, int) or accum_steps < 1:
        raise ValueError("accum_steps must be a positive integer")
    if len(loader) == 0:
        raise ValueError("training loader is empty")
    model.train()
    optimizer.zero_grad(set_to_none=True)
    local_count, windows = 0, []
    for i, (xb, yb) in enumerate(loader):
        last = (i + 1) % accum_steps == 0 or i + 1 == len(loader)
        context = nullcontext() if last else model.no_sync()
        with context:
            xb, yb = xb.to(device), yb.to(device)
            F.cross_entropy(model(xb), yb, reduction="sum").backward()
        local_count += yb.numel()
        if last:
            count = torch.tensor(local_count, dtype=torch.long, device=device)
            dist.all_reduce(count, op=dist.ReduceOp.SUM)
            global_count = count.item()
            factor = dist.get_world_size() / global_count
            for p in model.parameters():
                if p.grad is not None:
                    p.grad.mul_(factor)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            windows.append(global_count)
            local_count = 0
    return windows

The function sums unweighted cross-entropy losses with no ignored labels. After the final backward, DDP has averaged the accumulated sums across ranks. If two ranks contributed gradients whose sums are 12 and 18 over 9 and 9 examples, the buffer holds 15. Multiplying by world size divided by the total count gives 15 × 2 / 18, the required gradient mean 30 / 18. This also accounts for unequal example counts within a synchronized window. The scalar count all-reduce is additional communication, even though intermediate gradient reductions are suppressed.

Every rank must have the same number of micro-batches and use the same window boundaries in this function; example counts within those batches may differ. With ten examples per rank, local batches of 3, and accum_steps=3, global update counts are 18 and 2. Class weights, ignored targets, and token losses need a denominator matched to the objective. This float32 function does not implement AMP or clipping. When adding them, keep one scale within the window, unscale once after the final backward, normalize before clipping, and keep update/skip decisions and scheduler state consistent across ranks.

Normalization across devices

Ordinary BatchNorm uses each rank’s local batch statistics during training. For BatchNorm2d, each channel’s statistics include the local N × H × W values, so “batch 8” does not mean only eight scalar observations. DDP’s running-buffer broadcast does not combine those current-batch statistics across ranks. SyncBatchNorm aggregates training statistics over its process group; it adds communication and does not guarantee better validation quality.

For a CUDA model that needs synchronized BatchNorm, apply torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) before wrapping it in DDP and before constructing the optimizer. The supported DDP setup uses one GPU per process. GroupNorm and LayerNorm use statistics within each example over their specified axes, providing alternatives whose behavior should be evaluated for the architecture. The synthetic model above has no normalization layer and needs no conversion.

Configuration and checkpoint components are introduced in Reproducible Training Pipelines: Seeds, Configs, and Checkpoints. A distributed resume additionally needs a coordinated restore on every rank; that article’s single-process example does not establish exact multi-rank resumption.

When data parallelism is not enough

In the one-device-per-process setup with an ordinary optimizer used here, each device must hold a full replica plus its gradients, optimizer state, activations, and workspaces. DDP alone does not combine two GPUs into a larger memory pool. Sharding or model parallelism can address different parts of that memory requirement.

Strategy	What is divided	Constraint to examine
DDP	Input work; model and ordinary optimizer replicated	Per-device memory and gradient communication
ZeRO-1	Optimizer state	Optimizer-state memory
ZeRO-2	Optimizer state and gradients	State and gradient memory
FSDP / ZeRO-3	Parameters, gradients, and optimizer state	Temporary parameter gathers and activations still need space
Tensor parallelism	Operations within layers	Layer memory, compute, and frequent communication
Pipeline parallelism	Groups of layers across stages	Stage balance, pipeline idle time, and activation transfers

These approaches can be combined. Pipeline parallelism can place communication at stage boundaries, but a slow interconnect still has a cost; depth alone does not establish that it is the right choice. Choose a strategy from measured memory and communication requirements, then account for its scheduling and checkpoint complexity.

For throughput T on W GPUs and baseline throughput T₁ on one GPU, scaling efficiency is T / (W × T₁). A 3.4× speedup on four GPUs gives 85%; a 1.6× speedup gives 40%. Those numbers do not identify the bottleneck or predict what adding more GPUs will do. Compare synchronized, warmed-up measurements and state whether local or global batch size is held fixed. Profile communication, loading, small kernels, and rank imbalance, and also compare time to an acceptable validation result when the training recipe changes.

How large models are pretrained is covered in LLM Pretraining: Objectives, Data, and Scaling Laws.

References

• PyTorch. DistributedDataParallel.
• PyTorch. DistributedSampler and DataLoader.
• PyTorch. SyncBatchNorm.
• Huang et al. (2019). GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism. NeurIPS.
• Shoeybi et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv preprint.
• Li et al. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training. VLDB.
• Rajbhandari, Rasley, Ruwase, and He (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC.
• Zhao et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. VLDB.