# Copyright (c) 2025 NVIDIA CORPORATION. # Licensed under the MIT license. # Adapted from https://github.com/NVlabs/VILA/tree/main under the Apache 2.0 license. # LICENSE is in incl_licenses directory. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math import os from fnmatch import fnmatch from typing import List, Optional, Tuple, Union import torch import torch.nn.functional as F import torch.utils.checkpoint from torch import nn from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers import AutoConfig, AutoModel, AutoModelForCausalLM from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.generation import GenerationMixin from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_flash_attention_utils import _flash_attention_forward from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput, ) from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS from transformers.modeling_utils import PreTrainedModel from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import ( LlamaAttention, LlamaDynamicNTKScalingRotaryEmbedding, LlamaForCausalLM, LlamaLinearScalingRotaryEmbedding, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm, LlamaRotaryEmbedding, _prepare_4d_causal_attention_mask_with_cache_position, apply_rotary_pos_emb, repeat_kv, rotate_half, ) from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, is_flash_attn_greater_or_equal_2_10, is_torchdynamo_compiling, logging, replace_return_docstrings, ) from ..real_quantization import ( Coat_quantize_bgn, Coat_quantize_end, fp8_add_Ifp_Ifp_Ofp_Og16, fp8_add_Ifp_Ifp_Ofp_Opt, fp8_division, fp8_division_transpose, fp8_gelu_backward, fp8_gelu_forward, fp8_layernorm_noparam_backward, fp8_layernorm_noparam_forward, fp8_linear_backward, fp8_linear_forward, fp8_mul_backward, fp8_mul_forward, fp8_quantize, fp8_quantize_pertensor, fp8_quantize_pertensor_transpose, fp8_rmsnorm_backward, fp8_rmsnorm_forward, fp8_silu_backward, fp8_silu_forward, fp8_transpose, ) # FP8 related from ._fp8_quantization_config import QuantizationConfig from ._fp8_weightcache import FP8CacheWeightModule from ._fp8manager import FP8Manager logger = logging.get_logger(__name__) class CoatLlamaConfig(LlamaConfig): model_type = "fp8_llama" class CoatLlamaBeforeAttentionResidual(FP8CacheWeightModule): """ This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 1 * Linear layers """ def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_idx: Optional[int] = None): super().__init__(config, qargs, layer_idx) self.qargs = qargs self.fwobits = { "fabit": self.qargs.fabit, "fwbit": self.qargs.fwbit, "fobit": self.qargs.fobit, "babit": self.qargs.babit, "bwbit": self.qargs.bwbit, "bobit": self.qargs.bobit, } self.config = config self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " "when creating this class." ) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) def forward(self, re_x, x, s, rmsnorm_weight): if self.training: if self.qargs.weight_memory_efficient: # Prepare with torch.no_grad(): weight1_s = self.prepare_weight(self.q_proj.weight, "q_proj", FP8Manager.is_first_microbatch) weight2_s = self.prepare_weight(self.k_proj.weight, "k_proj", FP8Manager.is_first_microbatch) weight3_s = self.prepare_weight(self.v_proj.weight, "v_proj", FP8Manager.is_first_microbatch) return _CoatLlamaBeforeAttentionResidual.apply( re_x, x, s, self.q_proj.weight, None, None, weight1_s, self.k_proj.weight, None, None, weight2_s, self.v_proj.weight, None, None, weight3_s, rmsnorm_weight, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: # Prepare with torch.no_grad(): weight1, weight1_t, weight1_s = self.prepare_weight( self.q_proj.weight, "q_proj", FP8Manager.is_first_microbatch ) weight2, weight2_t, weight2_s = self.prepare_weight( self.k_proj.weight, "k_proj", FP8Manager.is_first_microbatch ) weight3, weight3_t, weight3_s = self.prepare_weight( self.v_proj.weight, "v_proj", FP8Manager.is_first_microbatch ) return _CoatLlamaBeforeAttentionResidual.apply( re_x, x, s, self.q_proj.weight, weight1, weight1_t, weight1_s, self.k_proj.weight, weight2, weight2_t, weight2_s, self.v_proj.weight, weight3, weight3_t, weight3_s, rmsnorm_weight, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: return re_x, self.att_proj(self.attn_norm(re_x)) class _CoatLlamaBeforeAttentionResidual(torch.autograd.Function): @staticmethod def forward( ctx, re_x, in_x, in_s, weight1_origin, weight1, weight1_t, weight1_s, weight2_origin, weight2, weight2_t, weight2_s, weight3_origin, weight3, weight3_t, weight3_s, rmsnorm_weight, group_size, fwobits, layer_id, config, qargs, eps=1e-5, ): # for autograd if fwobits["fabit"] == "E4M3": # in_x = in_x.to(torch.float8_e4m3fn) in_x = in_x.view(torch.float8_e4m3fn) else: raise ValueError("fabit should be E4M3") # LayerNorm ln_x, ln_s, ln_x_t, ln_utils = fp8_rmsnorm_forward( in_x, in_s, rmsnorm_weight, group_size, eps, transpose_output_2d=True ) # Linear Layer QKV Projection if qargs.weight_memory_efficient: assert weight1 is None # memory efficient weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s) weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s) weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s) fc1_x = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, False, group_size) # query states fc2_x = fp8_linear_forward(ln_x, ln_s, weight2, weight2_s, False, group_size) # key states fc3_x = fp8_linear_forward(ln_x, ln_s, weight3, weight3_s, False, group_size) # value states # ==================== save for backward ==================== ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s) if qargs.weight_memory_efficient: assert weight1_t is None and weight2_t is None and weight3_t is None ctx.weight = weight1_origin, weight1_s, weight2_origin, weight2_s, weight3_origin, weight3_s else: ctx.weight = weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s ctx.group_size = group_size ctx.ln_utils = ln_utils ctx.utils = fwobits, layer_id, config, qargs return re_x, fc1_x, fc2_x, fc3_x @staticmethod def backward(ctx, fp_grad, query_g, key_g, value_g): in_x, in_s, ln_x_t, ln_s = ctx.saved_tensors weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s = ctx.weight group_size = ctx.group_size rms_weight, rstd, num_warps = ctx.ln_utils fwobits, layer_id, config, qargs = ctx.utils # ==================== Begin backward ==================== # Quantize the RoPE and FlashAttention Output. grad_input and grad_weight requires different data layout. query_g, query_gs, query_g_t = fp8_quantize_pertensor_transpose( query_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False ) key_g, key_gs, key_g_t = fp8_quantize_pertensor_transpose( key_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False ) value_g, value_gs, value_g_t = fp8_quantize_pertensor_transpose( value_g, group_size, fwobits["babit"], transpose_output_2d=True, stochastic=False ) # Linear Layer QKV Projection if qargs.weight_memory_efficient: weight1_t, weight1_s = fp8_division_transpose( weight1_t, qargs.group_size, fwobits["fwbit"], weight1_s, only_transposed=True ) weight2_t, weight2_s = fp8_division_transpose( weight2_t, qargs.group_size, fwobits["fwbit"], weight2_s, only_transposed=True ) weight3_t, weight3_s = fp8_division_transpose( weight3_t, qargs.group_size, fwobits["fwbit"], weight3_s, only_transposed=True ) fc1_g1, att_q_wg = fp8_linear_backward( ln_x_t, ln_s, query_g, query_gs, query_g_t, weight1_t, weight1_s, group_size ) fc1_g2, att_k_wg = fp8_linear_backward(ln_x_t, ln_s, key_g, key_gs, key_g_t, weight2_t, weight2_s, group_size) fc1_g3, att_v_wg = fp8_linear_backward( ln_x_t, ln_s, value_g, value_gs, value_g_t, weight3_t, weight3_s, group_size ) fc1_g = fc1_g1 + fc1_g2 + fc1_g3 # LayerNorm in_g, rms_weight_grad = fp8_rmsnorm_backward(in_x, in_s, fc1_g, rms_weight, rstd, group_size, num_warps) # Add the gradient together, and prepare the input of the next layer. re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt( fp_grad, in_g, group_size, fwobits["babit"], stochastic=False ) # for autograd. forward's data type should be the same of backward tensor. this will not change the actual binary representation. in_g = in_g.view(torch.float8_e4m3fn) # Although the next operator is a linear layer in MLPResidual module, we return in_sg_g16 to make the size compatible with the forward. Otherwise it will not pass autograd. return ( re_g, in_g, in_sg_g16, att_q_wg, None, None, None, att_k_wg, None, None, None, att_v_wg, None, None, None, rms_weight_grad, None, None, None, None, None, None, ) class CoatLlamaAfterAttentionResidual(FP8CacheWeightModule): """ This is a typical transformer attention module that contains (1) Residual (2) 1 * Linear layers """ def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_id): super().__init__(config, qargs, layer_id) self.qargs = qargs self.fwobits = { "fabit": self.qargs.fabit, "fwbit": self.qargs.fwbit, "fobit": self.qargs.fobit, "babit": self.qargs.babit, "bwbit": self.qargs.bwbit, "bobit": self.qargs.bobit, } self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) def forward(self, re_x, in_x): if self.training: if self.qargs.weight_memory_efficient: # prepare for the weight with torch.no_grad(): weight4_s = self.prepare_weight(self.o_proj.weight, "o_proj", FP8Manager.is_first_microbatch) return _CoatLlamaAfterAttentionResidual.apply( re_x, in_x, self.o_proj.weight, None, None, weight4_s, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: # prepare for the weight with torch.no_grad(): weight4, weight4_t, weight4_s = self.prepare_weight( self.o_proj.weight, "o_proj", FP8Manager.is_first_microbatch ) return _CoatLlamaAfterAttentionResidual.apply( re_x, in_x, self.o_proj.weight, weight4, weight4_t, weight4_s, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: return re_x + self.attn_out(in_x), None, None class _CoatLlamaAfterAttentionResidual(torch.autograd.Function): @staticmethod def forward( ctx, re_x, flash_x, weight4_origin, weight4, weight4_t, weight4_s, group_size, fwobits, layer_id, config, qargs ): # Quantize the FlashAttention Output flash_qx, flash_s, _ = fp8_quantize_pertensor( flash_x, group_size, fwobits["fabit"] ) # Modified to make it memory efficient # # Attention Projection Linear Layer if qargs.weight_memory_efficient: assert weight4 is None # memory efficient weight4, weight4_s = fp8_division(weight4_origin, qargs.group_size, fwobits["fwbit"], weight4_s) fc4_x = fp8_linear_forward(flash_qx, flash_s, weight4, weight4_s, False, group_size) # # import IPython # IPython.embed() # Add the activations together fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc4_x, flash_qx.dtype, group_size) # ==================== save for backward ==================== ctx.save_for_backward(flash_x, flash_s) if qargs.weight_memory_efficient: assert weight4_t is None ctx.weight = weight4_origin, weight4_s else: ctx.weight = weight4_t, weight4_s ctx.group_size = group_size ctx.fwobits = fwobits ctx.utils = fwobits, layer_id, config, qargs # For autograd out_x = out_x.view(torch.float8_e4m3fn) return fp_x, out_x, out_s @staticmethod def backward(ctx, fp_grad, out_g, out_gs): flash_x, flash_s = ctx.saved_tensors weight4_t, weight4_s = ctx.weight group_size = ctx.group_size fwobits = ctx.fwobits fwobits, layer_id, config, qargs = ctx.utils # for autograd if fwobits["babit"] == "E5M2": # out_g = out_g.to(torch.float8_e5m2) out_g = out_g.view(torch.float8_e5m2) else: raise ValueError("babit should be E5M2") out_gs_max = out_gs.max() # ==================== Begin backward ==================== # Output Projection out_g_t = fp8_transpose(out_g, transpose_output_2d=True) # We do not save an extra flash_x to save the memory usage flash_x_t, flash_s = fp8_division_transpose( flash_x, group_size, fwobits["fabit"], flash_s, stochastic=False, only_transposed=True ) if qargs.weight_memory_efficient: weight4_t, weight4_s = fp8_division_transpose( weight4_t, qargs.group_size, fwobits["fwbit"], weight4_s, only_transposed=True ) fc4_g, attn_out_wg = fp8_linear_backward( flash_x_t, flash_s, out_g, out_gs_max, out_g_t, weight4_t, weight4_s, group_size ) return fp_grad, fc4_g, attn_out_wg, None, None, None, None, None, None, None, None class CoatLlamaMLPResidual(FP8CacheWeightModule): """ This is a typical transformer attention module that contains (1) Residual (2) LayerNorm / RMSNorm (3) 2 / 3 * Linear layers (4) GELU / Silu Activation """ def __init__(self, config: CoatLlamaConfig, qargs: QuantizationConfig, layer_id, hidden_size: int): super().__init__(config, qargs, layer_id) self.qargs = qargs self.fwobits = { "fabit": self.qargs.fabit, "fwbit": self.qargs.fwbit, "fobit": self.qargs.fobit, "babit": self.qargs.babit, "bwbit": self.qargs.bwbit, "bobit": self.qargs.bobit, } self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) self.training = True # below is only used when training = False assert config.hidden_act == "silu", "We only support silu activation currently" self.act_fn = ACT2FN[config.hidden_act] def forward(self, re_x, x, s, rmsnorm_weight): if self.training: if self.qargs.weight_memory_efficient: # prepare for the weight with torch.no_grad(): weight1_s = self.prepare_weight(self.gate_proj.weight, "gate_proj", FP8Manager.is_first_microbatch) weight2_s = self.prepare_weight(self.up_proj.weight, "up_proj", FP8Manager.is_first_microbatch) weight3_s = self.prepare_weight(self.down_proj.weight, "down_proj", FP8Manager.is_first_microbatch) return _CoatLlamaMLPResidual.apply( re_x, x, s, self.gate_proj.weight, None, None, weight1_s, self.up_proj.weight, None, None, weight2_s, self.down_proj.weight, None, None, weight3_s, rmsnorm_weight, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: # prepare for the weight with torch.no_grad(): weight1, weight1_t, weight1_s = self.prepare_weight( self.gate_proj.weight, "gate_proj", FP8Manager.is_first_microbatch ) weight2, weight2_t, weight2_s = self.prepare_weight( self.up_proj.weight, "up_proj", FP8Manager.is_first_microbatch ) weight3, weight3_t, weight3_s = self.prepare_weight( self.down_proj.weight, "down_proj", FP8Manager.is_first_microbatch ) return _CoatLlamaMLPResidual.apply( re_x, x, s, self.gate_proj.weight, weight1, weight1_t, weight1_s, self.up_proj.weight, weight2, weight2_t, weight2_s, self.down_proj.weight, weight3, weight3_t, weight3_s, rmsnorm_weight, self.qargs.group_size, self.fwobits, self.layer_id, self.config, self.qargs, ) else: raise NotImplementedError("Need TODO") og_x = re_x re_x = self.ff_norm(re_x) re_x = self.ff_proj(re_x) re_x = self.act(re_x) re_x = self.ff_out(re_x) re_x = og_x + re_x return re_x, None, None class _CoatLlamaMLPResidual(torch.autograd.Function): @staticmethod def forward( ctx, re_x, in_x, in_s, weight1_origin, weight1, weight1_t, weight1_s, weight2_origin, weight2, weight2_t, weight2_s, weight3_origin, weight3, weight3_t, weight3_s, rmsnorm_weight, group_size, fwobits, layer_id, config, qargs, eps=1e-5, ): # For autograd if fwobits["fabit"] == "E4M3": # in_x = in_x.to(torch.float8_e4m3fn) in_x = in_x.view(torch.float8_e4m3fn) else: raise ValueError("fabit should be E4M3") # LayerNorm ln_x, ln_s, ln_x_t, ln_utils = fp8_rmsnorm_forward( in_x, in_s, rmsnorm_weight, group_size, eps, transpose_output_2d=True ) # Linear Layer of Up Projection and Gate Projection. They are fused as one linear layer. if qargs.weight_memory_efficient: assert weight1 is None and weight2 is None and weight3 is None # memory efficient weight1, weight1_s = fp8_division(weight1_origin, qargs.group_size, fwobits["fwbit"], weight1_s) weight2, weight2_s = fp8_division(weight2_origin, qargs.group_size, fwobits["fwbit"], weight2_s) weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s) gate_x, gate_s = fp8_linear_forward(ln_x, ln_s, weight1, weight1_s, True, group_size) # Gate Proj up_x, up_s = fp8_linear_forward(ln_x, ln_s, weight2, weight2_s, True, group_size) # Up Proj # silu Activation silu_x, silu_s = fp8_silu_forward(gate_x, gate_s, group_size) # Element-wise Multiplication mul_x, mul_s, mul_x_t = fp8_mul_forward(silu_x, silu_s, up_x, up_s, group_size, transpose_output_2d=True) # Output Projection if weight3 is None: # memory efficient weight3, weight3_s = fp8_division(weight3_origin, qargs.group_size, fwobits["fwbit"], weight3_s) fc3_x = fp8_linear_forward(mul_x, mul_s, weight3, weight3_s, False, group_size) # Add the activation together fp_x, (out_x, out_s) = fp8_add_Ifp_Ifp_Ofp_Og16(re_x, fc3_x, mul_x.dtype, group_size) # ==================== save for backward ==================== ctx.save_for_backward(in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s) ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s) if ( qargs.weight_memory_efficient ): # Weight_1/2_origin will not be saved twice, so it will be more memory efficient. assert weight1_t is None and weight2_t is None and weight3_t is None ctx.weight = (weight1_origin, weight1_s, weight2_origin, weight2_s, weight3_origin, weight3_s) else: # Weight1/2_t is different from the origin weight, so saving it will consumes additional memory footprint. ctx.weight = (weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s) ctx.group_size = group_size ctx.ln_utils = ln_utils ctx.utils = fwobits, layer_id, config, qargs out_x = out_x.view(torch.float8_e4m3fn) return fp_x, out_x, out_s @staticmethod def backward(ctx, fp_grad, out_g, out_gs): fwobits, layer_id, config, qargs = ctx.utils in_x, in_s, ln_x_t, ln_s, gate_x, gate_s, up_x, up_s, silu_x, silu_s, mul_x_t, mul_s = ctx.saved_tensors (weight1_t, weight1_s, weight2_t, weight2_s, weight3_t, weight3_s) = ctx.weight group_size = ctx.group_size rms_weight, rstd, num_warps = ctx.ln_utils fwobits, layer_id, config, qargs = ctx.utils # For autograd if fwobits["babit"] == "E5M2": # out_g = out_g.to(torch.float8_e5m2) out_g = out_g.view(torch.float8_e5m2) else: raise ValueError("babit should be E5M2") out_gs_max = out_gs.max() # ==================== Begin backward ==================== # Output Projection out_gs = out_gs.max() out_g_t = fp8_transpose(out_g, transpose_output_2d=True) if qargs.weight_memory_efficient: weight3_t, weight3_s = fp8_division_transpose( weight3_t, qargs.group_size, fwobits["fwbit"], weight3_s, only_transposed=True ) fc3_g, weight3_grad = fp8_linear_backward( mul_x_t, mul_s, out_g, out_gs_max, out_g_t, weight3_t, weight3_s, group_size ) # [MEM TEST] del out_g, out_g_t, weight3_t # Element-wise Multiplication, 1 means gate, 2 means up mul_g1, (mul_g2, mul_gs2, mul_g2_t) = fp8_mul_backward( silu_x, silu_s, up_x, up_s, fc3_g, group_size, fwobits["babit"], output_quantized_transpose=True ) # Silu activation silu_g, silu_gs, silu_g_t = fp8_silu_backward( gate_x, gate_s, mul_g1, group_size, fwobits["babit"], output_quantized_transpose=True ) # Linear Layer of Up and Gate Projection if qargs.weight_memory_efficient: weight1_t, weight1_s = fp8_division_transpose( weight1_t, group_size, fwobits["fwbit"], weight1_s, only_transposed=True ) weight2_t, weight2_s = fp8_division_transpose( weight2_t, group_size, fwobits["fwbit"], weight2_s, only_transposed=True ) # Gate Proj fc1_g, weight1_grad = fp8_linear_backward( ln_x_t, ln_s, silu_g, silu_gs, silu_g_t, weight1_t, weight1_s, group_size ) fc2_g, weight2_grad = fp8_linear_backward( ln_x_t, ln_s, mul_g2, mul_gs2, mul_g2_t, weight2_t, weight2_s, group_size ) fc_g = fc1_g + fc2_g # layerNorm in_g, rms_weight_grad = fp8_rmsnorm_backward(in_x, in_s, fc_g, rms_weight, rstd, group_size, num_warps) # Add the gradient together re_g, (in_g, in_sg, in_sg_g16) = fp8_add_Ifp_Ifp_Ofp_Opt( fp_grad, in_g, group_size, fwobits["babit"], stochastic=False ) in_g = in_g.view(torch.float8_e4m3fn) return ( re_g, in_g, in_sg_g16, weight1_grad, None, None, None, weight2_grad, None, None, None, weight3_grad, None, None, None, rms_weight_grad, None, None, None, None, None, None, ) class LlamaAttentionWithoutLinear(nn.Module): """ Remove the Q/K/V/O projection layer in LlamaAttention module and only calculate the attention logic. The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module. """ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " "when creating this class." ) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = True # TODO (joao): remove in v4.46 (RoPE is computed in the model, not in the decoder layers) self.rotary_emb = LlamaRotaryEmbedding(config=self.config) def forward( self, query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = query_states.size() query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " "removed and `position_embeddings` will be mandatory." ) cos, sin = self.rotary_emb(value_states, position_ids) else: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: # no matter the length, we just slice it causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): raise ValueError( f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" f" {attn_output.size()}" ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class LlamaFlashAttention2WithoutLinear(LlamaAttentionWithoutLinear): """ Remove the Q/K/V/O projection layer in LlamaFlashAttention2 module and only calculate the attention logic. The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() def forward( self, query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if isinstance(past_key_value, StaticCache): raise ValueError( "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" ) output_attentions = False bsz, q_len, _ = query_states.size() # Flash attention requires the input to have the shape # batch_size x seq_length x head_dim x hidden_dim # therefore we just need to keep the original shape query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " "removed and `position_embeddings` will be mandatory." ) cos, sin = self.rotary_emb(value_states, position_ids) else: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache # to be able to avoid many of these transpose/reshape/view. query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) dropout_rate = self.attention_dropout if self.training else 0.0 # In PEFT, usually we cast the layer norms in float32 for training stability reasons # therefore the input hidden states gets silently casted in float32. Hence, we need # cast them back in the correct dtype just to be sure everything works as expected. # This might slowdown training & inference so it is recommended to not cast the LayerNorms # in fp32. (LlamaRMSNorm handles it correctly) input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, "_pre_quantization_dtype"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.q_proj.weight.dtype logger.warning_once( f"The input hidden states seems to be silently casted in float32, this might be related to" f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" f" {target_dtype}." ) query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) attn_output = _flash_attention_forward( query_states, key_states, value_states, attention_mask, q_len, position_ids=position_ids, dropout=dropout_rate, sliding_window=getattr(self, "sliding_window", None), use_top_left_mask=self._flash_attn_uses_top_left_mask, is_causal=self.is_causal, ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class LlamaSdpaAttentionWithoutLinear(LlamaAttentionWithoutLinear): """ Remove the Q/K/V/O projection layer in LlamaSdpaAttention module and only calculate the attention logic. The Q/K/V Projection is moved to BeforeAttention Module, and the O Projection is moved to AfterAttention Module. """ # Adapted from LlamaAttention.forward def forward( self, query_states: torch.Tensor, key_states: torch.Tensor, value_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: if output_attentions: # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented. logger.warning_once( "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, " 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' ) return super().forward( query_states=query_states, key_states=key_states, value_states=value_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, ) bsz, q_len, _ = query_states.size() query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) if position_embeddings is None: logger.warning_once( "The attention layers in this model are transitioning from computing the RoPE embeddings internally " "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " "removed and `position_embeddings` will be mandatory." ) cos, sin = self.rotary_emb(value_states, position_ids) else: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) causal_mask = attention_mask if attention_mask is not None: causal_mask = causal_mask[:, :, :, : key_states.shape[-2]] # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask, # Reference: https://github.com/pytorch/pytorch/issues/112577. if query_states.device.type == "cuda" and causal_mask is not None: query_states = query_states.contiguous() key_states = key_states.contiguous() value_states = value_states.contiguous() # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. is_causal = True if causal_mask is None and q_len > 1 else False attn_output = torch.nn.functional.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=causal_mask, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=is_causal, ) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, q_len, -1) return attn_output, None, past_key_value COAT_LLAMA_ATTENTION_CLASSES = { "eager": LlamaAttentionWithoutLinear, "flash_attention_2": LlamaFlashAttention2WithoutLinear, "sdpa": LlamaSdpaAttentionWithoutLinear, } class CoatLlamaDecoderLayer(nn.Module): def __init__(self, config: CoatLlamaConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.self_attn = COAT_LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) self.qargs = QuantizationConfig(**config.coat_fp8_args) self.BeforeAttention = CoatLlamaBeforeAttentionResidual(config, self.qargs, layer_idx) self.AfterAttention = CoatLlamaAfterAttentionResidual(config, self.qargs, layer_idx) self.MLPResidual = CoatLlamaMLPResidual(config, self.qargs, layer_idx, self.hidden_size) self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, quant_hidden_states: torch.Tensor, scale_hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: """ Args: hidden_states (`torch.FloatTensor`): BF16 input to the layer of shape `(batch, seq_len, embed_dim)` quant_hidden_states (`torch.float8_e4m3fn`): FP8 input to the layer of shape `(batch, seq_len, embed_dim)` scale_hidden_states (`torch.bfloat16`): BF16 scaling factor to the layer of shape `(batch, seq_len, embed_dim // group_size)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, query_sequence_length, key_sequence_length)` if default attention is used. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, with `head_dim` being the embedding dimension of each attention head. kwargs (`dict`, *optional*): Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code into the model """ # Coat: The residual, LayerNorm, and the Q/K/V Projection Linear Layer residual, query_states, key_states, value_states = self.BeforeAttention( hidden_states, quant_hidden_states, scale_hidden_states, self.input_layernorm.weight ) # Self Attention without any linear layer hidden_states, self_attn_weights, present_key_value = self.self_attn( query_states=query_states, key_states=key_states, value_states=value_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) # Coat: The Output Projection Linear Layer and Residual hidden_states, quant_hidden_states, scale_hidden_states = self.AfterAttention(residual, hidden_states) # Residual Connection, LayerNorm, and the whole MLP module hidden_states, quant_hidden_states, scale_hidden_states = self.MLPResidual( hidden_states, quant_hidden_states, scale_hidden_states, self.post_attention_layernorm.weight ) outputs = ((hidden_states, quant_hidden_states, scale_hidden_states),) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) return outputs class CoatLlamaPreTrainedModel(PreTrainedModel): config_class = CoatLlamaConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["LlamaDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_quantized_cache = True _supports_static_cache = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class CoatLlamaModel(CoatLlamaPreTrainedModel): """ Coat Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`CoatLlamaDecoderLayer`] Args: config: CoatLlamaConfig """ def __init__(self, config: CoatLlamaConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [CoatLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = LlamaRotaryEmbedding(config=config) self.gradient_checkpointing = False # Quantize self.qargs = QuantizationConfig(**config.coat_fp8_args) self.quantize_input_before_block = Coat_quantize_bgn(self.qargs) self.quantize_output_after_block = Coat_quantize_end(self.qargs) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = return_dict if return_dict is not None else self.config.use_return_dict if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) if self.gradient_checkpointing and self.training and use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." ) use_cache = False if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) # kept for BC (non `Cache` `past_key_values` inputs) return_legacy_cache = False if use_cache and not isinstance(past_key_values, Cache): return_legacy_cache = True if past_key_values is None: past_key_values = DynamicCache() else: past_key_values = DynamicCache.from_legacy_cache(past_key_values) logger.warning_once( "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and " "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class " "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)" ) if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device ) if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = self._update_causal_mask( attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions ) hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None next_decoder_cache = None # Prepare the input for Coat decoderlayer hidden_states, quant_hidden_states, scale_hidden_states = self.quantize_input_before_block(hidden_states) for decoder_layer in self.layers: if output_hidden_states: all_hidden_states += (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = self._gradient_checkpointing_func( decoder_layer.__call__, hidden_states, quant_hidden_states, scale_hidden_states, causal_mask, position_ids, past_key_values, output_attentions, use_cache, cache_position, position_embeddings, ) else: layer_outputs = decoder_layer( hidden_states, quant_hidden_states, scale_hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, ) hidden_states, quant_hidden_states, scale_hidden_states = layer_outputs[0] if use_cache: next_decoder_cache = layer_outputs[2 if output_attentions else 1] if output_attentions: all_self_attns += (layer_outputs[1],) # Summarize the output of the Decoder Layer hidden_states = self.quantize_output_after_block(hidden_states, quant_hidden_states, scale_hidden_states) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None if return_legacy_cache: next_cache = next_cache.to_legacy_cache() if not return_dict: return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) _update_causal_mask = LlamaModel._update_causal_mask class CoatLlamaForCausalLM(CoatLlamaPreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.model = CoatLlamaModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.model.embed_tokens def set_input_embeddings(self, value): self.model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model forward = LlamaForCausalLM.forward prepare_inputs_for_generation = LlamaForCausalLM.prepare_inputs_for_generation # TODO # class LlamaForSequenceClassification(LlamaPreTrainedModel): # class LlamaForQuestionAnswering(LlamaPreTrainedModel): # class LlamaForTokenClassification(LlamaPreTrainedModel): def make_state_dict_compatible(state_dict: dict[str, torch.Tensor]): compatible_state_dict = {} for key, value in state_dict.items(): if fnmatch(key, "*self_attn.q_proj*"): new_key = key.replace("self_attn.q_proj", "BeforeAttention.q_proj") elif fnmatch(key, "*self_attn.k_proj*"): new_key = key.replace("self_attn.k_proj", "BeforeAttention.k_proj") elif fnmatch(key, "*self_attn.v_proj*"): new_key = key.replace("self_attn.v_proj", "BeforeAttention.v_proj") elif fnmatch(key, "*self_attn.o_proj*"): new_key = key.replace("self_attn.o_proj", "AfterAttention.o_proj") elif fnmatch(key, "*mlp.gate_proj*"): new_key = key.replace("mlp.gate_proj", "MLPResidual.gate_proj") elif fnmatch(key, "*mlp.up_proj*"): new_key = key.replace("mlp.up_proj", "MLPResidual.up_proj") elif fnmatch(key, "*mlp.down_proj*"): new_key = key.replace("mlp.down_proj", "MLPResidual.down_proj") else: new_key = key compatible_state_dict[new_key] = value return compatible_state_dict AutoConfig.register("fp8_llama", CoatLlamaConfig) AutoModel.register(CoatLlamaConfig, CoatLlamaModel) AutoModelForCausalLM.register(CoatLlamaConfig, CoatLlamaForCausalLM)