How to run distributed training in PyTorch

Training one model across multiple processes only works when each worker receives a separate data shard and every backward pass contributes to the same parameter update. A two-rank CPU smoke run exposes launch, sampling, and synchronization errors before accelerator or cluster setup adds more variables.

The torchrun launcher starts both local workers and supplies the rank, local-rank, and world-size values through environment variables. Standalone rendezvous mode keeps this first run on one host, while the gloo backend performs the CPU gradient exchange without CUDA hardware.

Fixed inputs and a shared random seed make the two ranks easy to compare. Different sample lists show that DistributedSampler partitioned the dataset, while an identical post-step weight shows that DistributedDataParallel synchronized the gradients before each optimizer updated its replica.

1. Create ddp_smoke.py with the imports and shared training-state builder.

   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 build_training_state(rank, world_size):
       torch.manual_seed(17)
       model = DistributedDataParallel(nn.Linear(2, 1))
    
       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)
       optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
       return model, loader, optimizer

   DistributedSampler assigns different dataset indexes to each rank. A multi-epoch training job with shuffling must call sampler.set_epoch(epoch) before creating each epoch's iterator.
   Related: How to use a DataLoader in PyTorch

2. Add the one-step training function below build_training_state().

   def train_one_step(model, loader, optimizer):
       batch_features, batch_targets = next(iter(loader))
       optimizer.zero_grad()
       loss = nn.functional.mse_loss(model(batch_features), batch_targets)
       loss.backward()
       optimizer.step()
    
       sample_ids = [int(value) for value in batch_features[:, 0].tolist()]
       weight = model.module.weight.detach()[0, 0].item()
       return sample_ids, loss.item(), weight

3. Add the distributed entry function below train_one_step().

   def main():
       dist.init_process_group("gloo")
       rank = dist.get_rank()
       world_size = dist.get_world_size()
       local_rank = int(os.environ["LOCAL_RANK"])
    
       model, loader, optimizer = build_training_state(rank, world_size)
       sample_ids, loss, weight = train_one_step(model, loader, optimizer)
       result = (
           f"rank={rank} local_rank={local_rank} world_size={world_size} "
           f"samples={sample_ids} loss={loss:.4f} weight={weight:.4f}"
       )
    
       gathered = [None] * world_size if rank == 0 else None
       dist.gather_object(result, gathered, dst=0)
       if rank == 0:
           for line in gathered:
               print(line)
           print("ddp_step_complete=True")
    
       dist.destroy_process_group()

   torchrun sets RANK, LOCAL_RANK, and WORLD_SIZE before each worker imports the script. For CUDA training, use one process per GPU, select the device from LOCAL_RANK, and use the nccl backend.

4. Add the script entry point below main().

   if __name__ == "__main__":
       main()

5. Run two local workers with torchrun to confirm synchronized model weights.

   $ OMP_NUM_THREADS=1 torchrun --standalone --nproc-per-node=2 ddp_smoke.py
   ##### snipped #####
   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

   The different samples lists prove that the ranks consumed separate partitions. Matching weight=1.1972 values prove that the optimizer step left both model replicas aligned.