# this code is inspired by the DeepSpeed library and implemented with our own design from scratch import copy import math import warnings from typing import Any, Dict, Iterator, OrderedDict, Set, Tuple import torch import torch.distributed as dist from torch.nn import Parameter from torch.optim import Optimizer from colossalai.amp.naive_amp.mixed_precision_mixin import BF16MixedPrecisionMixin, FP16MixedPrecisionMixin from colossalai.checkpoint_io.utils import calculate_tensor_size from colossalai.logging import get_dist_logger from colossalai.nn.optimizer import ColossalaiOptimizer, CPUAdam, FusedAdam, HybridAdam from colossalai.tensor.d_tensor import is_distributed_tensor from colossalai.utils import disposable, get_current_device, is_ddp_ignored from .chunk import Chunk, ChunkManager from .gemini_ddp import ZeroDDP __all__ = ['ZeroOptimizer', 'GeminiAdamOptimizer'] _AVAIL_OPTIM_LIST = {FusedAdam, CPUAdam, HybridAdam} class GeminiFP16MixedPrecisionMixin(FP16MixedPrecisionMixin): def __init__(self, module: ZeroDDP, initial_scale: float = 2**16, min_scale: float = 1, growth_factor: float = 2, backoff_factor: float = 0.5, growth_interval: int = 1000, hysteresis: int = 2, max_scale: float = 2**32) -> None: super().__init__(initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, max_scale) self.module = module def check_local_overflow(self) -> bool: return self.module.overflow_counter > 0 def pre_zero_grad(self) -> None: self.module.overflow_counter = 0 class ZeroOptimizer(ColossalaiOptimizer): """A wrapper for optimizer. ``ZeroDDP`` and ``ZeroOptimizer`` implement Zero Redundancy Optimizer (ZeRO state-3). Note: You must use ``ZeroDDP`` with ``ZeroOptimizer``. Note: Make sure you set ``placement_policy`` of ``GeminiManager`` to `"auto"`, if you set ``gpu_margin_mem_ratio > 0``. Args: optim (Optimizer): An Optimizer instance. module (ZeroDDP): A ``ZeroDDP`` instance. gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward) which will be used when using hybrid CPU optimizer. This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto". Defaults to 0.0. initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32. min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1. growth_factor (float, optional): Growth_factor used by DynamicGradScaler. Defaults to 2. backoff_factor (float, optional): Backoff_factor used by DynamicGradScaler. Defaults to 0.5. growth_interval (float, optional): Growth_interval used by DynamicGradScaler. Defaults to 1000. hysteresis (float, optional): Hysteresis used by DynamicGradScaler. Defaults to 2. max_scale (int, optional): Max_scale used by DynamicGradScaler. Defaults to 2**32. clipping_norm (float, optional): The norm value used to clip gradient. Defaults to 0.0. norm_type (float, optional): The type of norm used for gradient clipping. Currently, only L2-norm (norm_type=2.0) is supported in ZeroOptimizer. Defaults to 2.0. verbose (bool, optional): Whether to print verbose information, including grad overflow info. Defaults to False. """ def __init__(self, optim: Optimizer, module: ZeroDDP, gpu_margin_mem_ratio: float = 0.0, initial_scale: float = 2**32, min_scale: float = 1, growth_factor: float = 2, backoff_factor: float = 0.5, growth_interval: int = 1000, hysteresis: int = 2, max_scale: float = 2**32, clipping_norm: float = 0.0, norm_type: float = 2.0, verbose: bool = False, **defaults: Any): super().__init__(optim) assert isinstance(module, ZeroDDP) assert type(optim) in _AVAIL_OPTIM_LIST, "You should use an optimizer in the available list:\n" \ f"{_AVAIL_OPTIM_LIST}" self.module = module self.gemini_manager = module.gemini_manager self.chunk_manager: ChunkManager = self.gemini_manager.chunk_manager self.param_to_range: Dict[Parameter, Tuple[int, int]] = dict() self.param_to_chunk32: Dict[Parameter, Chunk] = dict() self.chunk16_set: Set[Chunk] = set() self.clipping_flag = clipping_norm > 0.0 self.max_norm = clipping_norm self.verbose = verbose self.param_groups_backup = list() # Mapping from integer id to real/fake param tensor, used for checkpointing. self.id_to_real_params: Dict[int, Parameter] = dict() self.id_to_fake_params: Dict[int, Parameter] = dict() if self.clipping_flag: assert norm_type == 2.0, "ZeroOptimizer only supports L2 norm now" ddp_param_list = [] for name, param in module.named_parameters(): if is_ddp_ignored(param): if param.requires_grad: warnings.warn(f"Parameter `{name}` is ignored by DDP but requires gradient! " "You should handle its optimizer update by yourself!") else: ddp_param_list.append(param) for p, fp32_p in zip(ddp_param_list, module.fp32_params): chunk_16 = self.chunk_manager.get_chunk(p) if chunk_16 not in self.chunk16_set: chunk_16.l2_norm_flag = self.clipping_flag self.chunk16_set.add(chunk_16) self.__init__optimizer() if module.mixed_precision is torch.float16: self.mix_precision_mixin = GeminiFP16MixedPrecisionMixin(module, initial_scale=initial_scale, min_scale=min_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, hysteresis=hysteresis, max_scale=max_scale) elif module.mixed_precision is torch.bfloat16: self.mix_precision_mixin = BF16MixedPrecisionMixin() else: raise RuntimeError(f"Unsupported mixed precision type: {module.mixed_precision}") self._logger = get_dist_logger() self.gpu_margin_mem_ratio: float = float(gpu_margin_mem_ratio) assert 0.0 <= self.gpu_margin_mem_ratio <= 1.0, f'gpu_margin_mem_ratio must >=0.0 and <=1.0' # Only move fp32 shards from CPU to GPU when user allows and inner optimizer is valid # Inner optimizer must support optimizing hybrid (CPU and CUDA) tensors, # and it must set `num_fp32_shards_per_param` correctly self._should_move_fp32_params_h2d: bool = self.gemini_manager.is_cuda_margin_mem_avail and self.gpu_margin_mem_ratio > 0.0 and getattr( optim, 'num_fp32_shards_per_param', 0) >= 2 if self.gpu_margin_mem_ratio > 0.0 and not self.gemini_manager.is_cuda_margin_mem_avail: self._logger.warning(f'gpu_margin_mem_ratio is meaningless when placement_policy is not "auto"', ranks=[0]) self._register_states = disposable(self._register_states_) def _set_grad_ptr(self): for group in self.param_groups: for fake_param in group['params']: chunk32 = self.param_to_chunk32[fake_param] begin, end = self.param_to_range[fake_param] chunk16 = chunk32.paired_chunk fake_param.data = chunk16.payload[begin:end] fake_param.grad = fake_param.data fake_param.data = chunk32.payload[begin:end] def _update_fp16_params(self): none_tensor = torch.empty([0]) for group in self.param_groups: for fake_param in group['params']: assert fake_param.grad is None fake_param.data = none_tensor.to(fake_param.device) for chunk16 in self.chunk16_set: chunk16.optim_update() def _clear_global_norm(self) -> None: for c16 in self.chunk16_set: c16.l2_norm = None def _calc_global_norm(self) -> float: norm_sqr: float = 0.0 group_to_norm = dict() for c16 in self.chunk16_set: assert c16.l2_norm is not None if c16.is_gathered: norm_sqr += c16.l2_norm else: # this chunk is sharded, use communication to collect total norm if c16.torch_pg not in group_to_norm: group_to_norm[c16.torch_pg] = 0.0 group_to_norm[c16.torch_pg] += c16.l2_norm c16.l2_norm = None # clear l2 norm comm_buffer = torch.zeros(1, dtype=torch.float, device=get_current_device()) for group, part_norm in group_to_norm.items(): comm_buffer.fill_(part_norm) dist.all_reduce(comm_buffer, group=group) norm_sqr += comm_buffer.item() global_norm = math.sqrt(norm_sqr) return global_norm def _get_combined_scale(self): div_scale = self.mix_precision_mixin.get_grad_div_scale() if self.clipping_flag: total_norm = self._calc_global_norm() clip = ((total_norm / div_scale) + 1e-6) / self.max_norm if clip > 1: div_scale = clip * div_scale return -1 if div_scale == 1.0 else div_scale def zero_grad(self, *args, **kwargs): self.mix_precision_mixin.pre_zero_grad() return self.optim.zero_grad(set_to_none=True) def step(self, *args, **kwargs): self._maybe_move_fp32_params() self._set_grad_ptr() if self.mix_precision_mixin.should_skip_step(): if self.verbose: self._logger.info(f'Found overflow. Skip step') self._clear_global_norm() # clear recorded norm self.zero_grad() # reset all gradients self._update_fp16_params() return # get combined scale. combined scale = loss scale * clipping norm # so that gradient = gradient / combined scale combined_scale = self._get_combined_scale() ret = self.optim.step(div_scale=combined_scale, *args, **kwargs) self._register_states() self.zero_grad() self._update_fp16_params() return ret def clip_grad_norm(self, model: torch.nn.Module, max_norm: float, norm_type: float = 2.0): raise NotImplementedError def backward(self, loss: torch.Tensor): loss = self.mix_precision_mixin.pre_backward(loss) self.module.backward(loss) def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tensor): # This function is called except the last stage of pipeline parallel # It receives the scaled grad from the previous rank # No need to scale the grad again # Need to unscale when optimizing grad = self.mix_precision_mixin.pre_backward_by_grad(grad) self.module.backward_by_grad(tensor, grad) def _maybe_move_fp32_params(self): if self._should_move_fp32_params_h2d: self._should_move_fp32_params_h2d = False available_cuda_margin_mem = self.gemini_manager.cuda_margin_mem * self.gpu_margin_mem_ratio fp32_params_available_cuda_margin_mem = available_cuda_margin_mem / self.optim.num_fp32_shards_per_param fp32_params_used_cuda_margin_mem = 0 for group in self.param_groups: for fake_param in group['params']: chunk32 = self.param_to_chunk32[fake_param] chunk16 = chunk32.paired_chunk if chunk32.device_type == 'cuda': continue if fp32_params_used_cuda_margin_mem + chunk32.payload_mem < fp32_params_available_cuda_margin_mem: self.chunk_manager.move_chunk(chunk32, get_current_device()) # stores grad now self.chunk_manager.move_chunk(chunk16, get_current_device()) self.module.set_chunk_grad_device(chunk16, get_current_device()) fp32_params_used_cuda_margin_mem += chunk32.payload_mem for group in self.param_groups: for fake_param in group['params']: chunk32 = self.param_to_chunk32[fake_param] if chunk32.device_type == 'cuda': state = self.optim.state[fake_param] for k, v in state.items(): if isinstance(v, torch.Tensor): state[k] = v.to(get_current_device()) def _register_states_(self): for group in self.optim.param_groups: for p in group['params']: state = self.optim.state[p] for val in state.values(): if isinstance(val, torch.Tensor): self.chunk_manager.add_extern_static_tensor(val) def __init__optimizer(self): def get_range_pair(local_chunk: Chunk, local_param: Parameter): param_info = local_chunk.tensors_info[local_param] if local_chunk.keep_gathered: return param_info.offset, param_info.end begin = max(0, param_info.offset - local_chunk.shard_begin) end = min(local_chunk.shard_size, param_info.end - local_chunk.shard_begin) return begin, end param_id = -1 for group in self.optim.param_groups: fake_params_list = list() group_backup = {k: v for k, v in group.items() if k != 'params'} group_ids = [] for param in group['params']: # Record the mapping of id to current param. param_id += 1 self.id_to_real_params[param_id] = param group_ids.append(param_id) # If current param is controlled by current process, add it to fake_param. if is_ddp_ignored(param): continue chunk16 = self.chunk_manager.get_chunk(param) range_pair = get_range_pair(chunk16, param) if range_pair[0] >= range_pair[1]: continue grad_device = self.module.grads_device[param] fake_param = torch.nn.Parameter(torch.empty([0], device=grad_device)) self.param_to_chunk32[fake_param] = chunk16.paired_chunk self.param_to_range[fake_param] = range_pair self.id_to_fake_params[param_id] = fake_param fake_params_list.append(fake_param) # Update self.optim.param_groups as well as backup group. group['params'] = fake_params_list group_backup['params'] = group_ids self.param_groups_backup.append(group_backup) def get_offsets(self, param_id: int) -> tuple: ''' Args: param_id(int): The id of parameter. Returns: chunk_offset(int): Offset of parameter inside the chunk. shard_offset(int): Offset of its optimizer state shard relative to the whole optimizer state. shard_size(int): Length of parameter shard owned by current process. ''' if param_id not in self.id_to_fake_params: return -1, -1, -1 fake_param = self.id_to_fake_params[param_id] chunk = self.param_to_chunk32[fake_param].paired_chunk param = self.id_to_real_params[param_id] param_info = chunk.tensors_info[param] begin_in_chunk, end_in_chunk = self.param_to_range[fake_param] chunk_offset = begin_in_chunk if chunk.keep_gathered: shard_offset = 0 else: shard_offset = begin_in_chunk + chunk.shard_begin - param_info.offset shard_size = end_in_chunk - begin_in_chunk assert chunk_offset >= 0 and shard_offset >= 0 return chunk_offset, shard_offset, shard_size def collect_states(self, param_id: int, only_rank_0: bool = True) -> dict: """ Args: param_id (int): id of the parameter whose state is to be gathered at master rank. only_rank_0(bool): if True, states will be collected only on master rank, otherwise collected on every rank. Returns: collected_states(dict): the gathered optimzier state of parameter with given id if this method is called by master rank, otherwise an empty dict. This method can work only when called by all processes simultaneously. """ # Get param & chunk & process group. param = self.id_to_real_params[param_id] fake_param = self.id_to_fake_params.get(param_id, None) chunk = self.chunk_manager.get_chunk(param) process_group = chunk.torch_pg rank = dist.get_rank(process_group) master_rank = 0 collected_states = {} # Fetch names of states through all_gather. local_state_names = None if fake_param is not None: local_state_names = list(self.optim.state[fake_param].keys()) gathered_state_names = [None for _ in range(dist.get_world_size(process_group))] dist.barrier() dist.all_gather_object(gathered_state_names, local_state_names) state_names = None for names in gathered_state_names: if names is not None: # Assume different devices share the same set of state names if they have. state_names = copy.deepcopy(names) break # Directly return if this parameter doesn't have optimizer states. # e.g. parameter freezed/layer dropped if state_names is None: return collected_states # Boolean variable is_collector indicates that whether the current rank # needs to gather the whole optimizer states. # Only master rank is collector when only_rank_0 is True. # Every rank is collector when only_rank_0 is False. is_collector = (rank == master_rank) or (not only_rank_0) # If the chunk is kept gathered, # the parameteres are treated the same as that of those in strict DDP during training. # So states can be directly fetched from current device. if chunk.keep_gathered: assert param_id in self.id_to_fake_params if is_collector: states = self.optim.state[fake_param] for state_name in state_names: if state_name == 'step': # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32. collected_states[state_name] = torch.tensor(states['step'], dtype=torch.float32, requires_grad=False).cpu() else: state_tensor = states[state_name].detach().clone().to(torch.float32).cpu() collected_states[state_name] = torch.reshape(state_tensor, param.shape) return collected_states # Check whether the param with given id is managed by current process. own_param = param_id in self.id_to_fake_params # Collector gets prepared for state collecting. if is_collector: for state_name in state_names: if state_name == 'step': # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32. collected_states[state_name] = torch.tensor(0.0, dtype=torch.float32, requires_grad=False).cpu() else: collected_states[state_name] = torch.zeros(param.numel(), dtype=torch.float32, requires_grad=False).cpu() # Materials for gathering, including compacted state tensors, and the offset of shard inside each state. compacted_states = self.pack_optimizer_states_to_tensor(param_id, state_names) if own_param else None _, shard_offset, shard_size = self.get_offsets(param_id) # Collectors gather state shards through all_gathering. gathered_state_shards = [None for _ in range(dist.get_world_size(process_group))] dist.barrier() dist.all_gather_object(gathered_state_shards, [compacted_states, shard_offset, shard_size]) if is_collector: for state_shard in gathered_state_shards: compacted_states = state_shard[0] shard_offset = state_shard[1] shard_size = state_shard[2] if compacted_states is None: continue self.load_from_compacted_states(compacted_states, collected_states, state_names, shard_offset, shard_size) # Reshape tensors if is_collector: for state_name, state_tensor in collected_states.items(): if state_tensor.numel() == param.numel(): collected_states[state_name] = torch.reshape(state_tensor, param.shape) return collected_states def pack_optimizer_states_to_tensor(self, param_id: int, state_names: list, device: torch.device = torch.device('cuda'), dtype: torch.dtype = torch.float32) -> torch.Tensor: ''' With param id given, pack its optimizer states into a compact tensor and return. ''' if param_id not in self.id_to_fake_params: return None fake_param = self.id_to_fake_params[param_id] param_range = self.param_to_range[fake_param] states = self.optim.state[fake_param] shard_size = param_range[1] - param_range[0] compacted_size = 0 for name in state_names: if name == 'step': compacted_size += 1 else: compacted_size += shard_size compacted_states = torch.zeros(compacted_size, dtype=dtype, device=device, requires_grad=False) next_state_offset = 0 for state_name, state_tensor in states.items(): # State 'step' needs special operation. if state_name == 'step': if isinstance(state_tensor, torch.Tensor): compacted_states[next_state_offset] = state_tensor[0].item() else: assert isinstance(state_tensor, int) compacted_states[next_state_offset] = state_tensor next_state_offset += 1 else: assert state_tensor.numel() == shard_size compacted_states[next_state_offset:next_state_offset + shard_size].copy_(state_tensor) next_state_offset += shard_size return compacted_states def load_from_compacted_states(self, compacted_states: torch.Tensor, collected_states: dict, state_names: list, shard_start: int, shard_size: int): ''' Given a tensor carrying compacted optimizer states, update these states to collected_states. ''' shard_end = shard_start + shard_size next_state_offset = 0 for state_name in state_names: if state_name == 'step': collected_states['step'].data = torch.tensor(compacted_states[next_state_offset].item(), dtype=torch.float32, requires_grad=False).cpu() next_state_offset += 1 else: target_segment = collected_states[state_name][shard_start:shard_end] target_segment.copy_(compacted_states[next_state_offset:next_state_offset + shard_size]) next_state_offset += shard_size def get_param_groups_for_saving(self) -> list: ''' Return the param_groups in Pytorch format when saving to checkpoint. ''' param_groups = copy.deepcopy(self.param_groups_backup) # To be compatible with pytorch checkpointing, # store extra hyperparameters used by pytorch Adam optimizer. torch_special_hyperparameters = { 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'differentiable': False, 'fused': False } for group in param_groups: for k, v in torch_special_hyperparameters.items(): if k not in group: group[k] = v return param_groups def state_dict(self, only_rank_0: bool = True) -> dict: """ Args: only_rank_0 (bool): a boolean value indicating whether the state_dict is collected only on rank 0, dafault to True. Returns: The complete state of the optimizer as a :class:`dict`. It contains two entries: * state - a dict holding current optimization state. Its content differs between optimizer classes. * param_groups - a list containing all parameter groups where each parameter group is a dict. Warning: This method will gather and return the whole optimizer state_dict, so it should be called only when memory resources are abundant. """ state_dict = {} state_dict['param_groups'] = self.get_param_groups_for_saving() # Collect optimizer states. state_dict['state'] = dict() for param_id in self.id_to_real_params.keys(): dist.barrier() state_dict['state'][param_id] = self.collect_states(param_id=param_id, only_rank_0=only_rank_0) return state_dict def load_param_groups(self, saved_param_groups: list): """ Load saved_param_groups into self.param_groups and self.param_groups_backup """ self.param_groups_backup = copy.deepcopy(saved_param_groups) # discard the older param_groups self.optim.param_groups = [] for group in saved_param_groups: fake_params_list = list() updated_group = {k: v for k, v in group.items() if k != 'params'} for param_id in group['params']: if param_id not in self.id_to_fake_params: continue fake_param = self.id_to_fake_params[param_id] fake_params_list.append(fake_param) updated_group['params'] = fake_params_list self.optim.param_groups.append(updated_group) def load_single_param_states(self, param_id: int, saved_states: dict): """ Load saved optimizer states into parameter with given id. """ def cast(param, state_range, value, key=None): """ Make a copy of the needed segment of value and cast it to device of param. """ assert isinstance(value, torch.Tensor) ret_val = value if (key == "step"): assert value.numel() == 1 ret_val = int(value.item()) else: state_start, state_end = state_range ret_val = torch.zeros(state_end - state_start, dtype=torch.float32, device=param.device, requires_grad=False) ret_val.copy_(value.flatten()[state_start:state_end]) return ret_val assert param_id in self.id_to_fake_params fake_param = self.id_to_fake_params[param_id] _, state_offset, param_size = self.get_offsets(param_id) state_range = (state_offset, state_offset + param_size) # Copy states assigned to param (and cast tensors to appropriate types). updated_states = dict() for k, v in saved_states.items(): updated_states[k] = cast(fake_param, state_range, v, k) del v # clean loaded states self.optim.state[fake_param].update(updated_states) def load_param_states(self, param_states: dict): """Loads param states from a state_dict. The param_states can be complete or sharded. During loading, filter out the part of states not considered by current process. Args: param_states (dict): A mapping from param_id to its states. """ for param_id, states in param_states.items(): if param_id in self.id_to_fake_params: self.load_single_param_states(param_id, states) def optimizer_loading_epilogue(self): # Epilogue when loading state_dict to pytorch optimizer. self.optim._hook_for_profile() # To support multiprocessing pickle/unpickle. self.optim.defaults.setdefault('differentiable', False) def load_state_dict(self, state_dict: dict): """Loads optimizer state from complete optimizer state_dict. During loading, filter out the part of states not considered by current process. Args: state_dict (dict): optimizer state. Should be an object returned from a call to :meth:`state_dict`. """ assert 'param_groups' in state_dict assert 'state' in state_dict self.load_param_groups(state_dict['param_groups']) self.load_param_states(state_dict['state']) self.optimizer_loading_epilogue() def state_shard(self, prefix: str = '', max_shard_size: int = 1024, only_rank_0: bool = True) -> Iterator[Tuple[OrderedDict, int]]: """Returns dictionaries containing shards of optimizer states one by one. The max size of each dictionary shard is specified by ``max_shard_size``. Args: prefix (str, optional): the prefix for states. Default to ''. max_shard_size (int, optional): max size of state dict shard (in MB). Defaults to 1024. only_rank_0 (bool, optional): a boolean value indicating whether the state_dict is collected only on rank 0, dafault to True. Yields: Iterator[OrderedDict]: A generator of state dict shard of optimizer states. """ current_block = {} current_block_size = 0 for param_id in self.id_to_real_params.keys(): dist.barrier() state = self.collect_states(param_id=param_id, only_rank_0=only_rank_0) ret_block = None ret_block_size = 0 # A state might contain more than one tensors. # e.g. each Adam state includes: 'step', 'exp_avg', 'exp_avg_sq' state_size = 0 isDTensor = False for state_tensor in state.values(): # When state_tensor is not of Tensor class, # e.g., a SGD optimizer with momentum set to 0 can have None as state # The calculation of tensor size should be skipped to avoid error. if not isinstance(state_tensor, torch.Tensor): continue # If the states are stored as DTensors, mark isDTensor as true. if is_distributed_tensor(state_tensor): isDTensor = True state_size += calculate_tensor_size(state_tensor) if not isDTensor: if current_block_size + state_size > max_shard_size and current_block_size > 0: ret_block = current_block ret_block_size = current_block_size current_block = {} current_block_size = 0 current_block[param_id] = state current_block_size += state_size if ret_block != None: yield ret_block, ret_block_size yield current_block, current_block_size class GeminiAdamOptimizer(ZeroOptimizer): def __init__(self, model: torch.nn.Module, **defaults: Any) -> None: optimizer = HybridAdam(model.parameters(), **defaults) super().__init__(optimizer, model, **defaults)