RedHatAI/Kimi-K2-Instruct-quantized.w4a16
12396
1# coding=utf-82# Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20""" PyTorch DeepSeek model."""21import math22import warnings23from typing import List, Optional, Tuple, Union24 25import torch26import torch.nn.functional as F27import torch.utils.checkpoint28from torch import nn29from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss30 31from transformers.activations import ACT2FN32from transformers.cache_utils import Cache, DynamicCache33from transformers.modeling_attn_mask_utils import (34 AttentionMaskConverter,35 _prepare_4d_attention_mask,36 _prepare_4d_causal_attention_mask,37)38from transformers.modeling_outputs import (39 BaseModelOutputWithPast,40 CausalLMOutputWithPast,41 SequenceClassifierOutputWithPast,42)43from transformers.modeling_utils import PreTrainedModel44from transformers.pytorch_utils import (45 ALL_LAYERNORM_LAYERS,46 is_torch_greater_or_equal_than_1_13,47)48from transformers.utils import (49 add_start_docstrings,50 add_start_docstrings_to_model_forward,51 is_flash_attn_2_available,52 is_flash_attn_greater_or_equal_2_10,53 logging,54 replace_return_docstrings,55)56from transformers.utils.import_utils import is_torch_fx_available57from .configuration_deepseek import DeepseekV3Config58import torch.distributed as dist59import numpy as np60 61if is_flash_attn_2_available():62 from flash_attn import flash_attn_func, flash_attn_varlen_func63 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa64 65 66# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.67# It means that the function will not be traced through and simply appear as a node in the graph.68if is_torch_fx_available():69 if not is_torch_greater_or_equal_than_1_13:70 import torch.fx71 72 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)73 74 75logger = logging.get_logger(__name__)76 77_CONFIG_FOR_DOC = "DeepseekV3Config"78 79 80def _get_unpad_data(attention_mask):81 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)82 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()83 max_seqlen_in_batch = seqlens_in_batch.max().item()84 cu_seqlens = F.pad(85 torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)86 )87 return (88 indices,89 cu_seqlens,90 max_seqlen_in_batch,91 )92 93 94class DeepseekV3RMSNorm(nn.Module):95 def __init__(self, hidden_size, eps=1e-6):96 """97 DeepseekV3RMSNorm is equivalent to T5LayerNorm98 """99 super().__init__()100 self.weight = nn.Parameter(torch.ones(hidden_size))101 self.variance_epsilon = eps102 103 def forward(self, hidden_states):104 input_dtype = hidden_states.dtype105 hidden_states = hidden_states.to(torch.float32)106 variance = hidden_states.pow(2).mean(-1, keepdim=True)107 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)108 return self.weight * hidden_states.to(input_dtype)109 110 111ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm)112 113 114class DeepseekV3RotaryEmbedding(nn.Module):115 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):116 super().__init__()117 118 self.dim = dim119 self.max_position_embeddings = max_position_embeddings120 self.base = base121 inv_freq = 1.0 / (122 self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)123 )124 self.register_buffer("inv_freq", inv_freq, persistent=False)125 126 # Build here to make `torch.jit.trace` work.127 self._set_cos_sin_cache(128 seq_len=max_position_embeddings,129 device=self.inv_freq.device,130 dtype=torch.get_default_dtype(),131 )132 self.max_seq_len_cached = None133 134 def _set_cos_sin_cache(self, seq_len, device, dtype):135 self.max_seq_len_cached = seq_len136 t = torch.arange(137 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype138 )139 140 freqs = torch.outer(t, self.inv_freq.to(t.device))141 # Different from paper, but it uses a different permutation in order to obtain the same calculation142 emb = torch.cat((freqs, freqs), dim=-1)143 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)144 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)145 146 def forward(self, x, seq_len=None):147 # x: [bs, num_attention_heads, seq_len, head_size]148 if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:149 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)150 151 return (152 self.cos_cached[:seq_len].to(dtype=x.dtype),153 self.sin_cached[:seq_len].to(dtype=x.dtype),154 )155 156 157# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3158class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):159 """DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""160 161 def __init__(162 self,163 dim,164 max_position_embeddings=2048,165 base=10000,166 device=None,167 scaling_factor=1.0,168 ):169 self.scaling_factor = scaling_factor170 super().__init__(dim, max_position_embeddings, base, device)171 172 def _set_cos_sin_cache(self, seq_len, device, dtype):173 self.max_seq_len_cached = seq_len174 t = torch.arange(175 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype176 )177 t = t / self.scaling_factor178 179 freqs = torch.outer(t, self.inv_freq)180 # Different from paper, but it uses a different permutation in order to obtain the same calculation181 emb = torch.cat((freqs, freqs), dim=-1)182 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)183 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)184 185 186# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3187class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):188 """DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""189 190 def __init__(191 self,192 dim,193 max_position_embeddings=2048,194 base=10000,195 device=None,196 scaling_factor=1.0,197 ):198 self.scaling_factor = scaling_factor199 super().__init__(dim, max_position_embeddings, base, device)200 201 def _set_cos_sin_cache(self, seq_len, device, dtype):202 self.max_seq_len_cached = seq_len203 204 if seq_len > self.max_position_embeddings:205 base = self.base * (206 (self.scaling_factor * seq_len / self.max_position_embeddings)207 - (self.scaling_factor - 1)208 ) ** (self.dim / (self.dim - 2))209 inv_freq = 1.0 / (210 base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)211 )212 self.register_buffer("inv_freq", inv_freq, persistent=False)213 214 t = torch.arange(215 self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype216 )217 218 freqs = torch.outer(t, self.inv_freq)219 # Different from paper, but it uses a different permutation in order to obtain the same calculation220 emb = torch.cat((freqs, freqs), dim=-1)221 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)222 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)223 224 225# Inverse dim formula to find dim based on number of rotations226def yarn_find_correction_dim(227 num_rotations, dim, base=10000, max_position_embeddings=2048228):229 return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (230 2 * math.log(base)231 )232 233 234# Find dim range bounds based on rotations235def yarn_find_correction_range(236 low_rot, high_rot, dim, base=10000, max_position_embeddings=2048237):238 low = math.floor(239 yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)240 )241 high = math.ceil(242 yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)243 )244 return max(low, 0), min(high, dim - 1) # Clamp values just in case245 246 247def yarn_get_mscale(scale=1, mscale=1):248 if scale <= 1:249 return 1.0250 return 0.1 * mscale * math.log(scale) + 1.0251 252 253def yarn_linear_ramp_mask(min, max, dim):254 if min == max:255 max += 0.001 # Prevent singularity256 257 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)258 ramp_func = torch.clamp(linear_func, 0, 1)259 return ramp_func260 261 262class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):263 264 def __init__(265 self,266 dim,267 max_position_embeddings=2048,268 base=10000,269 device=None,270 scaling_factor=1.0,271 original_max_position_embeddings=4096,272 beta_fast=32,273 beta_slow=1,274 mscale=1,275 mscale_all_dim=0,276 ):277 self.scaling_factor = scaling_factor278 self.original_max_position_embeddings = original_max_position_embeddings279 self.beta_fast = beta_fast280 self.beta_slow = beta_slow281 self.mscale = mscale282 self.mscale_all_dim = mscale_all_dim283 super().__init__(dim, max_position_embeddings, base, device)284 285 def _set_cos_sin_cache(self, seq_len, device, dtype):286 self.max_seq_len_cached = seq_len287 dim = self.dim288 289 freq_extra = 1.0 / (290 self.base291 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)292 )293 freq_inter = 1.0 / (294 self.scaling_factor295 * self.base296 ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)297 )298 299 low, high = yarn_find_correction_range(300 self.beta_fast,301 self.beta_slow,302 dim,303 self.base,304 self.original_max_position_embeddings,305 )306 inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(307 device=device, dtype=torch.float32308 )309 inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask310 self.register_buffer("inv_freq", inv_freq, persistent=False)311 312 t = torch.arange(seq_len, device=device, dtype=torch.float32)313 314 freqs = torch.outer(t, inv_freq)315 316 _mscale = float(317 yarn_get_mscale(self.scaling_factor, self.mscale)318 / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim)319 )320 321 emb = torch.cat((freqs, freqs), dim=-1)322 self.register_buffer(323 "cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False324 )325 self.register_buffer(326 "sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False327 )328 329 330# Copied from transformers.models.llama.modeling_llama.rotate_half331def rotate_half(x):332 """Rotates half the hidden dims of the input."""333 x1 = x[..., : x.shape[-1] // 2]334 x2 = x[..., x.shape[-1] // 2 :]335 return torch.cat((-x2, x1), dim=-1)336 337 338# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb339def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):340 """Applies Rotary Position Embedding to the query and key tensors.341 342 Args:343 q (`torch.Tensor`): The query tensor.344 k (`torch.Tensor`): The key tensor.345 cos (`torch.Tensor`): The cosine part of the rotary embedding.346 sin (`torch.Tensor`): The sine part of the rotary embedding.347 position_ids (`torch.Tensor`):348 The position indices of the tokens corresponding to the query and key tensors. For example, this can be349 used to pass offsetted position ids when working with a KV-cache.350 unsqueeze_dim (`int`, *optional*, defaults to 1):351 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and352 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note353 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and354 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes355 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have356 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.357 Returns:358 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.359 """360 cos = cos[position_ids].unsqueeze(unsqueeze_dim)361 sin = sin[position_ids].unsqueeze(unsqueeze_dim)362 363 b, h, s, d = q.shape364 q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)365 366 b, h, s, d = k.shape367 k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)368 369 q_embed = (q * cos) + (rotate_half(q) * sin)370 k_embed = (k * cos) + (rotate_half(k) * sin)371 return q_embed, k_embed372 373 374class DeepseekV3MLP(nn.Module):375 def __init__(self, config, hidden_size=None, intermediate_size=None):376 super().__init__()377 self.config = config378 self.hidden_size = config.hidden_size if hidden_size is None else hidden_size379 self.intermediate_size = (380 config.intermediate_size if intermediate_size is None else intermediate_size381 )382 383 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)384 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)385 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)386 self.act_fn = ACT2FN[config.hidden_act]387 388 def forward(self, x):389 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))390 return down_proj391 392 393class MoEGate(nn.Module):394 def __init__(self, config):395 super().__init__()396 self.config = config397 self.top_k = config.num_experts_per_tok398 self.n_routed_experts = config.n_routed_experts399 self.routed_scaling_factor = config.routed_scaling_factor400 self.scoring_func = config.scoring_func401 self.seq_aux = config.seq_aux402 self.topk_method = config.topk_method403 self.n_group = config.n_group404 self.topk_group = config.topk_group405 406 # topk selection algorithm407 self.norm_topk_prob = config.norm_topk_prob408 self.gating_dim = config.hidden_size409 self.weight = nn.Parameter(410 torch.empty((self.n_routed_experts, self.gating_dim))411 )412 if self.topk_method == "noaux_tc":413 self.e_score_correction_bias = nn.Parameter(414 torch.empty((self.n_routed_experts))415 )416 self.reset_parameters()417 418 def reset_parameters(self) -> None:419 import torch.nn.init as init420 421 init.kaiming_uniform_(self.weight, a=math.sqrt(5))422 423 def forward(self, hidden_states):424 bsz, seq_len, h = hidden_states.shape425 ### compute gating score426 hidden_states = hidden_states.view(-1, h)427 logits = F.linear(428 hidden_states.type(torch.float32), self.weight.type(torch.float32), None429 )430 if self.scoring_func == "sigmoid":431 scores = logits.sigmoid()432 else:433 raise NotImplementedError(434 f"insupportable scoring function for MoE gating: {self.scoring_func}"435 )436 437 ### select top-k experts438 if self.topk_method == "noaux_tc":439 assert not self.training440 scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)441 group_scores = (442 scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim = -1)443 ) # [n, n_group]444 group_idx = torch.topk(445 group_scores, k=self.topk_group, dim=-1, sorted=False446 )[447 1448 ] # [n, top_k_group]449 group_mask = torch.zeros_like(group_scores) # [n, n_group]450 group_mask.scatter_(1, group_idx, 1) # [n, n_group]451 score_mask = (452 group_mask.unsqueeze(-1)453 .expand(454 bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group455 )456 .reshape(bsz * seq_len, -1)457 ) # [n, e]458 tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0) # [n, e]459 _, topk_idx = torch.topk(460 tmp_scores, k=self.top_k, dim=-1, sorted=False461 )462 topk_weight = scores.gather(1, topk_idx)463 else:464 raise NotImplementedError(465 f"insupportable TopK function for MoE gating: {self.topk_method}"466 )467 468 ### norm gate to sum 1469 if self.top_k > 1 and self.norm_topk_prob:470 denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20471 topk_weight = topk_weight / denominator472 topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor473 474 return topk_idx, topk_weight475 476class DeepseekV3MoE(nn.Module):477 """478 A mixed expert module containing shared experts.479 """480 481 def __init__(self, config):482 super().__init__()483 self.config = config484 self.num_experts_per_tok = config.num_experts_per_tok485 486 if hasattr(config, "ep_size") and config.ep_size > 1:487 assert config.ep_size == dist.get_world_size()488 self.ep_size = config.ep_size489 self.experts_per_rank = config.n_routed_experts // config.ep_size490 self.ep_rank = dist.get_rank()491 self.experts = nn.ModuleList(492 [493 (494 DeepseekV3MLP(495 config, intermediate_size=config.moe_intermediate_size496 )497 if i >= self.ep_rank * self.experts_per_rank498 and i < (self.ep_rank + 1) * self.experts_per_rank499 else None500 )501 for i in range(config.n_routed_experts)502 ]503 )504 else:505 self.ep_size = 1506 self.experts_per_rank = config.n_routed_experts507 self.ep_rank = 0508 self.experts = nn.ModuleList(509 [510 DeepseekV3MLP(511 config, intermediate_size=config.moe_intermediate_size512 )513 for i in range(config.n_routed_experts)514 ]515 )516 self.gate = MoEGate(config)517 if config.n_shared_experts is not None:518 intermediate_size = config.moe_intermediate_size * config.n_shared_experts519 self.shared_experts = DeepseekV3MLP(520 config=config, intermediate_size=intermediate_size521 )522 523 def forward(self, hidden_states):524 identity = hidden_states525 orig_shape = hidden_states.shape526 topk_idx, topk_weight = self.gate(hidden_states)527 hidden_states = hidden_states.view(-1, hidden_states.shape[-1])528 flat_topk_idx = topk_idx.view(-1)529 if not self.training:530 y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(*orig_shape)531 if self.config.n_shared_experts is not None:532 y = y + self.shared_experts(identity)533 return y534 535 @torch.no_grad()536 def moe_infer(self, x, topk_ids, topk_weight):537 cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))538 cnts.scatter_(1, topk_ids, 1)539 tokens_per_expert = cnts.sum(dim=0)540 idxs = topk_ids.view(-1).argsort()541 sorted_tokens = x[idxs // topk_ids.shape[1]]542 sorted_tokens_shape = sorted_tokens.shape543 if self.ep_size > 1:544 tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1)545 tokens_per_expert_group = tokens_per_expert.new_empty(546 tokens_per_expert.shape[0]547 )548 dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)549 output_splits = (550 tokens_per_expert_group.view(self.ep_size, -1)551 .sum(1)552 .cpu()553 .numpy()554 .tolist()555 )556 gathered_tokens = sorted_tokens.new_empty(557 tokens_per_expert_group.sum(dim=0).cpu().item(), sorted_tokens.shape[1]558 )559 input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()560 dist.all_to_all(561 list(gathered_tokens.split(output_splits)),562 list(sorted_tokens.split(input_split_sizes)),563 )564 tokens_per_expert_post_gather = tokens_per_expert_group.view(565 self.ep_size, self.experts_per_rank566 ).sum(dim=0)567 gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32)568 s = 0569 for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):570 gatherd_idxs[s : s + k] = i % self.experts_per_rank571 s += k572 gatherd_idxs = gatherd_idxs.argsort()573 sorted_tokens = gathered_tokens[gatherd_idxs]574 tokens_per_expert = tokens_per_expert_post_gather575 tokens_per_expert = tokens_per_expert.cpu().numpy()576 577 outputs = []578 start_idx = 0579 for i, num_tokens in enumerate(tokens_per_expert):580 end_idx = start_idx + num_tokens581 if num_tokens == 0:582 continue583 expert = self.experts[i + self.ep_rank * self.experts_per_rank]584 tokens_for_this_expert = sorted_tokens[start_idx:end_idx]585 expert_out = expert(tokens_for_this_expert)586 outputs.append(expert_out)587 start_idx = end_idx588 589 outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)590 if self.ep_size > 1:591 new_x = torch.empty_like(outs)592 new_x[gatherd_idxs] = outs593 gathered_tokens = new_x.new_empty(*sorted_tokens_shape)594 dist.all_to_all(595 list(gathered_tokens.split(input_split_sizes)),596 list(new_x.split(output_splits)),597 )598 outs = gathered_tokens599 600 new_x = torch.empty_like(outs)601 new_x[idxs] = outs602 final_out = (603 new_x.view(*topk_ids.shape, -1)604 .type(topk_weight.dtype)605 .mul_(topk_weight.unsqueeze(dim=-1))606 .sum(dim=1)607 .type(new_x.dtype)608 )609 return final_out610 611 612# Copied from transformers.models.llama.modeling_llama.repeat_kv613def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:614 """615 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,616 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)617 """618 batch, num_key_value_heads, slen, head_dim = hidden_states.shape619 if n_rep == 1:620 return hidden_states621 hidden_states = hidden_states[:, :, None, :, :].expand(622 batch, num_key_value_heads, n_rep, slen, head_dim623 )624 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)625 626 627# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3628class DeepseekV3Attention(nn.Module):629 """Multi-headed attention from 'Attention Is All You Need' paper"""630 631 def __init__(self, config: DeepseekV3Config, layer_idx: Optional[int] = None):632 super().__init__()633 self.config = config634 self.layer_idx = layer_idx635 if layer_idx is None:636 logger.warning_once(637 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "638 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "639 "when creating this class."640 )641 642 self.attention_dropout = config.attention_dropout643 self.hidden_size = config.hidden_size644 self.num_heads = config.num_attention_heads645 646 self.max_position_embeddings = config.max_position_embeddings647 self.rope_theta = config.rope_theta648 self.q_lora_rank = config.q_lora_rank649 self.qk_rope_head_dim = config.qk_rope_head_dim650 self.kv_lora_rank = config.kv_lora_rank651 self.v_head_dim = config.v_head_dim652 self.qk_nope_head_dim = config.qk_nope_head_dim653 self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim654 655 self.is_causal = True656 657 if self.q_lora_rank is None:658 self.q_proj = nn.Linear(659 self.hidden_size, self.num_heads * self.q_head_dim, bias=False660 )661 else:662 self.q_a_proj = nn.Linear(663 self.hidden_size, config.q_lora_rank, bias=config.attention_bias664 )665 self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)666 self.q_b_proj = nn.Linear(667 config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False668 )669 670 self.kv_a_proj_with_mqa = nn.Linear(671 self.hidden_size,672 config.kv_lora_rank + config.qk_rope_head_dim,673 bias=config.attention_bias,674 )675 self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)676 self.kv_b_proj = nn.Linear(677 config.kv_lora_rank,678 self.num_heads679 * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),680 bias=False,681 )682 683 self.o_proj = nn.Linear(684 self.num_heads * self.v_head_dim,685 self.hidden_size,686 bias=config.attention_bias,687 )688 self._init_rope()689 690 self.softmax_scale = self.q_head_dim ** (-0.5)691 if self.config.rope_scaling is not None:692 mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)693 scaling_factor = self.config.rope_scaling["factor"]694 if mscale_all_dim:695 mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)696 self.softmax_scale = self.softmax_scale * mscale * mscale697 698 def _init_rope(self):699 if self.config.rope_scaling is None:700 self.rotary_emb = DeepseekV3RotaryEmbedding(701 self.qk_rope_head_dim,702 max_position_embeddings=self.max_position_embeddings,703 base=self.rope_theta,704 )705 else:706 scaling_type = self.config.rope_scaling["type"]707 scaling_factor = self.config.rope_scaling["factor"]708 if scaling_type == "linear":709 self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding(710 self.qk_rope_head_dim,711 max_position_embeddings=self.max_position_embeddings,712 scaling_factor=scaling_factor,713 base=self.rope_theta,714 )715 elif scaling_type == "dynamic":716 self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding(717 self.qk_rope_head_dim,718 max_position_embeddings=self.max_position_embeddings,719 scaling_factor=scaling_factor,720 base=self.rope_theta,721 )722 elif scaling_type == "yarn":723 kwargs = {724 key: self.config.rope_scaling[key]725 for key in [726 "original_max_position_embeddings",727 "beta_fast",728 "beta_slow",729 "mscale",730 "mscale_all_dim",731 ]732 if key in self.config.rope_scaling733 }734 self.rotary_emb = DeepseekV3YarnRotaryEmbedding(735 self.qk_rope_head_dim,736 max_position_embeddings=self.max_position_embeddings,737 scaling_factor=scaling_factor,738 base=self.rope_theta,739 **kwargs,740 )741 else:742 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")743 744 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):745 return (746 tensor.view(bsz, seq_len, self.num_heads, self.v_head_dim)747 .transpose(1, 2)748 .contiguous()749 )750 751 def forward(752 self,753 hidden_states: torch.Tensor,754 attention_mask: Optional[torch.Tensor] = None,755 position_ids: Optional[torch.LongTensor] = None,756 past_key_value: Optional[Cache] = None,757 output_attentions: bool = False,758 use_cache: bool = False,759 **kwargs,760 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:761 if "padding_mask" in kwargs:762 warnings.warn(763 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"764 )765 bsz, q_len, _ = hidden_states.size()766 767 if self.q_lora_rank is None:768 q = self.q_proj(hidden_states)769 else:770 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))771 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)772 q_nope, q_pe = torch.split(773 q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1774 )775 776 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)777 compressed_kv, k_pe = torch.split(778 compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1779 )780 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)781 kv = (782 self.kv_b_proj(self.kv_a_layernorm(compressed_kv))783 .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)784 .transpose(1, 2)785 )786 787 k_nope, value_states = torch.split(788 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1789 )790 kv_seq_len = value_states.shape[-2]791 if past_key_value is not None:792 if self.layer_idx is None:793 raise ValueError(794 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "795 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "796 "with a layer index."797 )798 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)799 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)800 801 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)802 803 query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)804 query_states[:, :, :, : self.qk_nope_head_dim] = q_nope805 query_states[:, :, :, self.qk_nope_head_dim :] = q_pe806 807 key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)808 key_states[:, :, :, : self.qk_nope_head_dim] = k_nope809 key_states[:, :, :, self.qk_nope_head_dim :] = k_pe810 if past_key_value is not None:811 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models812 key_states, value_states = past_key_value.update(813 key_states, value_states, self.layer_idx, cache_kwargs814 )815 816 attn_weights = (817 torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale818 )819 820 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):821 raise ValueError(822 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"823 f" {attn_weights.size()}"824 )825 assert attention_mask is not None826 if attention_mask is not None:827 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):828 raise ValueError(829 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"830 )831 attn_weights = attn_weights + attention_mask832 833 # upcast attention to fp32834 attn_weights = nn.functional.softmax(835 attn_weights, dim=-1, dtype=torch.float32836 ).to(query_states.dtype)837 attn_weights = nn.functional.dropout(838 attn_weights, p=self.attention_dropout, training=self.training839 )840 attn_output = torch.matmul(attn_weights, value_states)841 842 if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim):843 raise ValueError(844 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is"845 f" {attn_output.size()}"846 )847 848 attn_output = attn_output.transpose(1, 2).contiguous()849 850 attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim)851 852 attn_output = self.o_proj(attn_output)853 854 if not output_attentions:855 attn_weights = None856 857 return attn_output, attn_weights, past_key_value858 859 860# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV3861class DeepseekV3FlashAttention2(DeepseekV3Attention):862 """863 DeepseekV3 flash attention module. This module inherits from `DeepseekV3Attention` as the weights of the module stays864 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of865 flash attention and deal with padding tokens in case the input contains any of them.866 """867 868 def __init__(self, *args, **kwargs):869 super().__init__(*args, **kwargs)870 871 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.872 # 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.873 # 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).874 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()875 876 def forward(877 self,878 hidden_states: torch.Tensor,879 attention_mask: Optional[torch.LongTensor] = None,880 position_ids: Optional[torch.LongTensor] = None,881 past_key_value: Optional[Cache] = None,882 output_attentions: bool = False,883 use_cache: bool = False,884 **kwargs,885 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:886 # DeepseekV3FlashAttention2 attention does not support output_attentions887 if "padding_mask" in kwargs:888 warnings.warn(889 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"890 )891 892 # overwrite attention_mask with padding_mask893 attention_mask = kwargs.pop("padding_mask")894 895 output_attentions = False896 897 bsz, q_len, _ = hidden_states.size()898 899 if self.q_lora_rank is None:900 q = self.q_proj(hidden_states)901 else:902 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))903 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)904 q_nope, q_pe = torch.split(905 q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1906 )907 908 # Flash attention requires the input to have the shape909 # batch_size x seq_length x head_dim x hidden_dim910 # therefore we just need to keep the original shape911 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)912 compressed_kv, k_pe = torch.split(913 compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1914 )915 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)916 kv = (917 self.kv_b_proj(self.kv_a_layernorm(compressed_kv))918 .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)919 .transpose(1, 2)920 )921 922 k_nope, value_states = torch.split(923 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1924 )925 kv_seq_len = value_states.shape[-2]926 927 kv_seq_len = value_states.shape[-2]928 if past_key_value is not None:929 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)930 931 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)932 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)933 934 query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)935 query_states[:, :, :, : self.qk_nope_head_dim] = q_nope936 query_states[:, :, :, self.qk_nope_head_dim :] = q_pe937 938 key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)939 key_states[:, :, :, : self.qk_nope_head_dim] = k_nope940 key_states[:, :, :, self.qk_nope_head_dim :] = k_pe941 942 if self.q_head_dim != self.v_head_dim:943 value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim])944 945 if past_key_value is not None:946 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models947 key_states, value_states = past_key_value.update(948 key_states, value_states, self.layer_idx, cache_kwargs949 )950 951 # 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 cache952 # to be able to avoid many of these transpose/reshape/view.953 query_states = query_states.transpose(1, 2)954 key_states = key_states.transpose(1, 2)955 value_states = value_states.transpose(1, 2)956 957 dropout_rate = self.attention_dropout if self.training else 0.0958 959 # In PEFT, usually we cast the layer norms in float32 for training stability reasons960 # therefore the input hidden states gets silently casted in float32. Hence, we need961 # cast them back in the correct dtype just to be sure everything works as expected.962 # This might slowdown training & inference so it is recommended to not cast the LayerNorms963 # in fp32. (DeepseekV3RMSNorm handles it correctly)964 965 input_dtype = query_states.dtype966 if input_dtype == torch.float32:967 # Handle the case where the model is quantized968 if hasattr(self.config, "_pre_quantization_dtype"):969 target_dtype = self.config._pre_quantization_dtype970 elif torch.is_autocast_enabled():971 target_dtype = torch.get_autocast_gpu_dtype()972 else:973 target_dtype = (974 self.q_proj.weight.dtype975 if self.q_lora_rank is None976 else self.q_a_proj.weight.dtype977 )978 979 logger.warning_once(980 f"The input hidden states seems to be silently casted in float32, this might be related to"981 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"982 f" {target_dtype}."983 )984 985 query_states = query_states.to(target_dtype)986 key_states = key_states.to(target_dtype)987 value_states = value_states.to(target_dtype)988 989 attn_output = self._flash_attention_forward(990 query_states,991 key_states,992 value_states,993 attention_mask,994 q_len,995 dropout=dropout_rate,996 softmax_scale=self.softmax_scale,997 )998 if self.q_head_dim != self.v_head_dim:999 attn_output = attn_output[:, :, :, : self.v_head_dim]1000 1001 attn_output = attn_output.reshape(1002 bsz, q_len, self.num_heads * self.v_head_dim1003 ).contiguous()1004 attn_output = self.o_proj(attn_output)1005 1006 if not output_attentions:1007 attn_weights = None1008 1009 return attn_output, attn_weights, past_key_value1010 1011 def _flash_attention_forward(1012 self,1013 query_states,1014 key_states,1015 value_states,1016 attention_mask,1017 query_length,1018 dropout=0.0,1019 softmax_scale=None,1020 ):1021 """1022 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1023 first unpad the input, then computes the attention scores and pad the final attention scores.1024 1025 Args:1026 query_states (`torch.Tensor`):1027 Input query states to be passed to Flash Attention API1028 key_states (`torch.Tensor`):1029 Input key states to be passed to Flash Attention API1030 value_states (`torch.Tensor`):1031 Input value states to be passed to Flash Attention API1032 attention_mask (`torch.Tensor`):1033 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1034 position of padding tokens and 1 for the position of non-padding tokens.1035 dropout (`int`, *optional*):1036 Attention dropout1037 softmax_scale (`float`, *optional*):1038 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1039 """1040 if not self._flash_attn_uses_top_left_mask:1041 causal = self.is_causal1042 else:1043 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in DeepseekV3FlashAttention2 __init__.1044 causal = self.is_causal and query_length != 11045 1046 # Contains at least one padding token in the sequence1047 if attention_mask is not None:1048 batch_size = query_states.shape[0]1049 (1050 query_states,1051 key_states,1052 value_states,1053 indices_q,1054 cu_seq_lens,1055 max_seq_lens,1056 ) = self._upad_input(1057 query_states, key_states, value_states, attention_mask, query_length1058 )1059 1060 cu_seqlens_q, cu_seqlens_k = cu_seq_lens1061 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1062 1063 attn_output_unpad = flash_attn_varlen_func(1064 query_states,1065 key_states,1066 value_states,1067 cu_seqlens_q=cu_seqlens_q,1068 cu_seqlens_k=cu_seqlens_k,1069 max_seqlen_q=max_seqlen_in_batch_q,1070 max_seqlen_k=max_seqlen_in_batch_k,1071 dropout_p=dropout,1072 softmax_scale=softmax_scale,1073 causal=causal,1074 )1075 1076 attn_output = pad_input(1077 attn_output_unpad, indices_q, batch_size, query_length1078 )1079 else:1080 attn_output = flash_attn_func(1081 query_states,1082 key_states,1083 value_states,1084 dropout,1085 softmax_scale=softmax_scale,1086 causal=causal,1087 )1088 1089 return attn_output1090 1091 def _upad_input(1092 self, query_layer, key_layer, value_layer, attention_mask, query_length1093 ):1094 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)1095 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape1096 1097 key_layer = index_first_axis(1098 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),1099 indices_k,1100 )1101 value_layer = index_first_axis(1102 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),1103 indices_k,1104 )1105 if query_length == kv_seq_len:1106 query_layer = index_first_axis(1107 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim),1108 indices_k,1109 )1110 cu_seqlens_q = cu_seqlens_k1111 max_seqlen_in_batch_q = max_seqlen_in_batch_k1112 indices_q = indices_k1113 elif query_length == 1:1114 max_seqlen_in_batch_q = 11115 cu_seqlens_q = torch.arange(1116 batch_size + 1, dtype=torch.int32, device=query_layer.device1117 ) # There is a memcpy here, that is very bad.1118 indices_q = cu_seqlens_q[:-1]1119 query_layer = query_layer.squeeze(1)1120 else:1121 # The -q_len: slice assumes left padding.1122 attention_mask = attention_mask[:, -query_length:]1123 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(1124 query_layer, attention_mask1125 )1126 1127 return (1128 query_layer,1129 key_layer,1130 value_layer,1131 indices_q,1132 (cu_seqlens_q, cu_seqlens_k),1133 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),1134 )1135 1136 1137ATTENTION_CLASSES = {1138 "eager": DeepseekV3Attention,1139 "flash_attention_2": DeepseekV3FlashAttention2,1140}1141 1142 1143class DeepseekV3DecoderLayer(nn.Module):1144 def __init__(self, config: DeepseekV3Config, layer_idx: int):1145 super().__init__()1146 self.hidden_size = config.hidden_size1147 1148 self.self_attn = ATTENTION_CLASSES[config._attn_implementation](1149 config=config, layer_idx=layer_idx1150 )1151 1152 self.mlp = (1153 DeepseekV3MoE(config)1154 if (1155 config.n_routed_experts is not None1156 and layer_idx >= config.first_k_dense_replace1157 and layer_idx % config.moe_layer_freq == 01158 )1159 else DeepseekV3MLP(config)1160 )1161 self.input_layernorm = DeepseekV3RMSNorm(1162 config.hidden_size, eps=config.rms_norm_eps1163 )1164 self.post_attention_layernorm = DeepseekV3RMSNorm(1165 config.hidden_size, eps=config.rms_norm_eps1166 )1167 1168 def forward(1169 self,1170 hidden_states: torch.Tensor,1171 attention_mask: Optional[torch.Tensor] = None,1172 position_ids: Optional[torch.LongTensor] = None,1173 past_key_value: Optional[Tuple[torch.Tensor]] = None,1174 output_attentions: Optional[bool] = False,1175 use_cache: Optional[bool] = False,1176 **kwargs,1177 ) -> Tuple[1178 torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]1179 ]:1180 """1181 Args:1182 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`1183 attention_mask (`torch.FloatTensor`, *optional*):1184 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,1185 query_sequence_length, key_sequence_length)` if default attention is used.1186 output_attentions (`bool`, *optional*):1187 Whether or not to return the attentions tensors of all attention layers. See `attentions` under1188 returned tensors for more detail.1189 use_cache (`bool`, *optional*):1190 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding1191 (see `past_key_values`).1192 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states1193 """1194 if "padding_mask" in kwargs:1195 warnings.warn(1196 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"1197 )1198 residual = hidden_states1199 1200 hidden_states = self.input_layernorm(hidden_states)