PyTorch DDP training loop, step by step

Losses, gradients and the single weight W are illustrative. The shuffle orders are the real torch.randperm(N) orders DistributedSampler draws with its default seed=0, and sharding, padding, batching and step counts follow DistributedSampler and DataLoader. Keys: ← → step, space plays.

world_size

drop_last=True dataset size batch_size

sample in current batch queued or used model weights gradients * padded repeat