jdopensource/JoyAI-LLM-Flash
175225
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/deepseek_v3/modular_deepseek_v3.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_deepseek_v3.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7import math8from functools import partial9from typing import Callable, Optional, Tuple, Union10 11import torch12import torch.nn.functional as F13from torch import nn14 15from transformers.activations import ACT2FN16from transformers.cache_utils import Cache, DynamicCache, StaticCache17from transformers.generation import GenerationMixin18from transformers.modeling_attn_mask_utils import AttentionMaskConverter19from transformers.modeling_flash_attention_utils import FlashAttentionKwargs20from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast21from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update22from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel23from transformers.processing_utils import Unpack24from transformers.utils import (25 add_start_docstrings,26 add_start_docstrings_to_model_forward,27 can_return_tuple,28 is_torch_flex_attn_available,29 logging,30 replace_return_docstrings,31)32from transformers.utils.deprecation import deprecate_kwarg33from .configuration_deepseek import DeepseekV3Config34 35 36if is_torch_flex_attn_available():37 from torch.nn.attention.flex_attention import BlockMask38 39 from transformers.integrations.flex_attention import make_flex_block_causal_mask40 41 42logger = logging.get_logger(__name__)43_CONFIG_FOR_DOC = "DeepseekV3Config"44 45 46class DeepseekV3RMSNorm(nn.Module):47 def __init__(self, hidden_size, eps=1e-6):48 """49 DeepseekV3RMSNorm is equivalent to T5LayerNorm50 """51 super().__init__()52 self.weight = nn.Parameter(torch.ones(hidden_size))53 self.variance_epsilon = eps54 55 def forward(self, hidden_states):56 input_dtype = hidden_states.dtype57 hidden_states = hidden_states.to(torch.float32)58 variance = hidden_states.pow(2).mean(-1, keepdim=True)59 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)60 return self.weight * hidden_states.to(input_dtype)61 62 def extra_repr(self):63 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"64 65 66class DeepseekV3RotaryEmbedding(nn.Module):67 def __init__(self, config: DeepseekV3Config, device=None):68 super().__init__()69 # BC: "rope_type" was originally "type"70 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:71 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))72 else:73 self.rope_type = "default"74 self.max_seq_len_cached = config.max_position_embeddings75 self.original_max_seq_len = config.max_position_embeddings76 77 self.config = config78 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]79 80 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)81 self.register_buffer("inv_freq", inv_freq, persistent=False)82 self.original_inv_freq = self.inv_freq83 84 @torch.no_grad()85 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)86 def forward(self, x, position_ids):87 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)88 position_ids_expanded = position_ids[:, None, :].float()89 90 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"91 with torch.autocast(device_type=device_type, enabled=False): # Force float3292 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)93 emb = torch.cat((freqs, freqs), dim=-1)94 cos = emb.cos() * self.attention_scaling95 sin = emb.sin() * self.attention_scaling96 97 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)98 99 100class DeepseekV3MLP(nn.Module):101 def __init__(self, config, hidden_size=None, intermediate_size=None):102 super().__init__()103 self.config = config104 self.hidden_size = config.hidden_size if hidden_size is None else hidden_size105 self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size106 107 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)108 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)109 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)110 self.act_fn = ACT2FN[config.hidden_act]111 112 def forward(self, x):113 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))114 return down_proj115 116 117class DeepseekV3TopkRouter(nn.Module):118 def __init__(self, config):119 super().__init__()120 self.config = config121 self.top_k = config.num_experts_per_tok122 self.n_routed_experts = config.n_routed_experts123 self.routed_scaling_factor = config.routed_scaling_factor124 self.n_group = config.n_group125 self.topk_group = config.topk_group126 self.norm_topk_prob = config.norm_topk_prob127 128 self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))129 self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts)))130 131 @torch.no_grad()132 def get_topk_indices(self, scores):133 scores_for_choice = scores.view(-1, self.n_routed_experts) + self.e_score_correction_bias.unsqueeze(0)134 group_scores = (135 scores_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)136 .topk(2, dim=-1)[0]137 .sum(dim=-1)138 )139 group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]140 group_mask = torch.zeros_like(group_scores)141 group_mask.scatter_(1, group_idx, 1)142 score_mask = (143 group_mask.unsqueeze(-1)144 .expand(-1, self.n_group, self.n_routed_experts // self.n_group)145 .reshape(-1, self.n_routed_experts)146 )147 scores_for_choice = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)148 topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]149 return topk_indices150 151 def forward(self, hidden_states):152 hidden_states = hidden_states.view(-1, self.config.hidden_size)153 router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))154 scores = router_logits.sigmoid()155 topk_indices = self.get_topk_indices(scores)156 topk_weights = scores.gather(1, topk_indices)157 if self.norm_topk_prob:158 denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20159 topk_weights /= denominator160 topk_weights = topk_weights * self.routed_scaling_factor161 return topk_indices, topk_weights162 163 164class DeepseekV3MoE(nn.Module):165 """166 A mixed expert module containing shared experts.167 """168 169 def __init__(self, config):170 super().__init__()171 self.config = config172 self.experts = nn.ModuleList(173 [174 DeepseekV3MLP(config, intermediate_size=config.moe_intermediate_size)175 for _ in range(config.n_routed_experts)176 ]177 )178 self.gate = DeepseekV3TopkRouter(config)179 self.shared_experts = DeepseekV3MLP(180 config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts181 )182 183 def moe(self, hidden_states: torch.Tensor, topk_indices: torch.Tensor, topk_weights: torch.Tensor):184 r"""185 CALL FOR CONTRIBUTION! I don't have time to optimise this right now, but expert weights need to be fused186 to not have to do a loop here (deepseek has 256 experts soooo yeah).187 """188 final_hidden_states = torch.zeros_like(hidden_states, dtype=topk_weights.dtype)189 expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=len(self.experts))190 expert_mask = expert_mask.permute(2, 0, 1)191 192 for expert_idx in range(len(self.experts)):193 expert = self.experts[expert_idx]194 mask = expert_mask[expert_idx]195 token_indices, weight_indices = torch.where(mask)196 197 if token_indices.numel() > 0:198 expert_weights = topk_weights[token_indices, weight_indices]199 expert_input = hidden_states[token_indices]200 expert_output = expert(expert_input)201 weighted_output = expert_output * expert_weights.unsqueeze(-1)202 final_hidden_states.index_add_(0, token_indices, weighted_output)203 204 # in original deepseek, the output of the experts are gathered once we leave this module205 # thus the moe module is itelsf an IsolatedParallel module206 # and all expert are "local" meaning we shard but we don't gather207 return final_hidden_states.type(hidden_states.dtype)208 209 def forward(self, hidden_states):210 residuals = hidden_states211 orig_shape = hidden_states.shape212 topk_indices, topk_weights = self.gate(hidden_states)213 hidden_states = hidden_states.view(-1, hidden_states.shape[-1])214 hidden_states = self.moe(hidden_states, topk_indices, topk_weights).view(*orig_shape)215 hidden_states = hidden_states + self.shared_experts(residuals)216 return hidden_states217 218 219def rotate_half(x):220 """Rotates half the hidden dims of the input."""221 x1 = x[..., : x.shape[-1] // 2]222 x2 = x[..., x.shape[-1] // 2 :]223 return torch.cat((-x2, x1), dim=-1)224 225 226def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):227 """Applies Rotary Position Embedding to the query and key tensors.228 229 Args:230 q (`torch.Tensor`): The query tensor.231 k (`torch.Tensor`): The key tensor.232 cos (`torch.Tensor`): The cosine part of the rotary embedding.233 sin (`torch.Tensor`): The sine part of the rotary embedding.234 position_ids (`torch.Tensor`, *optional*):235 Deprecated and unused.236 unsqueeze_dim (`int`, *optional*, defaults to 1):237 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and238 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note239 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and240 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes241 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have242 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.243 Returns:244 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.245 """246 cos = cos.unsqueeze(unsqueeze_dim)247 sin = sin.unsqueeze(unsqueeze_dim)248 q_embed = (q * cos) + (rotate_half(q) * sin)249 k_embed = (k * cos) + (rotate_half(k) * sin)250 return q_embed, k_embed251 252 253def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:254 """255 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,256 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)257 """258 batch, num_key_value_heads, slen, head_dim = hidden_states.shape259 if n_rep == 1:260 return hidden_states261 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)262 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)263 264 265def eager_attention_forward(266 module: nn.Module,267 query: torch.Tensor,268 key: torch.Tensor,269 value: torch.Tensor,270 attention_mask: Optional[torch.Tensor],271 scaling: float,272 dropout: float = 0.0,273 **kwargs,274):275 key_states = repeat_kv(key, module.num_key_value_groups)276 value_states = repeat_kv(value, module.num_key_value_groups)277 278 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling279 if attention_mask is not None:280 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]281 attn_weights = attn_weights + causal_mask282 283 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)284 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)285 attn_output = torch.matmul(attn_weights, value_states)286 attn_output = attn_output.transpose(1, 2).contiguous()287 288 return attn_output, attn_weights289 290 291def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):292 r"""293 TODO let's just use the original freqcis computation to not have the view294 transpose + reshape! This is not optimized!295 Applies Rotary Position Embedding to the query and key tensors.296 297 Args:298 q (`torch.Tensor`): The query tensor.299 k (`torch.Tensor`): The key tensor.300 cos (`torch.Tensor`): The cosine part of the rotary embedding.301 sin (`torch.Tensor`): The sine part of the rotary embedding.302 position_ids (`torch.Tensor`):303 The position indices of the tokens corresponding to the query and key tensors. For example, this can be304 used to pass offsetted position ids when working with a KV-cache.305 unsqueeze_dim (`int`, *optional*, defaults to 1):306 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and307 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note308 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and309 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes310 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have311 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.312 Returns:313 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.314 """315 cos = cos.unsqueeze(unsqueeze_dim)316 sin = sin.unsqueeze(unsqueeze_dim)317 318 b, h, s, d = q.shape319 q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)320 321 b, h, s, d = k.shape322 k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)323 324 q_embed = (q * cos) + (rotate_half(q) * sin)325 k_embed = (k * cos) + (rotate_half(k) * sin)326 return q_embed, k_embed327 328 329def yarn_get_mscale(scale=1, mscale=1):330 if scale <= 1:331 return 1.0332 return 0.1 * mscale * math.log(scale) + 1.0333 334 335class DeepseekV3Attention(nn.Module):336 """Multi-headed attention from 'Attention Is All You Need' paper"""337 338 def __init__(self, config: DeepseekV3Config, layer_idx: int):339 super().__init__()340 self.config = config341 self.layer_idx = layer_idx342 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads343 self.attention_dropout = config.attention_dropout344 self.num_heads = config.num_attention_heads345 self.rope_theta = config.rope_theta346 self.q_lora_rank = config.q_lora_rank347 self.qk_rope_head_dim = config.qk_rope_head_dim348 self.kv_lora_rank = config.kv_lora_rank349 self.v_head_dim = config.v_head_dim350 self.qk_nope_head_dim = config.qk_nope_head_dim351 self.qk_head_dim = config.qk_head_dim352 353 self.is_causal = True354 self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias)355 self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)356 self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)357 358 self.kv_a_proj_with_mqa = nn.Linear(359 config.hidden_size,360 self.kv_lora_rank + self.qk_rope_head_dim,361 bias=config.attention_bias,362 )363 self.kv_a_layernorm = DeepseekV3RMSNorm(self.kv_lora_rank)364 self.kv_b_proj = nn.Linear(365 self.kv_lora_rank,366 self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),367 bias=False,368 )369 370 self.o_proj = nn.Linear(371 self.num_heads * self.v_head_dim,372 config.hidden_size,373 bias=config.attention_bias,374 )375 376 self.scaling = self.qk_head_dim ** (-0.5)377 if self.config.rope_scaling is not None:378 mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)379 scaling_factor = self.config.rope_scaling["factor"]380 if mscale_all_dim:381 mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)382 self.scaling = self.scaling * mscale * mscale383 384 def forward(385 self,386 hidden_states: torch.Tensor,387 position_embeddings: Tuple[torch.Tensor, torch.Tensor],388 attention_mask: Optional[torch.Tensor],389 past_key_value: Optional[Cache] = None,390 cache_position: Optional[torch.LongTensor] = None,391 **kwargs: Unpack[FlashAttentionKwargs],392 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:393 batch_size, seq_length = hidden_states.shape[:-1]394 query_shape = (batch_size, seq_length, -1, self.qk_head_dim)395 key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)396 397 q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(query_shape).transpose(1, 2)398 q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)399 400 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)401 k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)402 403 k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)404 k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)405 406 k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)407 408 cos, sin = position_embeddings409 if self.config.rope_interleave: # support using interleaved weights for efficiency410 q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)411 else:412 q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)413 k_rot = k_rot.expand(*k_pass.shape[:-1], -1)414 415 query_states = torch.cat((q_pass, q_rot), dim=-1)416 key_states = torch.cat((k_pass, k_rot), dim=-1)417 418 if past_key_value is not None:419 # sin and cos are specific to RoPE models; cache_position needed for the static cache420 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}421 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)422 423 if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:424 value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])425 426 attention_interface: Callable = eager_attention_forward427 if self.config._attn_implementation != "eager":428 if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):429 logger.warning_once(430 "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "431 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'432 )433 else:434 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]435 436 attn_output, attn_weights = attention_interface(437 self,438 query_states,439 key_states,440 value_states,441 attention_mask,442 dropout=0.0 if not self.training else self.attention_dropout,443 scaling=self.scaling,444 **kwargs,445 )446 447 if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim:448 attn_output = attn_output[:, :, :, : self.v_head_dim]449 450 attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()451 attn_output = self.o_proj(attn_output)452 return attn_output, attn_weights453 454 455class DeepseekV3DecoderLayer(nn.Module):456 def __init__(self, config: DeepseekV3Config, layer_idx: int):457 super().__init__()458 self.hidden_size = config.hidden_size459 460 self.self_attn = DeepseekV3Attention(config=config, layer_idx=layer_idx)461 462 if layer_idx >= config.first_k_dense_replace:463 self.mlp = DeepseekV3MoE(config)464 else:465 self.mlp = DeepseekV3MLP(config)466 467 self.input_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)468 self.post_attention_layernorm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)469 470 def forward(471 self,472 hidden_states: torch.Tensor,473 attention_mask: Optional[torch.Tensor] = None,474 position_ids: Optional[torch.LongTensor] = None,475 past_key_value: Optional[Cache] = None,476 output_attentions: Optional[bool] = False,477 use_cache: Optional[bool] = False,478 cache_position: Optional[torch.LongTensor] = None,479 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC480 **kwargs: Unpack[FlashAttentionKwargs],481 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:482 residual = hidden_states483 484 hidden_states = self.input_layernorm(hidden_states)485 486 # Self Attention487 hidden_states, self_attn_weights = self.self_attn(488 hidden_states=hidden_states,489 attention_mask=attention_mask,490 position_ids=position_ids,491 past_key_value=past_key_value,492 output_attentions=output_attentions,493 use_cache=use_cache,494 cache_position=cache_position,495 position_embeddings=position_embeddings,496 **kwargs,497 )498 hidden_states = residual + hidden_states499 500 # Fully Connected501 residual = hidden_states502 hidden_states = self.post_attention_layernorm(hidden_states)503 hidden_states = self.mlp(hidden_states)504 hidden_states = residual + hidden_states505 506 outputs = (hidden_states,)507 if output_attentions:508 outputs += (self_attn_weights,)509 510 return outputs511 512 513DEEPSEEK_V3_START_DOCSTRING = r"""514 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the515 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads516 etc.)517 518 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.519 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage520 and behavior.521 522 Parameters:523 config ([`DeepseekV3Config`]):524 Model configuration class with all the parameters of the model. Initializing with a config file does not525 load the weights associated with the model, only the configuration. Check out the526 [`~PreTrainedModel.from_pretrained`] method to load the model weights.527"""528 529 530@add_start_docstrings(531 "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.",532 DEEPSEEK_V3_START_DOCSTRING,533)534class DeepseekV3PreTrainedModel(PreTrainedModel):535 config_class = DeepseekV3Config536 base_model_prefix = "model"537 supports_gradient_checkpointing = True538 _no_split_modules = ["DeepseekV3DecoderLayer"]539 _skip_keys_device_placement = ["past_key_values"]540 _supports_flash_attn_2 = True541 _supports_sdpa = True542 _supports_flex_attn = True543 _supports_cache_class = True544 _supports_quantized_cache = True545 _supports_static_cache = True546 _supports_attention_backend = True547 548 def _init_weights(self, module):549 std = self.config.initializer_range550 if isinstance(module, nn.Linear):551 module.weight.data.normal_(mean=0.0, std=std)552 if module.bias is not None:553 module.bias.data.zero_()554 elif isinstance(module, nn.Embedding):555 module.weight.data.normal_(mean=0.0, std=std)556 if module.padding_idx is not None:557 module.weight.data[module.padding_idx].zero_()558 elif isinstance(module, DeepseekV3TopkRouter):559 module.weight.data.normal_(mean=0.0, std=std)560 elif isinstance(module, nn.Parameter):561 module.weight.data.normal_(mean=0.0, std=std)562 563 564DEEPSEEK_V3_INPUTS_DOCSTRING = r"""565 Args:566 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):567 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide568 it.569 570 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and571 [`PreTrainedTokenizer.__call__`] for details.572 573 [What are input IDs?](../glossary#input-ids)574 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):575 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:576 577 - 1 for tokens that are **not masked**,578 - 0 for tokens that are **masked**.579 580 [What are attention masks?](../glossary#attention-mask)581 582 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and583 [`PreTrainedTokenizer.__call__`] for details.584 585 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see586 `past_key_values`).587 588 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]589 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more590 information on the default strategy.591 592 - 1 indicates the head is **not masked**,593 - 0 indicates the head is **masked**.594 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):595 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,596 config.n_positions - 1]`.597 598 [What are position IDs?](../glossary#position-ids)599 past_key_values (`Cache`, *optional*):600 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention601 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`602 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.603 604 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).605 606 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't607 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`608 of shape `(batch_size, sequence_length)`.609 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):610 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This611 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the612 model's internal embedding lookup matrix.613 use_cache (`bool`, *optional*):614 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see615 `past_key_values`).616 output_attentions (`bool`, *optional*):617 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned618 tensors for more detail.619 output_hidden_states (`bool`, *optional*):620 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for621 more detail.622 return_dict (`bool`, *optional*):623 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.624 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):625 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,626 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer627 the complete sequence length.628"""629 630 631@add_start_docstrings(632 "The bare DeepseekV3 Model outputting raw hidden-states without any specific head on top.",633 DEEPSEEK_V3_START_DOCSTRING,634)635class DeepseekV3Model(DeepseekV3PreTrainedModel):636 """637 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DeepseekV3DecoderLayer`]638 639 Args:640 config: DeepseekV3Config641 """642 643 _keys_to_ignore_on_load_unexpected = [r"model\.layers\.61.*"]644 645 def __init__(self, config: DeepseekV3Config):646 super().__init__(config)647 self.padding_idx = config.pad_token_id648 self.vocab_size = config.vocab_size649 650 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)651 self.layers = nn.ModuleList(652 [DeepseekV3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]653 )654 self.norm = DeepseekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)655 self.rotary_emb = DeepseekV3RotaryEmbedding(config=config)656 self.gradient_checkpointing = False657 658 # Initialize weights and apply final processing659 self.post_init()660 661 def get_input_embeddings(self):662 return self.embed_tokens663 664 def set_input_embeddings(self, value):665 self.embed_tokens = value666 667 @can_return_tuple668 @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING)669 def forward(670 self,671 input_ids: Optional[torch.LongTensor] = None,672 attention_mask: Optional[torch.Tensor] = None,673 position_ids: Optional[torch.LongTensor] = None,674 past_key_values: Optional[Cache] = None,675 inputs_embeds: Optional[torch.FloatTensor] = None,676 use_cache: Optional[bool] = None,677 output_attentions: Optional[bool] = None,678 output_hidden_states: Optional[bool] = None,679 cache_position: Optional[torch.LongTensor] = None,680 **flash_attn_kwargs: Unpack[FlashAttentionKwargs],681 ) -> BaseModelOutputWithPast:682 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions683 output_hidden_states = (684 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states685 )686 use_cache = use_cache if use_cache is not None else self.config.use_cache687 688 if (input_ids is None) ^ (inputs_embeds is not None):689 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")690 691 if self.gradient_checkpointing and self.training and use_cache:692 logger.warning_once(693 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."694 )695 use_cache = False696 697 # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache698 if not isinstance(past_key_values, (type(None), Cache)):699 raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")700 701 if inputs_embeds is None:702 inputs_embeds = self.embed_tokens(input_ids)703 704 if use_cache and past_key_values is None:705 past_key_values = DynamicCache()706 707 if cache_position is None:708 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0709 cache_position = torch.arange(710 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device711 )712 713 if position_ids is None:714 position_ids = cache_position.unsqueeze(0)715 716 causal_mask = self._update_causal_mask(717 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions718 )719 720 hidden_states = inputs_embeds721 722 # create position embeddings to be shared across the decoder layers723 position_embeddings = self.rotary_emb(hidden_states, position_ids)724 725 # decoder layers726 all_hidden_states = () if output_hidden_states else None727 all_self_attns = () if output_attentions else None728 729 for decoder_layer in self.layers[: self.config.num_hidden_layers]:730 if output_hidden_states:731 all_hidden_states += (hidden_states,)732 733 if self.gradient_checkpointing and self.training:734 layer_outputs = self._gradient_checkpointing_func(735 partial(decoder_layer.__call__, **flash_attn_kwargs),736 hidden_states,737 causal_mask,738 position_ids,739 past_key_values,740 output_attentions,741 use_cache,742 cache_position,743 position_embeddings,744 )745 else:746 layer_outputs = decoder_layer(747 hidden_states,748 attention_mask=causal_mask,749 position_ids=position_ids,750 past_key_value=past_key_values,751 output_attentions=output_attentions,752 use_cache=use_cache,753 cache_position=cache_position,754 position_embeddings=position_embeddings,755 **flash_attn_kwargs,756 )757 758 hidden_states = layer_outputs[0]759 760 if output_attentions:761 all_self_attns += (layer_outputs[1],)762 763 hidden_states = self.norm(hidden_states)764 765 # add hidden states from the last decoder layer766 if output_hidden_states:767 all_hidden_states += (hidden_states,)768 769 return BaseModelOutputWithPast(770 last_hidden_state=hidden_states,771 past_key_values=past_key_values if use_cache else None,772 hidden_states=all_hidden_states,773 attentions=all_self_attns,774 )775 776 def _update_causal_mask(777 self,778 attention_mask: torch.Tensor,779 input_tensor: torch.Tensor,780 cache_position: torch.Tensor,781 past_key_values: Cache,782 output_attentions: bool = False,783 ):784 if self.config._attn_implementation == "flash_attention_2":785 if attention_mask is not None and (attention_mask == 0.0).any():786 return attention_mask787 return None788 if self.config._attn_implementation == "flex_attention":789 if isinstance(attention_mask, torch.Tensor):790 attention_mask = make_flex_block_causal_mask(attention_mask)791 if isinstance(attention_mask, BlockMask):792 return attention_mask793 794 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in795 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail796 # to infer the attention mask.797 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0798 using_static_cache = isinstance(past_key_values, StaticCache)799 800 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward801 if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:802 if AttentionMaskConverter._ignore_causal_mask_sdpa(803 attention_mask,804 inputs_embeds=input_tensor,805 past_key_values_length=past_seen_tokens,806 is_training=self.training,807 ):808 return None809 810 dtype, device = input_tensor.dtype, input_tensor.device811 sequence_length = input_tensor.shape[1]812 if using_static_cache:813 target_length = past_key_values.get_max_cache_shape()814 else:815 target_length = (816 attention_mask.shape[-1]817 if isinstance(attention_mask, torch.Tensor)818 else past_seen_tokens + sequence_length + 1819 )820 821 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).822 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(823 attention_mask,824 sequence_length=sequence_length,825 target_length=target_length,826 dtype=dtype,827 device=device,828 cache_position=cache_position,829 batch_size=input_tensor.shape[0],830 )831 832 if (833 self.config._attn_implementation == "sdpa"834 and attention_mask is not None835 and attention_mask.device.type in ["cuda", "xpu"]836 and not output_attentions837 ):838 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when839 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.840 # Details: https://github.com/pytorch/pytorch/issues/110213841 min_dtype = torch.finfo(dtype).min842 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)843 844 return causal_mask845 846 @staticmethod847 def _prepare_4d_causal_attention_mask_with_cache_position(848 attention_mask: torch.Tensor,849 sequence_length: int,850 target_length: int,851 dtype: torch.dtype,852 device: torch.device,853 cache_position: torch.Tensor,854 batch_size: int,855 **kwargs,856 ):857 """858 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape859 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.860 861 Args:862 attention_mask (`torch.Tensor`):863 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape864 `(batch_size, 1, query_length, key_value_length)`.865 sequence_length (`int`):866 The sequence length being processed.867 target_length (`int`):868 The target length: when generating with static cache, the mask should be as long as the static cache,869 to account for the 0 padding, the part of the cache that is not filled yet.870 dtype (`torch.dtype`):871 The dtype to use for the 4D attention mask.872 device (`torch.device`):873 The device to place the 4D attention mask on.874 cache_position (`torch.Tensor`):875 Indices depicting the position of the input sequence tokens in the sequence.876 batch_size (`torch.Tensor`):877 Batch size.878 """879 if attention_mask is not None and attention_mask.dim() == 4:880 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.881 causal_mask = attention_mask882 else:883 min_dtype = torch.finfo(dtype).min884 causal_mask = torch.full(885 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device886 )887 if sequence_length != 1:888 causal_mask = torch.triu(causal_mask, diagonal=1)889 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)890 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)891 if attention_mask is not None:892 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit893 mask_length = attention_mask.shape[-1]894 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(895 causal_mask.device896 )897 padding_mask = padding_mask == 0898 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(899 padding_mask, min_dtype900 )901 902 return causal_mask903 904 905class DeepseekV3ForCausalLM(DeepseekV3PreTrainedModel, GenerationMixin):906 _tied_weights_keys = ["lm_head.weight"]907 _tp_plan = {"lm_head": "colwise_rep"}908 _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}909 910 def __init__(self, config):911 super().__init__(config)912 self.model = DeepseekV3Model(config)913 self.vocab_size = config.vocab_size914 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)915 916 # Initialize weights and apply final processing917 self.post_init()918 919 def get_input_embeddings(self):920 return self.model.embed_tokens921 922 def set_input_embeddings(self, value):923 self.model.embed_tokens = value924 925 def get_output_embeddings(self):926 return self.lm_head927 928 def set_output_embeddings(self, new_embeddings):929 self.lm_head = new_embeddings930 931 def set_decoder(self, decoder):932 self.model = decoder933 934 def get_decoder(self):935 return self.model936 937 @can_return_tuple938 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")939 @add_start_docstrings_to_model_forward(DEEPSEEK_V3_INPUTS_DOCSTRING)940 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)941 def forward(942 self,943 input_ids: Optional[torch.LongTensor] = None,944 attention_mask: Optional[torch.Tensor] = None,945 position_ids: Optional[torch.LongTensor] = None,946 past_key_values: Optional[Cache] = None,947 inputs_embeds: Optional[torch.FloatTensor] = None,948 labels: Optional[torch.LongTensor] = None,949 use_cache: Optional[bool] = None,950 output_attentions: Optional[bool] = None,951 output_hidden_states: Optional[bool] = None,952 cache_position: Optional[torch.LongTensor] = None,953 logits_to_keep: Union[int, torch.Tensor] = 0,954 **kwargs955 ) -> CausalLMOutputWithPast:956 r"""957 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):958 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,959 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored960 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.961 962 logits_to_keep (`int` or `torch.Tensor`, *optional*):963 If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all964 `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that965 token can save memory, which becomes pretty significant for long sequences or large vocabulary size.966 If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.967 This is useful when using packed tensor format (single dimension for batch and sequence length).968 969 Returns:970 971 Example:972 973 ```python974 >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM975 976 >>> model = DeepseekV3ForCausalLM.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf")977 >>> tokenizer = AutoTokenizer.from_pretrained("meta-deepseek_v3/DeepseekV3-2-7b-hf")978 979 >>> prompt = "Hey, are you conscious? Can you talk to me?"980 >>> inputs = tokenizer(prompt, return_tensors="pt")981 982 >>> # Generate983 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)984 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]985 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."986 ```"""987 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions988 output_hidden_states = (989 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states990 )991 992 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)993 outputs: BaseModelOutputWithPast = self.model(994 input_ids=input_ids,995 attention_mask=attention_mask,996 position_ids=position_ids,997 past_key_values=past_key_values,998 inputs_embeds=inputs_embeds,999 use_cache=use_cache,1000 output_attentions=output_attentions,1001 output_hidden_states=output_hidden_states,1002 cache_position=cache_position,1003 **kwargs,1004 )1005 1006 hidden_states = outputs.last_hidden_state1007 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss1008 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep1009 logits = self.lm_head(hidden_states[:, slice_indices, :])1010 1011 loss = None1012 if labels is not None:1013 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)1014 1015 return CausalLMOutputWithPast(1016 loss=loss,1017 logits=logits,1018 past_key_values=outputs.past_key_values,1019 hidden_states=outputs.hidden_states,1020 attentions=outputs.attentions,1021 )1022 1023 1024__all__ = ["DeepseekV3PreTrainedModel", "DeepseekV3Model", "DeepseekV3ForCausalLM"]1025