hvai/sdset
Stable Difusion store for learner get files
1596
1# coding=utf-8
2# 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 at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# 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 and
14# limitations under the License.
15
16""" PyTorch Phi-3 model."""
17
18import inspect
19import math
20import warnings
21from typing import List, Optional, Tuple, Union
22
23import torch
24import torch.nn.functional as F
25import torch.utils.checkpoint
26from torch import nn
27from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
28
29from transformers.activations import ACT2FN
30from transformers.cache_utils import Cache, DynamicCache
31from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
32from transformers.modeling_outputs import (
33 BaseModelOutputWithPast,
34 CausalLMOutputWithPast,
35 SequenceClassifierOutputWithPast,
36 TokenClassifierOutput,
37)
38from transformers.modeling_utils import PreTrainedModel
39from 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 Phi3Config
49
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 statements
54# if is_flash_attn_2_available():
55_flash_supports_window_size = False
56try:
57 from flash_attn import flash_attn_func, flash_attn_varlen_func
58 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
59
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-attention` 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-3
77]
78
79
80# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Phi3
81class Phi3RMSNorm(nn.Module):
82 def __init__(self, hidden_size, eps=1e-6):
83 """
84 Phi3RMSNorm is equivalent to T5LayerNorm
85 """
86 super().__init__()
87 self.weight = nn.Parameter(torch.ones(hidden_size))
88 self.variance_epsilon = eps
89
90 def forward(self, hidden_states):
91 input_dtype = hidden_states.dtype
92 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_data
99def _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->Phi3
112class Phi3RotaryEmbedding(nn.Module):
113 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
114 super().__init__()
115
116 self.dim = dim
117 self.max_position_embeddings = max_position_embeddings
118 self.base = base
119 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 contexts
131 # See https://github.com/huggingface/transformers/pull/29285
132 device_type = x.device.type
133 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 Phi3LongRoPEScaledRotaryEmbedding(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_embeddings
149
150 @torch.no_grad()
151 def forward(self, x, position_ids, seq_len=None):
152 seq_len = seq_len or torch.max(position_ids) + 1
153 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.dim
159 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 contexts
165 # See https://github.com/huggingface/transformers/pull/29285
166 device_type = x.device.type
167 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_embeddings
173 if scale <= 1.0:
174 scaling_factor = 1.0
175 else:
176 scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))
177
178 cos = emb.cos() * scaling_factor
179 sin = emb.sin() * scaling_factor
180 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
181
182
183# Copied from transformers.models.llama.modeling_llama.rotate_half
184def rotate_half(x):
185 """Rotates half the hidden dims of the input."""
186 x1 = x[..., : x.shape[-1] // 2]
187 x2 = x[..., x.shape[-1] // 2 :]
188 return torch.cat((-x2, x1), dim=-1)
189
190
191# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
192def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
193 """Applies Rotary Position Embedding to the query and key tensors.
194
195 Args:
196 q (`torch.Tensor`): The query tensor.
197 k (`torch.Tensor`): The key tensor.
198 cos (`torch.Tensor`): The cosine part of the rotary embedding.
199 sin (`torch.Tensor`): The sine part of the rotary embedding.
200 position_ids (`torch.Tensor`, *optional*):
201 Deprecated and unused.
202 unsqueeze_dim (`int`, *optional*, defaults to 1):
203 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
204 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
205 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
206 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
207 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
208 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
209 Returns:
210 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
211 """
212 cos = cos.unsqueeze(unsqueeze_dim)
213 sin = sin.unsqueeze(unsqueeze_dim)
214 q_embed = (q * cos) + (rotate_half(q) * sin)
215 k_embed = (k * cos) + (rotate_half(k) * sin)
216 return q_embed, k_embed
217
218
219class Phi3MLP(nn.Module):
220 def __init__(self, config):
221 super().__init__()
222
223 self.config = config
224 self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
225 self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
226
227 self.activation_fn = ACT2FN[config.hidden_act]
228
229 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
230 up_states = self.gate_up_proj(hidden_states)
231
232 gate, up_states = up_states.chunk(2, dim=-1)
233 up_states = up_states * self.activation_fn(gate)
234
235 return self.down_proj(up_states)
236
237
238# Copied from transformers.models.llama.modeling_llama.repeat_kv with llama->phi
239def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
240 """
241 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
242 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
243 """
244 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
245 if n_rep == 1:
246 return hidden_states
247 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
248 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
249
250
251class Phi3Attention(nn.Module):
252 """Multi-headed attention from 'Attention Is All You Need' paper"""
253
254 def __init__(self, config: Phi3Config, layer_idx: Optional[int] = None):
255 super().__init__()
256 self.config = config
257 self.layer_idx = layer_idx
258 if layer_idx is None:
259 logger.warning_once(
260 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
261 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
262 "when creating this class."
263 )
264
265 self.attention_dropout = config.attention_dropout
266 self.hidden_size = config.hidden_size
267 self.num_heads = config.num_attention_heads
268 self.head_dim = self.hidden_size // self.num_heads
269 self.num_key_value_heads = config.num_key_value_heads
270 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
271 self.max_position_embeddings = config.max_position_embeddings
272 self.original_max_position_embeddings = config.original_max_position_embeddings
273 self.rope_theta = config.rope_theta
274 self.rope_scaling = config.rope_scaling
275 self.is_causal = True
276
277 if (self.head_dim * self.num_heads) != self.hidden_size:
278 raise ValueError(
279 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
280 f" and `num_heads`: {self.num_heads})."
281 )
282
283 op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim)
284 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
285 self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False)
286 self._init_rope()
287
288 def _init_rope(self):
289 if self.rope_scaling is None:
290 self.rotary_emb = Phi3RotaryEmbedding(
291 self.head_dim,
292 max_position_embeddings=self.max_position_embeddings,
293 base=self.rope_theta,
294 )
295 else:
296 scaling_type = self.config.rope_scaling["type"]
297 if scaling_type == "longrope":
298 self.rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(self.head_dim, self.config)
299 else:
300 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
301
302 def forward(
303 self,
304 hidden_states: torch.Tensor,
305 attention_mask: Optional[torch.Tensor] = None,
306 position_ids: Optional[torch.LongTensor] = None,
307 past_key_value: Optional[Cache] = None,
308 output_attentions: bool = False,
309 use_cache: bool = False,
310 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
311 logger.warning_once("You are not running the flash-attention implementation, expect numerical differences.")
312
313 bsz, q_len, _ = hidden_states.size()
314
315 qkv = self.qkv_proj(hidden_states)
316 query_pos = self.num_heads * self.head_dim
317 query_states = qkv[..., :query_pos]
318 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
319 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
320
321 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
322 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
323 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
324
325 kv_seq_len = key_states.shape[-2]
326 if past_key_value is not None:
327 if self.layer_idx is None:
328 raise ValueError(
329 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
330 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
331 "with a layer index."
332 )
333 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
334 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
335
336 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
337
338 if past_key_value is not None:
339 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
340 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
341
342 # repeat k/v heads if n_kv_heads < n_heads
343 key_states = repeat_kv(key_states, self.num_key_value_groups)
344 value_states = repeat_kv(value_states, self.num_key_value_groups)
345
346 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
347
348 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
349 raise ValueError(
350 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
351 f" {attn_weights.size()}"
352 )
353
354 if attention_mask is not None:
355 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
356 raise ValueError(
357 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
358 )
359 attn_weights = attn_weights + attention_mask
360
361 # upcast attention to fp32
362 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
363 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
364
365 attn_output = torch.matmul(attn_weights, value_states)
366
367 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
368 raise ValueError(
369 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
370 f" {attn_output.size()}"
371 )
372
373 attn_output = attn_output.transpose(1, 2).contiguous()
374 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
375
376 attn_output = self.o_proj(attn_output)
377
378 if not output_attentions:
379 attn_weights = None
380
381 return attn_output, attn_weights, past_key_value
382
383
384class Phi3FlashAttention2(Phi3Attention):
385 """
386 Phi-3 flash attention module. This module inherits from `Phi3Attention` as the weights of the module stays
387 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
388 flash attention and deal with padding tokens in case the input contains any of them.
389 """
390
391 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
392 def __init__(self, *args, **kwargs):
393 super().__init__(*args, **kwargs)
394
395 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
396 # 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.
397 # 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).
398 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
399
400 def forward(
401 self,
402 hidden_states: torch.Tensor,
403 attention_mask: Optional[torch.LongTensor] = None,
404 position_ids: Optional[torch.LongTensor] = None,
405 past_key_value: Optional[Cache] = None,
406 output_attentions: bool = False,
407 use_cache: bool = False,
408 **kwargs,
409 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
410 # Phi3FlashAttention2 attention does not support output_attentions
411
412 if not _flash_supports_window_size:
413 logger.warning_once(
414 "The current flash attention version does not support sliding window attention. Please use `attn_implementation='eager'` or upgrade flash-attn library."
415 )
416 raise ValueError("The current flash attention version does not support sliding window attention.")
417
418 output_attentions = False
419
420 if "padding_mask" in kwargs:
421 warnings.warn(
422 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
423 )
424
425 # overwrite attention_mask with padding_mask
426 attention_mask = kwargs.pop("padding_mask")
427
428 bsz, q_len, _ = hidden_states.size()
429
430 qkv = self.qkv_proj(hidden_states)
431 query_pos = self.num_heads * self.head_dim
432 query_states = qkv[..., :query_pos]
433 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
434 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
435
436 # Flash attention requires the input to have the shape
437 # batch_size x seq_length x head_dim x hidden_dim
438 # therefore we just need to keep the original shape
439 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
440 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
441 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
442
443 kv_seq_len = key_states.shape[-2]
444 if past_key_value is not None:
445 if self.layer_idx is None:
446 raise ValueError(
447 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
448 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
449 "with a layer index."
450 )
451 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
452
453 # Because the input can be padded, the absolute sequence length depends on the max position id.
454 rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item() + 1)
455 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=rotary_seq_len)
456
457 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
458
459 use_sliding_windows = (
460 _flash_supports_window_size
461 and getattr(self.config, "sliding_window", None) is not None
462 and kv_seq_len > self.config.sliding_window
463 )
464
465 if past_key_value is not None:
466 # Activate slicing cache only if the config has a value `sliding_windows` attribute
467 cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
468 if (
469 getattr(self.config, "sliding_window", None) is not None
470 and kv_seq_len > self.config.sliding_window
471 and cache_has_contents
472 ):
473 slicing_tokens = 1 - self.config.sliding_window
474
475 past_key = past_key_value[self.layer_idx][0]
476 past_value = past_key_value[self.layer_idx][1]
477
478 past_key = past_key[:, :, slicing_tokens:, :].contiguous()
479 past_value = past_value[:, :, slicing_tokens:, :].contiguous()
480
481 if past_key.shape[-2] != self.config.sliding_window - 1:
482 raise ValueError(
483 f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
484 f" {past_key.shape}"
485 )
486
487 if attention_mask is not None:
488 attention_mask = attention_mask[:, slicing_tokens:]
489 attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
490
491 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
492 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
493
494 # repeat k/v heads if n_kv_heads < n_heads
495 key_states = repeat_kv(key_states, self.num_key_value_groups)
496 value_states = repeat_kv(value_states, self.num_key_value_groups)
497
498 attn_dropout = self.attention_dropout if self.training else 0.0
499
500 # In PEFT, usually we cast the layer norms in float32 for training stability reasons
501 # therefore the input hidden states gets silently casted in float32. Hence, we need
502 # cast them back in the correct dtype just to be sure everything works as expected.
503 # This might slowdown training & inference so it is recommended to not cast the LayerNorms
504 # in fp32.
505
506 if query_states.dtype == torch.float32:
507 if torch.is_autocast_enabled():
508 target_dtype = torch.get_autocast_gpu_dtype()
509 # Handle the case where the model is quantized
510 elif hasattr(self.config, "_pre_quantization_dtype"):
511 target_dtype = self.config._pre_quantization_dtype
512 else:
513 target_dtype = self.qkv_proj.weight.dtype
514
515 logger.warning_once(
516 f"The input hidden states seems to be silently casted in float32, this might be related to"
517 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
518 f" {target_dtype}."
519 )
520
521 query_states = query_states.to(target_dtype)
522 key_states = key_states.to(target_dtype)
523 value_states = value_states.to(target_dtype)
524
525 # Reashape to the expected shape for Flash Attention
526 query_states = query_states.transpose(1, 2)
527 key_states = key_states.transpose(1, 2)
528 value_states = value_states.transpose(1, 2)
529
530 attn_output = self._flash_attention_forward(
531 query_states,
532 key_states,
533 value_states,
534 attention_mask,
535 q_len,
536 dropout=attn_dropout,
537 use_sliding_windows=use_sliding_windows,
538 )
539
540 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
541 attn_output = self.o_proj(attn_output)
542
543 if not output_attentions:
544 attn_weights = None
545
546 return attn_output, attn_weights, past_key_value
547
548 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._flash_attention_forward
549 def _flash_attention_forward(
550 self,
551 query_states,
552 key_states,
553 value_states,
554 attention_mask,
555 query_length,
556 dropout=0.0,
557 softmax_scale=None,
558 use_sliding_windows=False,
559 ):
560 """
561 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
562 first unpad the input, then computes the attention scores and pad the final attention scores.
563
564 Args:
565 query_states (`torch.Tensor`):
566 Input query states to be passed to Flash Attention API
567 key_states (`torch.Tensor`):
568 Input key states to be passed to Flash Attention API
569 value_states (`torch.Tensor`):
570 Input value states to be passed to Flash Attention API
571 attention_mask (`torch.Tensor`):
572 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
573 position of padding tokens and 1 for the position of non-padding tokens.
574 dropout (`float`):
575 Attention dropout
576 softmax_scale (`float`, *optional*):
577 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
578 use_sliding_windows (`bool`, *optional*):
579 Whether to activate sliding window attention.
580 """
581 if not self._flash_attn_uses_top_left_mask:
582 causal = self.is_causal
583 else:
584 # 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__.
585 causal = self.is_causal and query_length != 1
586
587 # Contains at least one padding token in the sequence
588 if attention_mask is not None:
589 batch_size = query_states.shape[0]
590 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
591 query_states, key_states, value_states, attention_mask, query_length
592 )
593
594 cu_seqlens_q, cu_seqlens_k = cu_seq_lens
595 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
596
597 if not use_sliding_windows:
598 attn_output_unpad = flash_attn_varlen_func(
599 query_states,
600 key_states,
601 value_states,
602 cu_seqlens_q=cu_seqlens_q,
603 cu_seqlens_k=cu_seqlens_k,
604 max_seqlen_q=max_seqlen_in_batch_q,
605 max_seqlen_k=max_seqlen_in_batch_k,
606 dropout_p=dropout,
607 softmax_scale=softmax_scale,
608 causal=causal,
609 )
610 else:
611 attn_output_unpad = flash_attn_varlen_func(
612 query_states,
613 key_states,
614 value_states,
615 cu_seqlens_q=cu_seqlens_q,
616 cu_seqlens_k=cu_seqlens_k,
617 max_seqlen_q=max_seqlen_in_batch_q,
618 max_seqlen_k=max_seqlen_in_batch_k,
619 dropout_p=dropout,
620 softmax_scale=softmax_scale,
621 causal=causal,
622 window_size=(self.config.sliding_window, self.config.sliding_window),
623 )
624
625 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
626 else:
627 if not use_sliding_windows:
628 attn_output = flash_attn_func(
629 query_states,
630 key_states,
631 value_states,
632 dropout,
633 softmax_scale=softmax_scale,
634 causal=causal,
635 )
636 else:
637 attn_output = flash_attn_func(
638 query_states,
639 key_states,
640 value_states,
641 dropout,
642 softmax_scale=softmax_scale,
643 causal=causal,
644 window_size=(self.config.sliding_window, self.config.sliding_window),
645 )
646
647 return attn_output
648
649 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input
650 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
651 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
652
653 # On the first iteration we need to properly re-create the padding mask
654 # by slicing it on the proper place
655 if kv_seq_len != attention_mask.shape[-1]:
656 attention_mask_num_tokens = attention_mask.shape[-1]
657 attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
658
659 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
660
661 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
662 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
663
664 if query_length == kv_seq_len:
665 query_layer = index_first_axis(
666 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
667 )
668 cu_seqlens_q = cu_seqlens_k
669 max_seqlen_in_batch_q = max_seqlen_in_batch_k
670 indices_q = indices_k
671 elif query_length == 1:
672 max_seqlen_in_batch_q = 1
673 cu_seqlens_q = torch.arange(
674 batch_size + 1, dtype=torch.int32, device=query_layer.device
675 ) # There is a memcpy here, that is very bad.
676 indices_q = cu_seqlens_q[:-1]
677 query_layer = query_layer.squeeze(1)
678 else:
679 # The -q_len: slice assumes left padding.
680 attention_mask = attention_mask[:, -query_length:]
681 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
682
683 return (
684 query_layer,
685 key_layer,
686 value_layer,
687 indices_q,
688 (cu_seqlens_q, cu_seqlens_k),
689 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
690 )
691
692
693# copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Phi3
694# TODO @Arthur no longer copied from LLama after static cache
695class Phi3SdpaAttention(Phi3Attention):
696 """
697 Phi3 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
698 `Phi3Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
699 SDPA API.
700 """
701
702 # Adapted from Phi3Attention.forward
703 def forward(
704 self,
705 hidden_states: torch.Tensor,
706 attention_mask: Optional[torch.Tensor] = None,
707 position_ids: Optional[torch.LongTensor] = None,
708 past_key_value: Optional[Cache] = None,
709 output_attentions: bool = False,
710 use_cache: bool = False,
711 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
712 if output_attentions:
713 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
714 logger.warning_once(
715 "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, "
716 '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.'
717 )
718 return super().forward(
719 hidden_states=hidden_states,
720 attention_mask=attention_mask,
721 position_ids=position_ids,
722 past_key_value=past_key_value,
723 output_attentions=output_attentions,
724 use_cache=use_cache,
725 )
726
727 bsz, q_len, _ = hidden_states.size()
728
729 qkv = self.qkv_proj(hidden_states)
730 query_pos = self.num_heads * self.head_dim
731 query_states = qkv[..., :query_pos]
732 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
733 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
734
735 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
736 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
737 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
738
739 kv_seq_len = key_states.shape[-2]
740 if past_key_value is not None:
741 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
742 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
743
744 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
745
746 if past_key_value is not None:
747 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
748 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
749
750 key_states = repeat_kv(key_states, self.num_key_value_groups)
751 value_states = repeat_kv(value_states, self.num_key_value_groups)
752
753 if attention_mask is not None:
754 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
755 raise ValueError(
756 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
757 )
758
759 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
760 # Reference: https://github.com/pytorch/pytorch/issues/112577.
761 if query_states.device.type == "cuda" and attention_mask is not None:
762 query_states = query_states.contiguous()
763 key_states = key_states.contiguous()
764 value_states = value_states.contiguous()
765
766 attn_output = torch.nn.functional.scaled_dot_product_attention(
767 query_states,
768 key_states,
769 value_states,
770 attn_mask=attention_mask,
771 dropout_p=self.attention_dropout if self.training else 0.0,
772 # 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.
773 is_causal=self.is_causal and attention_mask is None and q_len > 1,
774 )
775
776 attn_output = attn_output.transpose(1, 2).contiguous()
777 attn_output = attn_output.view(bsz, q_len, self.hidden_size)
778
779 attn_output = self.o_proj(attn_output)
780
781 return attn_output, None, past_key_value
782
783
784PHI3_ATTENTION_CLASSES = {
785 "eager": Phi3Attention,
786 "flash_attention_2": Phi3FlashAttention2,
787 "sdpa": Phi3SdpaAttention,
788}
789
790
791class Phi3DecoderLayer(nn.Module):
792 def __init__(self, config: Phi3Config, layer_idx: int):
793 super().__init__()
794
795 self.config = config
796 self.self_attn = PHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
797
798 self.mlp = Phi3MLP(config)
799 self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
800
801 self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)
802 self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)
803 self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
804
805 def forward(
806 self,
807 hidden_states: torch.Tensor,
808 attention_mask: Optional[torch.Tensor] = None,
809 position_ids: Optional[torch.LongTensor] = None,
810 past_key_value: Optional[Tuple[torch.Tensor]] = None,
811 output_attentions: Optional[bool] = False,
812 use_cache: Optional[bool] = False,
813 **kwargs,
814 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
815 if "padding_mask" in kwargs:
816 warnings.warn(
817 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
818 )
819 """
820 Args:
821 hidden_states (`torch.FloatTensor`):
822 input to the layer of shape `(batch, seq_len, embed_dim)`
823 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
824 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
825 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
826 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
827 `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
828 output_attentions (`bool`, *optional*):
829 Whether or not to return the attentions tensors of all attention layers. See `attentions` under
830 returned tensors for more detail.
831 use_cache (`bool`, *optional*):
832 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
833 (see `past_key_values`).
834 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
835 """
836
837 residual = hidden_states
838
839 hidden_states = self.input_layernorm(hidden_states)
840
841 # Self Attention
842 attn_outputs, self_attn_weights, present_key_value = self.self_attn(
843 hidden_states=hidden_states,
844 attention_mask=attention_mask,
845 position_ids=position_ids,
846 past_key_value=past_key_value,
847 output_attentions=output_attentions,
848 use_cache=use_cache,
849 )
850
851 hidden_states = residual + self.resid_attn_dropout(attn_outputs)
852
853 residual = hidden_states
854 hidden_states = self.post_attention_layernorm(hidden_states)
855 hidden_states = self.mlp(hidden_states)
856 hidden_states = residual + self.resid_mlp_dropout(hidden_states)
857
858 outputs = (hidden_states,)
859
860 if output_attentions:
861 outputs += (self_attn_weights,)
862
863 if use_cache:
864 outputs += (present_key_value,)
865
866 return outputs
867
868
869PHI3_START_DOCSTRING = r"""
870 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
871 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
872 etc.)
873
874 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
875 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
876 and behavior.
877
878 Parameters:
879 config ([`Phi3Config`]):
880 Model configuration class with all the parameters of the model. Initializing with a config file does not
881 load the weights associated with the model, only the configuration. Check out the
882 [`~PreTrainedModel.from_pretrained`] method to load the model weights.
883"""
884
885
886@add_start_docstrings(
887 "The bare Phi-3 model outputting raw hidden-states without any specific head on top.",
888 PHI3_START_DOCSTRING,
889)
890class Phi3PreTrainedModel(PreTrainedModel):
891 config_class = Phi3Config
892 base_model_prefix = "model"
893 supports_gradient_checkpointing = True
894 _no_split_modules = ["Phi3DecoderLayer"]
895 _skip_keys_device_placement = "past_key_values"
896 _supports_flash_attn_2 = True
897 _supports_sdpa = False
898 _supports_cache_class = True
899
900 _version = "0.0.5"
901
902 def _init_weights(self, module):
903 std = self.config.initializer_range
904 if isinstance(module, nn.Linear):
905 module.weight.data.normal_(mean=0.0, std=std)
906 if module.bias is not None:
907 module.bias.data.zero_()
908 elif isinstance(module, nn.Embedding):
909 module.weight.data.normal_(mean=0.0, std=std)
910 if module.padding_idx is not None:
911 module.weight.data[module.padding_idx].zero_()
912
913
914PHI3_INPUTS_DOCSTRING = r"""
915 Args:
916 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
917 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
918 it.
919
920 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
921 [`PreTrainedTokenizer.__call__`] for details.
922
923 [What are input IDs?](../glossary#input-ids)
924 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
925 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
926
927 - 1 for tokens that are **not masked**,
928 - 0 for tokens that are **masked**.
929
930 [What are attention masks?](../glossary#attention-mask)
931
932 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
933 [`PreTrainedTokenizer.__call__`] for details.
934
935 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
936 `past_key_values`).
937
938 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
939 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
940 information on the default strategy.
941
942 - 1 indicates the head is **not masked**,
943 - 0 indicates the head is **masked**.
944 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
945 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
946 config.n_positions - 1]`.
947
948 [What are position IDs?](../glossary#position-ids)
949 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
950 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
951 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
952 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
953
954 Two formats are allowed:
955 - a [`~cache_utils.Cache`] instance;
956 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
957 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
958 cache format.
959
960 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
961 legacy cache format will be returned.
962
963 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
964 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
965 of shape `(batch_size, sequence_length)`.
966 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
967 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
968 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
969 model's internal embedding lookup matrix.
970 use_cache (`bool`, *optional*):
971 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
972 `past_key_values`).
973 output_attentions (`bool`, *optional*):
974 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
975 tensors for more detail.
976 output_hidden_states (`bool`, *optional*):
977 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
978 more detail.
979 return_dict (`bool`, *optional*):
980 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
981"""
982
983
984@add_start_docstrings(
985 "The bare Phi-3 model outputting raw hidden-states without any specific head on top.",
986 PHI3_START_DOCSTRING,
987)
988class Phi3Model(Phi3PreTrainedModel):
989 """
990 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
991
992 Args:
993 config: Phi3Config
994 """
995
996 def __init__(self, config: Phi3Config):
997 super().__init__(config)
998 self.padding_idx = config.pad_token_id
999 self.vocab_size = config.vocab_size
1000
1001 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
1002 self.embed_dropout = nn.Dropout(config.embd_pdrop)
1003 self.layers = nn.ModuleList(
1004 [Phi3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
1005 )
1006 self._attn_implementation = config._attn_implementation
1007 self.norm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1008
1009 self.gradient_checkpointing = False
1010 # Initialize weights and apply final processing
1011 self.post_init()
1012
1013 def get_input_embeddings(self):
1014 return self.embed_tokens
1015
1016 def set_input_embeddings(self, value):
1017 self.embed_tokens = value
1018
1019 @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
1020 def forward(
1021 self,
1022 input_ids: torch.LongTensor = None,
1023 attention_mask: Optional[torch.Tensor] = None,
1024 position_ids: Optional[torch.LongTensor] = None,
1025 past_key_values: Optional[List[torch.FloatTensor]] = None,
1026 inputs_embeds: Optional[torch.FloatTensor] = None,
1027 use_cache: Optional[bool] = None,
1028 output_attentions: Optional[bool] = None,
1029 output_hidden_states: Optional[bool] = None,
1030 return_dict: Optional[bool] = None,
1031 ) -> Union[Tuple, BaseModelOutputWithPast]:
1032 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1033 output_hidden_states = (
1034 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1035 )
1036 use_cache = use_cache if use_cache is not None else self.config.use_cache
1037
1038 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1039
1040 # retrieve input_ids and inputs_embeds
1041 if input_ids is not None and inputs_embeds is not None:
1042 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
1043 elif input_ids is not None:
1044 batch_size, seq_length = input_ids.shape[:2]
1045 elif inputs_embeds is not None:
1046 batch_size, seq_length = inputs_embeds.shape[:2]
1047 else:
1048 raise ValueError("You have to specify either input_ids or inputs_embeds")
1049
1050 past_key_values_length = 0
1051
1052 if self.gradient_checkpointing and self.training:
1053 if use_cache:
1054 logger.warning_once(
1055 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1056 )
1057 use_cache = False
1058
1059 if use_cache:
1060 use_legacy_cache = not isinstance(past_key_values, Cache)
1061 if use_legacy_cache:
1062 past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1063 past_key_values_length = past_key_values.get_usable_length(seq_length)
1064
1065 if position_ids is None:
1066 device = input_ids.device if input_ids is not None else inputs_embeds.device
1067 position_ids = torch.arange(
1068 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
1069 )
1070 position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
1071 else:
1072 position_ids = position_ids.view(-1, seq_length).long()
1073
1074 if inputs_embeds is None:
1075 inputs_embeds = self.embed_tokens(input_ids)
1076
1077 if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
1078 is_padding_right = attention_mask[:, -1].sum().item() != batch_size
1079 if is_padding_right:
1080 raise ValueError(
1081 "You are attempting to perform batched generation with padding_side='right'"
1082 " this may lead to unexpected behaviour for Flash Attention version of Phi3. Make sure to "
1083 " call `tokenizer.padding_side = 'left'` before tokenizing the input. "
1084 )
1085
1086 if self._attn_implementation == "flash_attention_2":
1087 # 2d mask is passed through the layers
1088 attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
1089 else:
1090 # 4d mask is passed through the layers
1091 attention_mask = _prepare_4d_causal_attention_mask(
1092 attention_mask,
1093 (batch_size, seq_length),
1094 inputs_embeds,
1095 past_key_values_length,
1096 sliding_window=self.config.sliding_window,
1097 )
1098
1099 hidden_states = inputs_embeds
1100
1101 # decoder layers
1102 all_hidden_states = () if output_hidden_states else None
1103 all_self_attns = () if output_attentions else None
1104 next_decoder_cache = None
1105
1106 for decoder_layer in self.layers:
1107 if output_hidden_states:
1108 all_hidden_states += (hidden_states,)
1109
1110 if self.gradient_checkpointing and self.training:
1111 layer_outputs = self._gradient_checkpointing_func(
1112 decoder_layer.__call__,
1113 hidden_states,
1114 attention_mask,
1115 position_ids,
1116 past_key_values,
1117 output_attentions,
1118 use_cache,
1119 )
1120 else:
1121 layer_outputs = decoder_layer(
1122 hidden_states,
1123 attention_mask=attention_mask,
1124 position_ids=position_ids,
1125 past_key_value=past_key_values,
1126 output_attentions=output_attentions,
1127 use_cache=use_cache,
1128 )
1129
1130 hidden_states = layer_outputs[0]
1131
1132 if use_cache:
1133 next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1134
1135 if output_attentions:
1136 all_self_attns += (layer_outputs[1],)
1137
1138 hidden_states = self.norm(hidden_states)
1139
1140 # add hidden states from the last decoder layer
1141 if output_hidden_states:
1142 all_hidden_states += (hidden_states,)
1143
1144 next_cache = None
1145 if use_cache:
1146 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1147 if not return_dict:
1148 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1149 return BaseModelOutputWithPast(
1150 last_hidden_state=hidden_states,
1151 past_key_values=next_cache,
1152 hidden_states=all_hidden_states,
1153 attentions=all_self_attns,
1154 )
1155
1156
1157class Phi3ForCausalLM(Phi3PreTrainedModel):
1158 _tied_weights_keys = ["lm_head.weight"]
1159
1160 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__ with Llama->Phi3
1161 def __init__(self, config):
1162 super().__init__(config)
1163 self.model = Phi3Model(config)
1164 self.vocab_size = config.vocab_size
1165 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1166
1167 # Initialize weights and apply final processing
1168 self.post_init()
1169
1170 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
1171 def get_input_embeddings(self):
1172 return self.model.embed_tokens
1173
1174 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
1175 def set_input_embeddings(self, value):
1176 self.model.embed_tokens = value
1177
1178 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
1179 def get_output_embeddings(self):
1180 return self.lm_head
1181
1182 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
1183 def set_output_embeddings(self, new_embeddings):
1184 self.lm_head = new_embeddings
1185
1186 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
1187 def set_decoder(self, decoder):
1188 self.model = decoder
1189
1190 # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
1191 def get_decoder(self):
1192 return self.model
1193
1194 # Ignore copy
1195 @add_start_docstrings_to_model_forward(PHI3_INPUTS_DOCSTRING)
1196 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1197 def forward(
1198 self,
1199 input_ids: torch.LongTensor = None,
1200 attention_mask: Optional[torch.Tensor] = None,