import dataclasses
import os
import socket
from typing import Optional
import torch
import torch.distributed
[docs]@dataclasses.dataclass
class DistributedOption:
# Enable distributed Training
distributed: bool = False
# torch.distributed.Backend: "nccl", "mpi", "gloo", or "tcp"
dist_backend: str = "nccl"
# if init_method="env://",
# env values of "MASTER_PORT", "MASTER_ADDR", "WORLD_SIZE", and "RANK" are referred.
dist_init_method: str = "env://"
dist_world_size: Optional[int] = None
dist_rank: Optional[int] = None
local_rank: Optional[int] = None
ngpu: int = 0
dist_master_addr: Optional[str] = None
dist_master_port: Optional[int] = None
dist_launcher: Optional[str] = None
multiprocessing_distributed: bool = True
[docs] def init_options(self):
if self.distributed:
if self.dist_init_method == "env://":
if get_master_addr(self.dist_master_addr, self.dist_launcher) is None:
raise RuntimeError(
"--dist_master_addr or MASTER_ADDR must be set "
"if --dist_init_method == 'env://'"
)
if get_master_port(self.dist_master_port) is None:
raise RuntimeError(
"--dist_master_port or MASTER_PORT must be set "
"if --dist_init_port == 'env://'"
)
# About priority order:
# If --dist_* is specified:
# Use the value of --dist_rank and overwrite it environ just in case.
# elif environ is set:
# Use the value of environ and set it to self
self.dist_rank = get_rank(self.dist_rank, self.dist_launcher)
self.dist_world_size = get_world_size(
self.dist_world_size, self.dist_launcher
)
self.local_rank = get_local_rank(self.local_rank, self.dist_launcher)
if self.local_rank is not None:
if self.ngpu > 1:
raise RuntimeError(f"Assuming 1GPU in this case: ngpu={self.ngpu}")
if "CUDA_VISIBLE_DEVICES" in os.environ:
cvd = os.environ["CUDA_VISIBLE_DEVICES"]
if self.local_rank >= len(cvd.split(",")):
raise RuntimeError(
f"LOCAL_RANK={self.local_rank} is bigger "
f"than the number of visible devices: {cvd}"
)
if (
self.dist_rank is not None
and self.dist_world_size is not None
and self.dist_rank >= self.dist_world_size
):
raise RuntimeError(
f"RANK >= WORLD_SIZE: {self.dist_rank} >= {self.dist_world_size}"
)
if self.dist_init_method == "env://":
self.dist_master_addr = get_master_addr(
self.dist_master_addr, self.dist_launcher
)
self.dist_master_port = get_master_port(self.dist_master_port)
if (
self.dist_master_addr is not None
and self.dist_master_port is not None
):
self.dist_init_method = (
f"tcp://{self.dist_master_addr}:{self.dist_master_port}"
)
[docs] def init_torch_distributed(self):
if self.distributed:
# See:
# https://docs.nvidia.com/deeplearning/sdk/nccl-developer-guide/docs/env.html
os.environ.setdefault("NCCL_DEBUG", "INFO")
# See:
# https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group
os.environ.setdefault("NCCL_BLOCKING_WAIT", "1")
torch.distributed.init_process_group(
backend=self.dist_backend,
init_method=self.dist_init_method,
world_size=self.dist_world_size,
rank=self.dist_rank,
)
# About distributed model:
# if self.local_rank is not None and ngpu == 1
# => Distributed with n-Process and n-GPU
# if self.local_rank is None and ngpu >= 1
# => Distributed with 1-Process and n-GPU
if self.local_rank is not None and self.ngpu > 0:
torch.cuda.set_device(self.local_rank)
[docs]def resolve_distributed_mode(args):
# Note that args.distributed is set by only this function.
# and ArgumentParser doesn't have such option
if args.multiprocessing_distributed:
num_nodes = get_num_nodes(args.dist_world_size, args.dist_launcher)
# a. multi-node
if num_nodes > 1:
args.distributed = True
# b. single-node and multi-gpu with multiprocessing_distributed mode
elif args.ngpu > 1:
args.distributed = True
# c. single-node and single-gpu
else:
args.distributed = False
if args.ngpu <= 1:
# Disable multiprocessing_distributed mode if 1process per node or cpu mode
args.multiprocessing_distributed = False
if args.ngpu == 1:
# If the number of GPUs equals to 1 with multiprocessing_distributed mode,
# LOCAL_RANK is always 0
args.local_rank = 0
if num_nodes > 1 and get_node_rank(args.dist_rank, args.dist_launcher) is None:
raise RuntimeError(
"--dist_rank or RANK must be set "
"if --multiprocessing_distributed == true"
)
# Note that RANK, LOCAL_RANK, and WORLD_SIZE is automatically set,
# so we don't need to check here
else:
# d. multiprocess and multi-gpu with external launcher
# e.g. torch.distributed.launch
if get_world_size(args.dist_world_size, args.dist_launcher) > 1:
args.distributed = True
# e. single-process
else:
args.distributed = False
if args.distributed and args.ngpu > 0:
if get_local_rank(args.local_rank, args.dist_launcher) is None:
raise RuntimeError(
"--local_rank or LOCAL_RANK must be set "
"if --multiprocessing_distributed == false"
)
if args.distributed:
if get_node_rank(args.dist_rank, args.dist_launcher) is None:
raise RuntimeError(
"--dist_rank or RANK must be set "
"if --multiprocessing_distributed == false"
)
if args.distributed and args.dist_launcher == "slurm" and not is_in_slurm_step():
raise RuntimeError("Launch by 'srun' command if --dist_launcher='slurm'")
[docs]def is_in_slurm_job() -> bool:
return "SLURM_PROCID" in os.environ and "SLURM_NTASKS" in os.environ
[docs]def is_in_slurm_step() -> bool:
return (
is_in_slurm_job()
and "SLURM_STEP_NUM_NODES" in os.environ
and "SLURM_STEP_NODELIST" in os.environ
)
def _int_or_none(x: Optional[str]) -> Optional[int]:
if x is None:
return x
return int(x)
[docs]def free_port():
"""Find free port using bind().
There are some interval between finding this port and using it
and the other process might catch the port by that time.
Thus it is not guaranteed that the port is really empty.
"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("", 0))
return sock.getsockname()[1]
[docs]def get_rank(prior=None, launcher: str = None) -> Optional[int]:
if prior is None:
if launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
prior = os.environ["SLURM_PROCID"]
elif launcher == "mpi":
raise RuntimeError(
"launcher=mpi is used for 'multiprocessing-distributed' mode"
)
elif launcher is not None:
raise RuntimeError(f"launcher='{launcher}' is not supported")
if prior is not None:
return int(prior)
else:
# prior is None and RANK is None -> RANK = None
return _int_or_none(os.environ.get("RANK"))
[docs]def get_world_size(prior=None, launcher: str = None) -> int:
if prior is None:
if launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
prior = int(os.environ["SLURM_NTASKS"])
elif launcher == "mpi":
raise RuntimeError(
"launcher=mpi is used for 'multiprocessing-distributed' mode"
)
elif launcher is not None:
raise RuntimeError(f"launcher='{launcher}' is not supported")
if prior is not None:
return int(prior)
else:
# prior is None and WORLD_SIZE is None -> WORLD_SIZE = 1
return int(os.environ.get("WORLD_SIZE", "1"))
[docs]def get_local_rank(prior=None, launcher: str = None) -> Optional[int]:
# LOCAL_RANK is same as GPU device id
if prior is None:
if launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
prior = int(os.environ["SLURM_LOCALID"])
elif launcher == "mpi":
raise RuntimeError(
"launcher=mpi is used for 'multiprocessing-distributed' mode"
)
elif launcher is not None:
raise RuntimeError(f"launcher='{launcher}' is not supported")
if prior is not None:
return int(prior)
elif "LOCAL_RANK" in os.environ:
return int(os.environ["LOCAL_RANK"])
elif "CUDA_VISIBLE_DEVICES" in os.environ:
# There are two possibility:
# - "CUDA_VISIBLE_DEVICES" is set to multiple GPU ids. e.g. "0.1,2"
# => This intends to specify multiple devices to to be used exactly
# and local_rank information is possibly insufficient.
# - "CUDA_VISIBLE_DEVICES" is set to an id. e.g. "1"
# => This could be used for LOCAL_RANK
cvd = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
if len(cvd) == 1 and "LOCAL_RANK" not in os.environ:
# If CUDA_VISIBLE_DEVICES is set and LOCAL_RANK is not set,
# then use it as LOCAL_RANK.
# Unset CUDA_VISIBLE_DEVICES
# because the other device must be visible to communicate
return int(os.environ.pop("CUDA_VISIBLE_DEVICES"))
else:
return None
else:
return None
[docs]def get_master_addr(prior=None, launcher: str = None) -> Optional[str]:
if prior is None:
if launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
# e.g nodelist = foo[1-10],bar[3-8] or foo4,bar[2-10]
nodelist = os.environ["SLURM_STEP_NODELIST"]
prior = nodelist.split(",")[0].split("-")[0].replace("[", "")
if prior is not None:
return str(prior)
else:
return os.environ.get("MASTER_ADDR")
[docs]def get_master_port(prior=None) -> Optional[int]:
if prior is not None:
return prior
else:
return _int_or_none(os.environ.get("MASTER_PORT"))
[docs]def get_node_rank(prior=None, launcher: str = None) -> Optional[int]:
"""Get Node Rank.
Use for "multiprocessing distributed" mode.
The initial RANK equals to the Node id in this case and
the real Rank is set as (nGPU * NodeID) + LOCAL_RANK in torch.distributed.
"""
if prior is not None:
return prior
elif launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
# Assume ntasks_per_node == 1
if os.environ["SLURM_STEP_NUM_NODES"] != os.environ["SLURM_NTASKS"]:
raise RuntimeError(
"Run with --ntasks_per_node=1 if mutliprocessing_distributed=true"
)
return int(os.environ["SLURM_NODEID"])
elif launcher == "mpi":
# Use mpi4py only for initialization and not using for communication
from mpi4py import MPI
comm = MPI.COMM_WORLD
# Assume ntasks_per_node == 1 (We can't check whether it is or not)
return comm.Get_rank()
elif launcher is not None:
raise RuntimeError(f"launcher='{launcher}' is not supported")
else:
return _int_or_none(os.environ.get("RANK"))
[docs]def get_num_nodes(prior=None, launcher: str = None) -> Optional[int]:
"""Get the number of nodes.
Use for "multiprocessing distributed" mode.
RANK equals to the Node id in this case and
the real Rank is set as (nGPU * NodeID) + LOCAL_RANK in torch.distributed.
"""
if prior is not None:
return prior
elif launcher == "slurm":
if not is_in_slurm_step():
raise RuntimeError("This process seems not to be launched by 'srun'")
# Assume ntasks_per_node == 1
if os.environ["SLURM_STEP_NUM_NODES"] != os.environ["SLURM_NTASKS"]:
raise RuntimeError(
"Run with --ntasks_per_node=1 if mutliprocessing_distributed=true"
)
return int(os.environ["SLURM_STEP_NUM_NODES"])
elif launcher == "mpi":
# Use mpi4py only for initialization and not using for communication
from mpi4py import MPI
comm = MPI.COMM_WORLD
# Assume ntasks_per_node == 1 (We can't check whether it is or not)
return comm.Get_size()
elif launcher is not None:
raise RuntimeError(f"launcher='{launcher}' is not supported")
else:
# prior is None -> NUM_NODES = 1
return int(os.environ.get("WORLD_SIZE", 1))