AbstractPerspective/phi-2_MoE_orange_RAG_hidden
019
1 2# Copyright (c) Microsoft Corporation.3# Licensed under the MIT license.4#5# Copyright (c) 2022, Tri Dao, trid@cs.stanford.edu.6# Licensed under the BSD 3-Clause License.7 8from __future__ import annotations9 10import math11from dataclasses import dataclass, field12from typing import Any, Dict, Optional, Tuple, Union13 14import torch15import torch.nn as nn16from einops import rearrange, repeat17from transformers import PretrainedConfig, PreTrainedModel18from transformers.activations import ACT2FN19from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast20 21from .configuration_phi import PhiConfig22 23try:24 from flash_attn.bert_padding import pad_input, unpad_input25 from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding26 from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention27 from flash_attn.ops.fused_dense import FusedDense28except:29 pad_input, unpad_input = None, None30 FlashRotaryEmbedding = None31 FlashSelfAttention, FlashCrossAttention = None, None32 FusedDense = None33 34 35@dataclass36class InferenceParams:37 #Inference parameters passed to model to efficiently calculate38 #and store context during inference.39 #Reference:40 # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.41 #Args:42 # max_seqlen: Maximum sequence length.43 # max_batch_size: Maximum batch size.44 # seqlen_offset: Sequence length offset.45 # batch_size_offset: Batch size offset.46 # key_value_memory_dict: Key value memory dictionary.47 # lengths_per_sample: Lengths per sample.48 49 max_seqlen: int = field(metadata={"help": "Maximum sequence length."})50 51 max_batch_size: int = field(metadata={"help": "Maximum batch size."})52 53 seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})54 55 batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})56 57 key_value_memory_dict: Dict[str, Any] = field(58 default_factory=dict, metadata={"help": "Key value memory dictionary."}59 )60 61 lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})62 63 64class Embedding(nn.Module):65 #Token embedding with dropout.66 67 def __init__(self, config: PretrainedConfig) -> None:68 super().__init__()69 70 self.wte = nn.Embedding(config.vocab_size, config.n_embd)71 self.drop = nn.Dropout(config.embd_pdrop)72 73 def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:74 input_shape = input_ids.size()75 input_ids = input_ids.view(-1, input_shape[-1])76 77 hidden_states = self.wte(input_ids)78 hidden_states = self.drop(hidden_states)79 80 return hidden_states81 82 83def _apply_rotary_emb(84 x: torch.FloatTensor,85 cos: torch.FloatTensor,86 sin: torch.FloatTensor,87) -> torch.FloatTensor:88 _, seqlen, _, _ = x.shape89 _, rotary_dim = cos.shape90 rotary_dim *= 291 92 x_rot = x[:, :, :, :rotary_dim]93 x_pass = x[:, :, :, rotary_dim:]94 95 x1, x2 = x_rot.chunk(2, dim=-1)96 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")97 x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]98 99 x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)100 101 return torch.cat([x_rot, x_pass], axis=-1)102 103 104def _apply_rotary_emb_kv(105 kv: torch.FloatTensor,106 cos: torch.FloatTensor,107 sin: torch.FloatTensor,108 cos_k: Optional[torch.FloatTensor] = None,109 sin_k: Optional[torch.FloatTensor] = None,110) -> torch.FloatTensor:111 _, seqlen, _, _, _ = kv.shape112 _, rotary_dim = cos.shape113 rotary_dim *= 2114 115 k_rot = kv[:, :, 0, :, :rotary_dim]116 k_pass = kv[:, :, 0, :, rotary_dim:]117 118 k1, k2 = k_rot.chunk(2, dim=-1)119 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")120 k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]121 122 k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)123 124 return torch.cat(125 [126 torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),127 kv[:, :, 1:2, :, :],128 ],129 axis=2,130 )131 132 133def _apply_rotary_emb_qkv(134 qkv: torch.FloatTensor,135 cos: torch.FloatTensor,136 sin: torch.FloatTensor,137 cos_k: Optional[torch.FloatTensor] = None,138 sin_k: Optional[torch.FloatTensor] = None,139) -> torch.FloatTensor:140 _, seqlen, _, _, _ = qkv.shape141 _, rotary_dim = cos.shape142 rotary_dim *= 2143 144 q_rot = qkv[:, :, 0, :, :rotary_dim]145 q_pass = qkv[:, :, 0, :, rotary_dim:]146 147 k_rot = qkv[:, :, 1, :, :rotary_dim]148 k_pass = qkv[:, :, 1, :, rotary_dim:]149 150 q1, q2 = q_rot.chunk(2, dim=-1)151 k1, k2 = k_rot.chunk(2, dim=-1)152 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")153 q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]154 155 q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)156 k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)157 158 return torch.cat(159 [160 torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),161 torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),162 qkv[:, :, 2:3, :, :],163 ],164 axis=2,165 )166 167 168class RotaryEmbedding(nn.Module):169 #Rotary positional embedding (RoPE).170 #Reference:171 # RoFormer: Enhanced Transformer with Rotary Position Embedding.172 # https://arxiv.org/pdf/2104.09864.pdf.173 174 def __init__(175 self,176 dim: int,177 base: int = 10000,178 scale_base: Optional[float] = None,179 pos_idx_in_fp32: bool = True,180 max_position_embeddings: int = 2048,181 device: Optional[str] = None,182 **kwargs,183 ) -> None:184 super().__init__()185 186 if scale_base is not None:187 raise NotImplementedError188 189 self.dim = dim190 self.base = float(base)191 self.scale_base = scale_base192 self.pos_idx_in_fp32 = pos_idx_in_fp32193 self.max_position_embeddings = max_position_embeddings194 self.device = device195 196 # Generate and save the inverse frequency buffer (non-trainable)197 inv_freq = self._compute_inv_freq(device)198 self.register_buffer("inv_freq", inv_freq, persistent=False)199 200 # Generate and save the scale buffer (non-trainable)201 scale = (202 (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)203 if scale_base is not None204 else None205 )206 self.register_buffer("scale", scale, persistent=False)207 208 # Initialize cached attributes since ONNX can't rely on dynamic initialization209 self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)210 211 def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:212 return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))213 214 def _update_cos_sin_cache(215 self,216 seqlen: int,217 device: Optional[str] = None,218 dtype: Optional[torch.dtype] = None,219 ) -> None:220 self._seq_len_cached = seqlen221 222 # fp32 is preferred since the output of `torch.arange` can be quite large223 # and bf16 would lose a lot of precision224 if self.pos_idx_in_fp32:225 t = torch.arange(seqlen, device=device, dtype=torch.float32)226 if self.inv_freq.dtype != torch.float32:227 inv_freq = self._compute_inv_freq(device=device)228 else:229 inv_freq = self.inv_freq230 else:231 t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)232 inv_freq = self.inv_freq233 234 # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP235 freqs = torch.outer(t, inv_freq)236 if self.scale is None:237 self._cos_cached = torch.cos(freqs).to(dtype)238 self._sin_cached = torch.sin(freqs).to(dtype)239 else:240 power = (241 torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2242 ) / self.scale_base243 scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")244 245 # Force the scale multiplication to happen in fp32246 self._cos_cached = (torch.cos(freqs) * scale).to(dtype)247 self._sin_cached = (torch.sin(freqs) * scale).to(dtype)248 self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)249 self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)250 251 def forward(252 self,253 qkv: torch.Tensor,254 kv: Optional[torch.Tensor] = None,255 seqlen_offset: int = 0,256 **kwargs,257 ) -> Tuple[torch.Tensor, torch.Tensor]:258 if (259 self._seq_len_cached < qkv.shape[1] + seqlen_offset260 or self._cos_cached.device != qkv.device261 or self._cos_cached.dtype != qkv.dtype262 or (self.training and self._cos_cached.is_inference())263 ):264 self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)265 266 if kv is None:267 return _apply_rotary_emb_qkv(268 qkv,269 self._cos_cached[seqlen_offset:],270 self._sin_cached[seqlen_offset:],271 )272 else:273 q = _apply_rotary_emb(274 qkv,275 self._cos_cached[seqlen_offset:],276 self._sin_cached[seqlen_offset:],277 )278 kv = _apply_rotary_emb_kv(279 kv,280 self._cos_cached[seqlen_offset:],281 self._sin_cached[seqlen_offset:],282 )283 284 return q, kv285 286 287class MoE(nn.Module):288 def __init__(289 self,290 config: PretrainedConfig,291 ):292 super().__init__()293 self.gate = nn.Linear(config.n_embd, config.num_local_experts, bias=False)294 self.mlp = nn.ModuleList([MLP(config) for i in range(config.num_local_experts)])295 self.num_experts_per_tok = config.num_experts_per_tok296 297 def forward(self, x):298 orig_shape = x.shape299 x = x.view(-1, x.shape[-1])300 301 scores = self.gate(x)302 expert_weights, expert_indices = torch.topk(scores, self.num_experts_per_tok, dim=-1)303 expert_weights = expert_weights.softmax(dim=-1)304 flat_expert_indices = expert_indices.view(-1)305 306 x = x.repeat_interleave(self.num_experts_per_tok, dim=0)307 y = torch.empty_like(x)308 for i, expert in enumerate(self.mlp):309 y[flat_expert_indices == i] = expert(x[flat_expert_indices == i])310 y = (y.view(*expert_weights.shape, -1) * expert_weights.unsqueeze(-1)).sum(dim=1)311 return y.view(*orig_shape)312 313 314class MLP(nn.Module):315 #Multi-Layer Perceptron.316 #Reference:317 # Attention Is All You Need.318 # https://arxiv.org/pdf/1706.03762.pdf.319 320 def __init__(321 self,322 config: PretrainedConfig,323 n_inner: Optional[int] = None,324 act_fn: Optional[str] = None,325 ) -> None:326 super().__init__()327 328 act_fn = config.activation_function if act_fn is None else act_fn329 330 n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner331 n_inner = n_inner if n_inner is not None else 4 * config.n_embd332 333 self.fc1 = nn.Linear(config.n_embd, n_inner)334 self.fc2 = nn.Linear(n_inner, config.n_embd)335 self.act = ACT2FN[act_fn]336 337 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:338 hidden_states = self.fc1(hidden_states)339 hidden_states = self.act(hidden_states)340 hidden_states = self.fc2(hidden_states)341 342 return hidden_states343 344 345class SelfAttention(nn.Module):346 #Self-attention layer (compatible with PyTorch).347 #Reference:348 # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.349 350 def __init__(351 self,352 causal: bool = True,353 softmax_scale: Optional[float] = None,354 attention_dropout: float = 0.0,355 ) -> None:356 super().__init__()357 358 self.causal = causal359 self.softmax_scale = softmax_scale360 self.drop = nn.Dropout(attention_dropout)361 362 @torch.autocast("cpu", enabled=False)363 @torch.autocast("cuda", enabled=False)364 def forward(365 self,366 qkv: torch.FloatTensor,367 causal: bool = None,368 key_padding_mask: Optional[torch.BoolTensor] = None,369 **kwargs,370 ) -> torch.FloatTensor:371 batch_size, seqlen = qkv.shape[0], qkv.shape[1]372 q, k, v = qkv.unbind(dim=2)373 374 q = q.to(torch.float32)375 k = k.to(torch.float32)376 377 causal = self.causal if causal is None else causal378 softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])379 380 # Autocast is manually disabled to avoid `torch.einsum` performing the operation381 # using float16, which might lead to overflow382 scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)383 384 if key_padding_mask is not None:385 padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)386 padding_mask.masked_fill_(key_padding_mask, 0.0)387 388 scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")389 390 if causal:391 causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)392 scores = scores + causal_mask.to(dtype=scores.dtype)393 394 attention = torch.softmax(scores, dim=-1).to(v.dtype)395 attention = self.drop(attention)396 397 output = torch.einsum("bhts,bshd->bthd", attention, v)398 399 return output400 401 402class CrossAttention(nn.Module):403 #Cross-attention layer (compatible with PyTorch).404 #Reference:405 # https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.406 407 def __init__(408 self,409 causal: bool = True,410 softmax_scale: Optional[float] = None,411 attention_dropout: float = 0.0,412 ) -> None:413 super().__init__()414 415 self.causal = causal416 self.softmax_scale = softmax_scale417 self.drop = nn.Dropout(attention_dropout)418 419 @torch.autocast("cpu", enabled=False)420 @torch.autocast("cuda", enabled=False)421 def forward(422 self,423 q: torch.FloatTensor,424 kv: torch.FloatTensor,425 causal: bool = None,426 key_padding_mask: Optional[torch.BoolTensor] = None,427 **kwargs,428 ) -> torch.FloatTensor:429 batch_size, seqlen_q = q.shape[0], q.shape[1]430 seqlen_k = kv.shape[1]431 432 if kv.shape[3] != q.shape[2]:433 kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])434 k, v = kv.unbind(dim=2)435 436 q = q.to(torch.float32)437 k = k.to(torch.float32)438 439 causal = self.causal if causal is None else causal440 softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])441 442 # Autocast is manually disabled to avoid `torch.einsum` performing the operation443 # using float16, which might lead to overflow444 scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)445 446 if key_padding_mask is not None:447 padding_mask = torch.full(448 (batch_size, seqlen_k),449 -10000.0,450 dtype=scores.dtype,451 device=scores.device,452 )453 padding_mask.masked_fill_(key_padding_mask, 0.0)454 455 scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")456 457 if causal:458 rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")459 cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)460 causal_mask = cols > rows + seqlen_k - seqlen_q461 462 scores = scores.masked_fill(causal_mask, -10000.0)463 464 attention = torch.softmax(scores, dim=-1).to(v.dtype)465 attention = self.drop(attention)466 467 output = torch.einsum("bhts,bshd->bthd", attention, v)468 469 return output470 471 472def _find_mha_dims(473 config: PretrainedConfig,474 n_head: Optional[int] = None,475 n_head_kv: Optional[int] = None,476 head_dim: Optional[int] = None,477) -> Tuple[int, int]:478 if n_head is None and head_dim is None:479 head_dim = config.n_embd // config.n_head480 n_head = config.n_head481 elif n_head is None or head_dim is None:482 raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")483 484 if n_head_kv is None:485 n_head_kv = getattr(config, "n_head_kv", None) or n_head486 487 return n_head, n_head_kv, head_dim488 489 490def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:491 num_heads, head_dim = kv.shape[-2:]492 493 if layer_idx not in inference_params.key_value_memory_dict:494 inference_params.key_value_memory_dict[layer_idx] = torch.empty(495 inference_params.max_batch_size,496 inference_params.max_seqlen,497 2,498 num_heads,499 head_dim,500 dtype=kv.dtype,501 device=kv.device,502 )503 504 batch_start = inference_params.batch_size_offset505 batch_end = batch_start + kv.shape[0]506 507 sequence_start = inference_params.seqlen_offset508 sequence_end = sequence_start + kv.shape[1]509 510 # When the current sequence length is equal to or larger than the maximum sequence length,511 # we need to concatenate the current `kv` with the cached `kv` to expand its length512 if sequence_end >= inference_params.max_seqlen:513 inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)514 515 inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv516 kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]517 518 return kv519 520 521class MHA(nn.Module):522 #Multi-head attention layer.523 524 def __init__(525 self,526 config: PretrainedConfig,527 dtype: Optional[torch.dtype] = None,528 device: Optional[str] = None,529 rotary_dim: Optional[int] = None,530 rotary_base: float = 10000.0,531 rotary_scale_base: Optional[float] = None,532 n_head: Optional[int] = None,533 n_head_kv: Optional[int] = None,534 head_dim: Optional[int] = None,535 bias: bool = True,536 causal: bool = True,537 softmax_scale: Optional[float] = None,538 layer_idx: Optional[int] = None,539 return_residual: bool = False,540 checkpointing: bool = False,541 ) -> None:542 super().__init__()543 544 # Rotary embedding545 self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)546 if self.rotary_dim > 0:547 rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding548 if rotary_cls is None:549 rotary_cls = RotaryEmbedding550 551 rotary_kwargs = {}552 if rotary_cls is RotaryEmbedding:553 rotary_kwargs["max_position_embeddings"] = config.n_positions554 555 self.rotary_emb = rotary_cls(556 self.rotary_dim,557 base=rotary_base,558 scale_base=rotary_scale_base,559 device=device,560 **rotary_kwargs,561 )562 563 # MLP564 self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(565 config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim566 )567 op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)568 hidden_size = config.n_embd569 570 linear_cls = FusedDense if config.fused_dense else nn.Linear571 if linear_cls is None:572 linear_cls = nn.Linear573 574 self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)575 self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)576 577 # Attention578 attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention579 if attn_cls is None:580 attn_cls = SelfAttention581 582 cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention583 if cross_attn_cls is None:584 cross_attn_cls = CrossAttention585 586 self.inner_attn = attn_cls(587 causal=causal,588 softmax_scale=softmax_scale,589 attention_dropout=config.attn_pdrop,590 )591 self.inner_cross_attn = cross_attn_cls(592 causal=causal,593 softmax_scale=softmax_scale,594 attention_dropout=config.attn_pdrop,595 )596 597 self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention598 self.layer_idx = layer_idx599 self.return_residual = return_residual600 self.checkpointing = checkpointing601 602 def _forward_self_attn(603 self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]604 ) -> torch.FloatTensor:605 qkv = self.Wqkv(x)606 qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)607 608 if self.rotary_dim > 0:609 qkv = self.rotary_emb(qkv)610 611 if self.flash_attn:612 batch_size, seqlen = qkv.shape[0], qkv.shape[1]613 614 cu_seqlens, max_seqlen = None, None615 if key_padding_mask is not None:616 # If `key_padding_mask` is supplied, we need to unpad the input and retrieve617 # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`618 qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)619 620 if self.checkpointing:621 attn_output = torch.utils.checkpoint.checkpoint(622 self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen623 )624 else:625 attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)626 627 # If `key_padding_mask` is supplied, we need to pad the output back to the original shape628 return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output629 630 if self.checkpointing:631 return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)632 633 return self.inner_attn(qkv, key_padding_mask=key_padding_mask)634 635 def _forward_cross_attn(636 self,637 x: torch.FloatTensor,638 past_key_values: Optional[InferenceParams],639 key_padding_mask: Optional[torch.BoolTensor],640 ) -> torch.FloatTensor:641 batch_size = x.shape[0]642 643 qkv = self.Wqkv(x)644 645 q = qkv[..., : self.n_head * self.head_dim]646 q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)647 648 kv = qkv[..., self.n_head * self.head_dim :]649 kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)650 651 seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0652 causal = None if seqlen_offset == 0 else False653 if self.rotary_dim > 0:654 q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)655 656 if past_key_values is not None:657 kv = _update_kv_cache(kv, past_key_values, self.layer_idx)658 659 if self.flash_attn:660 batch_size, seqlen_q = q.shape[0], q.shape[1]661 seqlen_k = kv.shape[1]662 663 cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (664 None,665 None,666 None,667 None,668 )669 if key_padding_mask is not None:670 kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)671 672 if seqlen_q == 1:673 key_padding_mask = torch.ones(batch_size, 1, device=q.device)674 elif seqlen_q != seqlen_k:675 key_padding_mask = key_padding_mask[:, -seqlen_q:]676 677 q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)678 679 if self.checkpointing:680 attn_output = torch.utils.checkpoint.checkpoint(681 self.inner_cross_attn,682 q,683 kv,684 causal=causal,685 cu_seqlens=cu_seqlens_q,686 max_seqlen=max_seqlen_q,687 cu_seqlens_k=cu_seqlens_k,688 max_seqlen_k=max_seqlen_k,689 )690 else:691 attn_output = self.inner_cross_attn(692 q,693 kv,694 causal=causal,695 cu_seqlens=cu_seqlens_q,696 max_seqlen=max_seqlen_q,697 cu_seqlens_k=cu_seqlens_k,698 max_seqlen_k=max_seqlen_k,699 )700 701 return (702 pad_input(attn_output, indices_q, batch_size, max_seqlen_q)703 if key_padding_mask is not None704 else attn_output705 )706 707 if self.checkpointing:708 return torch.utils.checkpoint.checkpoint(709 self.inner_cross_attn,710 q,711 kv,712 key_padding_mask=key_padding_mask,713 causal=causal,714 )715 716 return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)717 718 def forward(719 self,720 x: torch.FloatTensor,721 past_key_values: Optional[InferenceParams] = None,722 attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,723 **kwargs,724 ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:725 if attention_mask is not None:726 attention_mask = attention_mask.bool()727 else:728 attention_mask = None729 730 # MHA731 if self.n_head == self.n_head_kv:732 if past_key_values is None:733 # If `past_key_values` are not supplied, we run self-attention734 attn_output = self._forward_self_attn(x, attention_mask)735 else:736 # If `past_key_values` are supplied, it means that we might have cached values and737 # could take advantage of cross-attention738 attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)739 # MQA / GQA740 else:741 # Regardless of `past_key_values` being supplied or not, it always use cross-attention742 # because `q` and `kv` lengths might be different743 attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)744 745 output = rearrange(attn_output, "... h d -> ... (h d)")746 output = self.out_proj(output)747 748 return output if not self.return_residual else (output, x)749 750 751class ParallelBlock(nn.Module):752 #Parallel block.753 #This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).754 755 def __init__(756 self,757 config: PretrainedConfig,758 block_idx: Optional[int] = None,759 ) -> None:760 super().__init__()761 762 self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)763 self.resid_dropout = nn.Dropout(config.resid_pdrop)764 self.block_idx = block_idx765 766 self.mixer = MHA(config, layer_idx=block_idx)767 self.moe = MoE(config)768 769 def forward(770 self,771 hidden_states: torch.FloatTensor,772 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,773 attention_mask: Optional[torch.BoolTensor] = None,774 **kwargs,775 ) -> torch.FloatTensor:776 residual = hidden_states777 hidden_states = self.ln(hidden_states)778 779 attn_outputs = self.mixer(780 hidden_states,781 past_key_values=past_key_values,782 attention_mask=attention_mask,783 )784 if isinstance(attn_outputs, tuple):785 attn_outputs = attn_outputs[0]786 787 attn_outputs = self.resid_dropout(attn_outputs)788 feed_forward_hidden_states = self.resid_dropout(self.moe(hidden_states))789 790 hidden_states = attn_outputs + feed_forward_hidden_states + residual791 792 return hidden_states, attn_outputs793 794 795class CausalLMHead(nn.Module):796 #Causal Language Modeling head.797 #Reference:798 # Improving Language Understanding by Generative Pre-Training.799 # https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.800 801 def __init__(self, config: PretrainedConfig) -> None:802 super().__init__()803 804 self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)805 self.linear = nn.Linear(config.n_embd, config.vocab_size)806 807 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:808 hidden_states = self.ln(hidden_states)809 logits = self.linear(hidden_states).to(torch.float32)810 811 return logits812 813 814class CausalLMLoss(nn.Module):815 #Causal Language Modeling loss.816 #Reference:817 # Improving Language Understanding by Generative Pre-Training.818 # https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.819 820 def __init__(self, shift_labels: bool = True) -> None:821 super().__init__()822 823 self.shift_labels = shift_labels824 self.loss_fct = nn.CrossEntropyLoss()825 826 def forward(self, logits: torch.FloatTensor, labels: torch.LongTensor) -> torch.FloatTensor:827 if self.shift_labels:828 logits = logits[..., :-1, :].contiguous()829 labels = labels[..., 1:].contiguous()830 831 loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))832 833 return loss834 835 836class PhiPreTrainedModel(PreTrainedModel):837 #Phi pre-trained model.838 839 config_class = PhiConfig840 base_model_prefix = "transformer"841 supports_gradient_checkpointing = False842 _no_split_modules = ["ParallelBlock"]843 844 def __init__(self, *inputs, **kwargs) -> None:845 super().__init__(*inputs, **kwargs)846 847 def _init_weights(self, module: nn.Module) -> None:848 if isinstance(module, (nn.Linear,)):849 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)850 if module.bias is not None:851 module.bias.data.zero_()852 elif isinstance(module, nn.Embedding):853 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)854 if module.padding_idx is not None:855 module.weight.data[module.padding_idx].zero_()856 elif isinstance(module, nn.LayerNorm):857 if module.bias is not None:858 module.bias.data.zero_()859 module.weight.data.fill_(1.0)860 861 def prepare_inputs_for_generation(862 self,863 input_ids: torch.LongTensor,864 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,865 attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,866 **kwargs,867 ) -> Dict[str, Any]:868 if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):869 past_key_values = InferenceParams(870 max_seqlen=self.config.n_positions,871 max_batch_size=input_ids.shape[0],872 seqlen_offset=0,873 batch_size_offset=0,874 key_value_memory_dict={},875 lengths_per_sample=None,876 )877 else:878 # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`879 past_key_values.seqlen_offset = input_ids.shape[1] - 1880 input_ids = input_ids[:, -1].unsqueeze(-1)881 882 return {883 "input_ids": input_ids,884 "past_key_values": past_key_values,885 "attention_mask": attention_mask,886 }887 888 889class PhiModel(PhiPreTrainedModel):890 #Phi model.891 892 _keys_to_ignore_on_load_missing = [""]893 _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]894 895 def __init__(self, config: PhiConfig) -> None:896 super().__init__(config)897 898 self.embd = Embedding(config)899 self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])900 self.gradient_checkpointing = False901 self.post_init()902 903 def get_input_embeddings(self) -> nn.Embedding:904 return self.embd.wte905 906 def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:907 self.embd.wte = new_embeddings908 909 def forward(910 self,911 input_ids: torch.LongTensor,912 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,913 attention_mask: Optional[torch.BoolTensor] = None,914 ) -> torch.FloatTensor:915 hidden_states = self.embd(input_ids)916 917 all_self_attns = []918 all_hidden_states = [hidden_states]919 920 for layer in self.h:921 hidden_states, attn_outputs = layer_outputs = layer(922 hidden_states,923 past_key_values=past_key_values,924 attention_mask=attention_mask,925 )926 927 all_hidden_states.append(hidden_states)928 all_self_attns.append(attn_outputs)929 930 return BaseModelOutputWithPast(last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attns)931 932 933class PhiForCausalLM(PhiPreTrainedModel):934 #Phi for Causal Language Modeling.935 936 _keys_to_ignore_on_load_missing = [""]937 _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]938 939 def __init__(self, config: PhiConfig) -> None:940 super().__init__(config)941 942 self.transformer = PhiModel(config)943 self.lm_head = CausalLMHead(config)944 self.loss = CausalLMLoss()945 946 self.post_init()947 948 def get_output_embeddings(self) -> nn.Linear:949 return self.lm_head.linear950 951 def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:952 self.lm_head.linear = new_embeddings953 954 def forward(955 self,956 input_ids: torch.LongTensor,957 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,958 attention_mask: Optional[torch.BoolTensor] = None,959 labels: Optional[torch.LongTensor] = None,960 **kwargs,961 ) -> CausalLMOutputWithPast:962 outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask)963 lm_logits = self.lm_head(outputs.last_hidden_state)964 965 966 loss = None967 if labels is not None:968 loss = self.loss(lm_logits, labels)969 970 return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)971