ManishThota/CustomModel
1158
1# Copyright (c) MILVLG team.2# Licensed under the Apache 2.0 license.3#4# Some code here is copied from the project Phi-2 (https://huggingface.co/microsoft/phi-2),5# SigLIP@transformers==4.37.0.dev0 (https://huggingface.co/google/siglip-so400m-patch14-384),6# and Llava (https://github.com/haotian-liu/LLaVA), and modified by 7# Zhenwei Shao (shaozw@hdu.edu.cn) @ MILVLG. We thank them for their great works.8# And their original licenses and copyright should be inherited (see the statements9# in `configuration_imp.py` for more details).10 11 12# Be careful: The way how `past_key_values.seqlen_offset` is updated is modified from13# the implementation of original Phi-2. See the comments below for details.14 15from __future__ import annotations16import os17import math18import re19from dataclasses import dataclass, field20from typing import Any, Dict, Optional, Tuple, Union, List21from abc import ABC, abstractmethod22 23import torch24import torch.nn as nn25from einops import rearrange, repeat26from transformers import (27 PretrainedConfig, 28 PreTrainedModel,29 AutoConfig,30 AutoModelForCausalLM31) 32from transformers.activations import ACT2FN33from transformers.modeling_outputs import CausalLMOutputWithPast34import sys35from .configuration_imp import PhiConfig, ImpConfig36from .vision_encoder import VisionTower37 38try:39 from flash_attn.bert_padding import pad_input, unpad_input40 from flash_attn.layers.rotary import RotaryEmbedding as FlashRotaryEmbedding41 from flash_attn.modules.mha import FlashCrossAttention, FlashSelfAttention42 from flash_attn.ops.fused_dense import FusedDense43except:44 pad_input, unpad_input = None, None45 FlashRotaryEmbedding = None46 FlashSelfAttention, FlashCrossAttention = None, None47 FusedDense = None48 49 50@dataclass51class InferenceParams:52 """Inference parameters passed to model to efficiently calculate53 and store context during inference.54 55 Reference:56 https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/utils/generation.py.57 58 Args:59 max_seqlen: Maximum sequence length.60 max_batch_size: Maximum batch size.61 seqlen_offset: Sequence length offset.62 batch_size_offset: Batch size offset.63 key_value_memory_dict: Key value memory dictionary.64 lengths_per_sample: Lengths per sample.65 66 """67 68 max_seqlen: int = field(metadata={"help": "Maximum sequence length."})69 70 max_batch_size: int = field(metadata={"help": "Maximum batch size."})71 72 seqlen_offset: int = field(default=0, metadata={"help": "Sequence length offset."})73 74 batch_size_offset: int = field(default=0, metadata={"help": "Batch size offset."})75 76 key_value_memory_dict: Dict[str, Any] = field(77 default_factory=dict, metadata={"help": "Key value memory dictionary."}78 )79 80 lengths_per_sample: torch.Tensor = field(default=None, metadata={"help": "Lengths per sample."})81 82 83class Embedding(nn.Module):84 """Token embedding with dropout."""85 86 def __init__(self, config: PretrainedConfig) -> None:87 super().__init__()88 89 self.wte = nn.Embedding(config.vocab_size, config.n_embd)90 self.drop = nn.Dropout(config.embd_pdrop)91 92 def forward(self, input_ids: torch.LongTensor) -> torch.FloatTensor:93 input_shape = input_ids.size()94 input_ids = input_ids.view(-1, input_shape[-1])95 96 hidden_states = self.wte(input_ids)97 hidden_states = self.drop(hidden_states)98 99 return hidden_states100 101 102 103def _apply_rotary_emb(104 x: torch.FloatTensor,105 cos: torch.FloatTensor,106 sin: torch.FloatTensor,107) -> torch.FloatTensor:108 _, seqlen, _, _ = x.shape109 _, rotary_dim = cos.shape110 rotary_dim *= 2111 112 x_rot = x[:, :, :, :rotary_dim]113 x_pass = x[:, :, :, rotary_dim:]114 115 x1, x2 = x_rot.chunk(2, dim=-1)116 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")117 x1, x2, c, s = [t.to(dtype=torch.float32) for t in [x1, x2, c, s]]118 119 x_rot = torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1).to(x.dtype)120 121 return torch.cat([x_rot, x_pass], axis=-1)122 123 124def _apply_rotary_emb_kv(125 kv: torch.FloatTensor,126 cos: torch.FloatTensor,127 sin: torch.FloatTensor,128 cos_k: Optional[torch.FloatTensor] = None,129 sin_k: Optional[torch.FloatTensor] = None,130) -> torch.FloatTensor:131 _, seqlen, _, _, _ = kv.shape132 _, rotary_dim = cos.shape133 rotary_dim *= 2134 135 k_rot = kv[:, :, 0, :, :rotary_dim]136 k_pass = kv[:, :, 0, :, rotary_dim:]137 138 k1, k2 = k_rot.chunk(2, dim=-1)139 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")140 k1, k2, c, s = [t.to(dtype=torch.float32) for t in [k1, k2, c, s]]141 142 k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(kv.dtype)143 144 return torch.cat(145 [146 torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),147 kv[:, :, 1:2, :, :],148 ],149 axis=2,150 )151 152 153def _apply_rotary_emb_qkv(154 qkv: torch.FloatTensor,155 cos: torch.FloatTensor,156 sin: torch.FloatTensor,157 cos_k: Optional[torch.FloatTensor] = None,158 sin_k: Optional[torch.FloatTensor] = None,159) -> torch.FloatTensor:160 _, seqlen, _, _, _ = qkv.shape161 _, rotary_dim = cos.shape162 rotary_dim *= 2163 164 q_rot = qkv[:, :, 0, :, :rotary_dim]165 q_pass = qkv[:, :, 0, :, rotary_dim:]166 167 k_rot = qkv[:, :, 1, :, :rotary_dim]168 k_pass = qkv[:, :, 1, :, rotary_dim:]169 170 q1, q2 = q_rot.chunk(2, dim=-1)171 k1, k2 = k_rot.chunk(2, dim=-1)172 c, s = rearrange(cos[:seqlen], "s d -> s 1 d"), rearrange(sin[:seqlen], "s d -> s 1 d")173 q1, q2, k1, k2, c, s = [t.to(dtype=torch.float32) for t in [q1, q2, k1, k2, c, s]]174 175 q_rot = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], axis=-1).to(qkv.dtype)176 k_rot = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], axis=-1).to(qkv.dtype)177 178 return torch.cat(179 [180 torch.cat([q_rot, q_pass], axis=-1).unsqueeze(2),181 torch.cat([k_rot, k_pass], axis=-1).unsqueeze(2),182 qkv[:, :, 2:3, :, :],183 ],184 axis=2,185 )186 187 188class RotaryEmbedding(nn.Module):189 """Rotary positional embedding (RoPE).190 191 Reference:192 RoFormer: Enhanced Transformer with Rotary Position Embedding.193 https://arxiv.org/pdf/2104.09864.pdf.194 195 """196 197 def __init__(198 self,199 dim: int,200 base: int = 10000,201 scale_base: Optional[float] = None,202 pos_idx_in_fp32: bool = True,203 max_position_embeddings: int = 2048,204 device: Optional[str] = None,205 **kwargs,206 ) -> None:207 super().__init__()208 209 if scale_base is not None:210 raise NotImplementedError211 212 self.dim = dim213 self.base = float(base)214 self.scale_base = scale_base215 self.pos_idx_in_fp32 = pos_idx_in_fp32216 self.max_position_embeddings = max_position_embeddings217 self.device = device218 219 # Generate and save the inverse frequency buffer (non-trainable)220 inv_freq = self._compute_inv_freq(device)221 self.register_buffer("inv_freq", inv_freq, persistent=False)222 223 # Generate and save the scale buffer (non-trainable)224 scale = (225 (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)226 if scale_base is not None227 else None228 )229 self.register_buffer("scale", scale, persistent=False)230 231 # Initialize cached attributes since ONNX can't rely on dynamic initialization232 self._update_cos_sin_cache(max_position_embeddings, device=device, dtype=torch.float32)233 234 def _compute_inv_freq(self, device: Optional[str] = None) -> torch.FloatTensor:235 return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))236 237 def _update_cos_sin_cache(238 self,239 seqlen: int,240 device: Optional[str] = None,241 dtype: Optional[torch.dtype] = None,242 ) -> None:243 self._seq_len_cached = seqlen244 245 # fp32 is preferred since the output of `torch.arange` can be quite large246 # and bf16 would lose a lot of precision247 if self.pos_idx_in_fp32:248 t = torch.arange(seqlen, device=device, dtype=torch.float32)249 if self.inv_freq.dtype != torch.float32:250 inv_freq = self._compute_inv_freq(device=device)251 else:252 inv_freq = self.inv_freq253 else:254 t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)255 inv_freq = self.inv_freq256 257 # `torch.outer` is preferred since `torch.einsum` converts from fp32 to fp16 if used with AMP258 freqs = torch.outer(t, inv_freq)259 if self.scale is None:260 self._cos_cached = torch.cos(freqs).to(dtype)261 self._sin_cached = torch.sin(freqs).to(dtype)262 else:263 power = (264 torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2265 ) / self.scale_base266 scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")267 268 # Force the scale multiplication to happen in fp32269 self._cos_cached = (torch.cos(freqs) * scale).to(dtype)270 self._sin_cached = (torch.sin(freqs) * scale).to(dtype)271 self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)272 self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)273 274 def forward(275 self,276 qkv: torch.Tensor,277 kv: Optional[torch.Tensor] = None,278 seqlen_offset: int = 0,279 **kwargs,280 ) -> Tuple[torch.Tensor, torch.Tensor]:281 if (282 self._seq_len_cached < qkv.shape[1] + seqlen_offset283 or self._cos_cached.device != qkv.device284 or self._cos_cached.dtype != qkv.dtype285 or (self.training and self._cos_cached.is_inference())286 ):287 self._update_cos_sin_cache(qkv.shape[1] + seqlen_offset, device=qkv.device, dtype=qkv.dtype)288 289 if kv is None:290 return _apply_rotary_emb_qkv(291 qkv,292 self._cos_cached[seqlen_offset:],293 self._sin_cached[seqlen_offset:],294 )295 else:296 q = _apply_rotary_emb(297 qkv,298 self._cos_cached[seqlen_offset:],299 self._sin_cached[seqlen_offset:],300 )301 kv = _apply_rotary_emb_kv(302 kv,303 self._cos_cached[seqlen_offset:],304 self._sin_cached[seqlen_offset:],305 )306 307 return q, kv308 309 310class MLP(nn.Module):311 """Multi-Layer Perceptron.312 313 Reference:314 Attention Is All You Need.315 https://arxiv.org/pdf/1706.03762.pdf.316 317 """318 319 def __init__(320 self,321 config: PretrainedConfig,322 n_inner: Optional[int] = None,323 act_fn: Optional[str] = None,324 ) -> None:325 super().__init__()326 327 act_fn = config.activation_function if act_fn is None else act_fn328 329 n_inner = getattr(config, "n_inner", None) if n_inner is None else n_inner330 n_inner = n_inner if n_inner is not None else 4 * config.n_embd331 332 self.fc1 = nn.Linear(config.n_embd, n_inner)333 self.fc2 = nn.Linear(n_inner, config.n_embd)334 self.act = ACT2FN[act_fn]335 336 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:337 hidden_states = self.fc1(hidden_states)338 hidden_states = self.act(hidden_states)339 hidden_states = self.fc2(hidden_states)340 341 return hidden_states342 343 344class SelfAttention(nn.Module):345 """Self-attention layer (compatible with PyTorch).346 347 Reference:348 https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.349 350 """351 352 def __init__(353 self,354 causal: bool = True,355 softmax_scale: Optional[float] = None,356 attention_dropout: float = 0.0,357 ) -> None:358 super().__init__()359 360 self.causal = causal361 self.softmax_scale = softmax_scale362 self.drop = nn.Dropout(attention_dropout)363 364 @torch.autocast("cpu", enabled=False)365 @torch.autocast("cuda", enabled=False)366 def forward(367 self,368 qkv: torch.FloatTensor,369 causal: bool = None,370 key_padding_mask: Optional[torch.BoolTensor] = None,371 **kwargs,372 ) -> torch.FloatTensor:373 batch_size, seqlen = qkv.shape[0], qkv.shape[1]374 q, k, v = qkv.unbind(dim=2)375 376 q = q.to(torch.float32)377 k = k.to(torch.float32)378 379 causal = self.causal if causal is None else causal380 softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])381 382 # Autocast is manually disabled to avoid `torch.einsum` performing the operation383 # using float16, which might lead to overflow384 scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)385 386 if key_padding_mask is not None:387 padding_mask = torch.full((batch_size, seqlen), -10000.0, dtype=scores.dtype, device=scores.device)388 padding_mask.masked_fill_(key_padding_mask, 0.0)389 390 scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")391 392 if causal:393 causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1)394 scores = scores + causal_mask.to(dtype=scores.dtype)395 396 attention = torch.softmax(scores, dim=-1).to(v.dtype)397 attention = self.drop(attention)398 399 output = torch.einsum("bhts,bshd->bthd", attention, v)400 401 return output402 403 404class CrossAttention(nn.Module):405 """Cross-attention layer (compatible with PyTorch).406 407 Reference:408 https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py.409 410 """411 412 def __init__(413 self,414 causal: bool = True,415 softmax_scale: Optional[float] = None,416 attention_dropout: float = 0.0,417 ) -> None:418 super().__init__()419 420 self.causal = causal421 self.softmax_scale = softmax_scale422 self.drop = nn.Dropout(attention_dropout)423 424 @torch.autocast("cpu", enabled=False)425 @torch.autocast("cuda", enabled=False)426 def forward(427 self,428 q: torch.FloatTensor,429 kv: torch.FloatTensor,430 causal: bool = None,431 key_padding_mask: Optional[torch.BoolTensor] = None,432 **kwargs,433 ) -> torch.FloatTensor:434 batch_size, seqlen_q = q.shape[0], q.shape[1]435 seqlen_k = kv.shape[1]436 437 if kv.shape[3] != q.shape[2]:438 kv = repeat(kv, "... hkv d -> ... (hkv g) d", g=q.shape[2] // kv.shape[3])439 k, v = kv.unbind(dim=2)440 441 q = q.to(torch.float32)442 k = k.to(torch.float32)443 444 causal = self.causal if causal is None else causal445 softmax_scale = self.softmax_scale or 1.0 / math.sqrt(q.shape[-1])446 447 # Autocast is manually disabled to avoid `torch.einsum` performing the operation448 # using float16, which might lead to overflow449 scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)450 451 if key_padding_mask is not None:452 padding_mask = torch.full(453 (batch_size, seqlen_k),454 -10000.0,455 dtype=scores.dtype,456 device=scores.device,457 )458 padding_mask.masked_fill_(key_padding_mask, 0.0)459 460 scores = scores + rearrange(padding_mask, "b s -> b 1 1 s")461 462 if causal:463 rows = rearrange(torch.arange(seqlen_q, device=q.device, dtype=torch.long), "s -> s 1")464 cols = torch.arange(seqlen_k, device=k.device, dtype=torch.long)465 causal_mask = cols > rows + seqlen_k - seqlen_q466 467 scores = scores.masked_fill(causal_mask, -10000.0)468 469 attention = torch.softmax(scores, dim=-1).to(v.dtype)470 attention = self.drop(attention)471 472 output = torch.einsum("bhts,bshd->bthd", attention, v)473 474 return output475 476 477def _find_mha_dims(478 config: PretrainedConfig,479 n_head: Optional[int] = None,480 n_head_kv: Optional[int] = None,481 head_dim: Optional[int] = None,482) -> Tuple[int, int]:483 if n_head is None and head_dim is None:484 head_dim = config.n_embd // config.n_head485 n_head = config.n_head486 elif n_head is None or head_dim is None:487 raise ValueError("`n_head` and `head_dim` must be both specified or `None`.")488 489 if n_head_kv is None:490 n_head_kv = getattr(config, "n_head_kv", None) or n_head491 492 return n_head, n_head_kv, head_dim493 494 495def _update_kv_cache(kv: torch.FloatTensor, inference_params: InferenceParams, layer_idx: int) -> torch.FloatTensor:496 num_heads, head_dim = kv.shape[-2:]497 498 if layer_idx not in inference_params.key_value_memory_dict:499 inference_params.key_value_memory_dict[layer_idx] = torch.empty(500 inference_params.max_batch_size,501 inference_params.max_seqlen,502 2,503 num_heads,504 head_dim,505 dtype=kv.dtype,506 device=kv.device,507 )508 509 batch_start = inference_params.batch_size_offset510 batch_end = batch_start + kv.shape[0]511 512 sequence_start = inference_params.seqlen_offset513 sequence_end = sequence_start + kv.shape[1]514 515 # When the current sequence length is equal to or larger than the maximum sequence length,516 # we need to concatenate the current `kv` with the cached `kv` to expand its length517 if sequence_end >= inference_params.max_seqlen:518 inference_params.key_value_memory_dict[layer_idx] = torch.concatenate((inference_params.key_value_memory_dict[layer_idx], kv), dim=1)519 520 inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...] = kv521 kv = inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, :sequence_end, ...]522 523 return kv524 525 526class MHA(nn.Module):527 """Multi-head attention layer."""528 529 def __init__(530 self,531 config: PretrainedConfig,532 dtype: Optional[torch.dtype] = None,533 device: Optional[str] = None,534 rotary_dim: Optional[int] = None,535 rotary_base: float = 10000.0,536 rotary_scale_base: Optional[float] = None,537 n_head: Optional[int] = None,538 n_head_kv: Optional[int] = None,539 head_dim: Optional[int] = None,540 bias: bool = True,541 causal: bool = True,542 softmax_scale: Optional[float] = None,543 layer_idx: Optional[int] = None,544 return_residual: bool = False,545 checkpointing: bool = False,546 ) -> None:547 super().__init__()548 549 # Rotary embedding550 self.rotary_dim = rotary_dim if rotary_dim is not None else getattr(config, "rotary_dim", 0)551 if self.rotary_dim > 0:552 rotary_cls = FlashRotaryEmbedding if config.flash_rotary else RotaryEmbedding553 if rotary_cls is None:554 rotary_cls = RotaryEmbedding555 556 rotary_kwargs = {}557 if rotary_cls is RotaryEmbedding:558 rotary_kwargs["max_position_embeddings"] = config.n_positions559 560 self.rotary_emb = rotary_cls(561 self.rotary_dim,562 base=rotary_base,563 scale_base=rotary_scale_base,564 device=device,565 **rotary_kwargs,566 )567 568 # MLP569 self.n_head, self.n_head_kv, self.head_dim = _find_mha_dims(570 config, n_head=n_head, n_head_kv=n_head_kv, head_dim=head_dim571 )572 op_size = self.head_dim * (self.n_head + 2 * self.n_head_kv)573 hidden_size = config.n_embd574 575 linear_cls = FusedDense if config.fused_dense else nn.Linear576 if linear_cls is None:577 linear_cls = nn.Linear578 579 self.Wqkv = linear_cls(hidden_size, op_size, bias=bias, device=device, dtype=dtype)580 self.out_proj = linear_cls(hidden_size, hidden_size, bias=bias, device=device, dtype=dtype)581 582 # Attention583 attn_cls = FlashSelfAttention if config.flash_attn else SelfAttention584 if attn_cls is None:585 attn_cls = SelfAttention586 587 cross_attn_cls = FlashCrossAttention if config.flash_attn else CrossAttention588 if cross_attn_cls is None:589 cross_attn_cls = CrossAttention590 591 self.inner_attn = attn_cls(592 causal=causal,593 softmax_scale=softmax_scale,594 attention_dropout=config.attn_pdrop,595 )596 self.inner_cross_attn = cross_attn_cls(597 causal=causal,598 softmax_scale=softmax_scale,599 attention_dropout=config.attn_pdrop,600 )601 602 self.flash_attn = config.flash_attn and attn_cls is FlashSelfAttention603 self.layer_idx = layer_idx604 self.return_residual = return_residual605 self.checkpointing = checkpointing606 607 def _forward_self_attn(608 self, x: torch.FloatTensor, key_padding_mask: Optional[torch.BoolTensor]609 ) -> torch.FloatTensor:610 qkv = self.Wqkv(x)611 qkv = rearrange(qkv, "... (three h d) -> ... three h d", three=3, d=self.head_dim)612 613 if self.rotary_dim > 0:614 qkv = self.rotary_emb(qkv)615 616 if self.flash_attn:617 batch_size, seqlen = qkv.shape[0], qkv.shape[1]618 619 cu_seqlens, max_seqlen = None, None620 if key_padding_mask is not None:621 # If `key_padding_mask` is supplied, we need to unpad the input and retrieve622 # the `cu_seqlens` and `max_seqlen` to be used by `flash-attn`623 qkv, indices, cu_seqlens, max_seqlen = unpad_input(qkv, key_padding_mask)624 625 if self.checkpointing:626 attn_output = torch.utils.checkpoint.checkpoint(627 self.inner_attn, qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen628 )629 else:630 attn_output = self.inner_attn(qkv, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen).to(qkv.device)631 632 # If `key_padding_mask` is supplied, we need to pad the output back to the original shape633 return pad_input(attn_output, indices, batch_size, seqlen) if key_padding_mask is not None else attn_output634 635 if self.checkpointing:636 return torch.utils.checkpoint.checkpoint(self.inner_attn, qkv, key_padding_mask=key_padding_mask)637 638 return self.inner_attn(qkv, key_padding_mask=key_padding_mask)639 640 def _forward_cross_attn(641 self,642 x: torch.FloatTensor,643 past_key_values: Optional[InferenceParams],644 key_padding_mask: Optional[torch.BoolTensor],645 ) -> torch.FloatTensor:646 batch_size = x.shape[0]647 648 qkv = self.Wqkv(x)649 650 q = qkv[..., : self.n_head * self.head_dim]651 q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)652 653 kv = qkv[..., self.n_head * self.head_dim :]654 kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)655 656 seqlen_offset = past_key_values.seqlen_offset if past_key_values is not None else 0657 causal = None if seqlen_offset == 0 else False658 if self.rotary_dim > 0:659 q, kv = self.rotary_emb(q, kv=kv, seqlen_offset=seqlen_offset)660 661 if past_key_values is not None:662 kv = _update_kv_cache(kv, past_key_values, self.layer_idx)663 664 if self.flash_attn:665 batch_size, seqlen_q = q.shape[0], q.shape[1]666 seqlen_k = kv.shape[1]667 668 cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k = (669 None,670 None,671 None,672 None,673 )674 if key_padding_mask is not None:675 kv, _, cu_seqlens_k, max_seqlen_k = unpad_input(kv, key_padding_mask)676 677 if seqlen_q == 1:678 key_padding_mask = torch.ones(batch_size, 1, device=q.device)679 elif seqlen_q != seqlen_k:680 key_padding_mask = key_padding_mask[:, -seqlen_q:]681 682 q, indices_q, cu_seqlens_q, max_seqlen_q = unpad_input(q, key_padding_mask)683 684 if self.checkpointing:685 attn_output = torch.utils.checkpoint.checkpoint(686 self.inner_cross_attn,687 q,688 kv,689 causal=causal,690 cu_seqlens=cu_seqlens_q,691 max_seqlen=max_seqlen_q,692 cu_seqlens_k=cu_seqlens_k,693 max_seqlen_k=max_seqlen_k,694 )695 else:696 attn_output = self.inner_cross_attn(697 q,698 kv,699 causal=causal,700 cu_seqlens=cu_seqlens_q,701 max_seqlen=max_seqlen_q,702 cu_seqlens_k=cu_seqlens_k,703 max_seqlen_k=max_seqlen_k,704 )705 706 return (707 pad_input(attn_output, indices_q, batch_size, max_seqlen_q)708 if key_padding_mask is not None709 else attn_output710 )711 712 if self.checkpointing:713 return torch.utils.checkpoint.checkpoint(714 self.inner_cross_attn,715 q,716 kv,717 key_padding_mask=key_padding_mask,718 causal=causal,719 )720 721 return self.inner_cross_attn(q, kv, key_padding_mask=key_padding_mask, causal=causal)722 723 def forward(724 self,725 x: torch.FloatTensor,726 past_key_values: Optional[InferenceParams] = None,727 attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,728 **kwargs,729 ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:730 if attention_mask is not None:731 attention_mask = attention_mask.bool()732 else:733 attention_mask = None734 735 # MHA736 if self.n_head == self.n_head_kv:737 if past_key_values is None:738 # If `past_key_values` are not supplied, we run self-attention739 attn_output = self._forward_self_attn(x, attention_mask)740 else:741 # If `past_key_values` are supplied, it means that we might have cached values and742 # could take advantage of cross-attention743 attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)744 # MQA / GQA745 else:746 # Regardless of `past_key_values` being supplied or not, it always use cross-attention747 # because `q` and `kv` lengths might be different748 attn_output = self._forward_cross_attn(x, past_key_values, attention_mask)749 750 output = rearrange(attn_output, "... h d -> ... (h d)")751 output = self.out_proj(output)752 753 return output if not self.return_residual else (output, x)754 755 756class ParallelBlock(nn.Module):757 """Parallel block.758 759 This block applies parallel mixer and MLP layers to the input (used in GPT-J and CodeGen).760 761 """762 763 def __init__(764 self,765 config: PretrainedConfig,766 block_idx: Optional[int] = None,767 ) -> None:768 super().__init__()769 770 self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)771 self.resid_dropout = nn.Dropout(config.resid_pdrop)772 self.block_idx = block_idx773 774 self.mixer = MHA(config, layer_idx=block_idx)775 self.mlp = MLP(config)776 777 def forward(778 self,779 hidden_states: torch.FloatTensor,780 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,781 attention_mask: Optional[torch.BoolTensor] = None,782 **kwargs,783 ) -> torch.FloatTensor:784 residual = hidden_states785 hidden_states = self.ln(hidden_states)786 787 attn_outputs = self.mixer(788 hidden_states,789 past_key_values=past_key_values,790 attention_mask=attention_mask,791 )792 if isinstance(attn_outputs, tuple):793 attn_outputs = attn_outputs[0]794 795 attn_outputs = self.resid_dropout(attn_outputs)796 feed_forward_hidden_states = self.resid_dropout(self.mlp(hidden_states))797 798 hidden_states = attn_outputs + feed_forward_hidden_states + residual799 800 return hidden_states801 802 803class CausalLMHead(nn.Module):804 """Causal Language Modeling head.805 806 Reference:807 Improving Language Understanding by Generative Pre-Training.808 https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf.809 810 """811 812 def __init__(self, config: PretrainedConfig) -> None:813 super().__init__()814 815 self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)816 self.linear = nn.Linear(config.n_embd, config.vocab_size)817 818 def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:819 hidden_states = self.ln(hidden_states)820 logits = self.linear(hidden_states).to(torch.float32)821 822 return logits823 824 825class PhiPreTrainedModel(PreTrainedModel):826 """Phi pre-trained model."""827 828 config_class = PhiConfig829 base_model_prefix = "transformer"830 supports_gradient_checkpointing = True831 _no_split_modules = ["ParallelBlock", "CLIPEncoderLayer", "Block"]832 833 def __init__(self, *inputs, **kwargs) -> None:834 super().__init__(*inputs, **kwargs)835 836 def _init_weights(self, module: nn.Module) -> None:837 if isinstance(module, (nn.Linear,)):838 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)839 if module.bias is not None:840 module.bias.data.zero_()841 elif isinstance(module, nn.Embedding):842 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)843 if module.padding_idx is not None:844 module.weight.data[module.padding_idx].zero_()845 elif isinstance(module, nn.LayerNorm):846 if module.bias is not None:847 module.bias.data.zero_()848 module.weight.data.fill_(1.0)849 850 def prepare_inputs_for_generation(851 self,852 input_ids: torch.LongTensor,853 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,854 attention_mask: Optional[Union[torch.LongTensor, torch.BoolTensor]] = None,855 **kwargs,856 ) -> Dict[str, Any]:857 if past_key_values is None or not (isinstance(past_key_values, InferenceParams)):858 past_key_values = InferenceParams(859 max_seqlen=self.config.n_positions,860 max_batch_size=input_ids.shape[0],861 seqlen_offset=0,862 batch_size_offset=0,863 key_value_memory_dict={},864 lengths_per_sample=None,865 )866 else:867 # ======================================================================868 # Assume that `past_key_values` has cached all tokens up to the last token in `input_ids`869 # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]870 # past_key_values.seqlen_offset = input_ids.shape[1] - 1871 # ======================================================================872 # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.873 # [Edited by zhenwei - 2024-01-20 21:15]874 input_ids = input_ids[:, -1].unsqueeze(-1)875 876 return {877 "input_ids": input_ids,878 "past_key_values": past_key_values,879 "attention_mask": attention_mask,880 }881 882 883class LlavaMetaModel(ABC):884 """885 Define the APIs for building components that are related to image perceiving.886 This implementation is based on the implementation from the Llave project.887 """888 889 def get_vision_tower(self):890 vision_tower = getattr(self, 'vision_tower', None)891 if type(vision_tower) is list:892 vision_tower = vision_tower[0]893 return vision_tower894 895 def build_vision_tower(self, config):896 self.vision_tower = VisionTower(config.vision_tower_cfg)897 898 def build_vision_projector(self, config):899 projector_type = getattr(config, 'mm_projector_type', 'linear')900 901 if projector_type == 'linear':902 self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)903 return904 905 mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)906 if mlp_gelu_match:907 mlp_depth = int(mlp_gelu_match.group(1))908 modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]909 for _ in range(1, mlp_depth):910 modules.append(nn.GELU())911 modules.append(nn.Linear(config.hidden_size, config.hidden_size))912 self.mm_projector = nn.Sequential(*modules)913 return914 915 if projector_type == 'identity':916 self.mm_projector = nn.Identity()917 return918 919 raise ValueError(f'Unknown projector type: {projector_type}')920 921 922class ImpModel(PhiPreTrainedModel, LlavaMetaModel):923 """Imp model. This implementation is modified from the implementation of Phi-2"""924 925 config_class = ImpConfig926 # _keys_to_ignore_on_load_missing = [""]927 # _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]928 929 def __init__(self, config: ImpConfig) -> None:930 super().__init__(config)931 932 self.embd = Embedding(config)933 self.h = nn.ModuleList([ParallelBlock(config, block_idx=i) for i in range(config.n_layer)])934 self.gradient_checkpointing = False935 936 if hasattr(config, "mm_vision_tower"):937 self.build_vision_tower(config)938 self.build_vision_projector(config)939 940 self.post_init()941 942 def embed_tokens(self, input_ids: torch.LongTensor) -> torch.FloatTensor:943 return self.embd(input_ids)[0]944 945 def get_input_embeddings(self) -> nn.Embedding:946 return self.embd.wte947 948 def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:949 self.embd.wte = new_embeddings950 951 def forward(952 self,953 input_ids: torch.LongTensor,954 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,955 attention_mask: Optional[torch.BoolTensor] = None,956 inputs_embeds: Optional[torch.FloatTensor] = None957 ) -> torch.FloatTensor:958 959 if inputs_embeds is None:960 hidden_states = self.embd(input_ids)961 else:962 hidden_states = inputs_embeds963 964 for layer in self.h:965 if self.gradient_checkpointing and self.training:966 967 def create_custom_forward(module):968 def custom_forward(*inputs):969 # None for past_key_value970 return module(*inputs)971 972 return custom_forward973 974 hidden_states = torch.utils.checkpoint.checkpoint(975 create_custom_forward(layer),976 hidden_states,977 None,978 attention_mask,979 )980 else:981 hidden_states = layer(982 hidden_states,983 past_key_values=past_key_values,984 attention_mask=attention_mask,985 )986 987 # I change the way of updating `past_key_values.seqlen_offset` to make the inference of imp work.988 # [Edited by zhenwei - 2024-01-20 21:15]989 if past_key_values is not None: # FIXME: when multi-batch inference, it is a bug990 past_key_values.seqlen_offset += hidden_states.shape[1]991 992 return hidden_states993 994 995class LlavaMetaForCausalLM(ABC):996 """This implementation is based on the implementation from the Llave project."""997 998 def init_constants(self, config):999 self.IGNORE_INDEX = getattr(config, 'ignore_index', -100)1000 self.IMAGE_TOKEN_INDEX = getattr(config, 'image_token_index', 50296)1001 self.DEFAULT_IMAGE_TOKEN = getattr(config, 'image_token', "<image>")1002 1003 @abstractmethod1004 def get_model(self):1005 pass1006 1007 def get_vision_tower(self):1008 return self.get_model().get_vision_tower()1009 1010 def encode_images(self, images):1011 image_features = self.get_model().get_vision_tower()(images)1012 image_features = self.get_model().mm_projector(image_features)1013 return image_features1014 1015 def prepare_inputs_labels_for_multimodal(1016 self, input_ids, position_ids, attention_mask, past_key_values, labels, images1017 ):1018 vision_tower = self.get_vision_tower()1019 # if vision_tower is None or images is None or past_key_values.seqlen_offset != 0:1020 if vision_tower is None or images is None or input_ids.shape[1] == 1:1021 if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:1022 target_shape = past_key_values.seqlen_offset + 11023 # inference_params.key_value_memory_dict[layer_idx][batch_start:batch_end, sequence_start:sequence_end, ...]1024 attention_mask = torch.cat((attention_mask, torch.ones(1025 (attention_mask.shape[0], target_shape - attention_mask.shape[1]),1026 dtype=attention_mask.dtype,1027 device=attention_mask.device1028 )), dim=1)1029 position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 11030 return input_ids, position_ids, attention_mask, past_key_values, None, labels1031 1032 if type(images) is list or images.ndim == 5:1033 concat_images = torch.cat([image for image in images], dim=0)1034 concat_images = concat_images.to(device=self.device, dtype=vision_tower.dtype)1035 image_features = self.encode_images(concat_images)1036 split_sizes = [image.shape[0] for image in images]1037 image_features = torch.split(image_features, split_sizes, dim=0)1038 image_features = [x.flatten(0, 1).to(self.device) for x in image_features]1039 else:1040 images = images.to(device=self.device, dtype=vision_tower.dtype)1041 image_features = self.encode_images(images).to(self.device)1042 1043 # TODO: image start / end is not implemented here to support pretraining.1044 if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):1045 raise NotImplementedError1046 1047 # Let's just add dummy tensors if they do not exist,1048 # it is a headache to deal with None all the time.1049 # But it is not ideal, and if you have a better idea,1050 # please open an issue / submit a PR, thanks.1051 _labels = labels1052 _position_ids = position_ids1053 _attention_mask = attention_mask1054 if attention_mask is None:1055 attention_mask = torch.ones_like(input_ids, dtype=torch.bool)1056 else:1057 attention_mask = attention_mask.bool()1058 if position_ids is None:1059 position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)1060 if labels is None:1061 labels = torch.full_like(input_ids, self.IGNORE_INDEX)1062 1063 # remove the padding using attention_mask -- TODO: double check1064 input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]1065 labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]1066 1067 new_input_embeds = []1068 new_labels = []1069 cur_image_idx = 01070 for batch_idx, cur_input_ids in enumerate(input_ids):1071 num_images = (cur_input_ids == self.IMAGE_TOKEN_INDEX).sum()1072 if num_images == 0:1073 cur_image_features = image_features[cur_image_idx]1074 cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)1075 cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)1076 new_input_embeds.append(cur_input_embeds)1077 new_labels.append(labels[batch_idx])1078 cur_image_idx += 11079 continue1080 1081 image_token_indices = [-1] + torch.where(cur_input_ids == self.IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]1082 cur_input_ids_noim = []1083 cur_labels = labels[batch_idx]1084 cur_labels_noim = []1085 for i in range(len(image_token_indices) - 1):1086 cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])1087 cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])1088 split_sizes = [x.shape[0] for x in cur_labels_noim]1089 cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))1090 # print(cur_input_embeds.shape)1091 cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)1092 cur_new_input_embeds = []1093 cur_new_labels = []1094 1095 for i in range(num_images + 1):1096 cur_new_input_embeds.append(cur_input_embeds_no_im[i])1097 cur_new_labels.append(cur_labels_noim[i])1098 if i < num_images:1099 cur_image_features = image_features[cur_image_idx]1100 cur_image_idx += 11101 cur_new_input_embeds.append(cur_image_features)1102 cur_new_labels.append(torch.full((cur_image_features.shape[0],), self.IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))1103 1104 cur_new_input_embeds = torch.cat(cur_new_input_embeds)1105 cur_new_labels = torch.cat(cur_new_labels)1106 1107 new_input_embeds.append(cur_new_input_embeds)1108 new_labels.append(cur_new_labels)1109 1110 # Truncate sequences to max length as image embeddings can make the sequence longer1111 tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)1112 if tokenizer_model_max_length is not None:1113 new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]1114 new_labels = [x[:tokenizer_model_max_length] for x in new_labels]1115 1116 # Combine them1117 max_len = max(x.shape[0] for x in new_input_embeds)1118 batch_size = len(new_input_embeds)1119 1120 new_input_embeds_padded = []1121 new_labels_padded = torch.full((batch_size, max_len), self.IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)1122 attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)1123 position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)1124 1125 for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):1126 cur_len = cur_new_embed.shape[0]1127 if getattr(self.config, 'tokenizer_padding_side', 'right') == "left":1128 new_input_embeds_padded.append(torch.cat((1129 torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),1130 cur_new_embed1131 ), dim=0))1132 if cur_len > 0:1133 new_labels_padded[i, -cur_len:] = cur_new_labels1134 attention_mask[i, -cur_len:] = True1135 position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)1136 else:1137 new_input_embeds_padded.append(torch.cat((1138 cur_new_embed,1139 torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)1140 ), dim=0))1141 if cur_len > 0:1142 new_labels_padded[i, :cur_len] = cur_new_labels1143 attention_mask[i, :cur_len] = True1144 position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)1145 1146 new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)1147 1148 if _labels is None:1149 new_labels = None1150 else:1151 new_labels = new_labels_padded1152 1153 if _attention_mask is None:1154 attention_mask = None1155 else:1156 attention_mask = attention_mask.to(dtype=_attention_mask.dtype)1157 1158 if _position_ids is None:1159 position_ids = None1160 1161 return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels1162 1163 1164class ImpForCausalLM(PhiPreTrainedModel, LlavaMetaForCausalLM):1165 """Imp for Causal Language Modeling."""1166 1167 # _keys_to_ignore_on_load_missing = [""]1168 # _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]1169 config_class = ImpConfig1170 1171 def __init__(self, config: ImpConfig) -> None:1172 super().__init__(config)1173 1174 self.transformer = ImpModel(config)1175 self.lm_head = CausalLMHead(config)1176 1177 self.post_init()1178 self.init_constants(config)1179 1180 def get_output_embeddings(self) -> nn.Linear:1181 return self.lm_head.linear1182 1183 def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:1184 self.lm_head.linear = new_embeddings1185 1186 def get_model(self):1187 return self.transformer1188 1189 def image_preprocess(self, images):1190 return self.get_vision_tower().image_processor(images)['pixel_values']1191 1192 def backbone_forward(1193 self,1194 input_ids: torch.LongTensor,1195 past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,1196 attention_mask: Optional[torch.BoolTensor] = None,1197 labels: Optional[torch.LongTensor] = None,1198 inputs_embeds: Optional[torch.FloatTensor] = None,1199 **kwargs,1200 ) -> CausalLMOutputWithPast: