One world, four processes

Chapter 8 runs a collective on a world of one, which proves the call path and nothing about coordination. Four is where the ideas start having consequences, and you do not need four machines to get there. torch.multiprocessing.spawn starts four fresh Python processes on this laptop, hands each one its index, and waits for all of them to finish.

Every one of those processes runs the same worker function. The only thing that differs is the integer it was handed, and that integer is what it passes to init_process_group as its rank. Nobody hands out assignments. Rank 2 knows it is rank 2 because it was told at startup, and it works out its own share of the data from that number and the world size.

Read the output ordering before anything else. Rank 1 printed first, then 2, then 3, and rank 0 landed last even though rank 0 is the process that printed the three header lines. Four processes writing into one terminal have no ordering discipline between them, and a collective imposes none either. All four came out of the all_reduce holding the same [10.0, 10.0], which is 1 + 2 + 3 + 4, and they arrived there in whatever order the operating system felt like.

That number is the first thing worth checking on any new machine: gloo is available here, nccl and mpi are not, because this is a CPU-only build. Every run in these four lessons is gloo, and every one of them is real.

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

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29531"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    x = torch.full((2,), float(rank + 1))
    dist.all_reduce(x)
    if rank == 0:
        print("backend:", dist.get_backend(), "| world:", dist.get_world_size())
        print("available: gloo", dist.is_gloo_available(),
              "nccl", dist.is_nccl_available(), "mpi", dist.is_mpi_available())
    print(f"rank {rank}: {x.tolist()}")
    dist.barrier()
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(4,), nprocs=4, join=True)

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

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29531"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    x = torch.full((2,), float(rank + 1))
    dist.all_reduce(x)
    if rank == 0:
        print("backend:", dist.get_backend(), "| world:", dist.get_world_size())
        print("available: gloo", dist.is_gloo_available(),
              "nccl", dist.is_nccl_available(), "mpi", dist.is_mpi_available())
    print(f"rank {rank}: {x.tolist()}")
    dist.barrier()
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(4,), nprocs=4, join=True)

# rank 1: [10.0, 10.0]
# rank 2: [10.0, 10.0]
# rank 3: [10.0, 10.0]
# backend: gloo | world: 4
# available: gloo True nccl False mpi False
# rank 0: [10.0, 10.0]

MASTER_ADDR and MASTER_PORT are not the address the ranks talk to each other on. They are the address of a rendezvous, and underneath every init method c10d supports there is one object doing the work: a Store, a small key-value service that rank 0 hosts and every other rank connects to. Ranks write their own listening addresses into it, read back everybody else's, and only then does gloo build the pairwise connections that carry the actual tensor traffic.

The store is a public class, not an implementation detail you have to infer. dist.TCPStore takes a host, a port, a world size, and a flag saying whether this process is the server, and it gives you set, get, add, wait, and num_keys. add is atomic, which is what makes it usable as the counter a barrier needs. Standing one up by hand takes four lines and no distributed job at all.

After init_process_group returns, the store your process group holds is not the raw TCPStore. Print its type and it is a PrefixStore, a wrapper that prepends a namespace to every key. That wrapper is why a second process group, a subgroup, or a later library that wants its own coordination channel can all share one TCP server without their keys colliding.

The rendezvous is a key-value store. The connections are built from what the ranks read out of it.

from datetime import timedelta
import torch.distributed as dist

s = dist.TCPStore("127.0.0.1", 29550, 1, True, timedelta(seconds=5))
s.set("rank0_addr", "tcp://127.0.0.1:1234")
print("get :", s.get("rank0_addr"))    # b"tcp://127.0.0.1:1234"
print("keys:", s.num_keys())           # 2
print("add :", s.add("counter", 1), s.add("counter", 1))   # 1 2

The string you hand init_process_group picks the library that moves the bytes, and the choice is a device question rather than a preference. gloo is the CPU answer and works on any machine. nccl is the CUDA answer and is what a GPU job wants. mpi only exists if torch was built against an MPI installation, which the wheels are not.

A single job can use two of them at once. The backend string accepts a device map, "cpu:gloo,cuda:nccl", so a process group can route CPU tensors through gloo and CUDA tensors through nccl without you building two groups by hand. The distributed reference documents the mapping form and the per-backend support matrix.

What gloo can carry is worth measuring rather than assuming, because the reputation is older than the code. Five operations, tried on a world of two on this build, all returned: all_reduce, all_gather_into_tensor, reduce_scatter_tensor, all_to_all_single, and barrier. That last pair matters for the FSDP lesson, which needs reduce-scatter to exist before it can shard a gradient.

all_reduce               ok
all_gather_into_tensor   ok
reduce_scatter_tensor    ok
all_to_all_single        ok
barrier                  ok

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

def try_op(name, fn, out):
    try:
        fn()
        out.append((name, "ok"))
    except Exception as e:
        out.append((name, f"{type(e).__name__}: {str(e).splitlines()[0][:70]}"))

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29560"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    out = []
    try_op("all_reduce", lambda: dist.all_reduce(torch.ones(4)), out)
    try_op("all_gather_into_tensor",
           lambda: dist.all_gather_into_tensor(torch.zeros(8), torch.ones(4)), out)
    try_op("reduce_scatter_tensor",
           lambda: dist.reduce_scatter_tensor(torch.zeros(4), torch.ones(8)), out)
    try_op("all_to_all_single",
           lambda: dist.all_to_all_single(torch.zeros(4), torch.ones(4)), out)
    try_op("barrier", lambda: dist.barrier(), out)
    if rank == 0:
        for n, s in out:
            print(f"{n:24s} {s}")
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(2,), nprocs=2, join=True)

# all_reduce               ok
# all_gather_into_tensor   ok
# reduce_scatter_tensor    ok
# all_to_all_single        ok
# barrier                  ok

Once a job has more than one axis of parallelism it needs process groups smaller than the world, and dist.new_group([0, 1]) builds one. The group is a first-class object: pass it as group= to any collective and the reduction happens only among its members, and dist.get_rank(group) gives you the position inside it, which is not the same integer as the global rank.

new_group is itself a collective, and that is the part people get wrong. Every rank in the world has to call it, including ranks that will never use the resulting group, because building the group requires agreement across all of them. Skip the call on rank 3 because rank 3 is not in the group, and the job hangs on group creation rather than on anything that looks like communication.

The run below builds two groups on a world of four and each rank all_reduces inside its own. Ranks 0 and 1 come out with 3.0, which is 1 + 2. Ranks 2 and 3 come out with 7.0, which is 3 + 4. Rank 2 reports global rank 2 and local rank 0, and that second number is the one a sharding calculation should be using.

def worker(rank, world):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29532"
    dist.init_process_group("gloo", rank=rank, world_size=world)
    pair = dist.new_group([0, 1])      # every rank calls both,
    odd = dist.new_group([2, 3])       # including the ranks not in them
    mine = pair if rank < 2 else odd
    x = torch.full((2,), float(rank + 1))
    dist.all_reduce(x, group=mine)
    print(f"rank {rank}: global {dist.get_rank()} local {dist.get_rank(mine)} -> {x.tolist()}")
    dist.destroy_process_group()

# rank 2: global 2 local 0 -> [7.0, 7.0]
# rank 1: global 1 local 1 -> [3.0, 3.0]
# rank 3: global 3 local 1 -> [7.0, 7.0]
# rank 0: global 0 local 0 -> [3.0, 3.0]

The xla path states the rule at the HLO level in its collectives chapter: every participant issues the matching call, in the same order, or nothing completes. c10d is where you get to feel that rule, because the failure arrives as a specific error rather than as a paragraph.

Break it the first way, by having rank 1 skip an all_reduce that rank 0 issues, and rank 1 sails past its own code while rank 0 sits in gloo waiting for a message that is never sent. The default timeout is thirty minutes, which is why a real hang looks like a job doing nothing rather than a job failing. Set the timeout to five seconds and the error text is immediate and exact: a recv operation that timed out waiting.

Break it the second way, by having the two ranks call all_reduce with different tensor sizes, and there is no exception at all. gloo reads a preamble that does not match the bytes it was handed, raises a C++ enforcement failure that nothing catches, and the process aborts. Python sees ProcessExitedException: process 1 terminated with signal SIGABRT, which is a different debugging situation from a traceback: no rank tells you what the mismatch was, and your try block never runs.

A wrong shape does not raise. It aborts the process.

rank 1 skips the collective, five-second timeout:
  rank 1: returned
  rank 0: RuntimeError: [... gloo/transport/uv/unbound_buffer.cc:67]
          Timed out waiting 5000ms for recv operation to complete

ranks call all_reduce with tensors of 4 and 8 elements:
  libc++abi: terminating due to uncaught exception of type gloo::EnforceNotMet:
  [enforce fail at ... gloo/transport/uv/pair.cc:248] op.nread == op.preamble.nbytes.
  torch.multiprocessing.spawn.ProcessExitedException: process 1 terminated with signal SIGABRT

import os
import sys
from datetime import timedelta
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def worker(rank, world, mode):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "29543" if mode == "skip" else "29544"
    dist.init_process_group("gloo", rank=rank, world_size=world,
                            timeout=timedelta(seconds=5))
    try:
        if mode == "skip":
            if rank == 0:
                dist.all_reduce(torch.ones(4))
        else:
            n = 4 if rank == 0 else 8
            dist.all_reduce(torch.ones(n))
        print(f"rank {rank}: returned")
    except Exception as e:
        print(f"rank {rank}: {type(e).__name__}: {str(e).splitlines()[0]}")
    dist.destroy_process_group()

if __name__ == "__main__":
    mp.spawn(worker, args=(2, sys.argv[1]), nprocs=2, join=True)

# python mismatch.py skip   -> the timeout above
# python mismatch.py shape  -> the abort above