FriendliAI/Phi-tiny-MoE-instruct
018
1# coding=utf-82# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16""" PyTorch PhiMoE model."""17import inspect18import math19import warnings20from typing import List, Optional, Tuple, Union21 22import torch23import torch.nn.functional as F24import torch.utils.checkpoint25from torch import nn26from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss27 28from transformers.activations import ACT2FN29from transformers.cache_utils import Cache, DynamicCache30from transformers.modeling_attn_mask_utils import (31 _prepare_4d_causal_attention_mask,32 _prepare_4d_causal_attention_mask_for_sdpa,33)34from transformers.modeling_outputs import (35 MoeCausalLMOutputWithPast,36 MoeModelOutputWithPast,37 SequenceClassifierOutputWithPast,38)39from transformers.modeling_utils import PreTrainedModel40from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_1341from transformers.utils import (42 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,46 logging,47 replace_return_docstrings,48)49from transformers.utils.import_utils import is_torch_fx_available50from .configuration_slimmoe import PhiMoEConfig51 52from einops import rearrange53from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding54 55 56if is_flash_attn_2_available():57 from flash_attn import flash_attn_func, flash_attn_varlen_func58 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa59 60 _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)61 62# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.63# It means that the function will not be traced through and simply appear as a node in the graph.64if is_torch_fx_available():65 if not is_torch_greater_or_equal_than_1_13:66 import torch.fx67 68 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)69 70 71logger = logging.get_logger(__name__)72 73_CONFIG_FOR_DOC = "PhiMoEConfig"74 75 76def load_balancing_loss_func(77 gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2, attention_mask: Optional[torch.Tensor] = None78) -> float:79 r"""80 Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.81 82 See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss83 function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between84 experts is too unbalanced.85 86 Args:87 gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):88 Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of89 shape [batch_size X sequence_length, num_experts].90 attention_mask (`torch.Tensor`, None):91 The attention_mask used in forward function92 shape [batch_size X sequence_length] if not None.93 num_experts (`int`, *optional*):94 Number of experts95 96 Returns:97 The auxiliary loss.98 """99 if gate_logits is None or not isinstance(gate_logits, tuple):100 return 0101 102 if isinstance(gate_logits, tuple):103 compute_device = gate_logits[0].device104 concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)105 106 routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)107 108 _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)109 110 expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)111 112 if attention_mask is None:113 # Compute the percentage of tokens routed to each experts114 tokens_per_expert = torch.mean(expert_mask.float(), dim=0)115 116 # Compute the average probability of routing to these experts117 router_prob_per_expert = torch.mean(routing_weights, dim=0)118 else:119 batch_size, sequence_length = attention_mask.shape120 num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)121 122 # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask123 expert_attention_mask = (124 attention_mask[None, :, :, None, None]125 .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))126 .reshape(-1, top_k, num_experts)127 .to(compute_device)128 )129 130 # Compute the percentage of tokens routed to each experts131 tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(132 expert_attention_mask, dim=0133 )134 135 # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert136 router_per_expert_attention_mask = (137 attention_mask[None, :, :, None]138 .expand((num_hidden_layers, batch_size, sequence_length, num_experts))139 .reshape(-1, num_experts)140 .to(compute_device)141 )142 143 # Compute the average probability of routing to these experts144 router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(145 router_per_expert_attention_mask, dim=0146 )147 148 overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))149 return overall_loss * num_experts150 151 152# Copied from transformers.models.llama.modeling_llama._get_unpad_data153def _get_unpad_data(attention_mask):154 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)155 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()156 max_seqlen_in_batch = seqlens_in_batch.max().item()157 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))158 return (159 indices,160 cu_seqlens,161 max_seqlen_in_batch,162 )163 164 165# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->PhiMoE166##https://dl.acm.org/doi/pdf/10.5555/3454287.3455397 The following is the implementation of layernorm167 168 169class PhiMoERotaryEmbedding(nn.Module):170 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):171 super().__init__()172 173 self.dim = dim174 self.max_position_embeddings = max_position_embeddings175 self.base = base176 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))177 self.register_buffer("inv_freq", inv_freq, persistent=False)178 179 # Build here to make `torch.jit.trace` work.180 self._set_cos_sin_cache(181 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()182 )183 184 def _set_cos_sin_cache(self, seq_len, device, dtype):185 self.max_seq_len_cached = seq_len186 t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)187 188 freqs = torch.outer(t, self.inv_freq)189 # Different from paper, but it uses a different permutation in order to obtain the same calculation190 emb = torch.cat((freqs, freqs), dim=-1)191 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)192 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)193 194 def forward(self, x, seq_len=None):195 # x: [bs, num_attention_heads, seq_len, head_size]196 if seq_len > self.max_seq_len_cached:197 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)198 199 return (200 self.cos_cached[:seq_len].to(dtype=x.dtype),201 self.sin_cached[:seq_len].to(dtype=x.dtype),202 )203 204 205class Phi3LongRoPEScaledRotaryEmbedding(nn.Module):206 207 def __init__(self, dim, config):208 super().__init__()209 self.dim = dim210 self.max_position_embeddings = config.max_position_embeddings211 self.base = config.rope_theta212 self.short_factor = config.rope_scaling["short_factor"]213 self.long_factor = config.rope_scaling["long_factor"]214 self.short_mscale = config.rope_scaling["short_mscale"]215 self.long_mscale = config.rope_scaling["long_mscale"]216 self.original_max_position_embeddings = config.rope_scaling["original_max_position_embeddings"]217 218 def forward(self, x, seq_len=None):219 if seq_len is None:220 seq_len = x.shape[-2]221 222 if seq_len > self.original_max_position_embeddings:223 rescale_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)224 mscale = self.long_mscale225 else:226 rescale_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)227 mscale = self.short_mscale228 assert rescale_factors.shape == (self.dim // 2, ), \229 f"misaligned shape for LongRoPE rescale factors: {rescale_factors.shape}"230 231 inv_freq = 1.0 / (rescale_factors * (self.base ** (torch.arange(0, self.dim, 2).float().to(x.device) / self.dim)))232 233 t = torch.arange(seq_len, device=x.device, dtype=torch.float32)234 freqs = torch.outer(t, inv_freq)235 236 emb = torch.cat((freqs, freqs), dim=-1)237 return (emb.cos() * mscale).to(x.dtype), (emb.sin() * mscale).to(x.dtype)238 239 240# Copied from transformers.models.llama.modeling_llama.rotate_half241def rotate_half(x):242 """Rotates half the hidden dims of the input."""243 x1 = x[..., : x.shape[-1] // 2]244 x2 = x[..., x.shape[-1] // 2 :]245 return torch.cat((-x2, x1), dim=-1)246 247 248 249def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):250 """Applies Rotary Position Embedding to the query and key tensors.251 252 Args:253 q (`torch.Tensor`): The query tensor.254 k (`torch.Tensor`): The key tensor.255 cos (`torch.Tensor`): The cosine part of the rotary embedding.256 sin (`torch.Tensor`): The sine part of the rotary embedding.257 position_ids (`torch.Tensor`):258 The position indices of the tokens corresponding to the query and key tensors. For example, this can be259 used to pass offsetted position ids when working with a KV-cache.260 unsqueeze_dim (`int`, *optional*, defaults to 1):261 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and262 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note263 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and264 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes265 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have266 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.267 Returns:268 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.269 """270 cos = cos[position_ids].unsqueeze(unsqueeze_dim)271 sin = sin[position_ids].unsqueeze(unsqueeze_dim)272 q_embed = (q * cos) + (rotate_half(q) * sin)273 k_embed = (k * cos) + (rotate_half(k) * sin)274 return q_embed, k_embed275 276 277# Copied from transformers.models.llama.modeling_llama.repeat_kv278def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:279 """280 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,281 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)282 """283 batch, num_key_value_heads, slen, head_dim = hidden_states.shape284 if n_rep == 1:285 return hidden_states286 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)287 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)288 289 290 291class PhiMoEAttention(nn.Module):292 """293 Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer294 and "Generating Long Sequences with Sparse Transformers".295 """296 297 def __init__(self, config: PhiMoEConfig, layer_idx: Optional[int] = None):298 super().__init__()299 self.config = config300 self.layer_idx = layer_idx301 if layer_idx is None:302 logger.warning_once(303 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "304 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "305 "when creating this class."306 )307 308 self.hidden_size = config.hidden_size309 self.num_heads = config.num_attention_heads310 self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)311 self.num_key_value_heads = config.num_key_value_heads312 self.num_key_value_groups = self.num_heads // self.num_key_value_heads313 self.max_position_embeddings = config.max_position_embeddings314 self.rope_theta = config.rope_theta315 self.is_causal = True316 self.attention_dropout = config.attention_dropout317 318 # if (self.head_dim * self.num_heads) != self.hidden_size:319 # raise ValueError(320 # f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"321 # f" and `num_heads`: {self.num_heads})."322 # )323 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.config.attention_bias)324 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.config.attention_bias)325 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.config.attention_bias)326 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.config.attention_bias)327 328 if getattr(config, 'rope_scaling', None) is None:329 self.rotary_emb = PhiMoERotaryEmbedding(330 self.head_dim,331 max_position_embeddings=self.max_position_embeddings,332 base=self.rope_theta,333 )334 else:335 scaling_type = self.config.rope_scaling["type"]336 if scaling_type == "longrope":337 self.rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(self.head_dim, self.config)338 else:339 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")340 341 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):342 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()343 344 def forward(345 self,346 hidden_states: torch.Tensor,347 attention_mask: Optional[torch.Tensor] = None,348 position_ids: Optional[torch.LongTensor] = None,349 past_key_value: Optional[Cache] = None,350 output_attentions: bool = False,351 use_cache: bool = False,352 **kwargs,353 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:354 if "padding_mask" in kwargs:355 warnings.warn(356 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"357 )358 bsz, q_len, _ = hidden_states.size()359 360 query_states = self.q_proj(hidden_states)361 key_states = self.k_proj(hidden_states)362 value_states = self.v_proj(hidden_states)363 364 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)365 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)366 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)367 368 kv_seq_len = key_states.shape[-2]369 if past_key_value is not None:370 if self.layer_idx is None:371 raise ValueError(372 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "373 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "374 "with a layer index."375 )376 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)377 378 # print ("before apply rotary pos_emb", len(kv_seq_len),torch.norm(value_states).items(),\379 # torch.norm(query_states).items(), torch.norm(key_states).items(), position_ids)380 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)381 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)382 383 # print ('after pos emb', torch.norm(query_states).item(), torch.norm(key_states).items())384 if past_key_value is not None:385 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models386 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)387 388 # repeat k/v heads if n_kv_heads < n_heads389 key_states = repeat_kv(key_states, self.num_key_value_groups)390 value_states = repeat_kv(value_states, self.num_key_value_groups)391 392 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)393 394 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):395 raise ValueError(396 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"397 f" {attn_weights.size()}"398 )399 400 if attention_mask is not None:401 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):402 raise ValueError(403 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"404 )405 406 attn_weights = attn_weights + attention_mask407 408 # upcast attention to fp32409 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)410 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)411 attn_output = torch.matmul(attn_weights, value_states)412 413 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):414 raise ValueError(415 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"416 f" {attn_output.size()}"417 )418 419 attn_output = attn_output.transpose(1, 2).contiguous()420 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)421 422 attn_output = self.o_proj(attn_output)423 424 if not output_attentions:425 attn_weights = None426 427 return attn_output, attn_weights, past_key_value428 429 430 431class PhiMoEFlashAttention2(PhiMoEAttention):432 """433 PhiMoE flash attention module. This module inherits from `PhiMoEAttention` as the weights of the module stays434 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of435 flash attention and deal with padding tokens in case the input contains any of them.436 """437 438 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__439 def __init__(self, *args, **kwargs):440 super().__init__(*args, **kwargs)441 442 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.443 # 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.444 # 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).445 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()446 447 def forward(448 self,449 hidden_states: torch.Tensor,450 attention_mask: Optional[torch.Tensor] = None,451 position_ids: Optional[torch.LongTensor] = None,452 past_key_value: Optional[Cache] = None,453 output_attentions: bool = False,454 use_cache: bool = False,455 **kwargs,456 ):457 if "padding_mask" in kwargs:458 warnings.warn(459 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"460 )461 462 # overwrite attention_mask with padding_mask463 attention_mask = kwargs.pop("padding_mask")464 bsz, q_len, _ = hidden_states.size()465 466 query_states = self.q_proj(hidden_states)467 key_states = self.k_proj(hidden_states)468 value_states = self.v_proj(hidden_states)469 470 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)471 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)472 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)473 474 kv_seq_len = key_states.shape[-2]475 if past_key_value is not None:476 if self.layer_idx is None:477 raise ValueError(478 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "479 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "480 "with a layer index."481 )482 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)483 484 # Because the input can be padded, the absolute sequence length depends on the max position id.485 rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)486 cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)487 488 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)489 490 use_sliding_windows = (491 _flash_supports_window_size492 and getattr(self.config, "sliding_window", None) is not None493 and kv_seq_len > self.config.sliding_window494 )495 496 if not _flash_supports_window_size:497 logger.warning_once(498 "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"499 " make sure to upgrade flash-attn library."500 )501 502 if past_key_value is not None:503 # Activate slicing cache only if the config has a value `sliding_windows` attribute504 cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0505 if (506 getattr(self.config, "sliding_window", None) is not None507 and kv_seq_len > self.config.sliding_window508 and cache_has_contents509 ):510 slicing_tokens = 1 - self.config.sliding_window511 512 past_key = past_key_value[self.layer_idx][0]513 past_value = past_key_value[self.layer_idx][1]514 515 past_key = past_key[:, :, slicing_tokens:, :].contiguous()516 past_value = past_value[:, :, slicing_tokens:, :].contiguous()517 518 if past_key.shape[-2] != self.config.sliding_window - 1:519 raise ValueError(520 f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"521 f" {past_key.shape}"522 )523 524 if attention_mask is not None:525 attention_mask = attention_mask[:, slicing_tokens:]526 attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)527 528 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models529 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)530 531 # repeat k/v heads if n_kv_heads < n_heads532 key_states = repeat_kv(key_states, self.num_key_value_groups)533 value_states = repeat_kv(value_states, self.num_key_value_groups)534 dropout_rate = 0.0 if not self.training else self.attention_dropout535 536 # In PEFT, usually we cast the layer norms in float32 for training stability reasons537 # therefore the input hidden states gets silently casted in float32. Hence, we need538 # cast them back in float16 just to be sure everything works as expected.539 input_dtype = query_states.dtype540 if input_dtype == torch.float32:541 if torch.is_autocast_enabled():542 target_dtype = torch.get_autocast_gpu_dtype()543 # Handle the case where the model is quantized544 elif hasattr(self.config, "_pre_quantization_dtype"):545 target_dtype = self.config._pre_quantization_dtype546 else:547 target_dtype = self.q_proj.weight.dtype548 549 logger.warning_once(550 f"The input hidden states seems to be silently casted in float32, this might be related to"551 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"552 f" {target_dtype}."553 )554 555 query_states = query_states.to(target_dtype)556 key_states = key_states.to(target_dtype)557 value_states = value_states.to(target_dtype)558 559 # Reashape to the expected shape for Flash Attention560 query_states = query_states.transpose(1, 2)561 key_states = key_states.transpose(1, 2)562 value_states = value_states.transpose(1, 2)563 564 attn_output = self._flash_attention_forward(565 query_states,566 key_states,567 value_states,568 attention_mask,569 q_len,570 dropout=dropout_rate,571 use_sliding_windows=use_sliding_windows,572 )573 574 attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads).contiguous()575 attn_output = self.o_proj(attn_output)576 577 if not output_attentions:578 attn_weights = None579 580 return attn_output, attn_weights, past_key_value581 582 def _flash_attention_forward(583 self,584 query_states,585 key_states,586 value_states,587 attention_mask,588 query_length,589 dropout=0.0,590 softmax_scale=None,591 use_sliding_windows=False,592 ):593 """594 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token595 first unpad the input, then computes the attention scores and pad the final attention scores.596 597 Args:598 query_states (`torch.Tensor`):599 Input query states to be passed to Flash Attention API600 key_states (`torch.Tensor`):601 Input key states to be passed to Flash Attention API602 value_states (`torch.Tensor`):603 Input value states to be passed to Flash Attention API604 attention_mask (`torch.Tensor`):605 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the606 position of padding tokens and 1 for the position of non-padding tokens.607 dropout (`float`):608 Attention dropout609 softmax_scale (`float`, *optional*):610 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)611 use_sliding_windows (`bool`, *optional*):612 Whether to activate sliding window attention.613 """614 if not self._flash_attn_uses_top_left_mask:615 causal = self.is_causal616 else:617 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.618 causal = self.is_causal and query_length != 1619 620 # Contains at least one padding token in the sequence621 if attention_mask is not None:622 batch_size = query_states.shape[0]623 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(624 query_states, key_states, value_states, attention_mask, query_length625 )626 627 cu_seqlens_q, cu_seqlens_k = cu_seq_lens628 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens629 630 if not use_sliding_windows:631 attn_output_unpad = flash_attn_varlen_func(632 query_states,633 key_states,634 value_states,635 cu_seqlens_q=cu_seqlens_q,636 cu_seqlens_k=cu_seqlens_k,637 max_seqlen_q=max_seqlen_in_batch_q,638 max_seqlen_k=max_seqlen_in_batch_k,639 dropout_p=dropout,640 softmax_scale=softmax_scale,641 causal=causal,642 )643 else:644 attn_output_unpad = flash_attn_varlen_func(645 query_states,646 key_states,647 value_states,648 cu_seqlens_q=cu_seqlens_q,649 cu_seqlens_k=cu_seqlens_k,650 max_seqlen_q=max_seqlen_in_batch_q,651 max_seqlen_k=max_seqlen_in_batch_k,652 dropout_p=dropout,653 softmax_scale=softmax_scale,654 causal=causal,655 window_size=(self.config.sliding_window, 0),656 )657 658 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)659 else:660 if not use_sliding_windows:661 attn_output = flash_attn_func(662 query_states,663 key_states,664 value_states,665 dropout,666 softmax_scale=softmax_scale,667 causal=causal,668 )669 else:670 attn_output = flash_attn_func(671 query_states,672 key_states,673 value_states,674 dropout,675 softmax_scale=softmax_scale,676 causal=causal,677 window_size=(self.config.sliding_window, 0),678 )679 680 return attn_output681 682 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):683 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape684 685 # On the first iteration we need to properly re-create the padding mask686 # by slicing it on the proper place687 if kv_seq_len != attention_mask.shape[-1]:688 attention_mask_num_tokens = attention_mask.shape[-1]689 attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]690 691 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)692 693 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)694 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)695 696 if query_length == kv_seq_len:697 query_layer = index_first_axis(698 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k699 )700 cu_seqlens_q = cu_seqlens_k701 max_seqlen_in_batch_q = max_seqlen_in_batch_k702 indices_q = indices_k703 elif query_length == 1:704 max_seqlen_in_batch_q = 1705 cu_seqlens_q = torch.arange(706 batch_size + 1, dtype=torch.int32, device=query_layer.device707 ) # There is a memcpy here, that is very bad.708 indices_q = cu_seqlens_q[:-1]709 query_layer = query_layer.squeeze(1)710 else:711 # The -q_len: slice assumes left padding.712 attention_mask = attention_mask[:, -query_length:]713 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)714 715 return (716 query_layer,717 key_layer,718 value_layer,719 indices_q,720 (cu_seqlens_q, cu_seqlens_k),721 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),722 )723 724 725 726class PhiMoESdpaAttention(PhiMoEAttention):727 """728 PhiMoE attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from729 `PhiMoEAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to730 SDPA API.731 """732 733 # Adapted from PhiMoEAttention.forward734 def forward(735 self,736 hidden_states: torch.Tensor,737 attention_mask: Optional[torch.Tensor] = None,738 position_ids: Optional[torch.LongTensor] = None,739 past_key_value: Optional[Cache] = None,740 output_attentions: bool = False,741 use_cache: bool = False,742 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:743 if output_attentions:744 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.745 logger.warning_once(746 "PhiMoEModel is using PhiMoESdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "747 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'748 )749 return super().forward(750 hidden_states=hidden_states,751 attention_mask=attention_mask,752 position_ids=position_ids,753 past_key_value=past_key_value,754 output_attentions=output_attentions,755 use_cache=use_cache,756 )757 758 bsz, q_len, _ = hidden_states.size()759 760 query_states = self.q_proj(hidden_states)761 key_states = self.k_proj(hidden_states)762 value_states = self.v_proj(hidden_states)763 764 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)765 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)766 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)767 768 kv_seq_len = key_states.shape[-2]769 if past_key_value is not None:770 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)771 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)772 773 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)774 775 if past_key_value is not None:776 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models777 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)778 779 key_states = repeat_kv(key_states, self.num_key_value_groups)780 value_states = repeat_kv(value_states, self.num_key_value_groups)781 782 if attention_mask is not None:783 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):784 raise ValueError(785 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"786 )787 788 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,789 # Reference: https://github.com/pytorch/pytorch/issues/112577.790 if query_states.device.type == "cuda" and attention_mask is not None:791 query_states = query_states.contiguous()792 key_states = key_states.contiguous()793 value_states = value_states.contiguous()794 795 attn_output = torch.nn.functional.scaled_dot_product_attention(796 query_states,797 key_states,798 value_states,799 attn_mask=attention_mask,800 dropout_p=self.attention_dropout if self.training else 0.0,801 # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.802 is_causal=self.is_causal and attention_mask is None and q_len > 1,803 )804 805 attn_output = attn_output.transpose(1, 2).contiguous()806 attn_output = attn_output.view(bsz, q_len, self.head_dim * self.num_heads)807 808 attn_output = self.o_proj(attn_output)809 810 return attn_output, None, past_key_value811 812 813PHIMOE_ATTENTION_CLASSES = {814 "eager": PhiMoEAttention,815 "flash_attention_2": PhiMoEFlashAttention2,816 "sdpa": PhiMoESdpaAttention,817}818 819 820class PhiMoEBlockSparseTop2MLP(nn.Module):821 def __init__(self, config: PhiMoEConfig):822 super().__init__()823 self.ffn_dim = config.intermediate_size824 self.hidden_dim = config.hidden_size825 826 self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)827 self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)828 self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)829 830 self.act_fn = ACT2FN[config.hidden_act]831 832 def forward(self, hidden_states):833 current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)834 current_hidden_states = self.w2(current_hidden_states)835 return current_hidden_states836 837 838class PhiMoEBLockSparseTop2MLP(PhiMoEBlockSparseTop2MLP):839 def __init__(self, *args, **kwargs):840 logger.warning_once(841 "PhiMoEBLockSparseTop2MLP is deprecated by PhiMoEBlockSparseTop2MLP and will be removed in v4.40."842 )843 super().__init__(*args, **kwargs)844 845 846class mp(torch.autograd.Function):847 @staticmethod848 def forward(849 ctx, 850 scores: torch.Tensor, 851 multiplier: torch.Tensor, 852 selected_experts: torch.Tensor,853 masked_gates: torch.Tensor,854 mask_for_one: torch.Tensor,855 ):856 ctx.save_for_backward(multiplier, selected_experts, masked_gates)857 return multiplier * mask_for_one858 859 @staticmethod860 def backward(861 ctx, 862 grad_at_output: torch.Tensor, 863 ):864 multiplier, selected_experts, masked_gates = ctx.saved_tensors865 866 grad_at_output = grad_at_output * multiplier867 868 grad_at_scores_expaned = masked_gates * grad_at_output.mul(-1)869 grad_at_scores_expaned.scatter_add_(870 dim=-1,871 index=selected_experts,872 src=grad_at_output,873 )874 875 return (876 grad_at_scores_expaned, 877 None, 878 None, 879 None, 880 None, 881 )882 883def sparsemixer(scores, top_k, jitter_eps, training):884 assert top_k == 2885 886 ################ first expert ################887 888 with torch.no_grad():889 # compute mask for sparsity890 mask_logits_threshold, max_ind = scores.max(dim=-1, keepdim=True)891 factor = scores.abs().clamp(min=mask_logits_threshold)892 mask_logits_threshold = (893 (mask_logits_threshold - scores) / factor894 ) > (2 * jitter_eps)895 896 # apply mask 897 masked_gates = scores.masked_fill(mask_logits_threshold, float('-inf'))898 if training:899 selected_experts = (900 masked_gates - torch.empty_like(masked_gates, memory_format=torch.legacy_contiguous_format).exponential_().log()901 ).max(dim=-1)[1].unsqueeze(-1) # gumbel sampling, more robust than than the multinomial method902 else:903 selected_experts = max_ind904 905 # compute scores for gradients906 masked_gates = torch.softmax(masked_gates, dim=-1)907 multiplier_o = masked_gates.gather(dim=-1, index=selected_experts)908 909 if training:910 # compute midpoint mask 911 max_scores, max_ind = masked_gates.max(dim=-1, keepdim=True)912 mask_for_one = torch.logical_or(913 selected_experts == max_ind,914 torch.rand_like(max_scores) > 0.75 # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.)915 ) 916 # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5917 mask_for_one = torch.add(0.3333, mask_for_one, alpha=0.6667).type_as(masked_gates)918 919 multiplier = mp.apply(920 scores, 921 multiplier_o, 922 selected_experts, 923 masked_gates, 924 mask_for_one,925 )926 else:927 multiplier = multiplier_o928 929 # masked out first expert 930 masked_scores = torch.scatter(931 scores,932 -1,933 selected_experts,934 float('-inf'),935 )936 with torch.no_grad():937 # compute mask for sparsity938 mask_logits_threshold, max_ind = masked_scores.max(dim=-1, keepdim=True)939 factor = scores.abs().clamp(min=mask_logits_threshold)940 mask_logits_threshold = (941 (mask_logits_threshold - scores) / factor942 ) > (2 * jitter_eps)943 944 # apply mask 945 masked_gates_top2 = masked_scores.masked_fill(mask_logits_threshold, float('-inf'))946 if training:947 selected_experts_top2 = (948 masked_gates_top2 - torch.empty_like(masked_gates_top2, memory_format=torch.legacy_contiguous_format).exponential_().log()949 ).max(dim=-1)[1].unsqueeze(-1) # gumbel sampling, more robust than than the multinomial method950 else:951 selected_experts_top2 = max_ind952 # compute scores for gradients953 masked_gates_top2 = torch.softmax(masked_gates_top2, dim=-1)954 multiplier_top2_o = masked_gates_top2.gather(dim=-1, index=selected_experts_top2)955 956 if training: 957 # compute midpoint mask 958 max_scores, max_ind = masked_gates_top2.max(dim=-1, keepdim=True)959 mask_for_one_top2 = torch.logical_or(960 selected_experts_top2 == max_ind,961 torch.rand_like(max_scores).uniform_() > 0.75 # Heun's third-order method: f(x) - f(0) = .25 f'(x) + .75 f'(x/3.)962 ) 963 # 1 -> 1.0 & 0 -> 1./3: lambda x: (x + 0.5) / 1.5964 mask_for_one_top2 = torch.add(0.3333, mask_for_one_top2, alpha=0.6667).type_as(masked_gates_top2)965 966 multiplier_top2 = mp.apply(967 scores, 968 multiplier_top2_o, 969 selected_experts_top2, 970 masked_gates_top2, 971 mask_for_one_top2,972 )973 else:974 multiplier_top2 = multiplier_top2_o975 976 multiplier = torch.concat((multiplier, multiplier_top2), dim=-1)977 selected_experts = torch.concat((selected_experts, selected_experts_top2), dim=-1)978 979 return (980 multiplier, 981 selected_experts,982 )983 984iterations = 0985class PhiMoESparseMoeBlock(nn.Module):986 """987 This implementation is988 strictly equivalent to standard MoE with full capacity (no989 dropped tokens). It's faster since it formulates MoE operations990 in terms of block-sparse operations to accomodate imbalanced991 assignments of tokens to experts, whereas standard MoE either992 (1) drop tokens at the cost of reduced performance or (2) set993 capacity factor to number of experts and thus waste computation994 and memory on padding.995 """996 997 def __init__(self, config):998 super().__init__()999 self.hidden_dim = config.hidden_size1000 self.ffn_dim = config.intermediate_size1001 self.num_experts = config.num_local_experts1002 self.top_k = config.num_experts_per_tok1003 global iterations1004 iterations +=11005 self.iter = iterations1006 # gating1007 self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)1008 1009 self.experts = nn.ModuleList([PhiMoEBlockSparseTop2MLP(config) for _ in range(self.num_experts)])1010 1011 # Jitter parameters1012 self.router_jitter_noise = config.router_jitter_noise1013 self.input_jitter_noise = config.input_jitter_noise1014 1015 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:1016 """ """1017 batch_size, sequence_length, hidden_dim = hidden_states.shape1018 if self.training and self.input_jitter_noise > 0:1019 hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.input_jitter_noise, 1.0 + self.input_jitter_noise)1020 hidden_states = hidden_states.view(-1, hidden_dim)1021 # router_logits: (batch * sequence_length, n_experts)1022 # print ( 'moe', self.iter, torch.norm(hidden_states).item())1023 router_logits = self.gate(hidden_states)1024 1025 routing_weights, selected_experts = sparsemixer(1026 router_logits, 1027 top_k=2, 1028 jitter_eps=self.router_jitter_noise, 1029 training=self.training,1030 )1031 1032 final_hidden_states = torch.zeros(1033 (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device1034 )1035 1036 # One hot encode the selected experts to create an expert mask1037 # this will be used to easily index which expert is going to be sollicitated1038 expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)1039 1040 # Loop over all available experts in the model and perform the computation on each expert1041 for expert_idx in range(self.num_experts):1042 expert_layer = self.experts[expert_idx]1043 idx, top_x = torch.where(expert_mask[expert_idx])1044 1045 if top_x.shape[0] == 0:1046 continue1047 1048 # in torch it is faster to index using lists than torch tensors1049 top_x_list = top_x.tolist()1050 idx_list = idx.tolist()1051 1052 # Index the correct hidden states and compute the expert hidden state for1053 # the current expert. We need to make sure to multiply the output hidden1054 # states by `routing_weights` on the corresponding tokens (top-1 and top-2)1055 current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)1056 current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None]1057 1058 # However `index_add_` only support torch tensors for indexing so we'll use1059 # the `top_x` tensor here.1060 final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))1061 final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)1062 # print ( 'moe', self.iter, torch.norm(final_hidden_states).item())1063 return final_hidden_states, router_logits1064 1065 1066class PhiMoEDecoderLayer(nn.Module):1067 def __init__(self, config: PhiMoEConfig, layer_idx: int):1068 super().__init__()1069 self.hidden_size = config.hidden_size1070 1071 self.self_attn = PHIMOE_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)1072 1073 self.block_sparse_moe = PhiMoESparseMoeBlock(config)1074 self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=True)1075 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=True)1076 1077 def forward(1078 self,1079 hidden_states: torch.Tensor,1080 attention_mask: Optional[torch.Tensor] = None,1081 position_ids: Optional[torch.LongTensor] = None,1082 past_key_value: Optional[Tuple[torch.Tensor]] = None,1083 output_attentions: Optional[bool] = False,1084 output_router_logits: Optional[bool] = False,1085 use_cache: Optional[bool] = False,1086 **kwargs,1087 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:1088 if "padding_mask" in kwargs:1089 warnings.warn(1090 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"1091 )1092 """1093 Args:1094 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`1095 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size1096 `(batch, sequence_length)` where padding elements are indicated by 0.1097 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states1098 output_attentions (`bool`, *optional*):1099 Whether or not to return the attentions tensors of all attention layers. See `attentions` under1100 returned tensors for more detail.1101 output_router_logits (`bool`, *optional*):1102 Whether or not to return the logits of all the routers. They are useful for computing the router loss, and1103 should not be returned during inference.1104 use_cache (`bool`, *optional*):1105 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding1106 (see `past_key_values`).1107 """1108 1109 residual = hidden_states1110 1111 hidden_states = self.input_layernorm(hidden_states)1112 1113 # Self Attention1114 hidden_states, self_attn_weights, present_key_value = self.self_attn(1115 hidden_states=hidden_states,1116 attention_mask=attention_mask,1117 position_ids=position_ids,1118 past_key_value=past_key_value,1119 output_attentions=output_attentions,1120 use_cache=use_cache,1121 )1122 hidden_states = residual + hidden_states1123 1124 # Fully Connected1125 residual = hidden_states1126 hidden_states = self.post_attention_layernorm(hidden_states)1127 hidden_states, router_logits = self.block_sparse_moe(hidden_states)1128 hidden_states = residual + hidden_states1129 1130 outputs = (hidden_states,)1131 1132 if output_attentions:1133 outputs += (self_attn_weights,)1134 1135 if use_cache:1136 outputs += (present_key_value,)1137 1138 if output_router_logits:1139 outputs += (router_logits,)1140 1141 return outputs1142 1143 1144PHIMOE_START_DOCSTRING = r"""1145 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the1146 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads1147 etc.)1148 1149 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.1150 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage1151 and behavior.1152 1153 Parameters:1154 config ([`PhiMoEConfig`]):1155 Model configuration class with all the parameters of the model. Initializing with a config file does not1156 load the weights associated with the model, only the configuration. Check out the1157 [`~PreTrainedModel.from_pretrained`] method to load the model weights.1158"""1159 1160 1161@add_start_docstrings(1162 "The bare PhiMoE Model outputting raw hidden-states without any specific head on top.",1163 PHIMOE_START_DOCSTRING,1164)1165 1166class PhiMoEPreTrainedModel(PreTrainedModel):1167 config_class = PhiMoEConfig1168 base_model_prefix = "model"1169 supports_gradient_checkpointing = True1170 _no_split_modules = ["PhiMoEDecoderLayer"]1171 _skip_keys_device_placement = "past_key_values"1172 _supports_flash_attn_2 = True1173 _supports_sdpa = True1174 _supports_cache_class = True1175 1176 def _init_weights(self, module):1177 pass1178 # std = self.config.initializer_range1179 # if isinstance(module, nn.Linear):1180 # module.weight.data.normal_(mean=0.0, std=std)1181 # if module.bias is not None:1182 # module.bias.data.zero_()1183 # elif isinstance(module, nn.Embedding):1184 # module.weight.data.normal_(mean=0.0, std=std)1185 # if module.padding_idx is not None:1186 # module.weight.data[module.padding_idx].zero_()1187 1188 1189PHIMOE_INPUTS_DOCSTRING = r"""1190 Args:1191 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):1192 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide1193 it.1194 1195 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and1196 [`PreTrainedTokenizer.__call__`] for details.1197 1198 [What are input IDs?](../glossary#input-ids)1199 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1200 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: