How to run distributed training in PyTorch

Distributed training in PyTorch lets separate worker processes train matching model replicas while gradient synchronization keeps their parameter updates aligned. A local smoke run with DistributedDataParallel is the safest first check before moving a script to GPUs or multiple nodes.

torchrun supplies the rank, local-rank, and world-size environment values that each worker process needs. For a single-node CPU check, --standalone starts a local rendezvous and --nproc-per-node=2 launches two ranks from the same script.

The smoke script uses the gloo backend so it can run without CUDA hardware, then uses DistributedSampler so each rank sees a different part of the dataset. Matching final weights across the rank output show that DistributedDataParallel synchronized gradients before the optimizer step completed.

1. Create a DistributedDataParallel smoke script.

   ddp_smoke.py

   import os
    
   import torch
   import torch.distributed as dist
   from torch import nn
   from torch.nn.parallel import DistributedDataParallel
   from torch.utils.data import DataLoader, TensorDataset
   from torch.utils.data.distributed import DistributedSampler
    
    
   def main():
       dist.init_process_group("gloo")
    
       rank = dist.get_rank()
       world_size = dist.get_world_size()
       local_rank = int(os.environ["LOCAL_RANK"])
    
       torch.manual_seed(17)
       model = nn.Linear(2, 1)
       ddp_model = DistributedDataParallel(model)
    
       features = torch.tensor(
           [
               [0.0, 0.0],
               [1.0, 1.0],
               [2.0, 2.0],
               [3.0, 3.0],
           ],
           dtype=torch.float32,
       )
       targets = torch.tensor([[0.0], [2.0], [4.0], [6.0]], dtype=torch.float32)
       dataset = TensorDataset(features, targets)
       sampler = DistributedSampler(
           dataset,
           num_replicas=world_size,
           rank=rank,
           shuffle=False,
       )
       loader = DataLoader(dataset, batch_size=2, sampler=sampler)
    
       batch_features, batch_targets = next(iter(loader))
       sample_ids = [int(value) for value in batch_features[:, 0].tolist()]
    
       optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.1)
       loss_fn = nn.MSELoss()
    
       optimizer.zero_grad()
       loss = loss_fn(ddp_model(batch_features), batch_targets)
       loss.backward()
       optimizer.step()
    
       weight = ddp_model.module.weight.detach()[0, 0].item()
       message = (
           f"rank={rank} local_rank={local_rank} world_size={world_size} "
           f"samples={sample_ids} loss={loss.item():.4f} weight={weight:.4f}"
       )
    
       gathered = [None] * world_size if rank == 0 else None
       dist.gather_object(message, gathered, dst=0)
    
       if rank == 0:
           for line in gathered:
               print(line)
           print("ddp_step_complete=True")
    
       dist.destroy_process_group()
    
    
   if __name__ == "__main__":
       main()

   DistributedSampler divides dataset indexes by rank. gloo is suitable for a CPU smoke run; use the nccl backend and one process per GPU after the same rank wiring works on GPU hardware.
   Related: How to use a DataLoader in PyTorch

2. Launch two local workers with torchrun.

   $ OMP_NUM_THREADS=1 torchrun --standalone --nproc-per-node=2 ddp_smoke.py
   rank=0 local_rank=0 world_size=2 samples=[0, 2] loss=6.6602 weight=1.1972
   rank=1 local_rank=1 world_size=2 samples=[1, 3] loss=17.2775 weight=1.1972
   ddp_step_complete=True

   OMP_NUM_THREADS=1 keeps this local CPU smoke run from starting extra OpenMP threads per rank. torchrun sets RANK, LOCAL_RANK, and WORLD_SIZE before each worker imports the script.

3. Confirm that both ranks printed different sample lists and the same final weight.

   samples=[0, 2] and samples=[1, 3] show the sampler split the dataset. The shared weight=1.1972 value shows the DDP optimizer step left both model replicas aligned.

4. Remove the temporary smoke script.

   $ rm ddp_smoke.py