# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team from abc import ABC, abstractmethod from typing import List, Tuple import torch class HybridEngineContainer(ABC): """ This container identifies which methods need to be overridden in addition to the base container to enable use in the RLHF pipeline. These methods are not necessary for inference alone. NOTE: If you are using this feature with a container that also inherits from `MetaTensorContainer`, ensure that `MetaTensorContainer` is inherited before `HybridEngineContainer` in the class definition. """ def initialize_tensors(self, enable_training=False): """ Same purposes as the base container, but also grabs the hooks for any LoRA parameters. If it's necessary to override specific sub-components of the model, it's best to augment the specific `set_[component]` itself rather than modifying the `initialize_tensors` method. See the `HybridSplitQKVContainer` for an example. """ super().initialize_tensors(enable_training=enable_training) self.set_lora_params() def transform_for_training(self): """ If the views on certain parameters are largely incompatible, it may be necessary to do more substantial transformations to the parameters. This method should be overridden to transform the inference format to what is necessary for training. """ pass def transform_for_inference(self): """ If the views on certain parameters are largely incompatible, it may be necessary to do more substantial transformations to the parameters. This method should be overridden to transform the training format to what is necessary for inference. """ pass @abstractmethod def set_lora_params(self): """ If available, set the LoRA parameters for the module. An implementation for this would iterate over all parameters of the model and use the `maybe_get_lora` helper method to check if the parameter does in fact have any LoRA params. """ raise NotImplementedError("A set_lora_params() function must be defined for the relevant parameters.") @abstractmethod def get_lora_matched_pair(self): """Get the pair of lora params and its matched model parameters.""" raise NotImplementedError("get_lora_matched_pair() must be defined for the relevant parameters.") def fuse_lora(self): """Fuse the LoRA parameters for the inference mode.""" for maybe_lora_param, param in self.get_lora_matched_pair(): if len(maybe_lora_param) == 3: lora_right_weight, \ lora_left_weight, \ lora_scaling = maybe_lora_param param.data += lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) def unfuse_lora(self): """Unfuse the LoRA parameters for the training mode.""" for maybe_lora_param, param in self.get_lora_matched_pair(): if len(maybe_lora_param) == 3: lora_right_weight, \ lora_left_weight, \ lora_scaling = maybe_lora_param param.data -= lora_scaling * torch.matmul(lora_left_weight.t(), lora_right_weight.t()) def apply_tensor_parallelism(self, mp_replace, reversed_dim=False): """ Add support for reversed dim in tensor parallelism. If necessary, override the called methods to handle partitioned weights (i.e. if qkv is split, override the `attention_qkv_mp` method). If the model component is not split, it should be safe to use the default implementation. """ # Setup the new Attention module self.attention_qkv_mp(mp_replace, reversed_dim=reversed_dim) self.attention_o_mp(mp_replace, reversed_dim=reversed_dim) # Setup the new MLP module self.mlp_inter_mp(mp_replace, reversed_dim=reversed_dim) self.mlp_output_mp(mp_replace, reversed_dim=reversed_dim) # Apply weight quantization # TODO(cmikeh2): Re-enable this once verified #self.apply_weight_quantization() def _release_params(self, param_pairs: List[Tuple[torch.Tensor, torch.Tensor]]): """ Helper for `release_[component]` methods. Accepts a list of tuples where the first element is the module param that needs to be deleted, and the second is the reassignment from the container. """ for module_param, container_param in param_pairs: if module_param is not None: del module_param module_param = container_param def release_memory(self): """ Delete module parameters if they exist and point them back to the container. The primary purpose of this is for TP-inference with ZeRO-3. In this scenario, we need to delete the parameters we've created for inference to free their memory. """ general_params = [ (self.module.attention.attn_ow, self.dense_w), (self.module.attention.attn_ob, self.dense_b), (self.module.mlp.attn_nw, self.attn_nw), (self.module.mlp.attn_nb, self.attn_nb), (self.module.norm_w, self.input_nw), (self.module.norm_b, self.input_nb), ] self._release_params(general_params) self.release_qkv() self.release_mlp() def release_qkv(self): """ Release for QKV parameters (as well as any aliases). """ qkv_params = [ (self.module.attention.attn_qkvw, self.qkvw), (self.module.attention.attn_qkvb, self.qkvb), ] self._release_params(qkv_params) def release_mlp(self): """ Release for MLP parameters (as well as any aliases). """ mlp_params = [ (self.module.mlp.inter_w, self._h4h_w), (self.module.mlp.inter_b, self._h4h_b), (self.module.mlp.output_w, self._4hh_w), (self.module.mlp.output_b, self._4hh_b), ] self._release_params(mlp_params) def reset_params(self): """ The purpose of reset params is to get the weights from the FP16 training copy of the model and copy to them to contiguous inference view. This only needs to be performed when the container parameters cannot be used directly for inference. """ self.reset_qkv() self.reset_mlp() def reset_qkv(self): """ Perform any necessary resets of the model parameters for the QKV components. """ pass def reset_mlp(self): """ Perform any necessary resets of the model parameters for the MLP components. """ pass def get_lora_params(self): """ Return a list of all parameters that would have LoRA for the module. """ if not hasattr(self, "lora_params"): self.set_lora_params() return self.lora_params def set_params_wo_copy(self, Z3_enabled=False): """ Rather than copying into, set the parameters directly. This is necessary to provide an inexpensive (low-memory-overhead) view onto the FP16 forward weights. """ self.module.mlp.attn_nw = self.attn_nw self.module.mlp.attn_nb = self.attn_nb self.module.norm_w = self.input_nw self.module.norm_b = self.input_nb self.set_attn_params_wo_copy(Z3_enabled=Z3_enabled) self.set_mlp_params_wo_copy(Z3_enabled=Z3_enabled) def set_attn_params_wo_copy(self, **kwargs): """ Narrower sub-method for finer grained overriding. """ self.module.attention.attn_ow = self.dense_w self.module.attention.attn_ob = self.dense_b self.module.attention.attn_qkvw = self.qkvw self.module.attention.attn_qkvb = self.qkvb def set_mlp_params_wo_copy(self, **kwargs): """ Narrower sub-method for finer grained overriding. """ self.module.mlp.inter_w = self._h4h_w self.module.mlp.inter_b = self._h4h_b self.module.mlp.output_w = self._4hh_w self.module.mlp.output_b = self._4hh_b