torch.distributed.init_process_group initializes a distributed process group: the collection of processes that communicate during distributed training. It bootstraps communication so processes can exchange gradients, parameters, and other tensor data through collective operations. section 16.2 section 16.1.1
The call needs enough information for every participating process to join the same group:
- backend selects the communication implementation; examples include gloo for CPU-oriented training, nccl for GPUs, and mpi for multiple machines. section 16.2
- rank identifies the current process, while world_size gives the total number of processes. section 16.1.3 section 16.1.1
- A TCPStore can bootstrap communication using MASTER_ADDR and MASTER_PORT; passing store is optional because PyTorch can create a default store. section 16.1.3
Explicit initialization
import torch.distributed as dist import os os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "12355" def init_process(rank, world_size): store = dist.TCPStore( host_name=os.environ["MASTER_ADDR"], port=int(os.environ["MASTER_PORT"]), world_size=world_size, is_master=rank == 0 ) dist.init_process_group( backend="gloo", store=store, rank=rank, world_size=world_size ) section 16.1.3
Initialization with torchrun
torchrun --nproc-per-node=4 3_torchrun.py import torch.distributed as dist import os def main(): rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) dist.init_process_group( backend="gloo", rank=rank, world_size=world_size ) section 16.1.3
Listing 16.5 Using torchrun
# torchrun --nproc-per-node=4 3_torchrun.py #1
import torch.distributed as dist
import os
def main():
rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])
print(f"{rank=} initializing")
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
if __name__ == "__main__":
main()
See the compact torchrun listing for environment-driven initialization without manually creating a store or spawning processes. section 16.1.3
