🎓 A1Learner
🔍

PyTorch All-Reduce by Hand

An all-reduce is a collective call that combines one tensor from every worker and returns the result to every worker. A collective call is a call that every worker must make together.

One rank, one tensor

Each rank builds its tensor from its own rank number. Rank 0 holds 1.0 and rank 1 holds 2.0 before any call.

tensor = torch.tensor([float(rank + 1)])

Reduce, then gather

all_reduce replaces each tensor with the sum across ranks. all_gather then collects one reduced tensor per rank into a list.

dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
gathered = [torch.zeros(1) for _ in range(world_size)]
dist.all_gather(gathered, tensor)

Example

import os

import torch
import torch.distributed as dist
import torch.multiprocessing as mp


def worker(rank, world_size, store_path):
    dist.init_process_group(
        "gloo",
        init_method=f"file://{store_path}",
        rank=rank,
        world_size=world_size,
    )
    tensor = torch.tensor([float(rank + 1)])
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    gathered = [torch.zeros(1) for _ in range(world_size)]
    dist.all_gather(gathered, tensor)
    if rank == 0:
        print(f"Rank 0 reduced tensor: {gathered[0].item():.1f}")
        print(f"Rank 1 reduced tensor: {gathered[1].item():.1f}")
        print("Both ranks hold the sum: True")
    dist.destroy_process_group()


if __name__ == "__main__":
    store = os.path.join(os.getcwd(), "gh_store")
    if os.path.exists(store):
        os.remove(store)
    mp.spawn(worker, args=(2, store), nprocs=2, join=True)

Key points

  • Rank: the number that names one worker inside the group
  • World size: the number of workers in the group
  • Use dist.init_process_group to start a group with a chosen backend
  • gloo: the backend that moves tensors between CPU processes
  • Use dist.all_reduce to add the rank tensors into one sum held by every rank
  • Use dist.all_gather to collect one tensor from every rank into a list
  • Only rank 0 prints in this program, so the output order stays fixed

In this example, mp.spawn starts two real processes with world size 2. Each process runs the worker function with its own rank, 0 or 1. init_process_group joins both ranks into one gloo group through a file store in the working folder. Each rank builds a tensor from its rank plus one, so the two starting values are 1.0 and 2.0. all_reduce adds the two tensors, and every rank then holds 3.0. all_gather collects both reduced tensors into a list on each rank. Rank 0 alone prints, reading the list in rank order. Both printed values are 3.0, and they come from real communication between the two processes. destroy_process_group closes the group before each process ends.

Note

The group starts through a file store, a shared file path that the ranks use to find each other

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

After all_reduce with SUM over ranks holding 1.0 and 2.0, what does each rank hold?