torchfoo.distributed.setup_distributed#
- torchfoo.distributed.setup_distributed(rank, world_size, master_addr=None, master_port=None, backend='auto', force=False)[source]#
Initialize a distributed process group.
If world_size is 1 and force is False, setup is skipped.
- Parameters:
rank (int) – rank of the current process
world_size (int) – total number of processes
master_addr (str | None) – address of the master node. Falls back to MASTER_ADDR env var, then “localhost”.
master_port (str | int | None) – port of the master node. Falls back to MASTER_PORT env var.
backend (str) – distributed backend (default: “auto”). “auto” selects “nccl” if CUDA is available and there are enough GPUs for the world_size, “gloo” otherwise.
force (bool) – if True, initialize even when world_size is 1
- Returns:
False if setup was skipped, True otherwise
- Return type: