unsloth/Kimi-K2.7-Code
21141
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 numpy as np26import torch27import torch.distributed as dist28import torch.nn.functional as F29import torch.utils.checkpoint30from torch import nn31from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss32from transformers.activations import ACT2FN33from transformers.cache_utils import Cache, DynamicCache34from transformers.modeling_attn_mask_utils import \35 _prepare_4d_causal_attention_mask36from transformers.modeling_outputs import (BaseModelOutputWithPast,37 CausalLMOutputWithPast,38 SequenceClassifierOutputWithPast)39from transformers.modeling_utils import PreTrainedModel40from transformers.pytorch_utils import (ALL_LAYERNORM_LAYERS,41 is_torch_greater_or_equal_than_1_13)42from transformers.utils import (add_start_docstrings,43 add_start_docstrings_to_model_forward,44 is_flash_attn_2_available,45 is_flash_attn_greater_or_equal_2_10, logging,46 replace_return_docstrings)47 48try:49 from transformers.utils.import_utils import is_torch_fx_available50except ImportError:51 52 def is_torch_fx_available() -> bool:53 return hasattr(torch, "fx")54 55 56from .configuration_deepseek import DeepseekV3Config57 58if is_flash_attn_2_available():59 from flash_attn import flash_attn_func, flash_attn_varlen_func60 from flash_attn.bert_padding import pad_input # noqa61 from flash_attn.bert_padding import index_first_axis, unpad_input62 63# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.64# It means that the function will not be traced through and simply appear as a node in the graph.65if is_torch_fx_available():66 if not is_torch_greater_or_equal_than_1_13:67 import torch.fx68 69 _prepare_4d_causal_attention_mask = torch.fx.wrap(70 _prepare_4d_causal_attention_mask)71 72logger = logging.get_logger(__name__)73 74_CONFIG_FOR_DOC = "DeepseekV3Config"75 76 77def _get_unpad_data(attention_mask):78 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)79 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()80 max_seqlen_in_batch = seqlens_in_batch.max().item()81 cu_seqlens = F.pad(82 torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))83 return (84 indices,85 cu_seqlens,86 max_seqlen_in_batch,87 )88 89 90# code modified from transformers 4.48.3 to amend breaks in newer transformers versions91def get_usable_length(past_key_value,92 new_seq_length: int,93 layer_idx: Optional[int] = 0) -> int:94 max_length = past_key_value.get_max_cache_shape()95 previous_seq_length = past_key_value.get_seq_length(layer_idx)96 if max_length is not None and max_length > 0 and previous_seq_length + new_seq_length > max_length:97 return max_length - new_seq_length98 return previous_seq_length99 100 101class DeepseekV3RMSNorm(nn.Module):102 103 def __init__(self, hidden_size, eps=1e-6):104 """105 DeepseekV3RMSNorm is equivalent to T5LayerNorm106 """107 super().__init__()108 self.weight = nn.Parameter(torch.ones(hidden_size))109 self.variance_epsilon = eps110 111 def forward(self, hidden_states):112 input_dtype = hidden_states.dtype113 hidden_states = hidden_states.to(torch.float32)114 variance = hidden_states.pow(2).mean(-1, keepdim=True)115 hidden_states = hidden_states * torch.rsqrt(variance +116 self.variance_epsilon)117 return self.weight * hidden_states.to(input_dtype)118 119 120ALL_LAYERNORM_LAYERS.append(DeepseekV3RMSNorm)121 122 123class DeepseekV3RotaryEmbedding(nn.Module):124 125 def __init__(self,126 dim,127 max_position_embeddings=2048,128 base=10000,129 device=None):130 super().__init__()131 132 self.dim = dim133 self.max_position_embeddings = max_position_embeddings134 self.base = base135 inv_freq = 1.0 / (self.base**(136 torch.arange(0, self.dim, 2).float().to(device) / self.dim))137 self.register_buffer("inv_freq", inv_freq, persistent=False)138 139 # Build here to make `torch.jit.trace` work.140 self._set_cos_sin_cache(141 seq_len=max_position_embeddings,142 device=self.inv_freq.device,143 dtype=torch.get_default_dtype(),144 )145 self.max_seq_len_cached = None146 147 def _set_cos_sin_cache(self, seq_len, device, dtype):148 self.max_seq_len_cached = seq_len149 t = torch.arange(self.max_seq_len_cached,150 device=device,151 dtype=self.inv_freq.dtype)152 153 freqs = torch.outer(t, self.inv_freq.to(t.device))154 # Different from paper, but it uses a different permutation in order to obtain the same calculation155 emb = torch.cat((freqs, freqs), dim=-1)156 self.register_buffer("cos_cached",157 emb.cos().to(dtype),158 persistent=False)159 self.register_buffer("sin_cached",160 emb.sin().to(dtype),161 persistent=False)162 163 def forward(self, x, seq_len=None):164 # x: [bs, num_attention_heads, seq_len, head_size]165 if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:166 self._set_cos_sin_cache(seq_len=seq_len,167 device=x.device,168 dtype=x.dtype)169 170 return (171 self.cos_cached[:seq_len].to(dtype=x.dtype),172 self.sin_cached[:seq_len].to(dtype=x.dtype),173 )174 175 176# Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->DeepseekV3177class DeepseekV3LinearScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):178 """DeepseekV3RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""179 180 def __init__(181 self,182 dim,183 max_position_embeddings=2048,184 base=10000,185 device=None,186 scaling_factor=1.0,187 ):188 self.scaling_factor = scaling_factor189 super().__init__(dim, max_position_embeddings, base, device)190 191 def _set_cos_sin_cache(self, seq_len, device, dtype):192 self.max_seq_len_cached = seq_len193 t = torch.arange(self.max_seq_len_cached,194 device=device,195 dtype=self.inv_freq.dtype)196 t = t / self.scaling_factor197 198 freqs = torch.outer(t, self.inv_freq)199 # Different from paper, but it uses a different permutation in order to obtain the same calculation200 emb = torch.cat((freqs, freqs), dim=-1)201 self.register_buffer("cos_cached",202 emb.cos().to(dtype),203 persistent=False)204 self.register_buffer("sin_cached",205 emb.sin().to(dtype),206 persistent=False)207 208 209# Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->DeepseekV3210class DeepseekV3DynamicNTKScalingRotaryEmbedding(DeepseekV3RotaryEmbedding):211 """DeepseekV3RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""212 213 def __init__(214 self,215 dim,216 max_position_embeddings=2048,217 base=10000,218 device=None,219 scaling_factor=1.0,220 ):221 self.scaling_factor = scaling_factor222 super().__init__(dim, max_position_embeddings, base, device)223 224 def _set_cos_sin_cache(self, seq_len, device, dtype):225 self.max_seq_len_cached = seq_len226 227 if seq_len > self.max_position_embeddings:228 base = self.base * ((self.scaling_factor * seq_len /229 self.max_position_embeddings) -230 (self.scaling_factor - 1))**(self.dim /231 (self.dim - 2))232 inv_freq = 1.0 / (base**(233 torch.arange(0, self.dim, 2).float().to(device) / self.dim))234 self.register_buffer("inv_freq", inv_freq, persistent=False)235 236 t = torch.arange(self.max_seq_len_cached,237 device=device,238 dtype=self.inv_freq.dtype)239 240 freqs = torch.outer(t, self.inv_freq)241 # Different from paper, but it uses a different permutation in order to obtain the same calculation242 emb = torch.cat((freqs, freqs), dim=-1)243 self.register_buffer("cos_cached",244 emb.cos().to(dtype),245 persistent=False)246 self.register_buffer("sin_cached",247 emb.sin().to(dtype),248 persistent=False)249 250 251# Inverse dim formula to find dim based on number of rotations252def yarn_find_correction_dim(num_rotations,253 dim,254 base=10000,255 max_position_embeddings=2048):256 return (dim * math.log(max_position_embeddings /257 (num_rotations * 2 * math.pi))) / (2 *258 math.log(base))259 260 261# Find dim range bounds based on rotations262def yarn_find_correction_range(low_rot,263 high_rot,264 dim,265 base=10000,266 max_position_embeddings=2048):267 low = math.floor(268 yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))269 high = math.ceil(270 yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings))271 return max(low, 0), min(high, dim - 1) # Clamp values just in case272 273 274def yarn_get_mscale(scale=1, mscale=1):275 if scale <= 1:276 return 1.0277 return 0.1 * mscale * math.log(scale) + 1.0278 279 280def yarn_linear_ramp_mask(min, max, dim):281 if min == max:282 max += 0.001 # Prevent singularity283 284 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)285 ramp_func = torch.clamp(linear_func, 0, 1)286 return ramp_func287 288 289class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):290 291 def __init__(292 self,293 dim,294 max_position_embeddings=2048,295 base=10000,296 device=None,297 scaling_factor=1.0,298 original_max_position_embeddings=4096,299 beta_fast=32,300 beta_slow=1,301 mscale=1,302 mscale_all_dim=0,303 ):304 self.scaling_factor = scaling_factor305 self.original_max_position_embeddings = original_max_position_embeddings306 self.beta_fast = beta_fast307 self.beta_slow = beta_slow308 self.mscale = mscale309 self.mscale_all_dim = mscale_all_dim310 super().__init__(dim, max_position_embeddings, base, device)311 312 def _set_cos_sin_cache(self, seq_len, device, dtype):313 self.max_seq_len_cached = seq_len314 dim = self.dim315 316 freq_extra = 1.0 / (self.base**(317 torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))318 freq_inter = 1.0 / (self.scaling_factor * self.base**(319 torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))320 321 low, high = yarn_find_correction_range(322 self.beta_fast,323 self.beta_slow,324 dim,325 self.base,326 self.original_max_position_embeddings,327 )328 inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(329 device=device, dtype=torch.float32)330 inv_freq = freq_inter * (1 -331 inv_freq_mask) + freq_extra * inv_freq_mask332 self.register_buffer("inv_freq", inv_freq, persistent=False)333 334 t = torch.arange(seq_len, device=device, dtype=torch.float32)335 336 freqs = torch.outer(t, inv_freq)337 338 _mscale = float(339 yarn_get_mscale(self.scaling_factor, self.mscale) /340 yarn_get_mscale(self.scaling_factor, self.mscale_all_dim))341 342 emb = torch.cat((freqs, freqs), dim=-1)343 self.register_buffer("cos_cached", (emb.cos() * _mscale).to(dtype),344 persistent=False)345 self.register_buffer("sin_cached", (emb.sin() * _mscale).to(dtype),346 persistent=False)347 348 349# Copied from transformers.models.llama.modeling_llama.rotate_half350def rotate_half(x):351 """Rotates half the hidden dims of the input."""352 x1 = x[..., :x.shape[-1] // 2]353 x2 = x[..., x.shape[-1] // 2:]354 return torch.cat((-x2, x1), dim=-1)355 356 357# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb358def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):359 """Applies Rotary Position Embedding to the query and key tensors.360 361 Args:362 q (`torch.Tensor`): The query tensor.363 k (`torch.Tensor`): The key tensor.364 cos (`torch.Tensor`): The cosine part of the rotary embedding.365 sin (`torch.Tensor`): The sine part of the rotary embedding.366 position_ids (`torch.Tensor`):367 The position indices of the tokens corresponding to the query and key tensors. For example, this can be368 used to pass offsetted position ids when working with a KV-cache.369 unsqueeze_dim (`int`, *optional*, defaults to 1):370 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and371 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note372 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and373 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes374 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have375 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.376 Returns:377 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.378 """379 cos = cos[position_ids].unsqueeze(unsqueeze_dim)380 sin = sin[position_ids].unsqueeze(unsqueeze_dim)381 382 b, h, s, d = q.shape383 q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)384 385 b, h, s, d = k.shape386 k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)387 388 q_embed = (q * cos) + (rotate_half(q) * sin)389 k_embed = (k * cos) + (rotate_half(k) * sin)390 return q_embed, k_embed391 392 393class DeepseekV3MLP(nn.Module):394 395 def __init__(self, config, hidden_size=None, intermediate_size=None):396 super().__init__()397 self.config = config398 self.hidden_size = config.hidden_size if hidden_size is None else hidden_size399 self.intermediate_size = (config.intermediate_size if intermediate_size400 is None else intermediate_size)401 402 self.gate_proj = nn.Linear(self.hidden_size,403 self.intermediate_size,404 bias=False)405 self.up_proj = nn.Linear(self.hidden_size,406 self.intermediate_size,407 bias=False)408 self.down_proj = nn.Linear(self.intermediate_size,409 self.hidden_size,410 bias=False)411 self.act_fn = ACT2FN[config.hidden_act]412 413 def forward(self, x):414 down_proj = self.down_proj(415 self.act_fn(self.gate_proj(x)) * self.up_proj(x))416 return down_proj417 418 419class MoEGate(nn.Module):420 421 def __init__(self, config):422 super().__init__()423 self.config = config424 self.top_k = config.num_experts_per_tok425 self.n_routed_experts = config.n_routed_experts426 self.routed_scaling_factor = config.routed_scaling_factor427 self.scoring_func = config.scoring_func428 self.seq_aux = config.seq_aux429 self.topk_method = config.topk_method430 self.n_group = config.n_group431 self.topk_group = config.topk_group432 433 # topk selection algorithm434 self.norm_topk_prob = config.norm_topk_prob435 self.gating_dim = config.hidden_size436 self.weight = nn.Parameter(437 torch.empty((self.n_routed_experts, self.gating_dim)))438 if self.topk_method == "noaux_tc":439 self.e_score_correction_bias = nn.Parameter(440 torch.empty((self.n_routed_experts)))441 self.reset_parameters()442 443 def reset_parameters(self) -> None:444 import torch.nn.init as init445 446 init.kaiming_uniform_(self.weight, a=math.sqrt(5))447 448 def forward(self, hidden_states):449 bsz, seq_len, h = hidden_states.shape450 ### compute gating score451 hidden_states = hidden_states.view(-1, h)452 logits = F.linear(hidden_states.type(torch.float32),453 self.weight.type(torch.float32), None)454 if self.scoring_func == "sigmoid":455 scores = logits.sigmoid()456 else:457 raise NotImplementedError(458 f"insupportable scoring function for MoE gating: {self.scoring_func}"459 )460 461 ### select top-k experts462 if self.topk_method == "noaux_tc":463 assert not self.training464 scores_for_choice = scores.view(465 bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)466 group_scores = (scores_for_choice.view(467 bsz * seq_len, self.n_group,468 -1).topk(2, dim=-1)[0].sum(dim=-1)) # [n, n_group]469 group_idx = torch.topk(group_scores,470 k=self.topk_group,471 dim=-1,472 sorted=False)[1] # [n, top_k_group]473 group_mask = torch.zeros_like(group_scores) # [n, n_group]474 group_mask.scatter_(1, group_idx, 1) # [n, n_group]475 score_mask = (group_mask.unsqueeze(-1).expand(476 bsz * seq_len, self.n_group,477 self.n_routed_experts // self.n_group).reshape(478 bsz * seq_len, -1)) # [n, e]479 tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(),480 0.0) # [n, e]481 _, topk_idx = torch.topk(tmp_scores,482 k=self.top_k,483 dim=-1,484 sorted=False)485 topk_weight = scores.gather(1, topk_idx)486 else:487 raise NotImplementedError(488 f"insupportable TopK function for MoE gating: {self.topk_method}"489 )490 491 ### norm gate to sum 1492 if self.top_k > 1 and self.norm_topk_prob:493 denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20494 topk_weight = topk_weight / denominator495 topk_weight = topk_weight * self.routed_scaling_factor # must multiply the scaling factor496 497 return topk_idx, topk_weight498 499 500class DeepseekV3MoE(nn.Module):501 """502 A mixed expert module containing shared experts.503 """504 505 def __init__(self, config):506 super().__init__()507 self.config = config508 self.num_experts_per_tok = config.num_experts_per_tok509 510 if hasattr(config, "ep_size") and config.ep_size > 1:511 assert config.ep_size == dist.get_world_size()512 self.ep_size = config.ep_size513 self.experts_per_rank = config.n_routed_experts // config.ep_size514 self.ep_rank = dist.get_rank()515 self.experts = nn.ModuleList([516 (DeepseekV3MLP(config,517 intermediate_size=config.moe_intermediate_size)518 if i >= self.ep_rank * self.experts_per_rank519 and i < (self.ep_rank + 1) * self.experts_per_rank else None)520 for i in range(config.n_routed_experts)521 ])522 else:523 self.ep_size = 1524 self.experts_per_rank = config.n_routed_experts525 self.ep_rank = 0526 self.experts = nn.ModuleList([527 DeepseekV3MLP(config,528 intermediate_size=config.moe_intermediate_size)529 for i in range(config.n_routed_experts)530 ])531 self.gate = MoEGate(config)532 if config.n_shared_experts is not None:533 intermediate_size = config.moe_intermediate_size * config.n_shared_experts534 self.shared_experts = DeepseekV3MLP(535 config=config, intermediate_size=intermediate_size)536 537 def forward(self, hidden_states):538 identity = hidden_states539 orig_shape = hidden_states.shape540 topk_idx, topk_weight = self.gate(hidden_states)541 hidden_states = hidden_states.view(-1, hidden_states.shape[-1])542 # flat_topk_idx = topk_idx.view(-1)543 if not self.training:544 y = self.moe_infer(hidden_states, topk_idx,545 topk_weight).view(*orig_shape)546 if self.config.n_shared_experts is not None:547 y = y + self.shared_experts(identity)548 return y549 550 @torch.no_grad()551 def moe_infer(self, x, topk_ids, topk_weight):552 cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))553 cnts.scatter_(1, topk_ids, 1)554 tokens_per_expert = cnts.sum(dim=0)555 idxs = topk_ids.view(-1).argsort()556 sorted_tokens = x[idxs // topk_ids.shape[1]]557 sorted_tokens_shape = sorted_tokens.shape558 if self.ep_size > 1:559 tokens_per_ep_rank = tokens_per_expert.view(self.ep_size,560 -1).sum(dim=1)561 tokens_per_expert_group = tokens_per_expert.new_empty(562 tokens_per_expert.shape[0])563 dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert)564 output_splits = (tokens_per_expert_group.view(565 self.ep_size, -1).sum(1).cpu().numpy().tolist())566 gathered_tokens = sorted_tokens.new_empty(567 tokens_per_expert_group.sum(dim=0).cpu().item(),568 sorted_tokens.shape[1])569 input_split_sizes = tokens_per_ep_rank.cpu().numpy().tolist()570 dist.all_to_all(571 list(gathered_tokens.split(output_splits)),572 list(sorted_tokens.split(input_split_sizes)),573 )574 tokens_per_expert_post_gather = tokens_per_expert_group.view(575 self.ep_size, self.experts_per_rank).sum(dim=0)576 gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0], ),577 dtype=np.int32)578 s = 0579 for i, k in enumerate(tokens_per_expert_group.cpu().numpy()):580 gatherd_idxs[s:s + k] = i % self.experts_per_rank581 s += k582 gatherd_idxs = gatherd_idxs.argsort()583 sorted_tokens = gathered_tokens[gatherd_idxs]584 tokens_per_expert = tokens_per_expert_post_gather585 tokens_per_expert = tokens_per_expert.cpu().numpy()586 587 outputs = []588 start_idx = 0589 for i, num_tokens in enumerate(tokens_per_expert):590 end_idx = start_idx + num_tokens591 if num_tokens == 0:592 continue593 expert = self.experts[i + self.ep_rank * self.experts_per_rank]594 tokens_for_this_expert = sorted_tokens[start_idx:end_idx]595 expert_out = expert(tokens_for_this_expert)596 outputs.append(expert_out)597 start_idx = end_idx598 599 outs = torch.cat(outputs,600 dim=0) if len(outputs) else sorted_tokens.new_empty(0)601 if self.ep_size > 1:602 new_x = torch.empty_like(outs)603 new_x[gatherd_idxs] = outs604 gathered_tokens = new_x.new_empty(*sorted_tokens_shape)605 dist.all_to_all(606 list(gathered_tokens.split(input_split_sizes)),607 list(new_x.split(output_splits)),608 )609 outs = gathered_tokens610 611 new_x = torch.empty_like(outs)612 new_x[idxs] = outs613 final_out = (new_x.view(614 *topk_ids.shape, -1).type(topk_weight.dtype).mul_(615 topk_weight.unsqueeze(dim=-1)).sum(dim=1).type(new_x.dtype))616 return final_out617 618 619# Copied from transformers.models.llama.modeling_llama.repeat_kv620def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:621 """622 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,623 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)624 """625 batch, num_key_value_heads, slen, head_dim = hidden_states.shape626 if n_rep == 1:627 return hidden_states628 hidden_states = hidden_states[:, :,629 None, :, :].expand(batch,630 num_key_value_heads,631 n_rep, slen, head_dim)632 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen,633 head_dim)634 635 636# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->DeepseekV3637class DeepseekV3Attention(nn.Module):638 """Multi-headed attention from 'Attention Is All You Need' paper"""639 640 def __init__(self,641 config: DeepseekV3Config,642 layer_idx: Optional[int] = None):643 super().__init__()644 self.config = config645 self.layer_idx = layer_idx646 if layer_idx is None:647 logger.warning_once(648 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "649 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "650 "when creating this class.")651 652 self.attention_dropout = config.attention_dropout653 self.hidden_size = config.hidden_size654 self.num_heads = config.num_attention_heads655 656 self.max_position_embeddings = config.max_position_embeddings657 self.rope_theta = config.rope_theta658 self.q_lora_rank = config.q_lora_rank659 self.qk_rope_head_dim = config.qk_rope_head_dim660 self.kv_lora_rank = config.kv_lora_rank661 self.v_head_dim = config.v_head_dim662 self.qk_nope_head_dim = config.qk_nope_head_dim663 self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim664 665 self.is_causal = True666 667 if self.q_lora_rank is None:668 self.q_proj = nn.Linear(self.hidden_size,669 self.num_heads * self.q_head_dim,670 bias=False)671 else:672 self.q_a_proj = nn.Linear(self.hidden_size,673 config.q_lora_rank,674 bias=config.attention_bias)675 self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)676 self.q_b_proj = nn.Linear(config.q_lora_rank,677 self.num_heads * self.q_head_dim,678 bias=False)679 680 self.kv_a_proj_with_mqa = nn.Linear(681 self.hidden_size,682 config.kv_lora_rank + config.qk_rope_head_dim,683 bias=config.attention_bias,684 )685 self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)686 self.kv_b_proj = nn.Linear(687 config.kv_lora_rank,688 self.num_heads *689 (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),690 bias=False,691 )692 693 self.o_proj = nn.Linear(694 self.num_heads * self.v_head_dim,695 self.hidden_size,696 bias=config.attention_bias,697 )698 self._init_rope()699 700 self.softmax_scale = self.q_head_dim**(-0.5)701 if self.config.rope_scaling is not None:702 mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)703 scaling_factor = self.config.rope_scaling["factor"]704 if mscale_all_dim:705 mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)706 self.softmax_scale = self.softmax_scale * mscale * mscale707 708 def _init_rope(self):709 if self.config.rope_scaling is None:710 self.rotary_emb = DeepseekV3RotaryEmbedding(711 self.qk_rope_head_dim,712 max_position_embeddings=self.max_position_embeddings,713 base=self.rope_theta,714 )715 else:716 scaling_type = self.config.rope_scaling["type"]717 scaling_factor = self.config.rope_scaling["factor"]718 if scaling_type == "linear":719 self.rotary_emb = DeepseekV3LinearScalingRotaryEmbedding(720 self.qk_rope_head_dim,721 max_position_embeddings=self.max_position_embeddings,722 scaling_factor=scaling_factor,723 base=self.rope_theta,724 )725 elif scaling_type == "dynamic":726 self.rotary_emb = DeepseekV3DynamicNTKScalingRotaryEmbedding(727 self.qk_rope_head_dim,728 max_position_embeddings=self.max_position_embeddings,729 scaling_factor=scaling_factor,730 base=self.rope_theta,731 )732 elif scaling_type == "yarn":733 kwargs = {734 key: self.config.rope_scaling[key]735 for key in [736 "original_max_position_embeddings",737 "beta_fast",738 "beta_slow",739 "mscale",740 "mscale_all_dim",741 ] if key in self.config.rope_scaling742 }743 self.rotary_emb = DeepseekV3YarnRotaryEmbedding(744 self.qk_rope_head_dim,745 max_position_embeddings=self.max_position_embeddings,746 scaling_factor=scaling_factor,747 base=self.rope_theta,748 **kwargs,749 )750 else:751 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")752 753 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):754 return (tensor.view(bsz, seq_len, self.num_heads,755 self.v_head_dim).transpose(1, 2).contiguous())756 757 def forward(758 self,759 hidden_states: torch.Tensor,760 attention_mask: Optional[torch.Tensor] = None,761 position_ids: Optional[torch.LongTensor] = None,762 past_key_value: Optional[Cache] = None,763 output_attentions: bool = False,764 use_cache: bool = False,765 **kwargs,766 ) -> Tuple[torch.Tensor, Optional[torch.Tensor],767 Optional[Tuple[torch.Tensor]]]:768 if "padding_mask" in kwargs:769 warnings.warn(770 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"771 )772 bsz, q_len, _ = hidden_states.size()773 774 if self.q_lora_rank is None:775 q = self.q_proj(hidden_states)776 else:777 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))778 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)779 q_nope, q_pe = torch.split(780 q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)781 782 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)783 compressed_kv, k_pe = torch.split(784 compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)785 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)786 kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(787 bsz, q_len, self.num_heads,788 self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))789 790 k_nope, value_states = torch.split(791 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)792 kv_seq_len = value_states.shape[-2]793 if past_key_value is not None:794 if self.layer_idx is None:795 raise ValueError(796 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "797 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "798 "with a layer index.")799 kv_seq_len += get_usable_length(past_key_value, kv_seq_len,800 self.layer_idx)801 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)802 803 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)804 805 query_states = k_pe.new_empty(bsz, self.num_heads, q_len,806 self.q_head_dim)807 query_states[:, :, :, :self.qk_nope_head_dim] = q_nope808 query_states[:, :, :, self.qk_nope_head_dim:] = q_pe809 810 key_states = k_pe.new_empty(bsz, self.num_heads, q_len,811 self.q_head_dim)812 key_states[:, :, :, :self.qk_nope_head_dim] = k_nope813 key_states[:, :, :, self.qk_nope_head_dim:] = k_pe814 if past_key_value is not None:815 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models816 key_states, value_states = past_key_value.update(817 key_states, value_states, self.layer_idx, cache_kwargs)818 819 attn_weights = (820 torch.matmul(query_states, key_states.transpose(2, 3)) *821 self.softmax_scale)822 823 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):824 raise ValueError(825 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"826 f" {attn_weights.size()}")827 assert attention_mask is not None828 if attention_mask is not None:829 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):830 raise ValueError(831 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"832 )833 attn_weights = attn_weights + attention_mask834 835 # upcast attention to fp32836 attn_weights = nn.functional.softmax(attn_weights,837 dim=-1,838 dtype=torch.float32).to(839 query_states.dtype)840 attn_weights = nn.functional.dropout(attn_weights,841 p=self.attention_dropout,842 training=self.training)843 attn_output = torch.matmul(attn_weights, value_states)844 845 if attn_output.size() != (bsz, self.num_heads, q_len, self.v_head_dim):846 raise ValueError(847 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.v_head_dim)}, but is"848 f" {attn_output.size()}")849 850 attn_output = attn_output.transpose(1, 2).contiguous()851 852 attn_output = attn_output.reshape(bsz, q_len,853 self.num_heads * self.v_head_dim)854 855 attn_output = self.o_proj(attn_output)856 857 if not output_attentions:858 attn_weights = None859 860 return attn_output, attn_weights, past_key_value861 862 863# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2 with Llama->DeepseekV3864class DeepseekV3FlashAttention2(DeepseekV3Attention):865 """866 DeepseekV3 flash attention module. This module inherits from `DeepseekV3Attention` as the weights of the module stays867 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of868 flash attention and deal with padding tokens in case the input contains any of them.869 """870 871 def __init__(self, *args, **kwargs):872 super().__init__(*args, **kwargs)873 874 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.875 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, 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.876 # 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).877 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10(878 )879 880 def forward(881 self,882 hidden_states: torch.Tensor,883 attention_mask: Optional[torch.LongTensor] = None,884 position_ids: Optional[torch.LongTensor] = None,885 past_key_value: Optional[Cache] = None,886 output_attentions: bool = False,887 use_cache: bool = False,888 **kwargs,889 ) -> Tuple[torch.Tensor, Optional[torch.Tensor],890 Optional[Tuple[torch.Tensor]]]:891 # DeepseekV3FlashAttention2 attention does not support output_attentions892 if "padding_mask" in kwargs:893 warnings.warn(894 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"895 )896 897 # overwrite attention_mask with padding_mask898 attention_mask = kwargs.pop("padding_mask")899 900 output_attentions = False901 902 bsz, q_len, _ = hidden_states.size()903 904 if self.q_lora_rank is None:905 q = self.q_proj(hidden_states)906 else:907 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))908 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)909 q_nope, q_pe = torch.split(910 q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)911 912 # Flash attention requires the input to have the shape913 # batch_size x seq_length x head_dim x hidden_dim914 # therefore we just need to keep the original shape915 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)916 compressed_kv, k_pe = torch.split(917 compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)918 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)919 kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(920 bsz, q_len, self.num_heads,921 self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))922 923 k_nope, value_states = torch.split(924 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)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 += get_usable_length(past_key_value, kv_seq_len,930 self.layer_idx)931 932 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)933 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)934 935 query_states = k_pe.new_empty(bsz, self.num_heads, q_len,936 self.q_head_dim)937 query_states[:, :, :, :self.qk_nope_head_dim] = q_nope938 query_states[:, :, :, self.qk_nope_head_dim:] = q_pe939 940 key_states = k_pe.new_empty(bsz, self.num_heads, q_len,941 self.q_head_dim)942 key_states[:, :, :, :self.qk_nope_head_dim] = k_nope943 key_states[:, :, :, self.qk_nope_head_dim:] = k_pe944 945 if self.q_head_dim != self.v_head_dim:946 value_states = F.pad(value_states,947 [0, self.q_head_dim - self.v_head_dim])948 949 if past_key_value is not None:950 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models951 key_states, value_states = past_key_value.update(952 key_states, value_states, self.layer_idx, cache_kwargs)953 954 # 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 cache955 # to be able to avoid many of these transpose/reshape/view.956 query_states = query_states.transpose(1, 2)957 key_states = key_states.transpose(1, 2)958 value_states = value_states.transpose(1, 2)959 960 dropout_rate = self.attention_dropout if self.training else 0.0961 962 # In PEFT, usually we cast the layer norms in float32 for training stability reasons963 # therefore the input hidden states gets silently casted in float32. Hence, we need964 # cast them back in the correct dtype just to be sure everything works as expected.965 # This might slowdown training & inference so it is recommended to not cast the LayerNorms966 # in fp32. (DeepseekV3RMSNorm handles it correctly)967 968 input_dtype = query_states.dtype969 if input_dtype == torch.float32:970 # Handle the case where the model is quantized971 if hasattr(self.config, "_pre_quantization_dtype"):972 target_dtype = self.config._pre_quantization_dtype973 elif torch.is_autocast_enabled():974 target_dtype = torch.get_autocast_gpu_dtype()975 else:976 target_dtype = (self.q_proj.weight.dtype if self.q_lora_rank977 is None else self.q_a_proj.weight.dtype)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 query_states = query_states.to(target_dtype)985 key_states = key_states.to(target_dtype)986 value_states = value_states.to(target_dtype)987 988 attn_output = self._flash_attention_forward(989 query_states,990 key_states,991 value_states,992 attention_mask,993 q_len,994 dropout=dropout_rate,995 softmax_scale=self.softmax_scale,996 )997 if self.q_head_dim != self.v_head_dim:998 attn_output = attn_output[:, :, :, :self.v_head_dim]999 1000 attn_output = attn_output.reshape(bsz, q_len, self.num_heads *1001 self.v_head_dim).contiguous()1002 attn_output = self.o_proj(attn_output)1003 1004 if not output_attentions:1005 attn_weights = None1006 1007 return attn_output, attn_weights, past_key_value1008 1009 def _flash_attention_forward(1010 self,1011 query_states,1012 key_states,1013 value_states,1014 attention_mask,1015 query_length,1016 dropout=0.0,1017 softmax_scale=None,1018 ):1019 """1020 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1021 first unpad the input, then computes the attention scores and pad the final attention scores.1022 1023 Args:1024 query_states (`torch.Tensor`):1025 Input query states to be passed to Flash Attention API1026 key_states (`torch.Tensor`):1027 Input key states to be passed to Flash Attention API1028 value_states (`torch.Tensor`):1029 Input value states to be passed to Flash Attention API1030 attention_mask (`torch.Tensor`):1031 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1032 position of padding tokens and 1 for the position of non-padding tokens.1033 dropout (`int`, *optional*):1034 Attention dropout1035 softmax_scale (`float`, *optional*):1036 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1037 """1038 if not self._flash_attn_uses_top_left_mask:1039 causal = self.is_causal1040 else:1041 # 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__.1042 causal = self.is_causal and query_length != 11043 1044 # Contains at least one padding token in the sequence1045 if attention_mask is not None:1046 batch_size = query_states.shape[0]1047 (1048 query_states,1049 key_states,1050 value_states,1051 indices_q,1052 cu_seq_lens,1053 max_seq_lens,1054 ) = self._upad_input(query_states, key_states, value_states,1055 attention_mask, query_length)1056 1057 cu_seqlens_q, cu_seqlens_k = cu_seq_lens1058 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1059 1060 attn_output_unpad = flash_attn_varlen_func(1061 query_states,1062 key_states,1063 value_states,1064 cu_seqlens_q=cu_seqlens_q,1065 cu_seqlens_k=cu_seqlens_k,1066 max_seqlen_q=max_seqlen_in_batch_q,1067 max_seqlen_k=max_seqlen_in_batch_k,1068 dropout_p=dropout,1069 softmax_scale=softmax_scale,1070 causal=causal,1071 )1072 1073 attn_output = pad_input(attn_output_unpad, indices_q, batch_size,1074 query_length)1075 else:1076 attn_output = flash_attn_func(1077 query_states,1078 key_states,1079 value_states,1080 dropout,1081 softmax_scale=softmax_scale,1082 causal=causal,1083 )1084 1085 return attn_output1086 1087 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask,1088 query_length):1089 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(1090 attention_mask)1091 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape1092 1093 key_layer = index_first_axis(1094 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,1095 head_dim),1096 indices_k,1097 )1098 value_layer = index_first_axis(1099 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads,1100 head_dim),1101 indices_k,1102 )1103 if query_length == kv_seq_len:1104 query_layer = index_first_axis(1105 query_layer.reshape(batch_size * kv_seq_len, self.num_heads,1106 head_dim),1107 indices_k,1108 )1109 cu_seqlens_q = cu_seqlens_k1110 max_seqlen_in_batch_q = max_seqlen_in_batch_k1111 indices_q = indices_k1112 elif query_length == 1:1113 max_seqlen_in_batch_q = 11114 cu_seqlens_q = torch.arange(1115 batch_size + 1, dtype=torch.int32, device=query_layer.device1116 ) # There is a memcpy here, that is very bad.1117 indices_q = cu_seqlens_q[:-1]1118 query_layer = query_layer.squeeze(1)1119 else:1120 # The -q_len: slice assumes left padding.1121 attention_mask = attention_mask[:, -query_length:]1122 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(1123 query_layer, attention_mask)1124 1125 return (1126 query_layer,1127 key_layer,1128 value_layer,1129 indices_q,1130 (cu_seqlens_q, cu_seqlens_k),1131 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),1132 )1133 1134 1135ATTENTION_CLASSES = {1136 "eager": DeepseekV3Attention,1137 "flash_attention_2": DeepseekV3FlashAttention2,1138}1139 1140 1141class DeepseekV3DecoderLayer(nn.Module):1142 1143 def __init__(self, config: DeepseekV3Config, layer_idx: int):1144 super().__init__()1145 self.hidden_size = config.hidden_size1146 1147 self.self_attn = ATTENTION_CLASSES[config._attn_implementation](1148 config=config, layer_idx=layer_idx)1149 1150 self.mlp = (DeepseekV3MoE(config) if1151 (config.n_routed_experts is not None1152 and layer_idx >= config.first_k_dense_replace1153 and layer_idx % config.moe_layer_freq == 0) else1154 DeepseekV3MLP(config))1155 self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size,1156 eps=config.rms_norm_eps)1157 self.post_attention_layernorm = DeepseekV3RMSNorm(1158 config.hidden_size, eps=config.rms_norm_eps)1159 1160 def forward(1161 self,1162 hidden_states: torch.Tensor,1163 attention_mask: Optional[torch.Tensor] = None,1164 position_ids: Optional[torch.LongTensor] = None,1165 past_key_value: Optional[Tuple[torch.Tensor]] = None,1166 output_attentions: Optional[bool] = False,1167 use_cache: Optional[bool] = False,1168 **kwargs,1169 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor,1170 torch.FloatTensor]]]:1171 """1172 Args:1173 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`1174 attention_mask (`torch.FloatTensor`, *optional*):1175 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,1176 query_sequence_length, key_sequence_length)` if default attention is used.1177 output_attentions (`bool`, *optional*):1178 Whether or not to return the attentions tensors of all attention layers. See `attentions` under1179 returned tensors for more detail.1180 use_cache (`bool`, *optional*):1181 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding1182 (see `past_key_values`).1183 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states1184 """1185 if "padding_mask" in kwargs:1186 warnings.warn(1187 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"1188 )1189 residual = hidden_states1190 1191 hidden_states = self.input_layernorm(hidden_states)1192 1193 # Self Attention1194 hidden_states, self_attn_weights, present_key_value = self.self_attn(1195 hidden_states=hidden_states,1196 attention_mask=attention_mask,1197 position_ids=position_ids,1198 past_key_value=past_key_value,1199 output_attentions=output_attentions,1200 use_cache=use_cache,