Hey @annisat, that looks good to me. One way to speed it up a bit is to run the dist.all_reduce at the beginning of the loop and set async_op=True. Then only wait for it when you need the result. In this way, the comm and the forward/backward/opt.step computation can overlap. Please see the code in the following thread:

  
    
    
    Multiprocessing - Barrier Blocks all Processes? 
  
  
    The self-contained code below works for me. 
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_i…