"""This code is adapted from Alpa https://github.com/alpa-projects/alpa/ with some changes. """ import operator from functools import reduce from typing import List, Tuple import torch import torch.distributed as dist # modified from alpa LogicalDeviceMesh(https://github.com/alpa-projects/alpa/blob/main/alpa/shard_parallel/auto_sharding.py) class DeviceMesh: """A logical view of a physical cluster. For example, we could view a physical cluster with 16 devices as a device mesh with shape (2, 2, 4) or (4, 4). Arguments: physical_mesh_id (torch.Tensor): physical view of the devices in global rank. logical_mesh_id (torch.Tensor): logical view of the devices in global rank. mesh_shape (torch.Size, optional): shape of logical view. mesh_alpha (List[float], optional): coefficients used for computing communication cost (default: None) mesh_beta (List[float], optional): coefficients used for computing communication cost (default: None) init_process_group (bool, optional): initialize logical process group during initializing the DeviceMesh instance if the init_process_group set to True. Otherwise, users need to call create_process_groups_for_logical_mesh manually to init logical process group. (default: False) need_flatten(bool, optional): initialize flatten_device_mesh during initializing the DeviceMesh instance if the need_flatten set to True. """ def __init__(self, physical_mesh_id: torch.Tensor, mesh_shape: torch.Size = None, logical_mesh_id: torch.Tensor = None, mesh_alpha: List[float] = None, mesh_beta: List[float] = None, init_process_group: bool = False, need_flatten: bool = True): self.physical_mesh_id = physical_mesh_id if logical_mesh_id is None: self.mesh_shape = mesh_shape self._logical_mesh_id = self.physical_mesh_id.reshape(self.mesh_shape) else: self._logical_mesh_id = logical_mesh_id self.mesh_shape = self._logical_mesh_id.shape # map global rank into logical rank self.convert_map = {} self._global_rank_to_logical_rank_map(self._logical_mesh_id, []) # coefficient for alpha-beta communication model if mesh_alpha is None: mesh_alpha = [1] * len(self.mesh_shape) if mesh_beta is None: mesh_beta = [1] * len(self.mesh_shape) self.mesh_alpha = tuple(mesh_alpha) self.mesh_beta = tuple(mesh_beta) self.init_process_group = init_process_group self.need_flatten = need_flatten if self.init_process_group: self.process_groups_dict = self.create_process_groups_for_logical_mesh() if self.need_flatten and self._logical_mesh_id.dim() > 1: self.flatten_device_mesh = self.flatten() # Create a new member `flatten_device_meshes` to distinguish from original flatten methods (Because I'm not sure if there are functions that rely on the self.flatten()) # self.flatten_device_meshes = FlattenDeviceMesh(self.physical_mesh_id, self.mesh_shape, self.mesh_alpha, # self.mesh_beta) @property def shape(self): return self.mesh_shape @property def num_devices(self): return reduce(operator.mul, self.physical_mesh_id.shape, 1) @property def logical_mesh_id(self): return self._logical_mesh_id def __deepcopy__(self, memo): cls = self.__class__ result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): if k != 'process_groups_dict': setattr(result, k, __import__("copy").deepcopy(v, memo)) else: setattr(result, k, v) return result def flatten(self): """ Flatten the logical mesh into an effective 1d logical mesh, """ flatten_mesh_shape_size = len(self.mesh_shape) flatten_mesh_shape = [self.num_devices] return DeviceMesh(self.physical_mesh_id, tuple(flatten_mesh_shape), mesh_alpha=[max(self.mesh_alpha)] * (flatten_mesh_shape_size - 1), mesh_beta=[max(self.mesh_beta)] * (flatten_mesh_shape_size - 1), init_process_group=self.init_process_group, need_flatten=False) def _global_rank_to_logical_rank_map(self, tensor, index_list): ''' This method is a helper function to build convert_map recursively. ''' for index, inner_tensor in enumerate(tensor): if inner_tensor.numel() == 1: self.convert_map[int(inner_tensor)] = index_list + [index] else: self._global_rank_to_logical_rank_map(inner_tensor, index_list + [index]) def create_process_groups_for_logical_mesh(self): ''' This method is used to initialize the logical process groups which will be used in communications among logical device mesh. Note: if init_process_group set to False, you have to call this method manually. Otherwise, the communication related function, such as ShapeConsistencyManager.apply will raise errors. ''' process_groups_dict = {} check_duplicate_list = [] global_rank_flatten_list = self.physical_mesh_id.view(-1).tolist() for global_rank in global_rank_flatten_list: process_groups = self.global_rank_to_process_groups_with_global_rank(global_rank) for axis, process_group in process_groups.items(): if axis not in process_groups_dict: process_groups_dict[axis] = [] if process_group not in check_duplicate_list: check_duplicate_list.append(process_group) process_group_handler = dist.new_group(process_group) process_groups_dict[axis].append((process_group, process_group_handler)) return process_groups_dict def global_rank_to_logical_rank(self, rank): return self.convert_map[rank] def global_rank_to_process_groups_with_logical_rank(self, rank): ''' Give a global rank and return all logical process groups of this rank. for example: physical_mesh_id = torch.arange(0, 16).reshape(2, 8) mesh_shape = (4, 4) # [[0, 1, 2, 3], # [4, 5, 6, 7], # [8, 9, 10,11], # [12,13,14,15]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) print(device_mesh.global_rank_to_process_groups_with_logical_rank(0)) output: # key is axis name # value is a list of logical ranks in same axis with rank 0 {0: [[0, 0], [1, 0], [2, 0], [3, 0]], 1: [[0, 0], [0, 1], [0, 2], [0, 3]]} ''' process_groups = {} for d in range(self.logical_mesh_id.dim()): for replacer in range(self.logical_mesh_id.shape[d]): if d not in process_groups: process_groups[d] = [] process_group_member = self.convert_map[rank].copy() process_group_member[d] = replacer process_groups[d].append(process_group_member) return process_groups def global_rank_to_process_groups_with_global_rank(self, rank): ''' Give a global rank and return all process groups of this rank. for example: physical_mesh_id = torch.arange(0, 16).reshape(2, 8) mesh_shape = (4, 4) # [[0, 1, 2, 3], # [4, 5, 6, 7], # [8, 9, 10,11], # [12,13,14,15]] device_mesh = DeviceMesh(physical_mesh_id, mesh_shape) print(device_mesh.global_rank_to_process_groups_with_global_rank(0)) output: # key is axis name # value is a list of global ranks in same axis with rank 0 {0: [0, 4, 8, 12], 1: [0, 1, 2, 3]} ''' logical_process_groups = self.global_rank_to_process_groups_with_logical_rank(rank) process_groups = {} for dim, logical_ranks in logical_process_groups.items(): process_groups[dim] = [] for logical_rank in logical_ranks: for g_rank, l_rank in self.convert_map.items(): if l_rank == logical_rank: process_groups[dim].append(g_rank) return process_groups def all_gather_cost(self, num_bytes, mesh_dim): num_devices = self.logical_mesh_id.shape[mesh_dim] return (self.mesh_alpha[mesh_dim] + self.mesh_beta[mesh_dim] * (num_devices - 1) / num_devices * num_bytes + 0.1) def all_reduce_cost(self, num_bytes, mesh_dim): num_devices = self.logical_mesh_id.shape[mesh_dim] return (self.mesh_alpha[mesh_dim] + self.mesh_beta[mesh_dim] * 2 * (num_devices - 1) / num_devices * num_bytes + 0.01) def reduce_scatter_cost(self, num_bytes, mesh_dim): num_devices = self.logical_mesh_id.shape[mesh_dim] return (self.mesh_alpha[mesh_dim] + self.mesh_beta[mesh_dim] * (num_devices - 1) / num_devices * num_bytes + 0.001) def all_to_all_cost(self, num_bytes, mesh_dim): num_devices = self.logical_mesh_id.shape[mesh_dim] penalty_factor = num_devices / 2.0 return (self.mesh_alpha[mesh_dim] + self.mesh_beta[mesh_dim] * (num_devices - 1) / num_devices / num_devices * num_bytes * penalty_factor + 0.001) class FlattenDeviceMesh(DeviceMesh): def __init__(self, physical_mesh_id, mesh_shape, mesh_alpha=None, mesh_beta=None): super().__init__(physical_mesh_id, mesh_shape, mesh_alpha, mesh_beta, init_process_group=False, need_flatten=False) # Different from flatten(), mesh_shape leaves unchanged, mesh_alpha and mesh_beta are scalars self.mesh_alpha = max(self.mesh_alpha) self.mesh_beta = min(self.mesh_beta) # Different from original process_groups_dict, rank_list is not stored self.process_number_dict = self.create_process_numbers_for_logical_mesh() def create_process_numbers_for_logical_mesh(self): ''' Build 1d DeviceMesh in column-major(0) and row-major(1) for example: mesh_shape = (2,4) # [[0, 1, 2, 3], # [4, 5, 6, 7]] # return {0: [0, 4, 1, 5, 2, 6, 3, 7], 1: [0, 1, 2, 3, 4, 5, 6, 7]} ''' num_devices = reduce(operator.mul, self.mesh_shape, 1) process_numbers_dict = {} process_numbers_dict[0] = torch.arange(num_devices).reshape(self.mesh_shape).transpose(1, 0).flatten().tolist() process_numbers_dict[1] = torch.arange(num_devices).reshape(self.mesh_shape).flatten().tolist() return process_numbers_dict def mix_gather_cost(self, num_bytes): num_devices = reduce(operator.mul, self.mesh_shape, 1) return (self.mesh_alpha + self.mesh_beta * (num_devices - 1) / num_devices * num_bytes + 0.1)