# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team # Create a container object to save model-specific tensors using the policy file above. from abc import ABC import torch import deepspeed from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig from deepspeed.accelerator import get_accelerator # If the intermediate size attribute is set DEFAULT_INTERMEDIATE_SIZE # it is assumed the intermediate size is 4x the embedding dimension DEFAULT_INTERMEDIATE_SIZE = -1 class BaseConvolutionContainer(ABC): # not implemented def __init__(self): pass class BaseTransformerContainer(ABC): def __init__(self, policy, config, model_config, layer_id, child): self.policy = policy self.config = config self.model_config = model_config self.layer_id = layer_id self.child = child self.megatron_v2 = self.policy.is_megatron_v2 self.scale_attention = self.policy.scale_attention self.ckpt_load_enabled = False # configuration for models. todo: can this be moved to a pydantic model config? self.hidden_size = None self.intermediate_size = None self.num_attention_heads = None self.mp_size = self.config.tensor_parallel.tp_size self.pre_layer_norm = self.model_config.do_layer_norm_before if \ hasattr(self.model_config, 'do_layer_norm_before') else self.policy.pre_attn_norm self.dtype = self.config.dtype self.attn_linear_layer = self.policy.linear_layer self.mlp_linear_layer = self.policy.linear_layer self.return_tuple = self.config.return_tuple self.triangular_masking = True self.local_attention = ((self.model_config.attention_layers[self.layer_id] == "local") if hasattr( self.model_config, 'attention_layers') else False) self.window_size = getattr(self.model_config, "window_size", 1) self.mlp_act_func_type = self.policy.mlp_act_func_type self.norm_type = self.policy.norm_type self.training_mp_size = self.config.training_mp_size self.bigscience_bloom = False self.max_out_tokens = self.config.max_out_tokens self.min_out_tokens = self.config.min_out_tokens self.scale_attn_by_inverse_layer_idx = getattr(self.config, "scale_attn_by_inverse_layer_idx", False) self.use_mup = self.policy.use_mup self.return_single_tuple = False self.rotary_dim = self.get_rotary_dim() self.mlp_after_attn = (self.rotary_dim is None or self.rotary_dim < 0) # Attention tensors self.qkvw = None self.qkvb = None self.dense_w = None self.dense_b = None # MLP tensors self._h4h_w = None self._h4h_b = None self._4hh_w = None self._4hh_b = None # LayerNorm tensors self.attn_nw = None self.attn_nb = None self.input_nw = None self.input_nb = None self.mp_group = None self.use_triton = False # Triton self.use_triton = config.use_triton and deepspeed.HAS_TRITON def create_ds_model_config(self): self.set_hidden_heads(*self.policy.get_hidden_heads()) assert self.num_attention_heads % self.mp_size == 0,\ "To run the model parallel across the GPUs, the attention_heads require to be divisible by the world_size!" +\ "This is because the attention computation is partitioned evenly among the parallel GPUs." self.ds_model_config = DeepSpeedInferenceConfig( hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, heads=self.num_attention_heads, layer_norm_eps=self.layernorm_epsilon, dtype=self.dtype, pre_layer_norm=self.pre_layer_norm, norm_type=self.norm_type, mp_size=self.mp_size, return_tuple=self.return_tuple, triangular_masking=self.triangular_masking, local_attention=self.local_attention, window_size=self.window_size, rotary_dim=self.rotary_dim, mlp_after_attn=self.mlp_after_attn, mlp_act_func_type=self.mlp_act_func_type, training_mp_size=self.training_mp_size, bigscience_bloom=self.bigscience_bloom, max_out_tokens=self.max_out_tokens, min_out_tokens=self.min_out_tokens, scale_attn_by_inverse_layer_idx=self.scale_attn_by_inverse_layer_idx, use_mup=self.use_mup, return_single_tuple=self.return_single_tuple, set_empty_params=self.config.set_empty_params, transposed_mode=self.config.transposed_mode, use_triton=self.use_triton, triton_autotune=self.config.triton_autotune) if self.use_triton and deepspeed.HAS_TRITON: from .bert import DS_BERTContainer if not isinstance(self, DS_BERTContainer): raise NotImplementedError("Triton kernels are only for BERT-like models yet") if not self.config.triton_autotune: from deepspeed.ops.transformer.inference.triton.matmul_ext import fp16_matmul fp16_matmul.skip_autotune() return self.ds_model_config def check_meta_tensor_support(self): if hasattr(self.qkvw, 'is_meta'): if self.qkvw.is_meta: assert self.ckpt_load_enabled, "Meta tensors are not supported for this model currently." else: raise NotImplementedError("Meta tensor support is not available, please upgrade to torch 1.10+") def initialize_tensors(self, enable_training=False): # Set the tensors from policy (user module) to container (DS module) self.set_attention(*self.policy.attention(enable_training=enable_training)) self.set_mlp(*self.policy.mlp(enable_training=enable_training)) self.set_layernorm(*self.policy.layernorm()) #self.check_meta_tensor_support() def convert_to_required_dtype(self): # Note: converting tensors to fp16 requires that we do it in-place using self.__dict__ and not make a list/dict copy if self.dtype in [torch.half, torch.bfloat16]: for k, v in self.__dict__.items(): # The list comprehension is used for MoE tensor lists if isinstance(v, list) and all((isinstance(tensor, torch.Tensor) \ or isinstance(tensor, torch.nn.Parameter)) for tensor in v): self.__dict__[k] = [moe_tensor.to(self.dtype) for moe_tensor in v] if isinstance(v, torch.Tensor) or isinstance(v, torch.nn.Parameter): self.__dict__[k] = v.to(self.dtype) def get_rotary_dim(self): if hasattr(self.model_config, 'rotary_dim'): return self.model_config.rotary_dim if hasattr(self.child, 'attention') and hasattr(self.child.attention, 'rotary_ndims'): return self.child.attention.rotary_ndims return -1 def set_moe(self, moe=False): self.moe = moe def set_tensor_parallel_config(self, mp_size, mp_group): self.mp_size = mp_size self.mp_group = mp_group def set_quantization_config(self, quantizer): self.quantizer = quantizer def set_hidden_heads(self, hidden_size, num_attention_heads, epsilon, intermediate_size): """ Args: hidden_size: embedding dimension of the model num_attention_heads: number of attention heads in the model epsilon: epsilon value for layer norm (same value used for all norms) intermediate_size: Size of MLP projection. If `DEFAULT_INTERMEDIATE_SIZE` is passed it is assumed to be `4 * hidden_size` """ self.hidden_size = hidden_size if intermediate_size == DEFAULT_INTERMEDIATE_SIZE: self.intermediate_size = 4 * hidden_size else: self.intermediate_size = intermediate_size self.num_attention_heads = num_attention_heads self.layernorm_epsilon = epsilon def set_attention(self, qkvw, qkvb, dense_w, dense_b): self.qkvw = qkvw self.qkvb = qkvb self.dense_w = dense_w self.dense_b = dense_b def set_mlp(self, _h4h_w, _h4h_b, _4hh_w, _4hh_b): self._h4h_w = _h4h_w self._h4h_b = _h4h_b self._4hh_w = _4hh_w self._4hh_b = _4hh_b def set_layernorm(self, attn_nw, attn_nb, input_nw, input_nb): self.attn_nw = attn_nw self.attn_nb = attn_nb self.input_nw = input_nw self.input_nb = input_nb def apply_weight_quantization(self): # quantize attention weights self.attention_quantization() # quantize mlp weights self.mlp_quantization() def attention_quantization(self): self.module.attention.attn_qkvw = self.quantizer.quantize(self.module.attention.attn_qkvw) self.module.attention.attn_ow = self.quantizer.quantize(self.module.attention.attn_ow) def mlp_quantization(self): self.module.mlp.inter_w = self.quantizer.quantize(self.module.mlp.inter_w) self.module.mlp.output_w = self.quantizer.quantize(self.module.mlp.output_w) def apply_tensor_parallelism(self, mp_replace): # setup the new Attention module self.attention_qkv_mp(mp_replace) self.attention_o_mp(mp_replace) # setup the new MLP module self.mlp_inter_mp(mp_replace) self.mlp_output_mp(mp_replace) # Apply weight quantization # TODO(cmikeh2): Re-enable this once verified #self.apply_weight_quantization() def attention_qkv_mp(self, mp_replace, reversed_dim=False): self.module.attention.attn_qkvw = mp_replace.strided_copy(self.module.attention.attn_qkvw, self.qkvw, num_splits=3, int8=reversed_dim) self.module.attention.attn_qkvb = mp_replace.strided_copy(self.module.attention.attn_qkvb, self.qkvb, num_splits=3, int8=reversed_dim) def attention_o_mp(self, mp_replace, reversed_dim=False): self.module.attention.attn_ow = mp_replace.copy(self.module.attention.attn_ow, self.dense_w, int8=reversed_dim) self.module.attention.attn_ob = mp_replace.copy(self.module.attention.attn_ob, self.dense_b, int8=reversed_dim, allocate_tensor=reversed_dim) def mlp_inter_mp(self, mp_replace, reversed_dim=False): self.module.mlp.inter_w = mp_replace.copy(self.module.mlp.inter_w, self._h4h_w, int8=reversed_dim) self.module.mlp.inter_b = mp_replace.copy(self.module.mlp.inter_b, self._h4h_b, int8=reversed_dim) def mlp_output_mp(self, mp_replace, reversed_dim=False): self.module.mlp.output_w = mp_replace.copy(self.module.mlp.output_w, self._4hh_w, int8=reversed_dim) self.module.mlp.output_b = mp_replace.copy(self.module.mlp.output_b, self._4hh_b, int8=reversed_dim, allocate_tensor=reversed_dim) def copy_data_to_new_module(self): params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb} for key in params: if params[key] is None: setattr(self.module.mlp, key, None) else: setattr(self.module.mlp, key, torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) params = {'norm_w': self.input_nw, 'norm_b': self.input_nb} for key in params: if params[key] is None: setattr(self.module, key, None) else: setattr(self.module, key, torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name()))) def transpose(self): self.transpose_attention() self.transpose_mlp() def transpose_attention(self): if self.attn_linear_layer: self.qkvw = self.transpose_impl(self.qkvw.data) self.dense_w = self.transpose_impl(self.dense_w.data) def transpose_mlp(self): if self.mlp_linear_layer: self._h4h_w = self.transpose_impl(self._h4h_w.data) self._4hh_w = self.transpose_impl(self._4hh_w.data) def transpose_impl(self, data): data = data.contiguous() data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1)) data = data.reshape(data.shape[-1], data.shape[-2]) data.to(get_accelerator().current_device_name()) return data def get_all_params(self): params = [ self.attn_nw, self.attn_nb, self.input_nw, self.input_nb, ] params.extend(self.get_attn_params()) params.extend(self.get_mlp_params()) return params def get_attn_params(self): return [self.qkvw, self.qkvb, self.dense_w, self.dense_b] def get_mlp_params(self): return [self._h4h_w, self._h4h_b, self._4hh_w, self._4hh_b]