nold/phi-3-medium-4k-instruct-8bit
125
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 Phi-3 model."""17 18import inspect19import math20import warnings21from typing import List, Optional, Tuple, Union22 23import torch24import torch.nn.functional as F25import torch.utils.checkpoint26from torch import nn27from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss28 29from transformers.activations import ACT2FN30from transformers.cache_utils import Cache, DynamicCache31from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask32from transformers.modeling_outputs import (33 BaseModelOutputWithPast,34 CausalLMOutputWithPast,35 SequenceClassifierOutputWithPast,36 TokenClassifierOutput,37)38from transformers.modeling_utils import PreTrainedModel39from transformers.utils import (40 add_code_sample_docstrings,41 add_start_docstrings,42 add_start_docstrings_to_model_forward,43 is_flash_attn_2_available,44 is_flash_attn_greater_or_equal_2_10,45 logging,46 replace_return_docstrings,47)48from .configuration_phi3 import Phi3Config49 50 51logger = logging.get_logger(__name__)52 53# Transformers scans dependencies in the modeling file, causing issues on conditional loading. The regex only ignores try/catch blocks, but not if statements54# if is_flash_attn_2_available():55_flash_supports_window_size = False56try: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)61except ImportError as error:62 logger.warning(63 f"`flash-attention` package not found, consider installing for better performance: {error}."64 )65 if not _flash_supports_window_size:66 logger.warning(67 "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`."68 )69 70_CHECKPOINT_FOR_DOC = "microsoft/Phi-3-mini-4k-instruct"71_CONFIG_FOR_DOC = "Phi3Config"72 73PHI3_PRETRAINED_MODEL_ARCHIVE_LIST = [74 "microsoft/Phi-3-mini-4k-instruct",75 "microsoft/Phi-3-mini-128k-instruct",76 # See all Phi-3 models at https://huggingface.co/models?filter=Phi-377]78 79 80# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Phi381class Phi3RMSNorm(nn.Module):82 def __init__(self, hidden_size, eps=1e-6):83 """84 Phi3RMSNorm is equivalent to T5LayerNorm85 """86 super().__init__()87 self.weight = nn.Parameter(torch.ones(hidden_size))88 self.variance_epsilon = eps89 90 def forward(self, hidden_states):91 input_dtype = hidden_states.dtype92 hidden_states = hidden_states.to(torch.float32)93 variance = hidden_states.pow(2).mean(-1, keepdim=True)94 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)95 return self.weight * hidden_states.to(input_dtype)96 97 98# Copied from transformers.models.llama.modeling_llama._get_unpad_data99def _get_unpad_data(attention_mask):100 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)101 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()102 max_seqlen_in_batch = seqlens_in_batch.max().item()103 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))104 return (105 indices,106 cu_seqlens,107 max_seqlen_in_batch,108 )109 110 111# Copied from transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding with gemma->phi3, Gemma->Phi3112class Phi3RotaryEmbedding(nn.Module):113 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):114 super().__init__()115 116 self.dim = dim117 self.max_position_embeddings = max_position_embeddings118 self.base = base119 self.register_buffer("inv_freq", None, persistent=False)120 121 @torch.no_grad()122 def forward(self, x, position_ids, seq_len=None):123 # x: [bs, num_attention_heads, seq_len, head_size]124 if self.inv_freq is None:125 self.inv_freq = 1.0 / (126 self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)127 )128 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)129 position_ids_expanded = position_ids[:, None, :].float()130 # Force float32 since bfloat16 loses precision on long contexts131 # See https://github.com/huggingface/transformers/pull/29285132 device_type = x.device.type133 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"134 with torch.autocast(device_type=device_type, enabled=False):135 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)136 emb = torch.cat((freqs, freqs), dim=-1)137 cos = emb.cos()138 sin = emb.sin()139 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)140 141 142class Phi3SuScaledRotaryEmbedding(Phi3RotaryEmbedding):143 def __init__(self, dim, config, device=None):144 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)145 146 self.short_factor = config.rope_scaling["short_factor"]147 self.long_factor = config.rope_scaling["long_factor"]148 self.original_max_position_embeddings = config.original_max_position_embeddings149 150 @torch.no_grad()151 def forward(self, x, position_ids, seq_len=None):152 seq_len = torch.max(position_ids) + 1153 if seq_len > self.original_max_position_embeddings:154 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)155 else:156 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)157 158 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim159 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)160 161 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)162 position_ids_expanded = position_ids[:, None, :].float()163 164 # Force float32 since bfloat16 loses precision on long contexts165 # See https://github.com/huggingface/transformers/pull/29285166 device_type = x.device.type167 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"168 with torch.autocast(device_type=device_type, enabled=False):169 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)170 emb = torch.cat((freqs, freqs), dim=-1)171 172 scale = self.max_position_embeddings / self.original_max_position_embeddings173 if scale <= 1.0:174 scaling_factor = 1.0175 else:176 scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))177 178 cos = emb.cos() * scaling_factor179 sin = emb.sin() * scaling_factor180 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)181 182 183class Phi3YarnScaledRotaryEmbedding(Phi3RotaryEmbedding):184 def __init__(self, dim, config, device=None):185 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)186 187 self.short_factor = config.rope_scaling["short_factor"]188 self.long_factor = config.rope_scaling["long_factor"]189 self.original_max_position_embeddings = config.original_max_position_embeddings190 191 @torch.no_grad()192 def forward(self, x, position_ids, seq_len=None):193 seq_len = torch.max(position_ids) + 1194 if seq_len > self.original_max_position_embeddings:195 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)196 else:197 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)198 199 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim200 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)201 202 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)203 position_ids_expanded = position_ids[:, None, :].float()204 205 # Force float32 since bfloat16 loses precision on long contexts206 # See https://github.com/huggingface/transformers/pull/29285207 device_type = x.device.type208 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"209 with torch.autocast(device_type=device_type, enabled=False):210 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)211 emb = torch.cat((freqs, freqs), dim=-1)212 213 scale = self.max_position_embeddings / self.original_max_position_embeddings214 if scale <= 1.0:215 scaling_factor = 1.0216 else:217 scaling_factor = 0.1 * math.log(scale) + 1.0218 219 cos = emb.cos() * scaling_factor220 sin = emb.sin() * scaling_factor221 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)222 223 224# Copied from transformers.models.llama.modeling_llama.rotate_half225def rotate_half(x):226 """Rotates half the hidden dims of the input."""227 x1 = x[..., : x.shape[-1] // 2]228 x2 = x[..., x.shape[-1] // 2 :]229 return torch.cat((-x2, x1), dim=-1)230 231 232# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb233def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):234 """Applies Rotary Position Embedding to the query and key tensors.235 236 Args:237 q (`torch.Tensor`): The query tensor.238 k (`torch.Tensor`): The key tensor.239 cos (`torch.Tensor`): The cosine part of the rotary embedding.240 sin (`torch.Tensor`): The sine part of the rotary embedding.241 position_ids (`torch.Tensor`, *optional*):242 Deprecated and unused.243 unsqueeze_dim (`int`, *optional*, defaults to 1):244 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and245 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note246 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and247 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes248 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have249 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.250 Returns:251 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.252 """253 cos = cos.unsqueeze(unsqueeze_dim)254 sin = sin.unsqueeze(unsqueeze_dim)255 q_embed = (q * cos) + (rotate_half(q) * sin)256 k_embed = (k * cos) + (rotate_half(k) * sin)257 return q_embed, k_embed258 259 260class Phi3MLP(nn.Module):261 def __init__(self, config):262 super().__init__()263 264 self.config = config265 self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)266 self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)267 268 self.activation_fn = ACT2FN[config.hidden_act]269 270 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:271 up_states = self.gate_up_proj(hidden_states)272 273 gate, up_states = up_states.chunk(2, dim=-1)274 up_states = up_states * self.activation_fn(gate)275 276 return self.down_proj(up_states)277 278 279# Copied from transformers.models.llama.modeling_llama.repeat_kv with llama->phi280def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:281 """282 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,283 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)284 """285 batch, num_key_value_heads, slen, head_dim = hidden_states.shape286 if n_rep == 1:287 return hidden_states288 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)289 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)290 291 292class Phi3Attention(nn.Module):293 """Multi-headed attention from 'Attention Is All You Need' paper"""294 295 def __init__(self, config: Phi3Config, layer_idx: Optional[int] = None):296 super().__init__()297 self.config = config298 self.layer_idx = layer_idx299 if layer_idx is None:300 logger.warning_once(301 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "302 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "303 "when creating this class."304 )305 306 self.attention_dropout = config.attention_dropout307 self.hidden_size = config.hidden_size308 self.num_heads = config.num_attention_heads309 self.head_dim = self.hidden_size // self.num_heads310 self.num_key_value_heads = config.num_key_value_heads311 self.num_key_value_groups = self.num_heads // self.num_key_value_heads312 self.max_position_embeddings = config.max_position_embeddings313 self.original_max_position_embeddings = config.original_max_position_embeddings314 self.rope_theta = config.rope_theta315 self.rope_scaling = config.rope_scaling316 self.is_causal = True317 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 324 op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim)325 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)326 self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False)327 self._init_rope()328 329 def _init_rope(self):330 if self.rope_scaling is None:331 self.rotary_emb = Phi3RotaryEmbedding(332 self.head_dim,333 max_position_embeddings=self.max_position_embeddings,334 base=self.rope_theta,335 )336 else:337 scaling_type = self.config.rope_scaling["type"]338 if scaling_type == "su":339 self.rotary_emb = Phi3SuScaledRotaryEmbedding(self.head_dim, self.config)340 elif scaling_type == "yarn":341 self.rotary_emb = Phi3YarnScaledRotaryEmbedding(self.head_dim, self.config)342 else:343 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")344 345 def forward(346 self,347 hidden_states: torch.Tensor,348 attention_mask: Optional[torch.Tensor] = None,349 position_ids: Optional[torch.LongTensor] = None,350 past_key_value: Optional[Cache] = None,351 output_attentions: bool = False,352 use_cache: bool = False,353 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:354 logger.warning_once("You are not running the flash-attention implementation, expect numerical differences.")355 356 bsz, q_len, _ = hidden_states.size()357 358 qkv = self.qkv_proj(hidden_states)359 query_pos = self.num_heads * self.head_dim360 query_states = qkv[..., :query_pos]361 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]362 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]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 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)378 379 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)380 381 if past_key_value is not None:382 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models383 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)384 385 # repeat k/v heads if n_kv_heads < n_heads386 key_states = repeat_kv(key_states, self.num_key_value_groups)387 value_states = repeat_kv(value_states, self.num_key_value_groups)388 389 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)390 391 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):392 raise ValueError(393 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"394 f" {attn_weights.size()}"395 )396 397 if attention_mask is not None:398 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):399 raise ValueError(400 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"401 )402 attn_weights = attn_weights + attention_mask403 404 # upcast attention to fp32405 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)406 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)407 408 attn_output = torch.matmul(attn_weights, value_states)409 410 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):411 raise ValueError(412 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"413 f" {attn_output.size()}"414 )415 416 attn_output = attn_output.transpose(1, 2).contiguous()417 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)418 419 attn_output = self.o_proj(attn_output)420 421 if not output_attentions:422 attn_weights = None423 424 return attn_output, attn_weights, past_key_value425 426 427class Phi3FlashAttention2(Phi3Attention):428 """429 Phi-3 flash attention module. This module inherits from `Phi3Attention` as the weights of the module stays430 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of431 flash attention and deal with padding tokens in case the input contains any of them.432 """433 434 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__435 def __init__(self, *args, **kwargs):436 super().__init__(*args, **kwargs)437 438 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.439 # 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.440 # 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).441 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()442 443 def forward(444 self,445 hidden_states: torch.Tensor,446 attention_mask: Optional[torch.LongTensor] = None,447 position_ids: Optional[torch.LongTensor] = None,448 past_key_value: Optional[Cache] = None,449 output_attentions: bool = False,450 use_cache: bool = False,451 **kwargs,452 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:453 # Phi3FlashAttention2 attention does not support output_attentions454 455 if not _flash_supports_window_size:456 logger.warning_once(457 "The current flash attention version does not support sliding window attention. Please use `attn_implementation='eager'` or upgrade flash-attn library."458 )459 raise ValueError("The current flash attention version does not support sliding window attention.")460 461 output_attentions = False462 463 if "padding_mask" in kwargs:464 warnings.warn(465 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"466 )467 468 # overwrite attention_mask with padding_mask469 attention_mask = kwargs.pop("padding_mask")470 471 bsz, q_len, _ = hidden_states.size()472 473 qkv = self.qkv_proj(hidden_states)474 query_pos = self.num_heads * self.head_dim475 query_states = qkv[..., :query_pos]476 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]477 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]478 479 # Flash attention requires the input to have the shape480 # batch_size x seq_length x head_dim x hidden_dim481 # therefore we just need to keep the original shape482 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)483 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)484 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)485 486 kv_seq_len = key_states.shape[-2]487 if past_key_value is not None:488 if self.layer_idx is None:489 raise ValueError(490 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "491 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "492 "with a layer index."493 )494 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)495 496 # Because the input can be padded, the absolute sequence length depends on the max position id.497 rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1498 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=rotary_seq_len)499 500 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)501 502 use_sliding_windows = (503 _flash_supports_window_size504 and getattr(self.config, "sliding_window", None) is not None505 and kv_seq_len > self.config.sliding_window506 )507 508 if past_key_value is not None:509 # Activate slicing cache only if the config has a value `sliding_windows` attribute510 cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0511 if (512 getattr(self.config, "sliding_window", None) is not None513 and kv_seq_len > self.config.sliding_window514 and cache_has_contents515 ):516 slicing_tokens = 1 - self.config.sliding_window517 518 past_key = past_key_value[self.layer_idx][0]519 past_value = past_key_value[self.layer_idx][1]520 521 past_key = past_key[:, :, slicing_tokens:, :].contiguous()522 past_value = past_value[:, :, slicing_tokens:, :].contiguous()523 524 if past_key.shape[-2] != self.config.sliding_window - 1:525 raise ValueError(526 f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"527 f" {past_key.shape}"528 )529 530 if attention_mask is not None:531 attention_mask = attention_mask[:, slicing_tokens:]532 attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)533 534 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models535 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)536 537 # repeat k/v heads if n_kv_heads < n_heads538 key_states = repeat_kv(key_states, self.num_key_value_groups)539 value_states = repeat_kv(value_states, self.num_key_value_groups)540 541 attn_dropout = self.attention_dropout if self.training else 0.0542 543 # In PEFT, usually we cast the layer norms in float32 for training stability reasons544 # therefore the input hidden states gets silently casted in float32. Hence, we need545 # cast them back in the correct dtype just to be sure everything works as expected.546 # This might slowdown training & inference so it is recommended to not cast the LayerNorms547 # in fp32.548 549 if query_states.dtype == torch.float32:550 if torch.is_autocast_enabled():551 target_dtype = torch.get_autocast_gpu_dtype()552 # Handle the case where the model is quantized553 elif hasattr(self.config, "_pre_quantization_dtype"):554 target_dtype = self.config._pre_quantization_dtype555 else:556 target_dtype = self.qkv_proj.weight.dtype557 558 logger.warning_once(559 f"The input hidden states seems to be silently casted in float32, this might be related to"560 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"561 f" {target_dtype}."562 )563 564 query_states = query_states.to(target_dtype)565 key_states = key_states.to(target_dtype)566 value_states = value_states.to(target_dtype)567 568 # Reashape to the expected shape for Flash Attention569 query_states = query_states.transpose(1, 2)570 key_states = key_states.transpose(1, 2)571 value_states = value_states.transpose(1, 2)572 573 attn_output = self._flash_attention_forward(574 query_states,575 key_states,576 value_states,577 attention_mask,578 q_len,579 dropout=attn_dropout,580 use_sliding_windows=use_sliding_windows,581 )582 583 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()584 attn_output = self.o_proj(attn_output)585 586 if not output_attentions:587 attn_weights = None588 589 return attn_output, attn_weights, past_key_value590 591 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._flash_attention_forward592 def _flash_attention_forward(593 self,594 query_states,595 key_states,596 value_states,597 attention_mask,598 query_length,599 dropout=0.0,600 softmax_scale=None,601 use_sliding_windows=False,602 ):603 """604 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token605 first unpad the input, then computes the attention scores and pad the final attention scores.606 607 Args:608 query_states (`torch.Tensor`):609 Input query states to be passed to Flash Attention API610 key_states (`torch.Tensor`):611 Input key states to be passed to Flash Attention API612 value_states (`torch.Tensor`):613 Input value states to be passed to Flash Attention API614 attention_mask (`torch.Tensor`):615 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the616 position of padding tokens and 1 for the position of non-padding tokens.617 dropout (`float`):618 Attention dropout619 softmax_scale (`float`, *optional*):620 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)621 use_sliding_windows (`bool`, *optional*):622 Whether to activate sliding window attention.623 """624 if not self._flash_attn_uses_top_left_mask:625 causal = self.is_causal626 else:627 # 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__.628 causal = self.is_causal and query_length != 1629 630 # Contains at least one padding token in the sequence631 if attention_mask is not None:632 batch_size = query_states.shape[0]633 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(634 query_states, key_states, value_states, attention_mask, query_length635 )636 637 cu_seqlens_q, cu_seqlens_k = cu_seq_lens638 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens639 640 if not use_sliding_windows:641 attn_output_unpad = flash_attn_varlen_func(642 query_states,643 key_states,644 value_states,645 cu_seqlens_q=cu_seqlens_q,646 cu_seqlens_k=cu_seqlens_k,647 max_seqlen_q=max_seqlen_in_batch_q,648 max_seqlen_k=max_seqlen_in_batch_k,649 dropout_p=dropout,650 softmax_scale=softmax_scale,651 causal=causal,652 )653 else:654 attn_output_unpad = flash_attn_varlen_func(655 query_states,656 key_states,657 value_states,658 cu_seqlens_q=cu_seqlens_q,659 cu_seqlens_k=cu_seqlens_k,660 max_seqlen_q=max_seqlen_in_batch_q,661 max_seqlen_k=max_seqlen_in_batch_k,662 dropout_p=dropout,663 softmax_scale=softmax_scale,664 causal=causal,665 window_size=(self.config.sliding_window, self.config.sliding_window),666 )667 668 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)669 else:670 if not use_sliding_windows:671 attn_output = flash_attn_func(672 query_states,673 key_states,674 value_states,675 dropout,676 softmax_scale=softmax_scale,677 causal=causal,678 )679 else:680 attn_output = flash_attn_func(681 query_states,682 key_states,683 value_states,684 dropout,685 softmax_scale=softmax_scale,686 causal=causal,687 window_size=(self.config.sliding_window, self.config.sliding_window),688 )689 690 return attn_output691 692 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input693 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):694 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape695 696 # On the first iteration we need to properly re-create the padding mask697 # by slicing it on the proper place698 if kv_seq_len != attention_mask.shape[-1]:699 attention_mask_num_tokens = attention_mask.shape[-1]700 attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]701 702 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)703 704 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)705 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)706 707 if query_length == kv_seq_len:708 query_layer = index_first_axis(709 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k710 )711 cu_seqlens_q = cu_seqlens_k712 max_seqlen_in_batch_q = max_seqlen_in_batch_k713 indices_q = indices_k714 elif query_length == 1:715 max_seqlen_in_batch_q = 1716 cu_seqlens_q = torch.arange(717 batch_size + 1, dtype=torch.int32, device=query_layer.device718 ) # There is a memcpy here, that is very bad.719 indices_q = cu_seqlens_q[:-1]720 query_layer = query_layer.squeeze(1)721 else:722 # The -q_len: slice assumes left padding.723 attention_mask = attention_mask[:, -query_length:]724 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)725 726 return (727 query_layer,728 key_layer,729 value_layer,730 indices_q,731 (cu_seqlens_q, cu_seqlens_k),732 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),733 )734 735 736# copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Phi3737# TODO @Arthur no longer copied from LLama after static cache738class Phi3SdpaAttention(Phi3Attention):739 """740 Phi3 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from741 `Phi3Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to742 SDPA API.743 """744 745 # Adapted from Phi3Attention.forward746 def forward(747 self,748 hidden_states: torch.Tensor,749 attention_mask: Optional[torch.Tensor] = None,750 position_ids: Optional[torch.LongTensor] = None,751 past_key_value: Optional[Cache] = None,752 output_attentions: bool = False,753 use_cache: bool = False,754 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:755 if output_attentions:756 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.757 logger.warning_once(758 "Phi3Model is using Phi3SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "759 '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.'760 )761 return super().forward(762 hidden_states=hidden_states,763 attention_mask=attention_mask,764 position_ids=position_ids,765 past_key_value=past_key_value,766 output_attentions=output_attentions,767 use_cache=use_cache,768 )769 770 bsz, q_len, _ = hidden_states.size()771 772 qkv = self.qkv_proj(hidden_states)773 query_pos = self.num_heads * self.head_dim774 query_states = qkv[..., :query_pos]775 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]776 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]777 778 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)779 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)780 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)781 782 kv_seq_len = key_states.shape[-2]783 if past_key_value is not None:784 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)785 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)786 787 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)788 789 if past_key_value is not None:790 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models791 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)792 793 key_states = repeat_kv(key_states, self.num_key_value_groups)794 value_states = repeat_kv(value_states, self.num_key_value_groups)795 796 if attention_mask is not None:797 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):798 raise ValueError(799 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"800 )801 802 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,803 # Reference: https://github.com/pytorch/pytorch/issues/112577.804 if query_states.device.type == "cuda" and attention_mask is not None:805 query_states = query_states.contiguous()806 key_states = key_states.contiguous()807 value_states = value_states.contiguous()808 809 attn_output = torch.nn.functional.scaled_dot_product_attention(810 query_states,811 key_states,812 value_states,813 attn_mask=attention_mask,814 dropout_p=self.attention_dropout if self.training else 0.0,815 # 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.816 is_causal=self.is_causal and attention_mask is None and q_len > 1,817 )818 819 attn_output = attn_output.transpose(1, 2).contiguous()820 attn_output = attn_output.view(bsz, q_len, self.hidden_size)821 822 attn_output = self.o_proj(attn_output)823 824 return attn_output, None, past_key_value825 826 827PHI3_ATTENTION_CLASSES = {828 "eager": Phi3Attention,829 "flash_attention_2": Phi3FlashAttention2,830 "sdpa": Phi3SdpaAttention,831}832 833 834class Phi3DecoderLayer(nn.Module):835 def __init__(self, config: Phi3Config, layer_idx: int):836 super().__init__()837 838 self.config = config839 self.self_attn = PHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)840 841 self.mlp = Phi3MLP(config)842 self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)843 844 self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)845 self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)846 self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)847 848 def forward(849 self,850 hidden_states: torch.Tensor,851 attention_mask: Optional[torch.Tensor] = None,852 position_ids: Optional[torch.LongTensor] = None,853 past_key_value: Optional[Tuple[torch.Tensor]] = None,854 output_attentions: Optional[bool] = False,855 use_cache: Optional[bool] = False,856 **kwargs,857 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:858 if "padding_mask" in kwargs:859 warnings.warn(860 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"861 )862 """863 Args:864 hidden_states (`torch.FloatTensor`):865 input to the layer of shape `(batch, seq_len, embed_dim)`866 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size867 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.868 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):869 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range870 `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)871 output_attentions (`bool`, *optional*):872 Whether or not to return the attentions tensors of all attention layers. See `attentions` under873 returned tensors for more detail.874 use_cache (`bool`, *optional*):875 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding876 (see `past_key_values`).877 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states878 """879 880 residual = hidden_states881 882 hidden_states = self.input_layernorm(hidden_states)883 884 # Self Attention885 attn_outputs, self_attn_weights, present_key_value = self.self_attn(886 hidden_states=hidden_states,887 attention_mask=attention_mask,888 position_ids=position_ids,889 past_key_value=past_key_value,890 output_attentions=output_attentions,891 use_cache=use_cache,892 )893 894 hidden_states = residual + self.resid_attn_dropout(attn_outputs)895 896 residual = hidden_states897 hidden_states = self.post_attention_layernorm(hidden_states)898 hidden_states = self.mlp(hidden_states)899 hidden_states = residual + self.resid_mlp_dropout(hidden_states)900 901 outputs = (hidden_states,)902 903 if output_attentions:904 outputs += (self_attn_weights,)905 906 if use_cache:907 outputs += (present_key_value,)908 909 return outputs910 911 912PHI3_START_DOCSTRING = r"""913 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the914 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads915 etc.)916 917 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.918 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage919 and behavior.920 921 Parameters:922 config ([`Phi3Config`]):923 Model configuration class with all the parameters of the model. Initializing with a config file does not924 load the weights associated with the model, only the configuration. Check out the925 [`~PreTrainedModel.from_pretrained`] method to load the model weights.926"""927 928 929@add_start_docstrings(930 "The bare Phi-3 model outputting raw hidden-states without any specific head on top.",931 PHI3_START_DOCSTRING,932)933class Phi3PreTrainedModel(PreTrainedModel):934 config_class = Phi3Config935 base_model_prefix = "model"936 supports_gradient_checkpointing = True937 _no_split_modules = ["Phi3DecoderLayer"]938 _skip_keys_device_placement = "past_key_values"939 _supports_flash_attn_2 = True940 _supports_sdpa = False941 _supports_cache_class = True942 943 _version = "0.0.5"944 945 def _init_weights(self, module):946 std = self.config.initializer_range947 if isinstance(module, nn.Linear):948 module.weight.data.normal_(mean=0.0, std=std)949 if module.bias is not None:950 module.bias.data.zero_()951 elif isinstance(module, nn.Embedding):952 module.weight.data.normal_(mean=0.0, std=std)953 if module.padding_idx is not None:954 module.weight.data[module.padding_idx].zero_()955 956 957PHI3_INPUTS_DOCSTRING = r"""958 Args:959 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):960 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide961 it.962 963 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and964 [`PreTrainedTokenizer.__call__`] for details.965 966 [What are input IDs?](../glossary#input-ids)967 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):968 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:969 970 - 1 for tokens that are **not masked**,971 - 0 for tokens that are **masked**.972 973 [What are attention masks?](../glossary#attention-mask)974 975 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and976 [`PreTrainedTokenizer.__call__`] for details.977 978 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see979 `past_key_values`).980 981 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]982 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more983 information on the default strategy.984 985 - 1 indicates the head is **not masked**,986 - 0 indicates the head is **masked**.987 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):988 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,989 config.n_positions - 1]`.990 991 [What are position IDs?](../glossary#position-ids)992 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):993 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention994 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`995 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.996 997 Two formats are allowed:998 - a [`~cache_utils.Cache`] instance;999 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of1000 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy1001 cache format.1002 1003 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the1004 legacy cache format will be returned.1005 1006 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't1007 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`1008 of shape `(batch_size, sequence_length)`.1009 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):1010 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This1011 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the1012 model's internal embedding lookup matrix.1013 use_cache (`bool`, *optional*):1014 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see1015 `past_key_values`).1016 output_attentions (`bool`, *optional*):1017 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned1018 tensors for more detail.1019 output_hidden_states (`bool`, *optional*):1020 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for1021 more detail.1022 return_dict (`bool`, *optional*):1023 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.1024"""1025 1026 1027@add_start_docstrings(1028 "The bare Phi-3 model outputting raw hidden-states without any specific head on top.",1029 PHI3_START_DOCSTRING,1030)1031class Phi3Model(Phi3PreTrainedModel):1032 """1033 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]1034 1035 Args:1036 config: Phi3Config1037 """1038 1039 def __init__(self, config: Phi3Config):1040 super().__init__(config)1041 self.padding_idx = config.pad_token_id1042 self.vocab_size = config.vocab_size1043 1044 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)1045 self.embed_dropout = nn.Dropout(config.embd_pdrop)1046 self.layers = nn.ModuleList(1047 [Phi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]1048 )1049 self._attn_implementation = config._attn_implementation1050 self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)1051 1052 self.gradient_checkpointing = False1053 # Initialize weights and apply final processing1054 self.post_init()1055 1056 def get_input_embeddings(self):1057 return self.embed_tokens1058 1059 def set_input_embeddings(self, value):1060 self.embed_tokens = value1061 1062 @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)1063 def forward(1064 self,1065 input_ids: torch.LongTensor = None,1066 attention_mask: Optional[torch.Tensor] = None,1067 position_ids: Optional[torch.LongTensor] = None,1068 past_key_values: Optional[List[torch.FloatTensor]] = None,1069 inputs_embeds: Optional[torch.FloatTensor] = None,1070 use_cache: Optional[bool] = None,1071 output_attentions: Optional[bool] = None,1072 output_hidden_states: Optional[bool] = None,1073 return_dict: Optional[bool] = None,1074 ) -> Union[Tuple, BaseModelOutputWithPast]:1075 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1076 output_hidden_states = (1077 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1078 )1079 use_cache = use_cache if use_cache is not None else self.config.use_cache1080 1081 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1082 1083 # retrieve input_ids and inputs_embeds1084 if input_ids is not None and inputs_embeds is not None:1085 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")1086 elif input_ids is not None:1087 batch_size, seq_length = input_ids.shape[:2]1088 elif inputs_embeds is not None:1089 batch_size, seq_length = inputs_embeds.shape[:2]1090 else:1091 raise ValueError("You have to specify either input_ids or inputs_embeds")1092 1093 past_key_values_length = 01094 1095 if self.gradient_checkpointing and self.training:1096 if use_cache:1097 logger.warning_once(1098 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."1099 )1100 use_cache = False1101 1102 if use_cache:1103 use_legacy_cache = not isinstance(past_key_values, Cache)1104 if use_legacy_cache:1105 past_key_values = DynamicCache.from_legacy_cache(past_key_values)1106 past_key_values_length = past_key_values.get_usable_length(seq_length)1107 1108 if position_ids is None:1109 device = input_ids.device if input_ids is not None else inputs_embeds.device1110 position_ids = torch.arange(1111 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device1112 )1113 position_ids = position_ids.unsqueeze(0).view(-1, seq_length)1114 else:1115 position_ids = position_ids.view(-1, seq_length).long()1116 1117 if inputs_embeds is None:1118 inputs_embeds = self.embed_tokens(input_ids)1119 1120 if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:1121 is_padding_right = attention_mask[:, -1].sum().item() != batch_size1122 if is_padding_right:1123 raise ValueError(1124 "You are attempting to perform batched generation with padding_side='right'"1125 " this may lead to unexpected behaviour for Flash Attention version of Phi3. Make sure to "1126 " call `tokenizer.padding_side = 'left'` before tokenizing the input. "1127 )1128 1129 if self._attn_implementation == "flash_attention_2":1130 # 2d mask is passed through the layers1131 attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None1132 else:1133 # 4d mask is passed through the layers1134 attention_mask = _prepare_4d_causal_attention_mask(1135 attention_mask,1136 (batch_size, seq_length),1137 inputs_embeds,1138 past_key_values_length,1139 sliding_window=self.config.sliding_window,1140 )1141 1142 hidden_states = inputs_embeds1143 1144 # decoder layers1145 all_hidden_states = () if output_hidden_states else None1146 all_self_attns = () if output_attentions else None1147 next_decoder_cache = None1148 1149 for decoder_layer in self.layers:1150 if output_hidden_states:1151 all_hidden_states += (hidden_states,)1152 1153 if self.gradient_checkpointing and self.training:1154 layer_outputs = self._gradient_checkpointing_func(1155 decoder_layer.__call__,1156 hidden_states,1157 attention_mask,1158 position_ids,1159 past_key_values,1160 output_attentions,1161 use_cache,1162 )1163 else:1164 layer_outputs = decoder_layer(1165 hidden_states,1166 attention_mask=attention_mask,1167 position_ids=position_ids,1168 past_key_value=past_key_values,1169 output_attentions=output_attentions,1170 use_cache=use_cache,1171 )1172 1173 hidden_states = layer_outputs[0]1174 1175 if use_cache:1176 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1177 1178 if output_attentions:1179 all_self_attns += (layer_outputs[1],)1180 1181 hidden_states = self.norm(hidden_states)1182 1183 # add hidden states from the last decoder layer1184 if output_hidden_states:1185 all_hidden_states += (hidden_states,)1186 1187 next_cache = None1188 if use_cache:1189 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache1190 if not return_dict:1191 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1192 return BaseModelOutputWithPast(1193 last_hidden_state=hidden_states,1194 past_key_values=next_cache,1195 hidden_states=all_hidden_states,1196 attentions=all_self_attns,1197 )1198 1199 1200class Phi3ForCausalLM(Phi3PreTrainedModel):