ByteDance/Sa2VA-InternVL3-14B
1092
1# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15""" PyTorch Phi-3 model."""16 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, MSELoss27from transformers.activations import ACT2FN28from transformers.cache_utils import Cache, DynamicCache29from transformers.modeling_attn_mask_utils import \30 _prepare_4d_causal_attention_mask31from transformers.modeling_outputs import (BaseModelOutputWithPast,32 CausalLMOutputWithPast,33 SequenceClassifierOutputWithPast,34 TokenClassifierOutput)35from transformers.modeling_utils import PreTrainedModel36from transformers.utils import (add_code_sample_docstrings,37 add_start_docstrings,38 add_start_docstrings_to_model_forward,39 is_flash_attn_2_available,40 is_flash_attn_greater_or_equal_2_10, logging,41 replace_return_docstrings)42 43from transformers.models.phi3.configuration_phi3 import Phi3Config44 45logger = logging.get_logger(__name__)46 47# Transformers scans dependencies in the modeling file, causing issues on conditional loading. The regex only ignores try/catch blocks, but not if statements48# if is_flash_attn_2_available():49_flash_supports_window_size = False50try:51 from flash_attn import flash_attn_func, flash_attn_varlen_func52 from flash_attn.bert_padding import (index_first_axis, pad_input, # noqa53 unpad_input)54 55 _flash_supports_window_size = 'window_size' in list(inspect.signature(flash_attn_func).parameters)56 has_flash_attn = True57except ImportError as error:58 logger.warning(59 f'`flash-attention` package not found, consider installing for better performance: {error}.'60 )61 if not _flash_supports_window_size:62 logger.warning(63 "Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`."64 )65 has_flash_attn = False66 67_CHECKPOINT_FOR_DOC = 'microsoft/Phi-3-mini-4k-instruct'68_CONFIG_FOR_DOC = 'Phi3Config'69 70PHI3_PRETRAINED_MODEL_ARCHIVE_LIST = [71 'microsoft/Phi-3-mini-4k-instruct',72 'microsoft/Phi-3-mini-128k-instruct',73 # See all Phi-3 models at https://huggingface.co/models?filter=Phi-374]75 76 77# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Phi378class Phi3RMSNorm(nn.Module):79 def __init__(self, hidden_size, eps=1e-6):80 """81 Phi3RMSNorm is equivalent to T5LayerNorm82 """83 super().__init__()84 self.weight = nn.Parameter(torch.ones(hidden_size))85 self.variance_epsilon = eps86 87 def forward(self, hidden_states):88 input_dtype = hidden_states.dtype89 hidden_states = hidden_states.to(torch.float32)90 variance = hidden_states.pow(2).mean(-1, keepdim=True)91 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)92 return self.weight * hidden_states.to(input_dtype)93 94 95# Copied from transformers.models.llama.modeling_llama._get_unpad_data96def _get_unpad_data(attention_mask):97 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)98 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()99 max_seqlen_in_batch = seqlens_in_batch.max().item()100 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))101 return (102 indices,103 cu_seqlens,104 max_seqlen_in_batch,105 )106 107 108# Copied from transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding with gemma->phi3, Gemma->Phi3109class Phi3RotaryEmbedding(nn.Module):110 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):111 super().__init__()112 113 self.dim = dim114 self.max_position_embeddings = max_position_embeddings115 self.base = base116 self.register_buffer('inv_freq', None, persistent=False)117 118 @torch.no_grad()119 def forward(self, x, position_ids, seq_len=None):120 # x: [bs, num_attention_heads, seq_len, head_size]121 if self.inv_freq is None:122 self.inv_freq = 1.0 / (123 self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)124 )125 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)126 position_ids_expanded = position_ids[:, None, :].float()127 # Force float32 since bfloat16 loses precision on long contexts128 # See https://github.com/huggingface/transformers/pull/29285129 device_type = x.device.type130 device_type = device_type if isinstance(device_type, str) and device_type != 'mps' else 'cpu'131 with torch.autocast(device_type=device_type, enabled=False):132 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)133 emb = torch.cat((freqs, freqs), dim=-1)134 cos = emb.cos()135 sin = emb.sin()136 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)137 138 139class Phi3SuScaledRotaryEmbedding(Phi3RotaryEmbedding):140 def __init__(self, dim, config, device=None):141 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)142 143 self.short_factor = config.rope_scaling['short_factor']144 self.long_factor = config.rope_scaling['long_factor']145 self.original_max_position_embeddings = config.original_max_position_embeddings146 147 @torch.no_grad()148 def forward(self, x, position_ids, seq_len=None):149 seq_len = torch.max(position_ids) + 1150 if seq_len > self.original_max_position_embeddings:151 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)152 else:153 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)154 155 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim156 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)157 158 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)159 position_ids_expanded = position_ids[:, None, :].float()160 161 # Force float32 since bfloat16 loses precision on long contexts162 # See https://github.com/huggingface/transformers/pull/29285163 device_type = x.device.type164 device_type = device_type if isinstance(device_type, str) and device_type != 'mps' else 'cpu'165 with torch.autocast(device_type=device_type, enabled=False):166 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)167 emb = torch.cat((freqs, freqs), dim=-1)168 169 scale = self.max_position_embeddings / self.original_max_position_embeddings170 if scale <= 1.0:171 scaling_factor = 1.0172 else:173 scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))174 175 cos = emb.cos() * scaling_factor176 sin = emb.sin() * scaling_factor177 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)178 179 180class Phi3YarnScaledRotaryEmbedding(Phi3RotaryEmbedding):181 def __init__(self, dim, config, device=None):182 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)183 184 self.short_factor = config.rope_scaling['short_factor']185 self.long_factor = config.rope_scaling['long_factor']186 self.original_max_position_embeddings = config.original_max_position_embeddings187 188 @torch.no_grad()189 def forward(self, x, position_ids, seq_len=None):190 seq_len = torch.max(position_ids) + 1191 if seq_len > self.original_max_position_embeddings:192 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)193 else:194 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)195 196 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim197 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)198 199 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)200 position_ids_expanded = position_ids[:, None, :].float()201 202 # Force float32 since bfloat16 loses precision on long contexts203 # See https://github.com/huggingface/transformers/pull/29285204 device_type = x.device.type205 device_type = device_type if isinstance(device_type, str) and device_type != 'mps' else 'cpu'206 with torch.autocast(device_type=device_type, enabled=False):207 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)208 emb = torch.cat((freqs, freqs), dim=-1)209 210 scale = self.max_position_embeddings / self.original_max_position_embeddings211 if scale <= 1.0:212 scaling_factor = 1.0213 else:214 scaling_factor = 0.1 * math.log(scale) + 1.0215 216 cos = emb.cos() * scaling_factor217 sin = emb.sin() * scaling_factor218 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)219 220 221# Copied from transformers.models.llama.modeling_llama.rotate_half222def rotate_half(x):223 """Rotates half the hidden dims of the input."""224 x1 = x[..., : x.shape[-1] // 2]225 x2 = x[..., x.shape[-1] // 2 :]226 return torch.cat((-x2, x1), dim=-1)227 228 229# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb230def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):231 """Applies Rotary Position Embedding to the query and key tensors.232 233 Args:234 q (`torch.Tensor`): The query tensor.235 k (`torch.Tensor`): The key tensor.236 cos (`torch.Tensor`): The cosine part of the rotary embedding.237 sin (`torch.Tensor`): The sine part of the rotary embedding.238 position_ids (`torch.Tensor`, *optional*):239 Deprecated and unused.240 unsqueeze_dim (`int`, *optional*, defaults to 1):241 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and242 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note243 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and244 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes245 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have246 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.247 Returns:248 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.249 """250 cos = cos.unsqueeze(unsqueeze_dim)251 sin = sin.unsqueeze(unsqueeze_dim)252 q_embed = (q * cos) + (rotate_half(q) * sin)253 k_embed = (k * cos) + (rotate_half(k) * sin)254 return q_embed, k_embed255 256 257class Phi3MLP(nn.Module):258 def __init__(self, config):259 super().__init__()260 261 self.config = config262 self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)263 self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)264 265 self.activation_fn = ACT2FN[config.hidden_act]266 267 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:268 up_states = self.gate_up_proj(hidden_states)269 270 gate, up_states = up_states.chunk(2, dim=-1)271 up_states = up_states * self.activation_fn(gate)272 273 return self.down_proj(up_states)274 275 276# Copied from transformers.models.llama.modeling_llama.repeat_kv with llama->phi277def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:278 """279 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,280 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)281 """282 batch, num_key_value_heads, slen, head_dim = hidden_states.shape283 if n_rep == 1:284 return hidden_states285 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)286 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)287 288 289class Phi3Attention(nn.Module):290 """Multi-headed attention from 'Attention Is All You Need' paper"""291 292 def __init__(self, config: Phi3Config, layer_idx: Optional[int] = None):293 super().__init__()294 self.config = config295 self.layer_idx = layer_idx296 if layer_idx is None:297 logger.warning_once(298 f'Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will '299 'lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` '300 'when creating this class.'301 )302 303 self.attention_dropout = config.attention_dropout304 self.hidden_size = config.hidden_size305 self.num_heads = config.num_attention_heads306 self.head_dim = self.hidden_size // self.num_heads307 self.num_key_value_heads = config.num_key_value_heads308 self.num_key_value_groups = self.num_heads // self.num_key_value_heads309 self.max_position_embeddings = config.max_position_embeddings310 self.original_max_position_embeddings = config.original_max_position_embeddings311 self.rope_theta = config.rope_theta312 self.rope_scaling = config.rope_scaling313 self.is_causal = True314 315 if (self.head_dim * self.num_heads) != self.hidden_size:316 raise ValueError(317 f'hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}'318 f' and `num_heads`: {self.num_heads}).'319 )320 321 op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim)322 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)323 self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False)324 self._init_rope()325 326 def _init_rope(self):327 if self.rope_scaling is None:328 self.rotary_emb = Phi3RotaryEmbedding(329 self.head_dim,330 max_position_embeddings=self.max_position_embeddings,331 base=self.rope_theta,332 )333 else:334 scaling_type = self.config.rope_scaling['type']335 if scaling_type == 'su':336 self.rotary_emb = Phi3SuScaledRotaryEmbedding(self.head_dim, self.config)337 elif scaling_type == 'yarn':338 self.rotary_emb = Phi3YarnScaledRotaryEmbedding(self.head_dim, self.config)339 else:340 raise ValueError(f'Unknown RoPE scaling type {scaling_type}')341 342 def forward(343 self,344 hidden_states: torch.Tensor,345 attention_mask: Optional[torch.Tensor] = None,346 position_ids: Optional[torch.LongTensor] = None,347 past_key_value: Optional[Cache] = None,348 output_attentions: bool = False,349 use_cache: bool = False,350 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:351 logger.warning_once('You are not running the flash-attention implementation, expect numerical differences.')352 353 bsz, q_len, _ = hidden_states.size()354 355 qkv = self.qkv_proj(hidden_states)356 query_pos = self.num_heads * self.head_dim357 query_states = qkv[..., :query_pos]358 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]359 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]360 361 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)362 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)363 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)364 365 kv_seq_len = key_states.shape[-2]366 if past_key_value is not None:367 if self.layer_idx is None:368 raise ValueError(369 f'The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} '370 'for auto-regressive decoding with k/v caching, please make sure to initialize the attention class '371 'with a layer index.'372 )373 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)374 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)375 376 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)377 378 if past_key_value is not None:379 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models380 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)381 382 # repeat k/v heads if n_kv_heads < n_heads383 key_states = repeat_kv(key_states, self.num_key_value_groups)384 value_states = repeat_kv(value_states, self.num_key_value_groups)385 386 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)387 388 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):389 raise ValueError(390 f'Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is'391 f' {attn_weights.size()}'392 )393 394 if attention_mask is not None:395 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):396 raise ValueError(397 f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'398 )399 attn_weights = attn_weights + attention_mask400 401 # upcast attention to fp32402 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)403 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)404 405 attn_output = torch.matmul(attn_weights, value_states)406 407 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):408 raise ValueError(409 f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'410 f' {attn_output.size()}'411 )412 413 attn_output = attn_output.transpose(1, 2).contiguous()414 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)415 416 attn_output = self.o_proj(attn_output)417 418 if not output_attentions:419 attn_weights = None420 421 return attn_output, attn_weights, past_key_value422 423 424class Phi3FlashAttention2(Phi3Attention):425 """426 Phi-3 flash attention module. This module inherits from `Phi3Attention` as the weights of the module stays427 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of428 flash attention and deal with padding tokens in case the input contains any of them.429 """430 431 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__432 def __init__(self, *args, **kwargs):433 super().__init__(*args, **kwargs)434 435 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.436 # 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.437 # 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).438 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()439 440 def forward(441 self,442 hidden_states: torch.Tensor,443 attention_mask: Optional[torch.LongTensor] = None,444 position_ids: Optional[torch.LongTensor] = None,445 past_key_value: Optional[Cache] = None,446 output_attentions: bool = False,447 use_cache: bool = False,448 **kwargs,449 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:450 # Phi3FlashAttention2 attention does not support output_attentions451 452 if not _flash_supports_window_size:453 logger.warning_once(454 "The current flash attention version does not support sliding window attention. Please use `attn_implementation='eager'` or upgrade flash-attn library."455 )456 raise ValueError('The current flash attention version does not support sliding window attention.')457 458 output_attentions = False459 460 if 'padding_mask' in kwargs:461 warnings.warn(462 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'463 )464 465 # overwrite attention_mask with padding_mask466 attention_mask = kwargs.pop('padding_mask')467 468 bsz, q_len, _ = hidden_states.size()469 470 qkv = self.qkv_proj(hidden_states)471 query_pos = self.num_heads * self.head_dim472 query_states = qkv[..., :query_pos]473 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]474 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]475 476 # Flash attention requires the input to have the shape477 # batch_size x seq_length x head_dim x hidden_dim478 # therefore we just need to keep the original shape479 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)480 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)481 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)482 483 kv_seq_len = key_states.shape[-2]484 if past_key_value is not None:485 if self.layer_idx is None:486 raise ValueError(487 f'The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} '488 'for auto-regressive decoding with k/v caching, please make sure to initialize the attention class '489 'with a layer index.'490 )491 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)492 493 # Because the input can be padded, the absolute sequence length depends on the max position id.494 rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1495 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=rotary_seq_len)496 497 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)498 499 use_sliding_windows = (500 _flash_supports_window_size501 and getattr(self.config, 'sliding_window', None) is not None502 and kv_seq_len > self.config.sliding_window503 )504 505 if past_key_value is not None:506 # Activate slicing cache only if the config has a value `sliding_windows` attribute507 cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0508 if (509 getattr(self.config, 'sliding_window', None) is not None510 and kv_seq_len > self.config.sliding_window511 and cache_has_contents512 ):513 slicing_tokens = 1 - self.config.sliding_window514 515 past_key = past_key_value[self.layer_idx][0]516 past_value = past_key_value[self.layer_idx][1]517 518 past_key = past_key[:, :, slicing_tokens:, :].contiguous()519 past_value = past_value[:, :, slicing_tokens:, :].contiguous()520 521 if past_key.shape[-2] != self.config.sliding_window - 1:522 raise ValueError(523 f'past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got'524 f' {past_key.shape}'525 )526 527 if attention_mask is not None:528 attention_mask = attention_mask[:, slicing_tokens:]529 attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)530 531 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models532 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)533 534 # repeat k/v heads if n_kv_heads < n_heads535 key_states = repeat_kv(key_states, self.num_key_value_groups)536 value_states = repeat_kv(value_states, self.num_key_value_groups)537 538 attn_dropout = self.attention_dropout if self.training else 0.0539 540 # In PEFT, usually we cast the layer norms in float32 for training stability reasons541 # therefore the input hidden states gets silently casted in float32. Hence, we need542 # cast them back in the correct dtype just to be sure everything works as expected.543 # This might slowdown training & inference so it is recommended to not cast the LayerNorms544 # in fp32.545 546 if query_states.dtype == torch.float32:547 if torch.is_autocast_enabled():548 target_dtype = torch.get_autocast_gpu_dtype()549 # Handle the case where the model is quantized550 elif hasattr(self.config, '_pre_quantization_dtype'):551 target_dtype = self.config._pre_quantization_dtype552 else:553 target_dtype = self.qkv_proj.weight.dtype554 555 logger.warning_once(556 f'The input hidden states seems to be silently casted in float32, this might be related to'557 f' the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in'558 f' {target_dtype}.'559 )560 561 query_states = query_states.to(target_dtype)562 key_states = key_states.to(target_dtype)563 value_states = value_states.to(target_dtype)564 565 # Reashape to the expected shape for Flash Attention566 query_states = query_states.transpose(1, 2)567 key_states = key_states.transpose(1, 2)568 value_states = value_states.transpose(1, 2)569 570 attn_output = self._flash_attention_forward(571 query_states,572 key_states,573 value_states,574 attention_mask,575 q_len,576 dropout=attn_dropout,577 use_sliding_windows=use_sliding_windows,578 )579 580 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()581 attn_output = self.o_proj(attn_output)582 583 if not output_attentions:584 attn_weights = None585 586 return attn_output, attn_weights, past_key_value587 588 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._flash_attention_forward589 def _flash_attention_forward(590 self,591 query_states,592 key_states,593 value_states,594 attention_mask,595 query_length,596 dropout=0.0,597 softmax_scale=None,598 use_sliding_windows=False,599 ):600 """601 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token602 first unpad the input, then computes the attention scores and pad the final attention scores.603 604 Args:605 query_states (`torch.Tensor`):606 Input query states to be passed to Flash Attention API607 key_states (`torch.Tensor`):608 Input key states to be passed to Flash Attention API609 value_states (`torch.Tensor`):610 Input value states to be passed to Flash Attention API611 attention_mask (`torch.Tensor`):612 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the613 position of padding tokens and 1 for the position of non-padding tokens.614 dropout (`float`):615 Attention dropout616 softmax_scale (`float`, *optional*):617 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)618 use_sliding_windows (`bool`, *optional*):619 Whether to activate sliding window attention.620 """621 if not self._flash_attn_uses_top_left_mask:622 causal = self.is_causal623 else:624 # 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__.625 causal = self.is_causal and query_length != 1626 627 # Contains at least one padding token in the sequence628 if attention_mask is not None:629 batch_size = query_states.shape[0]630 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(631 query_states, key_states, value_states, attention_mask, query_length632 )633 634 cu_seqlens_q, cu_seqlens_k = cu_seq_lens635 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens636 637 if not use_sliding_windows:638 attn_output_unpad = flash_attn_varlen_func(639 query_states,640 key_states,641 value_states,642 cu_seqlens_q=cu_seqlens_q,643 cu_seqlens_k=cu_seqlens_k,644 max_seqlen_q=max_seqlen_in_batch_q,645 max_seqlen_k=max_seqlen_in_batch_k,646 dropout_p=dropout,647 softmax_scale=softmax_scale,648 causal=causal,649 )650 else:651 attn_output_unpad = flash_attn_varlen_func(652 query_states,653 key_states,654 value_states,655 cu_seqlens_q=cu_seqlens_q,656 cu_seqlens_k=cu_seqlens_k,657 max_seqlen_q=max_seqlen_in_batch_q,658 max_seqlen_k=max_seqlen_in_batch_k,659 dropout_p=dropout,660 softmax_scale=softmax_scale,661 causal=causal,662 window_size=(self.config.sliding_window, self.config.sliding_window),663 )664 665 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)666 else:667 if not use_sliding_windows:668 attn_output = flash_attn_func(669 query_states,670 key_states,671 value_states,672 dropout,673 softmax_scale=softmax_scale,674 causal=causal,675 )676 else:677 attn_output = flash_attn_func(678 query_states,679 key_states,680 value_states,681 dropout,682 softmax_scale=softmax_scale,683 causal=causal,684 window_size=(self.config.sliding_window, self.config.sliding_window),685 )686 687 return attn_output688 689 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input690 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):691 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape692 693 # On the first iteration we need to properly re-create the padding mask694 # by slicing it on the proper place695 if kv_seq_len != attention_mask.shape[-1]:696 attention_mask_num_tokens = attention_mask.shape[-1]697 attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]698 699 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)700 701 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)702 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)703 704 if query_length == kv_seq_len:705 query_layer = index_first_axis(706 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k707 )708 cu_seqlens_q = cu_seqlens_k709 max_seqlen_in_batch_q = max_seqlen_in_batch_k710 indices_q = indices_k711 elif query_length == 1:712 max_seqlen_in_batch_q = 1713 cu_seqlens_q = torch.arange(714 batch_size + 1, dtype=torch.int32, device=query_layer.device715 ) # There is a memcpy here, that is very bad.716 indices_q = cu_seqlens_q[:-1]717 query_layer = query_layer.squeeze(1)718 else:719 # The -q_len: slice assumes left padding.720 attention_mask = attention_mask[:, -query_length:]721 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)722 723 return (724 query_layer,725 key_layer,726 value_layer,727 indices_q,728 (cu_seqlens_q, cu_seqlens_k),729 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),730 )731 732 733# copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Phi3734# TODO @Arthur no longer copied from LLama after static cache735class Phi3SdpaAttention(Phi3Attention):736 """737 Phi3 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from738 `Phi3Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to739 SDPA API.740 """741 742 # Adapted from Phi3Attention.forward743 def forward(744 self,745 hidden_states: torch.Tensor,746 attention_mask: Optional[torch.Tensor] = None,747 position_ids: Optional[torch.LongTensor] = None,748 past_key_value: Optional[Cache] = None,749 output_attentions: bool = False,750 use_cache: bool = False,751 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:752 if output_attentions:753 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.754 logger.warning_once(755 '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, '756 '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.'757 )758 return super().forward(759 hidden_states=hidden_states,760 attention_mask=attention_mask,761 position_ids=position_ids,762 past_key_value=past_key_value,763 output_attentions=output_attentions,764 use_cache=use_cache,765 )766 767 bsz, q_len, _ = hidden_states.size()768 769 qkv = self.qkv_proj(hidden_states)770 query_pos = self.num_heads * self.head_dim771 query_states = qkv[..., :query_pos]772 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]773 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]774 775 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)776 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)777 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)778 779 kv_seq_len = key_states.shape[-2]780 if past_key_value is not None:781 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)782 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)783 784 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)785 786 if past_key_value is not None:787 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models788 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)789 790 key_states = repeat_kv(key_states, self.num_key_value_groups)791 value_states = repeat_kv(value_states, self.num_key_value_groups)792 793 if attention_mask is not None:794 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):795 raise ValueError(796 f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'797 )798 799 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,800 # Reference: https://github.com/pytorch/pytorch/issues/112577.801 if query_states.device.type == 'cuda' and attention_mask is not None:802 query_states = query_states.contiguous()803 key_states = key_states.contiguous()804 value_states = value_states.contiguous()805 806 attn_output = torch.nn.functional.scaled_dot_product_attention(807 query_states,808 key_states,809 value_states,810 attn_mask=attention_mask,811 dropout_p=self.attention_dropout if self.training else 0.0,812 # 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.813 is_causal=self.is_causal and attention_mask is None and q_len > 1,814 )815 816 attn_output = attn_output.transpose(1, 2).contiguous()817 attn_output = attn_output.view(bsz, q_len, self.hidden_size)818 819 attn_output = self.o_proj(attn_output)820 821 return attn_output, None, past_key_value822 823 824PHI3_ATTENTION_CLASSES = {825 'eager': Phi3Attention,826 'flash_attention_2': Phi3FlashAttention2,827 'sdpa': Phi3SdpaAttention,828}829 830 831class Phi3DecoderLayer(nn.Module):832 def __init__(self, config: Phi3Config, layer_idx: int):833 super().__init__()834 835 self.config = config836 self.self_attn = PHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)837 838 self.mlp = Phi3MLP(config)839 self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)840 841 self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)842 self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)843 self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)844 845 def forward(846 self,847 hidden_states: torch.Tensor,848 attention_mask: Optional[torch.Tensor] = None,849 position_ids: Optional[torch.LongTensor] = None,850 past_key_value: Optional[Tuple[torch.Tensor]] = None,851 output_attentions: Optional[bool] = False,852 use_cache: Optional[bool] = False,853 **kwargs,854 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:855 if 'padding_mask' in kwargs:856 warnings.warn(857 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'858 )859 """860 Args:861 hidden_states (`torch.FloatTensor`):862 input to the layer of shape `(batch, seq_len, embed_dim)`863 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size864 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.865 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):866 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range867 `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)868 output_attentions (`bool`, *optional*):869 Whether or not to return the attentions tensors of all attention layers. See `attentions` under870 returned tensors for more detail.871 use_cache (`bool`, *optional*):872 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding873 (see `past_key_values`).874 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states875 """876 877 residual = hidden_states878 879 hidden_states = self.input_layernorm(hidden_states)880 881 # Self Attention882 attn_outputs, self_attn_weights, present_key_value = self.self_attn(883 hidden_states=hidden_states,884 attention_mask=attention_mask,885 position_ids=position_ids,886 past_key_value=past_key_value,887 output_attentions=output_attentions,888 use_cache=use_cache,889 )890 891 hidden_states = residual + self.resid_attn_dropout(attn_outputs)892 893 residual = hidden_states894 hidden_states = self.post_attention_layernorm(hidden_states)895 hidden_states = self.mlp(hidden_states)896 hidden_states = residual + self.resid_mlp_dropout(hidden_states)897 898 outputs = (hidden_states,)899 900 if output_attentions:901 outputs += (self_attn_weights,)902 903 if use_cache:904 outputs += (present_key_value,)905 906 return outputs907 908 909PHI3_START_DOCSTRING = r"""910 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the911 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads912 etc.)913 914 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.915 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage916 and behavior.917 918 Parameters:919 config ([`Phi3Config`]):920 Model configuration class with all the parameters of the model. Initializing with a config file does not921 load the weights associated with the model, only the configuration. Check out the922 [`~PreTrainedModel.from_pretrained`] method to load the model weights.923"""924 925 926@add_start_docstrings(927 'The bare Phi-3 model outputting raw hidden-states without any specific head on top.',928 PHI3_START_DOCSTRING,929)930class Phi3PreTrainedModel(PreTrainedModel):931 config_class = Phi3Config932 base_model_prefix = 'model'933 supports_gradient_checkpointing = True934 _no_split_modules = ['Phi3DecoderLayer']935 _skip_keys_device_placement = 'past_key_values'936 _supports_flash_attn_2 = True937 _supports_sdpa = False938 _supports_cache_class = True939 940 _version = '0.0.5'941 942 def __init__(self, config: Phi3Config):943 if not has_flash_attn:944 config._attn_implementation = 'eager'945 print('Warning: Flash attention is not available, using eager attention instead.')946 super().__init__(config)947 948 def _init_weights(self, module):949 std = self.config.initializer_range950 if isinstance(module, nn.Linear):951 module.weight.data.normal_(mean=0.0, std=std)952 if module.bias is not None:953 module.bias.data.zero_()954 elif isinstance(module, nn.Embedding):955 module.weight.data.normal_(mean=0.0, std=std)956 if module.padding_idx is not None:957 module.weight.data[module.padding_idx].zero_()958 959 960PHI3_INPUTS_DOCSTRING = r"""961 Args:962 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):963 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide964 it.965 966 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and967 [`PreTrainedTokenizer.__call__`] for details.968 969 [What are input IDs?](../glossary#input-ids)970 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):971 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:972 973 - 1 for tokens that are **not masked**,974 - 0 for tokens that are **masked**.975 976 [What are attention masks?](../glossary#attention-mask)977 978 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and979 [`PreTrainedTokenizer.__call__`] for details.980 981 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see982 `past_key_values`).983 984 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]985 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more986 information on the default strategy.987 988 - 1 indicates the head is **not masked**,989 - 0 indicates the head is **masked**.990 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):991 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,992 config.n_positions - 1]`.993 994 [What are position IDs?](../glossary#position-ids)995 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):996 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention997 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`998 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.999 1000 Two formats are allowed:1001 - a [`~cache_utils.Cache`] instance;1002 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of1003 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy1004 cache format.1005 1006 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the1007 legacy cache format will be returned.1008 1009 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't1010 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`1011 of shape `(batch_size, sequence_length)`.1012 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):1013 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This1014 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the1015 model's internal embedding lookup matrix.1016 use_cache (`bool`, *optional*):1017 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see1018 `past_key_values`).1019 output_attentions (`bool`, *optional*):1020 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned1021 tensors for more detail.1022 output_hidden_states (`bool`, *optional*):1023 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for1024 more detail.1025 return_dict (`bool`, *optional*):1026 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.1027"""1028 1029 1030@add_start_docstrings(1031 'The bare Phi-3 model outputting raw hidden-states without any specific head on top.',1032 PHI3_START_DOCSTRING,1033)1034class Phi3Model(Phi3PreTrainedModel):1035 """1036 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]1037 1038 Args:1039 config: Phi3Config1040 """1041 1042 def __init__(self, config: Phi3Config):1043 super().__init__(config)1044 self.padding_idx = config.pad_token_id1045 self.vocab_size = config.vocab_size1046 1047 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)1048 self.embed_dropout = nn.Dropout(config.embd_pdrop)1049 self.layers = nn.ModuleList(1050 [Phi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]1051 )1052 self._attn_implementation = config._attn_implementation1053 1054 self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)1055 1056 self.gradient_checkpointing = False1057 # Initialize weights and apply final processing1058 self.post_init()1059 1060 def get_input_embeddings(self):1061 return self.embed_tokens1062 1063 def set_input_embeddings(self, value):1064 self.embed_tokens = value1065 1066 @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)1067 def forward(1068 self,1069 input_ids: torch.LongTensor = None,1070 attention_mask: Optional[torch.Tensor] = None,1071 position_ids: Optional[torch.LongTensor] = None,1072 past_key_values: Optional[List[torch.FloatTensor]] = None,1073 inputs_embeds: Optional[torch.FloatTensor] = None,1074 use_cache: Optional[bool] = None,1075 output_attentions: Optional[bool] = None,1076 output_hidden_states: Optional[bool] = None,1077 return_dict: Optional[bool] = None,1078 ) -> Union[Tuple, BaseModelOutputWithPast]:1079 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1080 output_hidden_states = (1081 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1082 )1083 use_cache = use_cache if use_cache is not None else self.config.use_cache1084 1085 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1086 1087 # retrieve input_ids and inputs_embeds1088 if input_ids is not None and inputs_embeds is not None:1089 raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')1090 elif input_ids is not None:1091 batch_size, seq_length = input_ids.shape[:2]1092 elif inputs_embeds is not None:1093 batch_size, seq_length = inputs_embeds.shape[:2]1094 else:1095 raise ValueError('You have to specify either input_ids or inputs_embeds')1096 1097 past_key_values_length = 01098 1099 if self.gradient_checkpointing and self.training:1100 if use_cache:1101 logger.warning_once(1102 '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...'1103 )1104 use_cache = False1105 1106 if use_cache:1107 use_legacy_cache = not isinstance(past_key_values, Cache)1108 if use_legacy_cache:1109 past_key_values = DynamicCache.from_legacy_cache(past_key_values)1110 past_key_values_length = past_key_values.get_usable_length(seq_length)1111 1112 if position_ids is None:1113 device = input_ids.device if input_ids is not None else inputs_embeds.device1114 position_ids = torch.arange(1115 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device1116 )1117 position_ids = position_ids.unsqueeze(0).view(-1, seq_length)1118 else:1119 position_ids = position_ids.view(-1, seq_length).long()1120 1121 if inputs_embeds is None:1122 inputs_embeds = self.embed_tokens(input_ids)1123 1124 if attention_mask is not None and self._attn_implementation == 'flash_attention_2' and use_cache:1125 is_padding_right = attention_mask[:, -1].sum().item() != batch_size1126 if is_padding_right:1127 raise ValueError(1128 "You are attempting to perform batched generation with padding_side='right'"1129 ' this may lead to unexpected behaviour for Flash Attention version of Phi3. Make sure to '1130 " call `tokenizer.padding_side = 'left'` before tokenizing the input. "1131 )1132 1133 if self._attn_implementation == 'flash_attention_2':1134 # 2d mask is passed through the layers1135 attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None1136 else:1137 # 4d mask is passed through the layers1138 attention_mask = _prepare_4d_causal_attention_mask(1139 attention_mask,1140 (batch_size, seq_length),1141 inputs_embeds,1142 past_key_values_length,1143 sliding_window=self.config.sliding_window,1144 )1145 1146 hidden_states = inputs_embeds1147 1148 # decoder layers1149 all_hidden_states = () if output_hidden_states else None1150 all_self_attns = () if output_attentions else None1151 next_decoder_cache = None1152 1153 for decoder_layer in self.layers:1154 if output_hidden_states:1155 all_hidden_states += (hidden_states,)1156 1157 if self.gradient_checkpointing and self.training:1158 layer_outputs = self._gradient_checkpointing_func(1159 decoder_layer.__call__,1160 hidden_states,1161 attention_mask,1162 position_ids,1163 past_key_values,1164 output_attentions,1165 use_cache,1166 )1167 else:1168 layer_outputs = decoder_layer(1169 hidden_states,1170 attention_mask=attention_mask,1171 position_ids=position_ids,1172 past_key_value=past_key_values,1173 output_attentions=output_attentions,1174 use_cache=use_cache,1175 )1176 1177 hidden_states = layer_outputs[0]1178 1179 if use_cache:1180 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1181 1182 if output_attentions:1183 all_self_attns += (layer_outputs[1],)1184 1185 hidden_states = self.norm(hidden_states)1186 1187 # add hidden states from the last decoder layer1188 if output_hidden_states:1189 all_hidden_states += (hidden_states,)1190 1191 next_cache = None1192 if use_cache:1193 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache1194 if not return_dict:1195 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1196 return BaseModelOutputWithPast(1197 last_hidden_state=hidden_states,1198 past_key_values=next_cache,1199 hidden_states=all_hidden_states,1200 attentions=all_self_attns,