#!/usr/bin/env python # -*- encoding: utf-8 -*- # adopted from https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/context import inspect import random import socket import sys from collections import Counter from importlib.machinery import SourceFileLoader from pathlib import Path from typing import Union import numpy as np import torch import torch.distributed as dist from internlm.utils.common import SingletonMeta from internlm.utils.logger import get_logger from . import process_group_initializer as pgroup_initializer from .process_group_initializer import ParallelMode from .random import add_seed, get_seeds, set_mode IS_TENSOR_PARALLEL = "is_tensor_parallel" logger = get_logger(__file__) class Config(dict): """This is a wrapper class for dict objects so that values of which can be accessed as attributes. Args: config (dict): The dict object to be wrapped. """ def __init__(self, config: dict = None): if config is not None: for k, v in config.items(): self._add_item(k, v) def __missing__(self, key): raise KeyError(key) def __getattr__(self, key): try: value = super().__getitem__(key) return value except KeyError: raise AttributeError(key) def __setattr__(self, key, value): super().__setitem__(key, value) def _add_item(self, key, value): if isinstance(value, dict): self.__setattr__(key, Config(value)) else: self.__setattr__(key, value) def update(self, config): assert isinstance(config, (Config, dict)), "can only update dictionary or Config objects." for k, v in config.items(): self._add_item(k, v) return self @staticmethod def from_file(filename: str): """Reads a python file and constructs a corresponding :class:`Config` object. Args: filename (str): Name of the file to construct the return object. Returns: :class:`Config`: A :class:`Config` object constructed with information in the file. Raises: AssertionError: Raises an AssertionError if the file does not exist, or the file is not .py file """ # check config path if isinstance(filename, str): filepath = Path(filename).absolute() elif isinstance(filename, Path): filepath = filename.absolute() assert filepath.exists(), f"{filename} is not found, please check your configuration path" # check extension extension = filepath.suffix assert extension == ".py", "only .py files are supported" # import the config as module remove_path = False if filepath.parent not in sys.path: sys.path.insert(0, (filepath)) remove_path = True module_name = filepath.stem source_file = SourceFileLoader(fullname=str(module_name), path=str(filepath)) module = source_file.load_module() # pylint: disable=W4902,E1120 # load into config config = Config() for k, v in module.__dict__.items(): if k.startswith("__") or inspect.ismodule(v) or inspect.isclass(v): continue else: config._add_item(k, v) # remove module del sys.modules[module_name] if remove_path: sys.path.pop(0) return config class ParallelContext(metaclass=SingletonMeta): """This class provides interface functions for users to get the parallel context, such as the global rank, the local rank, the world size, etc. of each device. """ def __init__(self): # distributed settings self._global_ranks = dict() self._local_ranks = dict() self._world_sizes = dict() self._groups = dict() self._cpu_groups = dict() self._ranks_in_group = dict() # load config from file self._config = None # default parallel args, will be overwritten during process group intialization self.world_size = 1 self.data_parallel_size = 1 self.pipeline_parallel_size = 1 self.tensor_parallel_size = 1 self.zero1_parallel_size = -1 self.num_processes_on_current_node = -1 self.virtual_pipeline_parallel_size = None self.virtual_pipeline_parallel_rank = None @property def config(self): return self._config def load_config(self, config: Union[dict, str]): """Loads the configuration from either a dict or a file. Args: config (dict or str): Either a dict containing the configuration information or the filename of a file containing the configuration information. Raises: TypeError: Raises a TypeError if `config` is neither a dict nor a str. """ if isinstance(config, str): self._config = Config.from_file(config) elif isinstance(config, dict): self._config = Config(config) else: raise TypeError("Invalid type for config, only dictionary or string is supported") def detect_num_processes_on_current_node(self): hostname = socket.gethostname() hostname_list = [None for _ in range(self.get_world_size(ParallelMode.GLOBAL))] dist.all_gather_object(hostname_list, hostname, group=self.get_group(ParallelMode.GLOBAL)) counter = Counter(hostname_list) self.num_processes_on_current_node = counter[hostname] @staticmethod def _check_parallel_mode(parallel_mode: ParallelMode): assert isinstance( parallel_mode, ParallelMode ), f"expected the argument parallel_mode to be of enum ParallelMode, but got {type(parallel_mode)}" def get_global_rank(self): """Returns the global rank of the current device. Returns: int: The global rank of the current device """ return self._global_ranks[ParallelMode.GLOBAL] def get_local_rank(self, parallel_mode: ParallelMode): """Returns the local rank of the current device. Args: parallel_mode: The parallel mode for the rank. Returns: int: The local rank of the current device for `parallel_mode`. """ self._check_parallel_mode(parallel_mode) return self._local_ranks.get(parallel_mode, 0) def get_next_global_rank(self, parallel_mode: ParallelMode): """Returns the global rank of the next device. Args: parallel_mode: The parallel mode for the rank. Returns: int: The global rank of the next device for `parallel_mode`. """ self._check_parallel_mode(parallel_mode) # get rank and world size local_rank = self.get_local_rank(parallel_mode) world_size = self.get_world_size(parallel_mode) ranks_in_group = self.get_ranks_in_group(parallel_mode) return ranks_in_group[(local_rank + 1) % world_size] def get_prev_global_rank(self, parallel_mode: ParallelMode): """Returns the global rank of the previous device. Args: parallel_mode: The chosen parallel mode. Returns: int: The global rank of the previous device for `parallel_mode`. """ self._check_parallel_mode(parallel_mode) # get rank and world size local_rank = self.get_local_rank(parallel_mode) world_size = self.get_world_size(parallel_mode) ranks_in_group = self.get_ranks_in_group(parallel_mode) return ranks_in_group[(local_rank - 1) % world_size] def is_using_dp(self): """Returns a boolean value indicating whether the current device is initilized with ParallelMode.DATA and its world_size is greater than 1. """ return self.is_initialized(ParallelMode.DATA) and self.get_world_size(ParallelMode.DATA) > 1 def is_using_tp(self): """Returns a boolean value indicating whether the current device is initilized with ParallelMode.TENSOR and its world_size is greater than 1. """ return self.is_initialized(ParallelMode.TENSOR) and self.get_world_size(ParallelMode.TENSOR) > 1 def is_using_pp(self): """Returns a boolean value indicating whether the current device is initilized with ParallelMode.PIPELINE and its world_size is greater than 1. """ return self.is_initialized(ParallelMode.PIPELINE) and self.get_world_size(ParallelMode.PIPELINE) > 1 def is_using_sequence(self): """Returns a boolean value indicating whether the current device is initilized with ParallelMode.SEQUENCE and its world_size is greater than 1. """ return False # return gpc.is_initialized(ParallelMode.SEQUENCE) and gpc.get_world_size(ParallelMode.SEQUENCE) > 1 def is_first_rank(self, parallel_mode: ParallelMode): """Returns a boolean value indicating whether the current device is the first one among its group for `parallel_mode`. Args: parallel_mode: The chosen parallel mode. Returns: bool: a boolean value indicating whether the current device is the first one among its group for `parallel_mode`. """ rank = 0 if self.is_initialized(parallel_mode): rank = self.get_local_rank(parallel_mode) return rank == 0 def is_rank_for_log(self): """Returns a boolean value indicating whether the current device should print log.""" is_log_rank = ( self.is_first_rank(ParallelMode.DATA) and self.is_first_rank(ParallelMode.TENSOR) and self.is_last_rank(ParallelMode.PIPELINE) ) return is_log_rank def is_last_rank(self, parallel_mode: ParallelMode): """Returns a boolean value indicating whether the current device is the last one among its group for `parallel_mode`. Args: parallel_mode: The chosen parallel mode. Returns: bool: a boolean value indicating whether the current device is the first one among its group for `parallel_mode`. """ rank = 0 world_size = 1 if self.is_initialized(parallel_mode): rank = self.get_local_rank(parallel_mode) world_size = self.get_world_size(parallel_mode) return rank == world_size - 1 def is_pipeline_first_stage(self, ignore_virtual=False): if not ignore_virtual: if self.virtual_pipeline_parallel_size is not None and self.virtual_pipeline_parallel_rank != 0: return False return self.is_first_rank(ParallelMode.PIPELINE) def is_pipeline_last_stage(self, ignore_virtual=False): if not ignore_virtual: if ( self.virtual_pipeline_parallel_size is not None and self.virtual_pipeline_parallel_rank != self.virtual_pipeline_parallel_size - 1 ): return False return self.is_last_rank(ParallelMode.PIPELINE) def get_world_size(self, parallel_mode: ParallelMode): """Returns the world size for `parallel_mode`. Args: parallel_mode: The chosen parallel mode. Returns: int: The world size for `parallel_mode`. """ self._check_parallel_mode(parallel_mode) return self._world_sizes.get(parallel_mode, 1) def get_group(self, parallel_mode: ParallelMode): """Returns the group of the current device for `parallel_mode`. Args: parallel_mode: The chosen parallel mode. Returns: torch.distributed.ProcessGroup: The group of the current device for `parallel_mode`. """ self._check_parallel_mode(parallel_mode) return self._groups[parallel_mode] def get_ranks_in_group(self, parallel_mode: ParallelMode): """Returns the rank of the current device for `parallel_mode` in the group. Args: parallel_mode: The chosen parallel mode. Returns: int: The rank of the current device for `parallel_mode` in the group. """ self._check_parallel_mode(parallel_mode) return self._ranks_in_group[parallel_mode] def get_cpu_group(self, parallel_mode: ParallelMode): self._check_parallel_mode(parallel_mode) return self._cpu_groups[parallel_mode] def init_global_dist(self, rank: int, world_size: int, backend: str, host: str, port: int, use_cpu: bool = False): """Initializes the global distributed environment Args: rank (int): rank for the default process group. world_size (int): world size of the default process group. backend (str): backend for ``torch.distributed`` host (str): the master address for distributed training. port (str): the master port for distributed training. use_cpu (bool): whether to set up cpu process group. """ # initialize the default process group init_method = f"tcp://[{host}]:{port}" dist.init_process_group(rank=rank, world_size=world_size, backend=backend, init_method=init_method) # None will give the default global process group for pytorch dist operations ranks = list(range(world_size)) if use_cpu: cpu_group = dist.new_group(ranks, backend="gloo") if dist.get_backend() != "gloo" else None else: cpu_group = None self._register_dist(rank, world_size, dist.GroupMember.WORLD, cpu_group, ranks, ParallelMode.GLOBAL) self._global_ranks[ParallelMode.GLOBAL] = rank def _register_dist(self, local_rank, world_size, process_group, cpu_group, ranks_in_group, mode): self._check_parallel_mode(mode) self._local_ranks[mode] = local_rank self._world_sizes[mode] = world_size self._groups[mode] = process_group self._cpu_groups[mode] = cpu_group self._ranks_in_group[mode] = ranks_in_group def check_sanity(self): """Checks sanity of the parallel context. Raises: AssertionError: Raises an AssertionError if the world size does not equal to the product of data parallel size, pipeline parallel size and tensor parallel size. """ dps = self.data_parallel_size pps = self.pipeline_parallel_size tps = self.tensor_parallel_size ws = self.world_size assert ws == dps * pps * tps, ( f"Expected the world size {ws} to be equal to data" f" parallel size ({dps}) * pipeline parallel size " f"({pps}) * tensor parallel size ({tps})" ) assert self.zero1_parallel_size > 0 assert self.data_parallel_size % self.zero1_parallel_size == 0 def _set_parallel_size_from_config(self, config: dict, key: str, attr_name: str): if key in config: ele = config[key] if isinstance(ele, int): setattr(self, attr_name, ele) elif isinstance(ele, dict): setattr(self, attr_name, ele["size"]) else: raise NotImplementedError( f'{"Parallel configuration does not support this kind of argument, please use int or dict"}' ) def init_parallel_groups(self): """Initializes the parallel groups.""" # get rank and world size rank = self.get_global_rank() world_size = self.get_world_size(ParallelMode.GLOBAL) self.world_size = world_size # set parallel size as attributes for global context parallel_config = self.config.get("parallel", None) if parallel_config is not None: self._set_parallel_size_from_config(parallel_config, "pipeline", "pipeline_parallel_size") self._set_parallel_size_from_config(parallel_config, "tensor", "tensor_parallel_size") self._set_parallel_size_from_config(parallel_config, "zero1", "zero1_parallel_size") # the user should not set the data parallel size manually # instead, it should be calculated based on other parallel config self.data_parallel_size = self.world_size // (self.pipeline_parallel_size * self.tensor_parallel_size) if self.zero1_parallel_size <= 0: self.zero1_parallel_size = self.data_parallel_size self.check_sanity() initializer_args = [ rank, world_size, self.data_parallel_size, self.pipeline_parallel_size, self.tensor_parallel_size, self.zero1_parallel_size, ] # run initialization of different process groups initializers = [] initializers.append(pgroup_initializer.Initializer_Data(*initializer_args)) initializers.append(pgroup_initializer.Initializer_Model(*initializer_args)) initializers.append(pgroup_initializer.Initializer_Tensor(*initializer_args)) initializers.append(pgroup_initializer.Initializer_Zero1(*initializer_args)) if self.pipeline_parallel_size > 1: initializers.append(pgroup_initializer.Initializer_Pipeline(*initializer_args)) for initializer in initializers: parallel_setting = initializer.init_dist_group() if isinstance(parallel_setting, list): for args in parallel_setting: self._register_dist(*args) else: self._register_dist(*parallel_setting) def is_initialized(self, parallel_mode: ParallelMode): """Returns a boolean value indicating whether `parallel_mode` is initialized in the current system. """ return parallel_mode in self._groups def destroy(self): """Destroys the current distributed parallel environment.""" for mode, group in self._groups.items(): if mode is not ParallelMode.GLOBAL: dist.destroy_process_group(group) # destroy global process group dist.destroy_process_group() self._groups.clear() def set_device(self, device_ordinal: int = None): """Sets distributed processes to be bound to devices. Args: device_ordinal (int, optional): the device id to be bound to """ global_rank = self.get_global_rank() if device_ordinal is None: devices_per_node = torch.cuda.device_count() device_ordinal = global_rank % devices_per_node torch.cuda.set_device(device_ordinal) logger.info(f"process rank {global_rank} is bound to host:{socket.gethostname()} device: {device_ordinal}") def set_seed(self, seed: int, dpseed_with_tpoffset: bool = False): """Sets seeds for all random libraries. Args: seed (int): seed for random states """ pipeline_offset = self._local_ranks.get(ParallelMode.PIPELINE, 0) global_rank = self.get_global_rank() random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) assert torch.cuda.is_available() # data parallel seed are kept the same in the same pipeline stage dp_seed = seed if dpseed_with_tpoffset: dp_seed = seed + pipeline_offset * 1024 add_seed(ParallelMode.DATA, dp_seed) # model parallel seeds are different across ranks if self.is_initialized(ParallelMode.TENSOR): tp_rank = self.get_local_rank(ParallelMode.TENSOR) tp_seed = seed + tp_rank + pipeline_offset * 1024 add_seed(ParallelMode.TENSOR, tp_seed) set_mode(ParallelMode.DATA) seeds = get_seeds() seed_str = ", ".join([f"{k}: {v}" for k, v in seeds.items()]) logger.info( f"initialized seed on rank {global_rank}, " f"numpy: {seed}, python random: {seed}, {seed_str}," f"the default parallel seed is {ParallelMode.DATA}." ) def set_virtual_pipeline_parallel_size(self, size): self.virtual_pipeline_parallel_size = size def set_virtual_pipeline_parallel_rank(self, rank): self.virtual_pipeline_parallel_rank = rank global_context = ParallelContext()