PyTorch DDP: Multi-GPU Training When One GPU Fits

Problem and takeaway

Your model fits on one GPU, but a single training run takes too long. You need higher throughput without redesigning the model for model parallelism. The practical takeaway is to use PyTorch DistributedDataParallel (DDP) with one process per GPU, a disjoint data sampler, and checkpoint writes restricted to rank 0. This is the smallest suitable design for synchronous multi-GPU training when the model and optimizer state fit on a single device.

Requirements

• The model, optimizer state, and a single batch must fit in one GPU's memory.
• The dataset can be divided into independent shards; samples do not depend on each other.
• The environment supports one process per GPU and a collective communication backend, typically NCCL.
• All ranks can reach the same master address and port for process group initialization.

Smallest suitable design

Each GPU runs a separate copy of the Python process. The process group is initialized with a rank and world size. The model is moved to the local GPU before it is wrapped in DDP. DDP then handles gradient all-reduce automatically: every rank computes gradients on its own shard, and the collectives average them across ranks so all replicas stay in sync.

A DistributedSampler assigns each rank a non-overlapping subset of the dataset. Set drop_last=True when using synchronized batch normalization or when batch shapes must match across ranks.

Only rank 0 writes checkpoints and logs. Offset the random seed per rank so ranks do not sample the same data in the same order.

Minimal configuration example

Run on a GPU node where your user account has CUDA device access. No root is normally required. Replace train_script.py with your entry point.

torchrun --nproc_per_node=4 train_script.py

torchrun spawns one process per visible GPU and sets the rank and world size environment variables. Inside the script:

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup(rank, world_size):
    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
    os.environ.setdefault('MASTER_PORT', '29500')
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

# model = MyModel().to(rank)
# model = DDP(model, device_ids=[rank])
# sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, drop_last=True)

Do not move the model to the GPU after wrapping it with DDP; the wrapper expects the module to already reside on the correct local device.

Trust and data boundaries

Each process owns its input shard and its gradient buckets. There is no shared memory between ranks except the collective communication buffers managed by the backend. Checkpoint files are a shared resource, so restrict writes to rank 0 to prevent corruption. Random seeds must differ by rank; otherwise every rank can draw the same samples and waste compute.

If a checkpoint write fails or produces a corrupt file, restore the previous rank-0 checkpoint. Other ranks do not write checkpoints, so there is no distributed state to roll back.

Operational checks

• Confirm world size equals the number of visible GPUs.
• Confirm NCCL initialization completes without warnings or timeouts.
• Compare loss trajectories across ranks; they should track closely, not diverge.
• Profile peak GPU memory with torch.cuda.max_memory_allocated to verify gradient bucket overhead stays inside the device budget.
• Verify all ranks terminate cleanly and only rank 0 writes checkpoint files.

Failure modes

• Mismatched collective calls or network misconfiguration can cause a silent hang where some ranks wait for messages that never arrive.
• Omitting DistributedSampler duplicates data across ranks, producing unstable loss and wasted computation.
• Using DistributedSampler without drop_last=True can leave a smaller final batch on some ranks, which breaks synchronized batch normalization or loss reduction.
• Large world sizes can trigger out-of-memory errors from gradient bucket overhead, even when the model itself fits.

Conditions that would change the design

• If the model no longer fits on one GPU, switch to model parallelism or pipeline parallelism instead of DDP.
• If multi-node networking is unreliable, add fault-tolerant checkpointing and restart logic; DDP alone does not recover from node failure.
• If the workload is inference only, DDP is unnecessary overhead; use a single process with batching or model replication behind a serving layer.

Limitations and verification

DDP behavior and backend defaults can shift across PyTorch and CUDA versions. Validate the exact behavior in your target environment before production rollout. A practical verification is to run a two-process DDP job on one node with a synthetic dataset and confirm gradients match a single-process baseline within numerical tolerance. This does not prove production stability, but it confirms the communication path is wired correctly.

Version assumption: current PyTorch releases with CUDA and the NCCL backend. Confirm defaults in your target environment.