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-V 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_greater_or_equal_2_10,
44 logging,
45 replace_return_docstrings,
46)
47from .configuration_phi3_v import Phi3VConfig
48
49try:
50 from flash_attn import flash_attn_func, flash_attn_varlen_func
51 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
52
53 _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
54except ImportError:
55 pass
56
57import torch
58from torch import nn
59from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig
60from transformers.models.clip.modeling_clip import CLIPAttention
61from transformers.utils import logging
62
63logger = logging.get_logger(__name__)
64
65
66MAX_INPUT_ID = int(1e9)
67
68CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(
69 attention_dropout=0.0,
70 dropout=0.0,
71 hidden_act="quick_gelu",
72 hidden_size=1024,
73 image_size=336,
74 initializer_factor=1.0,
75 initializer_range=0.02,
76 intermediate_size=4096,
77 layer_norm_eps=1e-05,
78 num_attention_heads=16,
79 num_channels=3,
80 num_hidden_layers=24,
81 patch_size=14,
82 projection_dim=768
83)
84
85class CLIPAttentionFA2(CLIPAttention):
86 """Add flash attention 2 to CLIPAttention. (This is only used in the vision encoder)"""
87
88 def forward(self,
89 hidden_states,
90 attention_mask=None,
91 causal_attention_mask=None,
92 output_attentions=False,
93 ):
94 """Input shape: Batch x Time x Channel"""
95
96 assert attention_mask is None, "CLIPAttentionFA2 does not support attention_mask"
97 assert causal_attention_mask is None, "CLIPAttentionFA2 does not support causal_attention_mask"
98 assert output_attentions is False, "CLIPAttentionFA2 does not support output_attentions"
99
100 bsz, tgt_len, embed_dim = hidden_states.size()
101 query_states = self.q_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
102 key_states = self.k_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
103 value_states = self.v_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
104
105 attn_output = flash_attn_func(
106 query_states,
107 key_states,
108 value_states,
109 dropout_p=self.dropout if self.training else 0.0,
110 softmax_scale=self.scale,
111 causal=False,
112 ).reshape(bsz, tgt_len, embed_dim)
113
114 attn_output = self.out_proj(attn_output)
115 return attn_output, None
116
117
118class Phi3ImageEmbedding(nn.Module):
119 """Phi3 Image embedding."""
120
121 def __init__(self, config: PretrainedConfig, wte=None, **kwargs) -> None:
122 super().__init__()
123
124 # n_embed or hidden_size
125 hidden_size = config.n_embd if hasattr(config, 'n_embd') else config.hidden_size
126 if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):
127 embd_drop = config.embd_pdrop if hasattr(config, 'embd_pdrop') else config.embed_pdrop
128 self.drop = nn.Dropout(embd_drop)
129 else:
130 self.drop = None
131
132 self.wte = wte
133
134 if isinstance(config.img_processor, dict) and config.img_processor.get('name', None) == 'clip_vision_model':
135 assert 'model_name' in config.img_processor, 'model_name must be provided for CLIPVisionModel'
136 assert 'image_dim_out' in config.img_processor, 'image_dim_out must be provided for CLIPVisionModel'
137 assert 'num_img_tokens' in config.img_processor, 'num_img_tokens must be provided for CLIPVisionModel'
138 assert config.img_processor['model_name'] == 'openai/clip-vit-large-patch14-336'
139 clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG
140 self.img_processor = CLIPVisionModel(clip_config)
141 image_dim_out = config.img_processor['image_dim_out']
142 self.num_img_tokens = config.img_processor['num_img_tokens']
143
144 # FA2 in CLIP
145 if config._attn_implementation == 'flash_attention_2':
146 for layer in self.img_processor.vision_model.encoder.layers:
147 clip_fa2 = CLIPAttentionFA2(clip_config)
148 del layer.self_attn
149 layer.self_attn = clip_fa2
150 else:
151 raise NotImplementedError(f'img_processor = {config.img_processor}, not implemented')
152
153 self.image_dim_out = image_dim_out
154 self.img_sizes = None
155
156 # global_gn and sub_gn for hd transform, serves as line separator
157 self.use_hd_transform = kwargs.get('use_hd_transform', False)
158 self.with_learnable_separator = kwargs.get('with_learnable_separator', False)
159 self.hd_transform_order = kwargs.get('hd_transform_order', 'glb_sub')
160 # with_hd_transform and with_learnable_separator should have same value
161 assert self.use_hd_transform == self.with_learnable_separator, 'use_hd_transform and with_learnable_separator should have same value'
162 if self.with_learnable_separator:
163 assert self.use_hd_transform, 'learnable separator is only for hd transform'
164 # 1024 * 4, merge spatial to channel dimension
165 self.glb_GN = nn.Parameter(torch.zeros([1, 1, self.image_dim_out * 4]))
166 self.sub_GN = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out * 4]))
167 logger.info(f'learnable separator enabled for hd transform, hd_transform_order = {self.hd_transform_order}')
168
169 projection_cls = kwargs.get('projection_cls', 'linear')
170 if projection_cls == 'linear':
171 self.img_projection = nn.Linear(image_dim_out, hidden_size)
172 elif projection_cls == 'mlp' and self.use_hd_transform:
173 dim_projection = hidden_size
174 depth = 2
175 layers = [nn.Linear(image_dim_out * 4, dim_projection)]
176 for _ in range(1, depth):
177 layers.extend([nn.GELU(),
178 nn.Linear(dim_projection, dim_projection)])
179 self.img_projection = nn.Sequential(*layers)
180 elif projection_cls == 'mlp':
181 dim_projection = hidden_size
182 depth = 2
183 layers = [nn.Linear(image_dim_out, dim_projection)]
184 for _ in range(1, depth):
185 layers.extend([nn.GELU(),
186 nn.Linear(dim_projection, dim_projection)])
187 self.img_projection = nn.Sequential(*layers)
188 else:
189 raise NotImplementedError(f'projection_cls = {projection_cls}, not implemented')
190
191 self.vocab_size = config.vocab_size
192 self.img_features = None
193
194 if isinstance(config.img_processor, dict):
195 self.layer_idx = config.img_processor.get('layer_idx', -2)
196 self.type_feature = config.img_processor.get('type_feature', 'patch')
197 else:
198 self.layer_idx = -2
199 self.type_feature = 'patch'
200
201
202 def set_img_features(self, img_features: torch.FloatTensor) -> None:
203 self.img_features = img_features
204
205 def set_img_sizes(self, img_sizes: torch.LongTensor) -> None:
206 self.img_sizes = img_sizes
207
208 def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:
209 LAYER_IDX = self.layer_idx
210 TYPE_FEATURE = self.type_feature
211
212 img_processor_output = self.img_processor(img_embeds, output_hidden_states=True)
213 img_feature = img_processor_output.hidden_states[LAYER_IDX]
214
215 if TYPE_FEATURE == "patch":
216 patch_feature = img_feature[:, 1:]
217 return patch_feature
218
219 raise NotImplementedError
220
221 def forward(
222 self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None
223 ) -> torch.FloatTensor:
224 input_shape = input_ids.size()
225 input_ids = input_ids.view(-1, input_shape[-1])
226
227 # positions for image tokens
228 positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)
229 has_image = len(positions[0].tolist()) > 0
230 input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()
231 hidden_states = self.wte(input_ids)
232
233 if has_image:
234 assert self.use_hd_transform
235 num_images, num_crops, c, h, w = pixel_values.shape
236 assert c == 3 and h == w == 336
237 img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(
238 num_images, num_crops, -1, self.image_dim_out
239 )
240 image_features_proj = self.hd_feature_transform(img_features, image_sizes)
241 hidden_states = hidden_states.index_put(
242 positions, image_features_proj, accumulate=False
243 )
244
245 if self.drop is not None:
246 hidden_states = self.drop(hidden_states)
247
248 return hidden_states
249
250 def hd_feature_transform(self, image_features, image_sizes):
251 """
252 image_features: (num_images, num_crops+1, 24*24, 1024)
253 """
254 assert (
255 self.hd_transform_order == 'sub_glb'
256 ), f'hd_transform_order `{self.hd_transform_order}` not implemented'
257 if isinstance(self.img_projection, nn.Sequential):
258 target_device = self.img_projection[0].bias.device
259 target_dtype = self.img_projection[0].bias.dtype
260 else: # It's a single nn.Linear layer
261 target_device = self.img_projection.bias.device
262 target_dtype = self.img_projection.bias.dtype
263
264 global_image_features = image_features[:, 0] # (num_images, 24*24, 1024)
265 # global feature can be viewed as a special HD case with num_crops 1x1
266 global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
267 global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
268
269 all_image_embeddings = []
270 # need a for loop to process each image because of different image sizes
271 # (patch arrangement is different for each image)
272 for i, img_size in enumerate(image_sizes):
273 h, w = img_size
274 h_crop = h // 336
275 w_crop = w // 336
276 num_crops = h_crop * w_crop
277
278 # NOTE: real num_crops is padded
279 # (num_crops, 24*24, 1024)
280 sub_image_features = image_features[i, 1 : 1 + num_crops]
281 sub_image_features_hd = self.reshape_hd_patches_2x2merge(
282 sub_image_features, h_crop, w_crop
283 )
284 sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
285
286 # [sub features, separator, global features]
287 all_image_embeddings.extend(
288 [
289 sub_image_features_hd_newline.squeeze(0), # (h_crop*12*(w_crop*12+1), 4096)
290 self.glb_GN.squeeze(0),
291 global_image_features_hd_newline[i],
292 ]
293 )
294
295 image_features_proj = self.img_projection(
296 torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)
297 )
298
299 return image_features_proj
300
301 def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
302 """
303 image_features: (num_images*num_crops, 24*24, 1024)
304 output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops
305 """
306 N, L, C = image_features.shape
307 assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0
308 num_images = N // (h_crop * w_crop)
309 H = int(L**0.5)
310 image_features_hd = (
311 image_features.reshape(N, H, H, C) # N, 24, 24, 1024
312 .reshape(N, H // 2, 2, H // 2, 2, C) # N, 12, 2, 12, 2, 1024
313 .permute(0, 1, 3, 2, 4, 5) # N, 12, 12, 2, 2, 1024
314 .reshape(N, -1, 4 * C) # N, 144, 4096
315 .reshape(
316 num_images, h_crop, w_crop, H // 2, H // 2, -1
317 ) # n_img, h_crop, w_crop, 12, 12, 4096
318 .permute(0, 1, 3, 2, 4, 5) # n_img, h_crop, 12, w_crop, 12, 4096
319 .reshape(
320 num_images, h_crop * H // 2, w_crop * H // 2, 4 * C
321 ) # n_img, h_crop*12, w_crop*12, 4096
322 )
323
324 # alternative implementation using einops
325 # from einops import rearrange
326 # image_features_nhwc = rearrange(
327 # image_features,
328 # 'N (H W) c -> N H W c',
329 # H=H,
330 # W=H,
331 # )
332 # image_features_2x2merge = rearrange(
333 # image_features_nhwc,
334 # 'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',
335 # h_pool=2,
336 # w_pool=2,
337 # )
338 # image_features_hd = rearrange(
339 # image_features_2x2merge,
340 # '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',
341 # h_crop=h_crop,
342 # w_crop=w_crop,
343 # )
344
345 return image_features_hd
346
347 def add_image_newline(self, image_features_hd):
348 """
349 image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
350 output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
351 """
352 num_images, h, w, hid_dim = image_features_hd.shape
353 # add the newline token to the HD image feature patches
354 newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1) # (n_img, h, 1, hid_dim)
355 image_features_hd_newline = torch.cat(
356 [image_features_hd, newline_embeddings], dim=2
357 ).reshape(num_images, -1, hid_dim)
358 return image_features_hd_newline
359
360
361logger = logging.get_logger(__name__)
362
363_CHECKPOINT_FOR_DOC = "microsoft/Phi-3-vision-128k-instruct"
364_CONFIG_FOR_DOC = "Phi3VConfig"
365
366PHI3V_PRETRAINED_MODEL_ARCHIVE_LIST = [
367 "microsoft/Phi-3-vision-128k-instruct",
368 # See all Phi-3 models at https://huggingface.co/models?filter=Phi-3
369]
370
371
372# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Phi3
373class Phi3RMSNorm(nn.Module):
374 def __init__(self, hidden_size, eps=1e-6):
375 """
376 Phi3RMSNorm is equivalent to T5LayerNorm
377 """
378 super().__init__()
379 self.weight = nn.Parameter(torch.ones(hidden_size))
380 self.variance_epsilon = eps
381
382 def forward(self, hidden_states):
383 input_dtype = hidden_states.dtype
384 hidden_states = hidden_states.to(torch.float32)
385 variance = hidden_states.pow(2).mean(-1, keepdim=True)
386 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
387 return self.weight * hidden_states.to(input_dtype)
388
389
390# Copied from transformers.models.llama.modeling_llama._get_unpad_data
391def _get_unpad_data(attention_mask):
392 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
393 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
394 max_seqlen_in_batch = seqlens_in_batch.max().item()
395 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
396 return (
397 indices,
398 cu_seqlens,
399 max_seqlen_in_batch,
400 )
401
402
403# Copied from transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding with gemma->phi3, Gemma->Phi3
404class Phi3RotaryEmbedding(nn.Module):
405 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
406 super().__init__()
407
408 self.dim = dim
409 self.max_position_embeddings = max_position_embeddings
410 self.base = base
411 self.register_buffer("inv_freq", None, persistent=False)
412
413 @torch.no_grad()
414 def forward(self, x, position_ids, seq_len=None):
415 # x: [bs, num_attention_heads, seq_len, head_size]
416 if self.inv_freq is None:
417 self.inv_freq = 1.0 / (
418 self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim)
419 )
420 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
421 position_ids_expanded = position_ids[:, None, :].float()
422 # Force float32 since bfloat16 loses precision on long contexts
423 # See https://github.com/huggingface/transformers/pull/29285
424 device_type = x.device.type
425 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
426 with torch.autocast(device_type=device_type, enabled=False):
427 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
428 emb = torch.cat((freqs, freqs), dim=-1)
429 cos = emb.cos()
430 sin = emb.sin()
431 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
432
433
434class Phi3SuScaledRotaryEmbedding(Phi3RotaryEmbedding):
435 def __init__(self, dim, config, device=None):
436 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)
437
438 self.short_factor = config.rope_scaling["short_factor"]
439 self.long_factor = config.rope_scaling["long_factor"]
440 self.original_max_position_embeddings = config.original_max_position_embeddings
441
442 @torch.no_grad()
443 def forward(self, x, position_ids, seq_len=None):
444 seq_len = torch.max(position_ids) + 1
445 if seq_len > self.original_max_position_embeddings:
446 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)
447 else:
448 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)
449
450 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim
451 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)
452
453 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
454 position_ids_expanded = position_ids[:, None, :].float()
455
456 # Force float32 since bfloat16 loses precision on long contexts
457 # See https://github.com/huggingface/transformers/pull/29285
458 device_type = x.device.type
459 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
460 with torch.autocast(device_type=device_type, enabled=False):
461 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
462 emb = torch.cat((freqs, freqs), dim=-1)
463
464 scale = self.max_position_embeddings / self.original_max_position_embeddings
465 if scale <= 1.0:
466 scaling_factor = 1.0
467 else:
468 scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))
469
470 cos = emb.cos() * scaling_factor
471 sin = emb.sin() * scaling_factor
472 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
473
474
475class Phi3YarnScaledRotaryEmbedding(Phi3RotaryEmbedding):
476 def __init__(self, dim, config, device=None):
477 super().__init__(dim, config.max_position_embeddings, config.rope_theta, device)
478
479 self.short_factor = config.rope_scaling["short_factor"]
480 self.long_factor = config.rope_scaling["long_factor"]
481 self.original_max_position_embeddings = config.original_max_position_embeddings
482
483 @torch.no_grad()
484 def forward(self, x, position_ids, seq_len=None):
485 seq_len = torch.max(position_ids) + 1
486 if seq_len > self.original_max_position_embeddings:
487 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=x.device)
488 else:
489 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=x.device)
490
491 inv_freq_shape = torch.arange(0, self.dim, 2, dtype=torch.int64, device=x.device).float() / self.dim
492 self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)
493
494 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
495 position_ids_expanded = position_ids[:, None, :].float()
496
497 # Force float32 since bfloat16 loses precision on long contexts
498 # See https://github.com/huggingface/transformers/pull/29285
499 device_type = x.device.type
500 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
501 with torch.autocast(device_type=device_type, enabled=False):
502 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
503 emb = torch.cat((freqs, freqs), dim=-1)
504
505 scale = self.max_position_embeddings / self.original_max_position_embeddings
506 if scale <= 1.0:
507 scaling_factor = 1.0
508 else:
509 scaling_factor = 0.1 * math.log(scale) + 1.0
510
511 cos = emb.cos() * scaling_factor
512 sin = emb.sin() * scaling_factor
513 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
514
515
516# Copied from transformers.models.llama.modeling_llama.rotate_half
517def rotate_half(x):
518 """Rotates half the hidden dims of the input."""
519 x1 = x[..., : x.shape[-1] // 2]
520 x2 = x[..., x.shape[-1] // 2 :]
521 return torch.cat((-x2, x1), dim=-1)
522
523
524# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
525def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
526 """Applies Rotary Position Embedding to the query and key tensors.
527
528 Args:
529 q (`torch.Tensor`): The query tensor.
530 k (`torch.Tensor`): The key tensor.
531 cos (`torch.Tensor`): The cosine part of the rotary embedding.
532 sin (`torch.Tensor`): The sine part of the rotary embedding.
533 position_ids (`torch.Tensor`, *optional*):
534 Deprecated and unused.
535 unsqueeze_dim (`int`, *optional*, defaults to 1):
536 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
537 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
538 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
539 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
540 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
541 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
542 Returns:
543 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
544 """
545 cos = cos.unsqueeze(unsqueeze_dim)
546 sin = sin.unsqueeze(unsqueeze_dim)
547 q_embed = (q * cos) + (rotate_half(q) * sin)
548 k_embed = (k * cos) + (rotate_half(k) * sin)
549 return q_embed, k_embed
550
551
552class Phi3MLP(nn.Module):
553 def __init__(self, config):
554 super().__init__()
555
556 self.config = config
557 self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
558 self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
559
560 self.activation_fn = ACT2FN[config.hidden_act]
561
562 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
563 up_states = self.gate_up_proj(hidden_states)
564
565 gate, up_states = up_states.chunk(2, dim=-1)
566 up_states = up_states * self.activation_fn(gate)
567
568 return self.down_proj(up_states)
569
570
571# Copied from transformers.models.llama.modeling_llama.repeat_kv with llama->phi
572def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
573 """
574 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
575 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
576 """
577 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
578 if n_rep == 1:
579 return hidden_states
580 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
581 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
582
583
584class Phi3Attention(nn.Module):
585 """Multi-headed attention from 'Attention Is All You Need' paper"""
586
587 def __init__(self, config: Phi3VConfig, layer_idx: Optional[int] = None):
588 super().__init__()
589 self.config = config
590 self.layer_idx = layer_idx
591 if layer_idx is None:
592 logger.warning_once(
593 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
594 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
595 "when creating this class."
596 )
597
598 self.attention_dropout = config.attention_dropout
599 self.hidden_size = config.hidden_size
600 self.num_heads = config.num_attention_heads
601 self.head_dim = self.hidden_size // self.num_heads
602 self.num_key_value_heads = config.num_key_value_heads
603 self.num_key_value_groups = self.num_heads // self.num_key_value_heads
604 self.max_position_embeddings = config.max_position_embeddings
605 self.original_max_position_embeddings = config.original_max_position_embeddings
606 self.rope_theta = config.rope_theta
607 self.rope_scaling = config.rope_scaling
608 self.is_causal = True
609
610 if (self.head_dim * self.num_heads) != self.hidden_size:
611 raise ValueError(
612 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
613 f" and `num_heads`: {self.num_heads})."
614 )
615
616 op_size = self.num_heads * self.head_dim + 2 * (self.num_key_value_heads * self.head_dim)
617 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
618 self.qkv_proj = nn.Linear(self.hidden_size, op_size, bias=False)
619 self._init_rope()
620
621 def _init_rope(self):
622 if self.rope_scaling is None:
623 self.rotary_emb = Phi3RotaryEmbedding(
624 self.head_dim,
625 max_position_embeddings=self.max_position_embeddings,
626 base=self.rope_theta,
627 )
628 else:
629 scaling_type = self.config.rope_scaling["type"]
630 if scaling_type == "su":
631 self.rotary_emb = Phi3SuScaledRotaryEmbedding(self.head_dim, self.config)
632 elif scaling_type == "yarn":
633 self.rotary_emb = Phi3YarnScaledRotaryEmbedding(self.head_dim, self.config)
634 else:
635 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
636
637 def forward(
638 self,
639 hidden_states: torch.Tensor,
640 attention_mask: Optional[torch.Tensor] = None,
641 position_ids: Optional[torch.LongTensor] = None,
642 past_key_value: Optional[Cache] = None,
643 output_attentions: bool = False,
644 use_cache: bool = False,
645 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
646 logger.warning_once("You are not running the flash-attention implementation, expect numerical differences.")
647
648 bsz, q_len, _ = hidden_states.size()
649
650 qkv = self.qkv_proj(hidden_states)
651 query_pos = self.num_heads * self.head_dim
652 query_states = qkv[..., :query_pos]
653 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
654 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
655
656 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
657 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
658 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
659
660 kv_seq_len = key_states.shape[-2]
661 if past_key_value is not None:
662 if self.layer_idx is None:
663 raise ValueError(
664 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
665 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
666 "with a layer index."
667 )
668 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
669 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
670
671 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
672
673 if past_key_value is not None:
674 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
675 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
676
677 # repeat k/v heads if n_kv_heads < n_heads
678 key_states = repeat_kv(key_states, self.num_key_value_groups)
679 value_states = repeat_kv(value_states, self.num_key_value_groups)
680
681 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
682
683 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
684 raise ValueError(
685 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
686 f" {attn_weights.size()}"
687 )
688
689 if attention_mask is not None:
690 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
691 raise ValueError(
692 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
693 )
694 attn_weights = attn_weights + attention_mask
695
696 # upcast attention to fp32
697 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype)
698 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
699
700 attn_output = torch.matmul(attn_weights, value_states)
701
702 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
703 raise ValueError(
704 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
705 f" {attn_output.size()}"
706 )
707
708 attn_output = attn_output.transpose(1, 2).contiguous()
709 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
710
711 attn_output = self.o_proj(attn_output)
712
713 if not output_attentions:
714 attn_weights = None
715
716 return attn_output, attn_weights, past_key_value
717
718
719class Phi3FlashAttention2(Phi3Attention):
720 """
721 Phi-3 flash attention module. This module inherits from `Phi3Attention` as the weights of the module stays
722 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
723 flash attention and deal with padding tokens in case the input contains any of them.
724 """
725
726 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
727 def __init__(self, *args, **kwargs):
728 super().__init__(*args, **kwargs)
729
730 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
731 # 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.
732 # 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).
733 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
734
735 def forward(
736 self,
737 hidden_states: torch.Tensor,
738 attention_mask: Optional[torch.LongTensor] = None,
739 position_ids: Optional[torch.LongTensor] = None,
740 past_key_value: Optional[Cache] = None,
741 output_attentions: bool = False,
742 use_cache: bool = False,
743 **kwargs,
744 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
745 # Phi3FlashAttention2 attention does not support output_attentions
746
747 if not _flash_supports_window_size:
748 logger.warning_once(
749 "The current flash attention version does not support sliding window attention. Please use `attn_implementation='eager'` or upgrade flash-attn library."
750 )
751 raise ValueError("The current flash attention version does not support sliding window attention.")
752
753 output_attentions = False
754
755 if "padding_mask" in kwargs:
756 warnings.warn(
757 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
758 )
759
760 # overwrite attention_mask with padding_mask
761 attention_mask = kwargs.pop("padding_mask")
762
763 bsz, q_len, _ = hidden_states.size()
764
765 qkv = self.qkv_proj(hidden_states)
766 query_pos = self.num_heads * self.head_dim
767 query_states = qkv[..., :query_pos]
768 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
769 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
770
771 # Flash attention requires the input to have the shape
772 # batch_size x seq_length x head_dim x hidden_dim
773 # therefore we just need to keep the original shape
774 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
775 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
776 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
777
778 kv_seq_len = key_states.shape[-2]
779 if past_key_value is not None:
780 if self.layer_idx is None:
781 raise ValueError(
782 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
783 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
784 "with a layer index."
785 )
786 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
787
788 # Because the input can be padded, the absolute sequence length depends on the max position id.
789 rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
790 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=rotary_seq_len)
791
792 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
793
794 use_sliding_windows = (
795 _flash_supports_window_size
796 and getattr(self.config, "sliding_window", None) is not None
797 and kv_seq_len > self.config.sliding_window
798 )
799
800 if past_key_value is not None:
801 # Activate slicing cache only if the config has a value `sliding_windows` attribute
802 cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
803 if (
804 getattr(self.config, "sliding_window", None) is not None
805 and kv_seq_len > self.config.sliding_window
806 and cache_has_contents
807 ):
808 slicing_tokens = 1 - self.config.sliding_window
809
810 past_key = past_key_value[self.layer_idx][0]
811 past_value = past_key_value[self.layer_idx][1]
812
813 past_key = past_key[:, :, slicing_tokens:, :].contiguous()
814 past_value = past_value[:, :, slicing_tokens:, :].contiguous()
815
816 if past_key.shape[-2] != self.config.sliding_window - 1:
817 raise ValueError(
818 f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
819 f" {past_key.shape}"
820 )
821
822 if attention_mask is not None:
823 attention_mask = attention_mask[:, slicing_tokens:]
824 attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
825
826 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
827 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
828
829 # repeat k/v heads if n_kv_heads < n_heads
830 key_states = repeat_kv(key_states, self.num_key_value_groups)
831 value_states = repeat_kv(value_states, self.num_key_value_groups)
832
833 attn_dropout = self.attention_dropout if self.training else 0.0
834
835 # In PEFT, usually we cast the layer norms in float32 for training stability reasons
836 # therefore the input hidden states gets silently casted in float32. Hence, we need
837 # cast them back in the correct dtype just to be sure everything works as expected.
838 # This might slowdown training & inference so it is recommended to not cast the LayerNorms
839 # in fp32.
840
841 if query_states.dtype == torch.float32:
842 if torch.is_autocast_enabled():
843 target_dtype = torch.get_autocast_gpu_dtype()
844 # Handle the case where the model is quantized
845 elif hasattr(self.config, "_pre_quantization_dtype"):
846 target_dtype = self.config._pre_quantization_dtype
847 else:
848 target_dtype = self.qkv_proj.weight.dtype
849
850 logger.warning_once(
851 f"The input hidden states seems to be silently casted in float32, this might be related to"
852 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
853 f" {target_dtype}."
854 )
855
856 query_states = query_states.to(target_dtype)
857 key_states = key_states.to(target_dtype)
858 value_states = value_states.to(target_dtype)
859
860 # Reashape to the expected shape for Flash Attention
861 query_states = query_states.transpose(1, 2)
862 key_states = key_states.transpose(1, 2)
863 value_states = value_states.transpose(1, 2)
864
865 attn_output = self._flash_attention_forward(
866 query_states,
867 key_states,
868 value_states,
869 attention_mask,
870 q_len,
871 dropout=attn_dropout,
872 use_sliding_windows=use_sliding_windows,
873 )
874
875 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
876 attn_output = self.o_proj(attn_output)
877
878 if not output_attentions:
879 attn_weights = None
880
881 return attn_output, attn_weights, past_key_value
882
883 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._flash_attention_forward
884 def _flash_attention_forward(
885 self,
886 query_states,
887 key_states,
888 value_states,
889 attention_mask,
890 query_length,
891 dropout=0.0,
892 softmax_scale=None,
893 use_sliding_windows=False,
894 ):
895 """
896 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
897 first unpad the input, then computes the attention scores and pad the final attention scores.
898
899 Args:
900 query_states (`torch.Tensor`):
901 Input query states to be passed to Flash Attention API
902 key_states (`torch.Tensor`):
903 Input key states to be passed to Flash Attention API
904 value_states (`torch.Tensor`):
905 Input value states to be passed to Flash Attention API
906 attention_mask (`torch.Tensor`):
907 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
908 position of padding tokens and 1 for the position of non-padding tokens.
909 dropout (`float`):
910 Attention dropout
911 softmax_scale (`float`, *optional*):
912 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
913 use_sliding_windows (`bool`, *optional*):
914 Whether to activate sliding window attention.
915 """
916 if not self._flash_attn_uses_top_left_mask:
917 causal = self.is_causal
918 else:
919 # 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__.
920 causal = self.is_causal and query_length != 1
921
922 # Contains at least one padding token in the sequence
923 if attention_mask is not None:
924 batch_size = query_states.shape[0]
925 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
926 query_states, key_states, value_states, attention_mask, query_length
927 )
928
929 cu_seqlens_q, cu_seqlens_k = cu_seq_lens
930 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
931
932 if not use_sliding_windows:
933 attn_output_unpad = flash_attn_varlen_func(
934 query_states,
935 key_states,
936 value_states,
937 cu_seqlens_q=cu_seqlens_q,
938 cu_seqlens_k=cu_seqlens_k,
939 max_seqlen_q=max_seqlen_in_batch_q,
940 max_seqlen_k=max_seqlen_in_batch_k,
941 dropout_p=dropout,
942 softmax_scale=softmax_scale,
943 causal=causal,
944 )
945 else:
946 attn_output_unpad = flash_attn_varlen_func(
947 query_states,
948 key_states,
949 value_states,
950 cu_seqlens_q=cu_seqlens_q,
951 cu_seqlens_k=cu_seqlens_k,
952 max_seqlen_q=max_seqlen_in_batch_q,
953 max_seqlen_k=max_seqlen_in_batch_k,
954 dropout_p=dropout,
955 softmax_scale=softmax_scale,
956 causal=causal,
957 window_size=(self.config.sliding_window, self.config.sliding_window),
958 )
959
960 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
961 else:
962 if not use_sliding_windows:
963 attn_output = flash_attn_func(
964 query_states,
965 key_states,
966 value_states,
967 dropout,
968 softmax_scale=softmax_scale,
969 causal=causal,
970 )
971 else:
972 attn_output = flash_attn_func(
973 query_states,
974 key_states,
975 value_states,
976 dropout,
977 softmax_scale=softmax_scale,
978 causal=causal,
979 window_size=(self.config.sliding_window, self.config.sliding_window),
980 )
981
982 return attn_output
983
984 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input
985 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
986 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
987
988 # On the first iteration we need to properly re-create the padding mask
989 # by slicing it on the proper place
990 if kv_seq_len != attention_mask.shape[-1]:
991 attention_mask_num_tokens = attention_mask.shape[-1]
992 attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
993
994 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
995
996 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
997 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
998
999 if query_length == kv_seq_len:
1000 query_layer = index_first_axis(
1001 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
1002 )
1003 cu_seqlens_q = cu_seqlens_k
1004 max_seqlen_in_batch_q = max_seqlen_in_batch_k
1005 indices_q = indices_k
1006 elif query_length == 1:
1007 max_seqlen_in_batch_q = 1
1008 cu_seqlens_q = torch.arange(
1009 batch_size + 1, dtype=torch.int32, device=query_layer.device
1010 ) # There is a memcpy here, that is very bad.
1011 indices_q = cu_seqlens_q[:-1]
1012 query_layer = query_layer.squeeze(1)
1013 else:
1014 # The -q_len: slice assumes left padding.
1015 attention_mask = attention_mask[:, -query_length:]
1016 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
1017
1018 return (
1019 query_layer,
1020 key_layer,
1021 value_layer,
1022 indices_q,
1023 (cu_seqlens_q, cu_seqlens_k),
1024 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
1025 )
1026
1027
1028# copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Phi3
1029# TODO @Arthur no longer copied from LLama after static cache
1030class Phi3SdpaAttention(Phi3Attention):
1031 """
1032 Phi3 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
1033 `Phi3Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
1034 SDPA API.
1035 """
1036
1037 # Adapted from Phi3Attention.forward
1038 def forward(
1039 self,
1040 hidden_states: torch.Tensor,
1041 attention_mask: Optional[torch.Tensor] = None,
1042 position_ids: Optional[torch.LongTensor] = None,
1043 past_key_value: Optional[Cache] = None,
1044 output_attentions: bool = False,
1045 use_cache: bool = False,
1046 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
1047 if output_attentions:
1048 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
1049 logger.warning_once(
1050 "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, "
1051 '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.'
1052 )
1053 return super().forward(
1054 hidden_states=hidden_states,
1055 attention_mask=attention_mask,
1056 position_ids=position_ids,
1057 past_key_value=past_key_value,
1058 output_attentions=output_attentions,
1059 use_cache=use_cache,
1060 )
1061
1062 bsz, q_len, _ = hidden_states.size()
1063
1064 qkv = self.qkv_proj(hidden_states)
1065 query_pos = self.num_heads * self.head_dim
1066 query_states = qkv[..., :query_pos]
1067 key_states = qkv[..., query_pos : query_pos + self.num_key_value_heads * self.head_dim]
1068 value_states = qkv[..., query_pos + self.num_key_value_heads * self.head_dim :]
1069
1070 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
1071 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
1072 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
1073
1074 kv_seq_len = key_states.shape[-2]
1075 if past_key_value is not None:
1076 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
1077 cos, sin = self.rotary_emb(value_states, position_ids, seq_len=kv_seq_len)
1078
1079 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
1080
1081 if past_key_value is not None:
1082 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
1083 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
1084
1085 key_states = repeat_kv(key_states, self.num_key_value_groups)
1086 value_states = repeat_kv(value_states, self.num_key_value_groups)
1087
1088 if attention_mask is not None:
1089 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
1090 raise ValueError(
1091 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
1092 )
1093
1094 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
1095 # Reference: https://github.com/pytorch/pytorch/issues/112577.
1096 if query_states.device.type == "cuda" and attention_mask is not None:
1097 query_states = query_states.contiguous()
1098 key_states = key_states.contiguous()
1099 value_states = value_states.contiguous()
1100
1101 attn_output = torch.nn.functional.scaled_dot_product_attention(
1102 query_states,
1103 key_states,
1104 value_states,
1105 attn_mask=attention_mask,
1106 dropout_p=self.attention_dropout if self.training else 0.0,
1107 # 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.
1108 is_causal=self.is_causal and attention_mask is None and q_len > 1,
1109 )
1110
1111 attn_output = attn_output.transpose(1, 2).contiguous()
1112 attn_output = attn_output.view(bsz, q_len, self.hidden_size)
1113
1114 attn_output = self.o_proj(attn_output)
1115
1116 return attn_output, None, past_key_value
1117
1118
1119PHI3_ATTENTION_CLASSES = {
1120 "eager": Phi3Attention,
1121 "flash_attention_2": Phi3FlashAttention2,
1122 "sdpa": Phi3SdpaAttention,
1123}
1124
1125
1126class Phi3DecoderLayer(nn.Module):
1127 def __init__(self, config: Phi3VConfig, layer_idx: int):
1128 super().__init__()
1129
1130 self.config = config
1131 self.self_attn = PHI3_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
1132
1133 self.mlp = Phi3MLP(config)
1134 self.input_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1135
1136 self.resid_attn_dropout = nn.Dropout(config.resid_pdrop)
1137 self.resid_mlp_dropout = nn.Dropout(config.resid_pdrop)
1138 self.post_attention_layernorm = Phi3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1139
1140 def forward(
1141 self,
1142 hidden_states: torch.Tensor,
1143 attention_mask: Optional[torch.Tensor] = None,
1144 position_ids: Optional[torch.LongTensor] = None,
1145 past_key_value: Optional[Tuple[torch.Tensor]] = None,
1146 output_attentions: Optional[bool] = False,
1147 use_cache: Optional[bool] = False,
1148 **kwargs,
1149 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
1150 if "padding_mask" in kwargs:
1151 warnings.warn(
1152 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
1153 )
1154 """
1155 Args:
1156 hidden_states (`torch.FloatTensor`):
1157 input to the layer of shape `(batch, seq_len, embed_dim)`
1158 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
1159 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
1160 position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
1161 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
1162 `[0, config.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)
1163 output_attentions (`bool`, *optional*):
1164 Whether or not to return the attentions tensors of all attention layers. See `attentions` under
1165 returned tensors for more detail.
1166 use_cache (`bool`, *optional*):
1167 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
1168 (see `past_key_values`).
1169 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
1170 """
1171
1172 residual = hidden_states
1173
1174 hidden_states = self.input_layernorm(hidden_states)
1175
1176 # Self Attention
1177 attn_outputs, self_attn_weights, present_key_value = self.self_attn(
1178 hidden_states=hidden_states,
1179 attention_mask=attention_mask,
1180 position_ids=position_ids,
1181 past_key_value=past_key_value,
1182 output_attentions=output_attentions,
1183 use_cache=use_cache,
1184 )
1185
1186 hidden_states = residual + self.resid_attn_dropout(attn_outputs)
1187
1188 residual = hidden_states
1189 hidden_states = self.post_attention_layernorm(hidden_states)
1190 hidden_states = self.mlp(hidden_states)
1191 hidden_states = residual + self.resid_mlp_dropout(hidden_states)
1192
1193 outputs = (hidden_states,)
1194
1195 if output_attentions:
1196 outputs += (self_attn_weights,)
1197
1198 if use_cache:
1199 outputs += (present_key_value,)
1200