import dataclasses import os import socket from typing import Optional import torch import torch.distributed @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 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}" ) 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) 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'") def is_in_slurm_job() -> bool: return "SLURM_PROCID" in os.environ and "SLURM_NTASKS" in os.environ 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) 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] 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")) 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")) 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 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") 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")) 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")) 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))