numind/NuExtract-large
122152
1import math2from typing import Any, Dict, Optional, List, Tuple, Union3import torch4import torch.nn as nn5import torch.nn.functional as F6 7 8from einops import rearrange9 10from transformers.modeling_outputs import SequenceClassifierOutputWithPast, CausalLMOutputWithPast, BaseModelOutputWithPast11from transformers.modeling_utils import PreTrainedModel12from transformers.utils import logging13 14from transformers.cache_utils import Cache, DynamicCache15 16from .triton_flash_blocksparse_attn import BlockSparseParams17from .triton_blocksparse_attention_layer import BlockSparseAttentionLayer18from .positional_embedding import RotaryEmbedding19 20from .configuration_phi3_small import Phi3SmallConfig21 22# Flash Attention Related Imports23is_flash_attention_available = False24try:25 import flash_attn26 if int(flash_attn.__version__.split('.')[0]) < 2:27 from flash_attn.flash_attn_interface import (28 flash_attn_func,29 flash_attn_unpadded_kvpacked_func as flash_attn_varlen_kvpacked_func,30 )31 32 # rename `max_seqlen`33 def flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens, max_seqlen, dropout_p=0.0, **kwargs):34 return flash_attn_func(qkv, cu_seqlens, dropout_p=dropout_p, max_s=max_seqlen, **kwargs)35 36 else:37 from flash_attn.flash_attn_interface import (38 flash_attn_varlen_kvpacked_func,39 )40 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input41 is_flash_attention_available = True42except ImportError:43 pass44 45logger = logging.get_logger(__name__)46 47LegacyCache = Tuple[Tuple[torch.FloatTensor]]48 49# Taken from https://github.com/allenai/allennlp/blob/main/allennlp/nn/util.py50def info_value_of_dtype(dtype: torch.dtype):51 """52 Returns the `finfo` or `iinfo` object of a given PyTorch data type. Does not allow torch.bool.53 """54 if dtype == torch.bool:55 raise TypeError("Does not support torch.bool")56 elif dtype.is_floating_point:57 return torch.finfo(dtype)58 else:59 return torch.iinfo(dtype)60 61 62# Taken from https://github.com/allenai/allennlp/blob/main/allennlp/nn/util.py63def min_value_of_dtype(dtype: torch.dtype):64 """65 Returns the minimum value of a given PyTorch data type. Does not allow torch.bool.66 """67 return info_value_of_dtype(dtype).min68 69# Copied from transformers.models.llama.modeling_llama._get_unpad_data70def _get_unpad_data(attention_mask):71 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)72 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()73 max_seqlen_in_batch = seqlens_in_batch.max().item()74 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))75 return (76 indices,77 cu_seqlens,78 max_seqlen_in_batch,79 )80 81 82@torch.jit.script83def quick_gelu(x):84 return x * torch.sigmoid(1.702 * x)85 86 87@torch.jit.script88def gegelu(input, limit: Optional[float] = None):89 a_gelu, a_linear = input[..., ::2], input[..., 1::2]90 if limit is not None:91 a_gelu = torch.where(92 torch.isinf(a_gelu), a_gelu, a_gelu.clamp(min=None, max=limit)93 )94 a_linear = torch.where(95 torch.isinf(a_linear), a_linear, a_linear.clamp(min=-limit, max=limit)96 )97 out_gelu = quick_gelu(a_gelu)98 return out_gelu * (a_linear + 1)99 100def collapse_first_n_dims(x: torch.Tensor, n: int) -> torch.Tensor:101 """102 Collapse the first `n` dimensions of a tensor into a single dimension.103 104 Args:105 x (torch.Tensor): The input tensor.106 n (int): The number of dimensions to collapse.107 108 Returns:109 torch.Tensor: The output tensor.110 """111 return x.view(-1, *x.shape[n:])112 113def pad_tensor_to_next_mult_of(114 tensor: torch.Tensor,115 dim: int,116 n: int,117) -> Tuple[torch.Tensor, int]:118 """119 Pads a tensor along a specified dimension to the next multiple of a given number.120 121 Args:122 tensor (torch.Tensor): The input tensor.123 dim (int): The dimension along which to pad the tensor.124 n (int): The number to pad the tensor to the next multiple of.125 126 Returns:127 Tuple[torch.Tensor, int]: A tuple containing the padded tensor and the amount of padding added.128 """129 residual = tensor.size(dim) % n130 if residual == 0:131 return tensor, 0132 padding = n - residual133 padding_tensor = torch.zeros((*tensor.size()[:dim], padding, *tensor.size()[dim + 1:]), device=tensor.device, dtype=tensor.dtype)134 return torch.cat([tensor, padding_tensor], dim=dim), padding135 136def strip_padding_from_tensor(137 tensor: torch.Tensor,138 dim: int,139 residual: int,140) -> torch.Tensor:141 """142 Removes padding from a tensor along a specified dimension.143 144 Args:145 tensor (torch.Tensor): The input tensor.146 dim (int): The dimension along which to remove padding.147 residual (int): The amount of padding to remove.148 149 Returns:150 torch.Tensor: The tensor with padding removed along the specified dimension.151 """152 return torch.narrow(tensor, dim, 0, tensor.size(dim) - residual)153 154class Phi3SmallMLP(nn.Module):155 def __init__(self, config: Phi3SmallConfig):156 super().__init__()157 self.config = config158 assert self.config.hidden_act == "gegelu", "Only `gegelu` is supported for the Phi-3-small model .."159 self.hidden_size = config.hidden_size160 self.gegelu_limit = config.gegelu_limit161 self.intermediate_size = config.intermediate_size162 163 self.up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size)164 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size)165 self.dropout = nn.Dropout(config.ffn_dropout_prob)166 167 def forward(self, x: torch.Tensor) -> torch.Tensor:168 return self.dropout(169 self.down_proj(170 gegelu(self.up_proj(x), limit=self.gegelu_limit)171 )172 )173 174 175class Phi3SmallSelfAttention(nn.Module):176 def __init__(self, config: Phi3SmallConfig, layer_idx: Optional[int] = None) -> None:177 super().__init__()178 self.config = config179 self.layer_idx = layer_idx180 if layer_idx is None:181 logger.warning_once(182 f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "183 "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "184 "when creating this class."185 )186 187 self.hidden_size = config.hidden_size188 # Number of Query Heads189 self.num_heads = config.num_attention_heads190 self.head_dim = self.hidden_size // self.num_heads191 # Number of Key Value Heads192 self.num_key_value_heads = config.num_key_value_heads193 self.num_q_per_kv = self.num_heads // self.num_key_value_heads194 self.max_position_embeddings = config.max_position_embeddings195 self.rope_embedding_base = config.rope_embedding_base196 self.rope_position_scale = config.rope_position_scale197 self.is_causal = True198 199 self.attention_dropout_rate = config.attention_dropout_prob200 201 norm_factor = None202 if config.mup_use_scaling:203 norm_factor = self.head_dim / config.mup_attn_multiplier204 else:205 norm_factor = math.sqrt(self.head_dim)206 self.softmax_scale = 1.0 / norm_factor207 208 self.query_key_value = nn.Linear(self.hidden_size, (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim)209 self.dense = nn.Linear(self.hidden_size, self.hidden_size)210 211 self.blocksparse_params = None212 # layer_idx is 0 indexed because that's what the KV Cache expects.213 if self.config.dense_attention_every_n_layers and ((self.layer_idx + 1) % self.config.dense_attention_every_n_layers == 0):214 logger.info(215 f"Layer {layer_idx + 1} is using dense attention since it is divisible by "216 f"{self.config.dense_attention_every_n_layers}"217 )218 assert is_flash_attention_available, "Flash Attention is not available, but is needed for dense attention"219 else:220 # BlockSparse related Parameters221 self.blocksparse_params = BlockSparseParams.from_config(config)222 223 if self.blocksparse:224 active_head_range = None225 """226 ... note(bapatra)::227 228 In case of tensor parallelism and while using the heterogeneous head patterns,229 the active head range needs to be modified based on the tensor parallel rank230 and the tensor parallel world size.231 232 This is because in the case of heterogeneous head patterns, the kernel needs to know233 which head is on which device, so that it can pick the corresponding blocksparse head234 pattern correctly.235 236 Example:237 ```python238 239 if not self.blocksparse_params.homo_head_pattern:240 tp_rank = torch.distributed.get_rank() % tp_world_size241 num_heads_per_partition = num_heads // tp_world_size242 active_head_range = (tp_rank * num_heads_per_partition, (tp_rank + 1) * num_heads_per_partition)243 244 ```245 246 """247 248 self._blocksparse_layer = BlockSparseAttentionLayer(249 n_heads=self.num_heads,250 max_seq_len=self.max_position_embeddings,251 sparse_block_size=self.blocksparse_params.block_size,252 local_blocks=self.blocksparse_params.num_local_blocks,253 vert_stride=self.blocksparse_params.vert_stride,254 kernel_block_size=self.blocksparse_params.kernel_block_size,255 homo_head=self.blocksparse_params.homo_head_pattern,256 active_head_range=active_head_range,257 )258 self.rotary_emb = RotaryEmbedding.from_config(config)259 260 261 @property262 def blocksparse(self):263 return self.blocksparse_params is not None264 265 def _split_heads(self, mixed_x_layer: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:266 bs, sq, _ = mixed_x_layer.size()267 r"""268 The main idea is that we group tensors as269 [bs, sq, (q00, q01, ... q0m, k0, v0), (q10, q11, ... q1m, k1, v1), ... (qn0, qn1, ... qnm, kn, vn)]270 That ways, when the MP column sharding happens, this tensor will be sharded keeping all the271 queries and keys intact. In order to get the correct qkv, we first break into groups, and then272 index into the groups.273 """274 275 intermediate_shape = (bs, sq, -1, (self.num_q_per_kv + 2), self.head_dim)276 mixed_x_layer = mixed_x_layer.view(*intermediate_shape)277 q = mixed_x_layer[:, :, :, :-2]278 k = mixed_x_layer[:, :, :, [-2]]279 v = mixed_x_layer[:, :, :, [-1]]280 q, k, v = [281 rearrange(282 x,283 "bs sq group nh hn -> bs sq (group nh) hn"284 ) for x in (q, k, v)285 ]286 return q, k, v287 288 # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._unpad_input289 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):290 batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape291 292 293 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)294 295 key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)296 value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)297 298 if query_length == kv_seq_len:299 query_layer = index_first_axis(300 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k301 )302 cu_seqlens_q = cu_seqlens_k303 max_seqlen_in_batch_q = max_seqlen_in_batch_k304 indices_q = indices_k305 elif query_length == 1:306 max_seqlen_in_batch_q = 1307 cu_seqlens_q = torch.arange(308 batch_size + 1, dtype=torch.int32, device=query_layer.device309 ) # There is a memcpy here, that is very bad.310 indices_q = cu_seqlens_q[:-1]311 query_layer = query_layer.squeeze(1)312 else:313 # The -q_len: slice assumes left padding.314 attention_mask = attention_mask[:, -query_length:]315 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)316 317 return (318 query_layer,319 key_layer,320 value_layer,321 indices_q,322 (cu_seqlens_q, cu_seqlens_k),323 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),324 )325 326 def _apply_blocksparse_attention(327 self,328 q: torch.Tensor,329 k: torch.Tensor,330 v: torch.Tensor,331 attention_mask: Optional[torch.LongTensor],332 return_attention_probs: bool = False,333 ) -> torch.Tensor:334 """335 Applies blocksparse attention to the input tensors.336 337 Args:338 q (torch.Tensor): The query tensor of shape (bs, nqp, seq_len, hn).339 k (torch.Tensor): The key tensor of shape (bs, nkp, seq_len, hn).340 v (torch.Tensor): The value tensor of shape (bs, nkp, seq_len, hn).341 attention_mask (Optional[torch.LongTensor]): The attention mask tensor of shape (bs, seq_len).342 return_attention_probs (bool, optional): Whether to return attention probabilities. Defaults to False.343 344 Returns:345 torch.Tensor: The context layer tensor of shape (bs, nqp, seq_len, hn).346 """347 assert not return_attention_probs, "return_attention_probs is not supported for blocksparse attention"348 q, k, v = q.contiguous(), k.contiguous(), v.contiguous()349 # shape: (bs, nqp, seq_len, hn)350 if torch.is_grad_enabled():351 # Training or non-batched inference352 context_layer = self._blocksparse_layer(353 q=q, k=k, v=v, sm_scale=self.softmax_scale354 )355 elif attention_mask is None:356 if q.size(0) != 1:357 logger.warning_once(358 "You are attempting to do batched inference without passing the attention mask.\n"359 "This is okay if you are running loglikelihood requests. However, if you want to do generation, "360 "this probably won't work as expected. Please pass the attention mask to the forward function."361 )362 context_layer = self._blocksparse_layer(363 q=q, k=k, v=v, sm_scale=self.softmax_scale364 )365 else:366 """367 Shapes of tensors are as follows:368 q: (bs, nqp, seq_len, hdim)369 k: (bs, nkp, seq_len, hdim)370 v: (bs, nkp, seq_len, hdim)371 We first need to transpose the shapes to fit what the372 kernel needs, and the reinvert it back at the end of the operations373 """374 assert attention_mask.ndim == 2, "The kernel, like flash-attention-2, only supports 2d attention masks ..."375 left_paddings = attention_mask.shape[1] - attention_mask.sum(dim=-1)376 # shape: (bs, seq_len, nqp, hdim)377 q = q.transpose(1, 2).contiguous()378 # shape: (bs, seq_len, nkp, hdim)379 k = k.transpose(1, 2).contiguous()380 # shape: (bs, seq_len, nkp, hdim)381 v = v.transpose(1, 2).contiguous()382 context_layer = self._blocksparse_layer(383 q=q, k=k, v=v, sm_scale=self.softmax_scale, left_paddings=left_paddings.to(torch.int32)384 )385 # shape: (bs, nqp, seq_len, hdim)386 context_layer = context_layer.transpose(1, 2).contiguous()387 return context_layer388 389 def _apply_dense_attention(390 self,391 q: torch.Tensor,392 k: torch.Tensor,393 v: torch.Tensor,394 attention_mask: torch.Tensor,395 return_attention_probs: bool = False,396 ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:397 """398 Apply dense attention399 400 Args:401 q (torch.Tensor):402 The query tensor, shape: (bs, num_query_heads, seq_len, head_size)403 k (torch.Tensor):404 The key tensor, shape: (bs, num_query_heads, seq_len, head_size)405 v (torch.Tensor):406 The value tensor, shape: (bs, num_query_heads, seq_len, head_size)407 408 return_attention_probs (bool, optional):409 Return the attention probabilities. Defaults to False.410 411 Returns:412 Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:413 Return the output of the attention aggregation. If `return_attention_probs` is True, then414 also return the attention probabilities415 416 .. note::417 Right now, am assuming the expansion for the query key values is already done418 outside. But ideally, since Flash attention handles the GQA correctly, we can419 avoid doing that.420 421 """422 attention_dropout_prob = self.attention_dropout_rate if self.training else 0.0423 # Get into the correct shape for the Flash Attention API424 # shape: (bs, seq_len, nqp, hn)425 q = q.transpose(1, 2).contiguous()426 query_length = q.size(1)427 # shape: (bs, seq_len, npq, hn)428 k = k.transpose(1, 2).contiguous()429 # shape: (bs, seq_len, npq, hn)430 v = v.transpose(1, 2).contiguous()431 432 if attention_mask is not None:433 causal = q.size(2) == k.size(2)434 batch_size = q.shape[0]435 flat_q, flat_k, flat_v, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(436 q, k, v, attention_mask, query_length437 )438 cu_seqlens_q, cu_seqlens_k = cu_seq_lens439 max_seqlen_q, max_seqlen_k = max_seq_lens440 flat_kv = torch.cat((flat_k.unsqueeze(1), flat_v.unsqueeze(1)), dim=1)441 attn_output_unpad = flash_attn_varlen_kvpacked_func(442 q=flat_q,443 kv=flat_kv,444 cu_seqlens_q=cu_seqlens_q,445 cu_seqlens_k=cu_seqlens_k,446 max_seqlen_q=max_seqlen_q,447 max_seqlen_k=max_seqlen_k,448 dropout_p=attention_dropout_prob,449 softmax_scale=self.softmax_scale,450 causal=causal,451 return_attn_probs=return_attention_probs452 )453 attention_output = pad_input(454 attn_output_unpad, indices_q, batch_size, query_length455 )456 else:457 kv = torch.cat((k.unsqueeze(2), v.unsqueeze(2)), dim=2)458 cu_seqlens_q = torch.arange(459 0, (q.size(0) + 1), device=q.device, dtype=torch.int32460 ) * q.size(1)461 cu_seqlens_kv = torch.arange(462 0, (kv.size(0) + 1), device=kv.device, dtype=torch.int32463 ) * kv.size(1)464 max_seqlen_q = q.size(1)465 max_seqlen_k = kv.size(1)466 attention_output = flash_attn_varlen_kvpacked_func(467 q=collapse_first_n_dims(q, 2),468 kv=collapse_first_n_dims(kv, 2),469 cu_seqlens_q=cu_seqlens_q,470 cu_seqlens_k=cu_seqlens_kv,471 max_seqlen_q=max_seqlen_q,472 max_seqlen_k=max_seqlen_k,473 dropout_p=attention_dropout_prob,474 softmax_scale=self.softmax_scale,475 causal=q.size(1) == kv.size(1),476 return_attn_probs=return_attention_probs477 )478 if return_attention_probs:479 (context_layer, attn_probs) = attention_output480 context_layer = context_layer.view(q.size(0), q.size(1), -1, q.size(3)).transpose(1, 2).contiguous()481 return (context_layer, attn_probs)482 context_layer = attention_output483 context_layer = context_layer.view(q.size(0), q.size(1), -1, q.size(3)).transpose(1, 2).contiguous()484 return context_layer485 486 487 def expand_kv_to_q_size(self, kv: torch.Tensor, num_q_per_kv: int) -> torch.Tensor:488 """489 Expand the key-value tensor to match the size of the query tensor.490 491 Args:492 kv (torch.Tensor): The key-value tensor of shape (bsz, nkp, 2, seq_len, hdim).493 num_q_per_kv (int): The number of queries per key-value.494 495 Returns:496 torch.Tensor: The expanded key-value tensor of shape (bsz, nqp, 2, seq_len, hdim).497 Where nqp = num_q_per_kv * nkp498 499 .. note(bapatra)::500 Right now, I am using a repeat_interleave to expand the kv to the size of q.501 This incurs a memory penalty, since the tensors are actually copied.502 TODO: If this does yield benefits, then potentially we can use the re-written503 flash attention kernel that can handle GQA.504 """505 506 repeats = torch.tensor([num_q_per_kv] * kv.size(1)).to(kv.device)507 total = repeats.sum()508 expanded_kv = torch.repeat_interleave(509 kv,510 repeats=repeats,511 dim=1,512 output_size=total513 )514 return expanded_kv515 516 def forward(517 self,518 hidden_states: torch.Tensor,519 attention_mask: Optional[torch.Tensor] = None,520 position_ids: Optional[torch.LongTensor] = None,521 past_key_values: Optional[Cache] = None,522 output_attentions: bool = False,523 use_cache: bool = False,524 **kwargs,525 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:526 """527 The forward function of the Self Attention Layer.528 529 Args:530 hidden_states (torch.Tensor):531 The input tensor of shape (bs, q_len, h).532 attention_mask (Optional[torch.Tensor], optional):533 The attention mask tensor of shape (bs, seq_len). This is the 2D attention mask tensor as is standard in the flash-attention534 kernel.535 Defaults to None.536 position_ids (Optional[torch.LongTensor], optional):537 The position ids tensor of shape (bs, q_len). Defaults to None. Unused by the function.538 past_key_value (Optional[Cache], optional): 539 The previous kv cache values. Defaults to None.540 output_attentions (bool, optional): 541 Whether to return the attention scores. Defaults to False.542 .. note::543 For the blocksparse attention kernel, we do not support returning the attention scores.544 use_cache (bool, optional): 545 Whether to use the cache for storing the kv. Defaults to False.546 547 Returns:548 Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:549 The output tensor of shape (bs, q_len, h), 550 the attention scores tensor of shape (bs, nqp, q_len, seq_len) if `output_attentions` is True, 551 and the updated cache values if `use_cache` is True.552 553 Notations:554 ------------555 bs: batch size556 sq_len: sequence length of the entire sequence557 q_len: sequence length of the query558 cache_sq: sequence length in the cache559 If there is no cache then cache_sq = 0560 and sq_len = q_len561 otherwise sq_len = q_len + cache_sq562 h: hidden size563 nq: number of query heads564 nkv: number of key heads565 hn: hidden size per head566 hn = h // nq567 nqp: number of query heads (per MP partition)568 nqp = nq // (num mp partitions)569 nkvp: number of key-value heads (per MP partition)570 nkvp = nk // (num mp partitions)571 572 """573 # shape: (bs, q_len, h)574 bsz, q_len, _ = hidden_states.size()575 576 # shape: (bs, q_len, (nqp + 2 * nkvp) * hn)577 mixed_x_layer = self.query_key_value(hidden_states)578 # shape: (bs, q_len, nqp, hn), shape: (bs, q_len, nkvp, hn), shape: (bs, q_len, nkvp, hn)579 q, k, v = self._split_heads(mixed_x_layer)580 581 # shape: (bs, qnp, q_len, hn)582 query_states = q.permute(0, 2, 1, 3).contiguous()583 # shape: (bs, nkvp, q_len, hn)584 key_states = k.permute(0, 2, 1, 3).contiguous()585 # shape: (bs, nkvp, q_len, hn)586 value_states = v.permute(0, 2, 1, 3).contiguous()587 588 kv_seq_len = key_states.shape[-2]589 if past_key_values is not None:590 if self.layer_idx is None:591 raise ValueError(592 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "593 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "594 "with a layer index."595 )596 if self.rotary_emb is not None:597 seqlen_offset = past_key_values.get_usable_length(kv_seq_len, layer_idx=self.layer_idx)598 # shape: (bs, nqp, q_len, hn), shape: (bs, nkvp, q_len, hn)599 query_states, key_states = self.rotary_emb(600 query_states, key_states, seq_dimension=2, seqlen_offset=seqlen_offset601 )602 key_states, value_states = past_key_values.update(key_states=key_states, value_states=value_states, layer_idx=self.layer_idx)603 else:604 # In this case seq_len = q_len and cache_sq = 0605 if self.rotary_emb is not None:606 # shape: (bs, nqp, seq_len, hn), shape: (bs, nkvp, seq_len, hn)607 query_states, key_states = self.rotary_emb(query_states, key_states, seq_dimension=2)608 609 # shape: (bs, nkvp, 2, seq_len, hn)610 kv_states = torch.cat((key_states.unsqueeze(2), value_states.unsqueeze(2)), dim=2)611 # shape: (bs, nqp, 2, seq_len, hn)612 expanded_kv_states = self.expand_kv_to_q_size(kv_states, num_q_per_kv=self.num_q_per_kv)613 # shape: (bs, nqp, seq_len, hn), shape: (bs, nqp, seq_len, hn)614 expanded_key_states, expanded_value_states = expanded_kv_states[:, :, 0], expanded_kv_states[:, :, 1]615 if self.blocksparse:616 attn_function_output = self._apply_blocksparse_attention(617 q=query_states,618 k=expanded_key_states,619 v=expanded_value_states,620 attention_mask=attention_mask,621 return_attention_probs=output_attentions622 )623 else:624 attn_function_output = self._apply_dense_attention(625 q=query_states,626 k=expanded_key_states,627 v=expanded_value_states,628 attention_mask=attention_mask,629 return_attention_probs=output_attentions630 )631 632 attn_weights = None633 if output_attentions:634 attn_output, attn_weights = attn_function_output635 else:636 # shape: (bs, nqp, seq_len, hn)637 attn_output = attn_function_output638 # shape: (bs, seq_len, nqp, hn)639 attn_output = attn_output.transpose(1, 2).contiguous()640 641 # shape: (bs, seq_len, h)642 attn_output = attn_output.view(bsz, q_len, -1)643 attn_output = self.dense(attn_output)644 return attn_output, attn_weights, past_key_values645 646 647class Phi3SmallDecoderLayer(nn.Module):648 def __init__(self, config: Phi3SmallConfig, layer_idx: int):649 super().__init__()650 self.hidden_size = config.hidden_size651 self.self_attn = Phi3SmallSelfAttention(config, layer_idx)652 self.mlp = Phi3SmallMLP(config)653 654 self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)655 self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)656 657 def forward(658 self,659 hidden_states: torch.Tensor,660 attention_mask: Optional[torch.Tensor] = None,661 position_ids: Optional[torch.LongTensor] = None,662 past_key_values: Optional[Cache] = None,663 output_attentions: Optional[bool] = None,664 use_cache: Optional[bool] = None,665 **kwargs,666 ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Cache]]:667 residual = hidden_states668 hidden_states = self.input_layernorm(hidden_states)669 670 # Self Attention671 hidden_states, self_attn_weights, present_key_values = self.self_attn(672 hidden_states=hidden_states,673 attention_mask=attention_mask,674 position_ids=position_ids,675 past_key_values=past_key_values,676 output_attentions=output_attentions,677 use_cache=use_cache,678 )679 hidden_states = residual + hidden_states680 681 # Fully Connected682 residual = hidden_states683 hidden_states = self.post_attention_layernorm(hidden_states)684 hidden_states = self.mlp(hidden_states)685 hidden_states = residual + hidden_states686 687 outputs = (hidden_states,)688 689 if output_attentions:690 outputs += (self_attn_weights,)691 692 if use_cache:693 outputs += (present_key_values,)694 695 return outputs696 697 698 699class Phi3SmallPreTrainedModel(PreTrainedModel):700 config_class = Phi3SmallConfig701 base_model_prefix = "model"702 supports_gradient_checkpointing = True703 _no_split_modules = ["Phi3SmallDecoderLayer"]704 skip_keys_device_placement = "past_key_values"705 _supports_flash_attn_2 = True706 _supports_sdpa = False707 _supports_cache_class = True708 709 def _init_weights(self, module: nn.Module):710 std = self.config.initializer_range711 if isinstance(module, nn.Linear):712 # Slightly different from the TF version which uses truncated_normal for initialization713 # cf https://github.com/pytorch/pytorch/pull/5617714 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)715 elif isinstance(module, nn.Embedding):716 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)717 if module.padding_idx is not None:718 module.weight.data[module.padding_idx].zero_()719 elif isinstance(module, nn.LayerNorm):720 module.bias.data.zero_()721 module.weight.data.fill_(1.0)722 723 # The output projection on the decoder attention layer as well as the down_proj in the MLP are scaled724 # differently (dubbed `output_layer_init_method` in the Megatron code). This is replicated here725 for name, p in module.named_parameters():726 if any(x in name for x in ("c_proj.weight", "down_proj.weight", "o_proj.weight")):727 # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block728 p.data.normal_(mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)))729 730 731class Phi3SmallModel(Phi3SmallPreTrainedModel):732 733 def __init__(self, config):734 super().__init__(config)735 self.config = config736 737 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)738 739 # Embedding Dropout740 self.embedding_dropout = nn.Dropout(config.embedding_dropout_prob)741 742 # MuP Embedding scaling743 self.mup_embedding_multiplier = config.mup_embedding_multiplier744 745 self.layers = nn.ModuleList([Phi3SmallDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])746 747 self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)748 749 self.gradient_checkpointing = False750 751 # Initialize weights and apply final processing752 self.post_init()753 754 def get_input_embeddings(self):755 return self.embed_tokens756 757 def set_input_embeddings(self, value):758 self.embed_tokens = value759 760 @property761 def pad_sequence_to_multiple_of_64(self):762 # We only need to do this for the backward pass. So only required763 # when we are in the context of generating gradients764 return self.config.pad_sequence_to_multiple_of_64 and torch.is_grad_enabled()765 766 def forward(767 self,768 input_ids: torch.LongTensor = None,769 attention_mask: Optional[torch.Tensor] = None,770 position_ids: Optional[torch.LongTensor] = None,771 past_key_values: Optional[Union[Cache, LegacyCache]] = None,772 inputs_embeds: Optional[torch.FloatTensor] = None,773 use_cache: Optional[bool] = None,774 output_attentions: Optional[bool] = None,775 output_hidden_states: Optional[bool] = None,776 return_dict: Optional[bool] = None,777 ) -> Union[Tuple, BaseModelOutputWithPast]:778 779 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions780 output_hidden_states = (781 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states782 )783 use_cache = use_cache if use_cache is not None else self.config.use_cache784 785 return_dict = return_dict if return_dict is not None else self.config.use_return_dict786 787 if input_ids is not None and inputs_embeds is not None:788 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")789 elif input_ids is not None:790 batch_size, seq_length = input_ids.shape791 elif inputs_embeds is not None:792 batch_size, seq_length, _ = inputs_embeds.shape793 else:794 raise ValueError("You have to specify either input_ids or inputs_embeds")795 796 if self.gradient_checkpointing and self.training:797 if use_cache:798 logger.warning_once(799 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."800 )801 use_cache = False802 803 past_key_values_length = 0804 805 if use_cache:806 use_legacy_cache = not isinstance(past_key_values, Cache)807 if use_legacy_cache:808 past_key_values = DynamicCache.from_legacy_cache(past_key_values)809 past_key_values_length = past_key_values.get_usable_length(seq_length)810 811 if position_ids is None:812 device = input_ids.device if input_ids is not None else inputs_embeds.device813 position_ids = torch.arange(814 past_key_values_length, past_key_values_length + seq_length, dtype=torch.long, device=device815 )816 position_ids = position_ids.unsqueeze(0).view(-1, seq_length)817 else:818 position_ids = position_ids.view(-1, seq_length).long()819 820 if attention_mask is not None:821 if batch_size <= 0:822 raise ValueError("batch_size has to be defined and > 0")823 824 if inputs_embeds is None:825 inputs_embeds = self.embed_tokens(input_ids)826 inputs_embeds = self.embedding_dropout(inputs_embeds)827 828 if self.mup_embedding_multiplier is not None and self.mup_embedding_multiplier > 0.0:829 inputs_embeds = inputs_embeds * self.mup_embedding_multiplier830 831 residual = 0832 if self.pad_sequence_to_multiple_of_64:833 # note(bapatra): Since we don't particularly use the position_ids and the attention mask834 # we don't need to pad them835 inputs_embeds, residual = pad_tensor_to_next_mult_of(tensor=inputs_embeds, dim=1, n=64)836 837 hidden_states = inputs_embeds838 839 # decoder layers840 all_hidden_states = () if output_hidden_states else None841 all_self_attns = () if output_attentions else None842 next_decoder_cache = None843 844 for decoder_layer in self.layers:845 if output_hidden_states:846 all_hidden_states += (hidden_states,)847 848 if self.gradient_checkpointing and self.training:849 layer_outputs = self._gradient_checkpointing_func(850 decoder_layer.__call__,851 hidden_states,852 attention_mask,853 position_ids,854 past_key_values,855 output_attentions,856 use_cache,857 )858 else:859 layer_outputs = decoder_layer(860 hidden_states,861 attention_mask=attention_mask,862 position_ids=position_ids,863 past_key_values=past_key_values,864 output_attentions=output_attentions,865 use_cache=use_cache,866 )867 hidden_states = layer_outputs[0]868 869 if use_cache:870 # Following the Mistral schema for layer return values871 next_decoder_cache = layer_outputs[2 if output_attentions else 1]872 if output_attentions:873 all_self_attns += (layer_outputs[1],)874 875 hidden_states = self.final_layernorm(hidden_states)876 877 if residual > 0:878 hidden_states = strip_padding_from_tensor(tensor=hidden_states, dim=1, residual=residual)879 880 # add hidden states from the last decoder layer881 if output_hidden_states:882 all_hidden_states += (hidden_states,)883 884 next_cache = None885 if use_cache:886 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache887 888 if not return_dict:889 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)890 return BaseModelOutputWithPast(891 last_hidden_state=hidden_states,892 past_key_values=next_cache,893 hidden_states=all_hidden_states,894 attentions=all_self_attns,895 )896 897 898class Phi3SmallForCausalLM(Phi3SmallPreTrainedModel):899 _tied_weights_keys = ["lm_head.weight"]900 901 def __init__(self, config):902 super().__init__(config)903 self.model = Phi3SmallModel(config)904 self.vocab_size = config.vocab_size905 self.lm_head = nn.Linear(config.hidden_size, self.vocab_size, bias=False)906 self.mup_width_multiplier = config.mup_width_multiplier907 908 # Create the mask for the dummy tokens in the vocabulary909 dummy_token_indices = config.dummy_token_indices910 dummy_tokens_mask = torch.zeros(self.vocab_size).bool()911 dummy_tokens_mask[dummy_token_indices] = True912 # shape: (vocab_size,)913 self.register_buffer("dummy_tokens_mask", dummy_tokens_mask, persistent=False)914 915 # Initialize weights and apply final processing916 self.post_init()917 918 def get_input_embeddings(self):919 return self.model.embed_tokens920 921 def set_input_embeddings(self, value):922 self.model.embed_tokens = value923 924 def get_output_embeddings(self):925 return self.lm_head926 927 def set_output_embeddings(self, value):928 self.lm_head = value929 930 def set_decoder(self, decoder):931 self.model = decoder932 933 def get_decoder(self):934 return self.model935 936 def forward(937 self,938 input_ids: torch.LongTensor = None,939 attention_mask: Optional[torch.Tensor] = None,940 position_ids: Optional[torch.LongTensor] = None,941 past_key_values: Optional[List[torch.FloatTensor]] = None,942 inputs_embeds: Optional[torch.FloatTensor] = None,943 labels: Optional[torch.LongTensor] = None,944 use_cache: Optional[bool] = None,945 output_attentions: Optional[bool] = None,946 output_hidden_states: Optional[bool] = None,947 return_dict: Optional[bool] = None, 948 ) -> Union[Tuple, CausalLMOutputWithPast]:949 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions950 output_hidden_states = (951 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states952 )953 return_dict = return_dict if return_dict is not None else self.config.use_return_dict954 955 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)956 outputs = self.model(957 input_ids=input_ids,958 attention_mask=attention_mask,959 position_ids=position_ids,960 past_key_values=past_key_values,961 inputs_embeds=inputs_embeds,962 use_cache=use_cache,963 output_attentions=output_attentions,964 output_hidden_states=output_hidden_states,965 return_dict=return_dict,966 )967 968 hidden_states = outputs[0]969 logits = self.lm_head(hidden_states)970 logits = logits.float()971 if self.mup_width_multiplier:972 logits = logits / self.mup_width_multiplier973 logits = logits.masked_fill(self.dummy_tokens_mask, min_value_of_dtype(logits.dtype))974 975 loss = None976 if labels is not None:977 # Shift so that tokens < n predict n978 shift_logits = logits[..., :-1, :].contiguous()979 shift_labels = labels[..., 1:].contiguous()980 # Flatten the tokens981 loss_fct = nn.CrossEntropyLoss()982 shift_logits = shift_logits.view(-1, self.config.vocab_size)983 shift_labels = shift_labels.view(-1)984 # Enable model parallelism985 shift_labels = shift_labels.to(shift_logits.device)986 loss = loss_fct(shift_logits, shift_labels)987 988 if not return_dict:989 output = (logits,) + outputs[1:]990 return (loss,) + output if loss is not None else output991 992 return CausalLMOutputWithPast(993 loss=loss,994 logits=logits,995 past_key_values=outputs.past_key_values,996 hidden_states=outputs.hidden_states,997 attentions=outputs.attentions,998 )999 1000 def prepare_inputs_for_generation(1001 self, 1002 input_ids: torch.LongTensor,1003 past_key_values: Optional[List[torch.FloatTensor]] = None,1004 attention_mask: Optional[torch.FloatTensor] = None,1005 inputs_embeds: Optional[torch.FloatTensor] = None,1006 **kwargs1007 ) -> Dict[str, Any]:1008 # only last token for inputs_ids if past is defined in kwargs1009 if past_key_values:1010 input_ids = input_ids[:, -1].unsqueeze(-1)1011 1012 position_ids = kwargs.get("position_ids", None)1013 1014 if attention_mask is not None and position_ids is None:1015 # create position_ids on the fly for batch generation1016 position_ids = attention_mask.long().cumsum(-1) - 11017 position_ids.masked_fill_(attention_mask == 0, 1)1018 if past_key_values:1019 position_ids = position_ids[:, -1].unsqueeze(-1)1020 else:1021 position_ids = None1022 1023 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step1024 if inputs_embeds is not None and past_key_values is None:1025 model_inputs = {"inputs_embeds": inputs_embeds}1026 else:1027 model_inputs = {"input_ids": input_ids}1028 1029 model_inputs.update(1030 {1031 "past_key_values": past_key_values,1032 "use_cache": kwargs.get("use_cache"),1033 "position_ids": position_ids,1034 "attention_mask": attention_mask,1035 }1036 )1037 return model_inputs1038 1039 1040# Copied from transformers.models.mistral.modeling_mistral.MistralForSequenceClassification with Mistral -> Phi3Small1041class Phi3SmallForSequenceClassification(Phi3SmallPreTrainedModel):1042 def __init__(self, config):1043 super().__init__(config)1044 self.num_labels = config.num_labels1045 self.model = Phi3SmallModel(config)1046 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)1047 1048 # Initialize weights and apply final processing1049 self.post_init()1050 1051 def get_input_embeddings(self):1052 return self.model.embed_tokens1053 1054 def set_input_embeddings(self, value):1055 self.model.embed_tokens = value1056 1057 1058 def forward(1059 self,1060 input_ids: torch.LongTensor = None,1061 attention_mask: Optional[torch.Tensor] = None,1062 position_ids: Optional[torch.LongTensor] = None,1063 past_key_values: Optional[List[torch.FloatTensor]] = None,1064 inputs_embeds: Optional[torch.FloatTensor] = None,1065 labels: Optional[torch.LongTensor] = None,1066 use_cache: Optional[bool] = None,1067 output_attentions: Optional[bool] = None,1068 output_hidden_states: Optional[bool] = None,1069 return_dict: Optional[bool] = None,1070 ) -> Union[Tuple, SequenceClassifierOutputWithPast]:1071 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1072 1073 transformer_outputs = self.model(1074 input_ids,1075 attention_mask=attention_mask,1076 position_ids=position_ids,1077 past_key_values=past_key_values,1078 inputs_embeds=inputs_embeds,1079 use_cache=use_cache,1080 output_attentions=output_attentions,1081 output_hidden_states=output_hidden_states,1082 return_dict=return_dict,1083 )1084 hidden_states = transformer_outputs[0]1085 logits = self.score(hidden_states)1086 1087 if input_ids is not None:1088 batch_size = input_ids.shape[0]1089 else:1090 batch_size = inputs_embeds.shape[0]1091 1092 if self.config.pad_token_id is None and batch_size != 1:1093 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")1094 if self.config.pad_token_id is None:1095 sequence_lengths = -11096 else:1097 if input_ids is not None:1098 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility1099 sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 11100 sequence_lengths = sequence_lengths % input_ids.shape[-1]1101 sequence_lengths = sequence_lengths.to(logits.device)1102 else:1103 sequence_lengths = -11104 1105 pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]1106 1107 loss = None1108 if labels is not None:1109 labels = labels.to(logits.device)1110 if self.config.problem_type is None:1111 if self.num_labels == 1:1112 self.config.problem_type = "regression"1113 elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):1114 self.config.problem_type = "single_label_classification"1115 else:1116 self.config.problem_type = "multi_label_classification"1117 1118 if self.config.problem_type == "regression":1119 loss_fct = nn.MSELoss()1120 if self.num_labels == 1:1121 loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())1122 else:1123 loss = loss_fct(pooled_logits, labels)1124 elif self.config.problem_type == "single_label_classification":1125 loss_fct = nn.CrossEntropyLoss()1126 loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))1127 elif self.config.problem_type == "multi_label_classification":1128 loss_fct = nn.BCEWithLogitsLoss()1129 loss = loss_fct(pooled_logits, labels)1130 if not return_dict:1131 output = (pooled_logits,) + transformer_outputs[1:]1132 return ((loss,) + output) if loss is not None else output1133 1134 return SequenceClassifierOutputWithPast(1135 loss=loss,1136 logits=pooled_logits,1137 past_key_values=transformer_outputs.past_key_values,1138 hidden_states=transformer_outputs.hidden_states,1139 attentions=transformer_outputs.attentions,1140 )1141