Multi-GPU Setup with PyTorch: Data Parallel and Distributed

DistributedDataParallel (DDP) normally launches one process per GPU and synchronizes gradients. Each process retains a model replica; two GPUs do not automatically combine their VRAM into one larger device. If the model itself cannot fit on one GPU, investigate sharded training or model/tensor parallelism rather than assuming DDP will solve it.

Prerequisites for the two-GPU example

Use a single Linux host with two allocated CUDA GPUs, a validated PyTorch build with NCCL support, and a small disposable work directory. Confirm that the devices are actually in one host. Two independently rented instances require multi-node networking and rendezvous setup beyond this example. No NVLink, peer-to-peer, or high-speed network capability is implied by GPU count.

nvidia-smi -L
nvidia-smi topo -m
python -c 'import torch; import torch.distributed as d; print(torch.cuda.device_count()); print(d.is_nccl_available())'

Topology output may be restricted in virtualized/container environments. Do not infer missing capabilities from hidden telemetry alone. Check the offering and operator documentation before paying for a topology-dependent workload.

Run a complete DDP smoke test

Save as ddp_smoke.py. The synthetic dataset avoids downloads; each rank receives a distributed shard. This is an execution check, not a useful training benchmark.

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

def main():
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl", timeout=timedelta(seconds=120))
    try:
        torch.manual_seed(7)
        data = TensorDataset(torch.randn(64, 8), torch.randn(64, 1))
        sampler = DistributedSampler(data, shuffle=True)
        loader = DataLoader(data, batch_size=8, sampler=sampler)
        device = torch.device("cuda", local_rank)
        model = DDP(torch.nn.Linear(8, 1).to(device), device_ids=[local_rank])
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
        for epoch in range(2):
            sampler.set_epoch(epoch)
            for x, y in loader:
                optimizer.zero_grad(set_to_none=True)
                loss = torch.nn.functional.mse_loss(model(x.to(device)), y.to(device))
                loss.backward()
                optimizer.step()
        torch.cuda.synchronize()
        print("rank", dist.get_rank(), "device", local_rank, "completed", flush=True)
    finally:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()

torchrun --standalone --nnodes=1 --nproc-per-node=2 ddp_smoke.py

Expect completion from ranks 0 and 1 on distinct local devices and a successful launcher exit. When restricting devices with CUDA_VISIBLE_DEVICES, local rank indexes the visible list, not necessarily the host's physical index. Never launch more processes than your allocated visible GPUs.

Adapt the pattern to real training

Keep DistributedSampler and call set_epoch each epoch for changed shuffling. Do not also enable DataLoader shuffle. Samplers can pad or drop records when dataset size is uneven, so account for duplicates in validation metrics. All ranks must execute matching collective operations. Save checkpoints from global rank zero, including model, optimizer, scheduler, scaler when used, and progress; coordinate recovery to avoid partial or conflicting files.

Global batch is approximately per-rank batch × world size × accumulation steps. Keeping per-rank batch fixed while adding GPUs changes the optimization problem. Compare either a fixed global batch or explicitly document that change when assessing speed and convergence.

Diagnose hangs and poor scaling

Inspect the earliest error across all ranks: one rank's OOM often appears as a collective timeout elsewhere. For a short reproduction, NCCL_DEBUG=INFO and TORCH_DISTRIBUTED_DEBUG=DETAIL can expose initialization and collective issues. Logs can contain topology and network details; keep them private. Do not disable transports or security controls indiscriminately to silence a timeout.

If the smoke test passes but real training is slow, profile input loading, communication, per-rank balance, and synchronization. Larger GPU count can increase communication and rental cost without improving time to a useful checkpoint. Stop all worker processes after a failed run, save artifacts, and terminate unused allocations. Next: tensor parallelism or performance diagnosis.

Technical references

• torchrun process and rank semantics
• PyTorch DDP design
• Official multi-GPU training example