SparseLLM/ProSparse-MiniCPM-1B-sft
3561
1# coding=utf-82# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20""" PyTorch MiniCPM model."""21import math22import warnings23from typing import List, Optional, Tuple, Union, Dict24 25import torch26import torch.nn.functional as F27import torch.utils.checkpoint28from torch import nn29from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss30 31from transformers.activations import ACT2FN32from transformers.cache_utils import Cache, DynamicCache33from transformers.modeling_attn_mask_utils import (34 AttentionMaskConverter,35 _prepare_4d_attention_mask,36 _prepare_4d_causal_attention_mask,37 _prepare_4d_causal_attention_mask_for_sdpa,38)39from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast40from transformers.modeling_utils import PreTrainedModel41from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_1342from transformers.utils import (43 add_start_docstrings,44 add_start_docstrings_to_model_forward,45 is_flash_attn_2_available,46 is_flash_attn_greater_or_equal_2_10,47 logging,48 replace_return_docstrings,49)50from transformers.utils.import_utils import is_torch_fx_available51from .configuration_minicpm import MiniCPMConfig52import re53 54try:55 from flash_attn import flash_attn_func, flash_attn_varlen_func56 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa57except:58 pass59 60 61# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.62# It means that the function will not be traced through and simply appear as a node in the graph.63if is_torch_fx_available():64 if not is_torch_greater_or_equal_than_1_13:65 import torch.fx66 67 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)68 69 70logger = logging.get_logger(__name__)71 72_CONFIG_FOR_DOC = "MiniCPMConfig"73 74 75def _get_unpad_data(attention_mask):76 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)77 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()78 max_seqlen_in_batch = seqlens_in_batch.max().item()79 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))80 return (81 indices,82 cu_seqlens,83 max_seqlen_in_batch,84 )85 86 87def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):88 warnings.warn(89 "Calling `transformers.models.minicpm.modeling_minicpm._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"90 )91 return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)92 93 94def _make_causal_mask(95 input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 096):97 warnings.warn(98 "Calling `transformers.models.minicpm.modeling_minicpm._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.minicpm.modeling_minicpm.AttentionMaskConverter._make_causal_mask"99 )100 return AttentionMaskConverter._make_causal_mask(101 input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length102 )103 104# @torch.jit.script # type: ignore105def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):106 old_dtype = hidden.dtype107 variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)108 hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)109 return hidden * weight110 111 112class MiniCPMRMSNorm(nn.Module):113 def __init__(self, hidden_size, eps=1e-6):114 """115 MiniCPMRMSNorm is equivalent to T5LayerNorm116 """117 super().__init__()118 self.weight = nn.Parameter(torch.ones(hidden_size))119 self.variance_epsilon = eps120 121 def forward(self, hidden_states):122 return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)123 124 125ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)126 127 128class MiniCPMRotaryEmbedding(nn.Module):129 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):130 super().__init__()131 132 self.dim = dim133 self.max_position_embeddings = max_position_embeddings134 self.base = base135 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))136 self.register_buffer("inv_freq", inv_freq, persistent=False)137 138 # Build here to make `torch.jit.trace` work.139 self._set_cos_sin_cache(140 # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()141 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32142 )143 144 def _set_cos_sin_cache(self, seq_len, device, dtype):145 self.max_seq_len_cached = seq_len146 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)147 freqs = torch.outer(t, self.inv_freq)148 # Different from paper, but it uses a different permutation in order to obtain the same calculation149 emb = torch.cat((freqs, freqs), dim=-1)150 151 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)152 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)153 154 def forward(self, x, seq_len=None):155 # x: [bs, num_attention_heads, seq_len, head_size]156 if seq_len > self.max_seq_len_cached:157 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)158 159 return (160 self.cos_cached[:seq_len].to(dtype=x.dtype),161 self.sin_cached[:seq_len].to(dtype=x.dtype),162 )163 164 165class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):166 """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""167 168 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):169 self.scaling_factor = scaling_factor170 super().__init__(dim, max_position_embeddings, base, device)171 172 def _set_cos_sin_cache(self, seq_len, device, dtype):173 self.max_seq_len_cached = seq_len174 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)175 t = t / self.scaling_factor176 177 freqs = torch.outer(t, self.inv_freq)178 # Different from paper, but it uses a different permutation in order to obtain the same calculation179 emb = torch.cat((freqs, freqs), dim=-1)180 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)181 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)182 183 184class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):185 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""186 187 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):188 self.scaling_factor = scaling_factor189 super().__init__(dim, max_position_embeddings, base, device)190 191 def _set_cos_sin_cache(self, seq_len, device, dtype):192 self.max_seq_len_cached = seq_len193 194 if seq_len > self.max_position_embeddings:195 base = self.base * (196 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)197 ) ** (self.dim / (self.dim - 2))198 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))199 self.register_buffer("inv_freq", inv_freq, persistent=False)200 201 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)202 203 freqs = torch.outer(t, self.inv_freq)204 # Different from paper, but it uses a different permutation in order to obtain the same calculation205 emb = torch.cat((freqs, freqs), dim=-1)206 207 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)208 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)209 210 211def rotate_half(x):212 """Rotates half the hidden dims of the input."""213 x1 = x[..., : x.shape[-1] // 2]214 x2 = x[..., x.shape[-1] // 2 :]215 return torch.cat((-x2, x1), dim=-1)216 217 218def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):219 """Applies Rotary Position Embedding to the query and key tensors.220 221 Args:222 q (`torch.Tensor`): The query tensor.223 k (`torch.Tensor`): The key tensor.224 cos (`torch.Tensor`): The cosine part of the rotary embedding.225 sin (`torch.Tensor`): The sine part of the rotary embedding.226 position_ids (`torch.Tensor`):227 The position indices of the tokens corresponding to the query and key tensors. For example, this can be228 used to pass offsetted position ids when working with a KV-cache.229 unsqueeze_dim (`int`, *optional*, defaults to 1):230 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and231 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note232 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and233 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes234 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have235 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.236 Returns:237 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.238 """239 # cos = cos[position_ids].unsqueeze(unsqueeze_dim)240 # sin = sin[position_ids].unsqueeze(unsqueeze_dim)241 # q_embed = (q * cos) + (rotate_half(q) * sin)242 # k_embed = (k * cos) + (rotate_half(k) * sin)243 orig_dtype = k.dtype244 cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]245 sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]246 q_fp32 = q.to(dtype=torch.float32, device=q.device)247 k_fp32 = k.to(dtype=torch.float32, device=k.device)248 q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)249 k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)250 return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)251 252class MiniCPMMLP(nn.Module):253 def __init__(self, config):254 super().__init__()255 self.config = config256 self.hidden_size = config.hidden_size257 self.intermediate_size = config.intermediate_size258 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)259 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)260 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)261 if config.hidden_act in ACT2FN:262 self.act_fn = ACT2FN[config.hidden_act]263 elif config.hidden_act == "shiftrelu":264 def shifted_relu(x):265 return torch.nn.functional.relu(x - config.hidden_act_param)266 self.act_fn = shifted_relu267 elif config.hidden_act == "fatrelu":268 def fat_relu(x):269 new_x = torch.zeros_like(x)270 mask = torch.ge(x, config.hidden_act_param)271 new_x[mask] = x[mask]272 return new_x273 self.act_fn = fat_relu274 else:275 raise NotImplementedError(f"Unsupported activation function: {config.hidden_act}")276 277 def forward(self, x):278 if self.config.pretraining_tp > 1:279 slice = self.intermediate_size // self.config.pretraining_tp280 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)281 up_proj_slices = self.up_proj.weight.split(slice, dim=0)282 down_proj_slices = self.down_proj.weight.split(slice, dim=1)283 284 gate_proj = torch.cat(285 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1286 )287 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)288 289 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)290 down_proj = [291 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)292 ]293 down_proj = sum(down_proj)294 else:295 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))296 297 return down_proj298 299 300def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:301 """302 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,303 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)304 """305 batch, num_key_value_heads, slen, head_dim = hidden_states.shape306 if n_rep == 1:307 return hidden_states308 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)309 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)310 311 312 313class MiniCPMAttention(nn.Module):314 """Multi-headed attention from 'Attention Is All You Need' paper"""315 316 def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):317 super().__init__()318 self.config = config319 self.layer_idx = layer_idx320 if layer_idx is None:321 logger.warning_once(322 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "323 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "324 "when creating this class."325 )326 327 self.attention_dropout = config.attention_dropout328 self.hidden_size = config.hidden_size329 self.num_heads = config.num_attention_heads330 self.head_dim = self.hidden_size // self.num_heads331 self.num_key_value_heads = config.num_key_value_heads332 self.num_key_value_groups = self.num_heads // self.num_key_value_heads333 self.max_position_embeddings = config.max_position_embeddings334 self.rope_theta = config.rope_theta335 self.is_causal = True336 337 if (self.head_dim * self.num_heads) != self.hidden_size:338 raise ValueError(339 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"340 f" and `num_heads`: {self.num_heads})."341 )342 343 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)344 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)345 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)346 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)347 self._init_rope()348 349 def _init_rope(self):350 if self.config.rope_scaling is None:351 self.rotary_emb = MiniCPMRotaryEmbedding(352 self.head_dim,353 max_position_embeddings=self.max_position_embeddings,354 base=self.rope_theta,355 )356 else:357 scaling_type = self.config.rope_scaling["type"]358 scaling_factor = self.config.rope_scaling["factor"]359 if scaling_type == "linear":360 self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(361 self.head_dim,362 max_position_embeddings=self.max_position_embeddings,363 scaling_factor=scaling_factor,364 base=self.rope_theta,365 )366 elif scaling_type == "dynamic":367 self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(368 self.head_dim,369 max_position_embeddings=self.max_position_embeddings,370 scaling_factor=scaling_factor,371 base=self.rope_theta,372 )373 else:374 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")375 376 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):377 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()378 379 def forward(380 self,381 hidden_states: torch.Tensor,382 attention_mask: Optional[torch.Tensor] = None,383 position_ids: Optional[torch.LongTensor] = None,384 past_key_value: Optional[Cache] = None,385 output_attentions: bool = False,386 use_cache: bool = False,387 **kwargs,388 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:389 if "padding_mask" in kwargs:390 warnings.warn(391 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"392 )393 394 bsz, q_len, _ = hidden_states.size()395 396 if self.config.pretraining_tp > 1:397 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp398 query_slices = self.q_proj.weight.split(399 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0400 )401 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)402 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)403 404 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]405 query_states = torch.cat(query_states, dim=-1)406 407 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]408 key_states = torch.cat(key_states, dim=-1)409 410 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]411 value_states = torch.cat(value_states, dim=-1)412 413 else:414 query_states = self.q_proj(hidden_states)415 key_states = self.k_proj(hidden_states)416 value_states = self.v_proj(hidden_states)417 418 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)419 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)420 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)421 422 kv_seq_len = key_states.shape[-2]423 if past_key_value is not None:424 if self.layer_idx is None:425 raise ValueError(426 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "427 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "428 "with a layer index."429 )430 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)431 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)432 433 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)434 435 if past_key_value is not None:436 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models437 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)438 439 key_states = repeat_kv(key_states, self.num_key_value_groups)440 value_states = repeat_kv(value_states, self.num_key_value_groups)441 442 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)443 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):444 raise ValueError(445 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"446 f" {attn_weights.size()}"447 )448 449 if attention_mask is not None:450 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):451 raise ValueError(452 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"453 )454 attn_weights = attn_weights + attention_mask455 456 # upcast attention to fp32457 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)458 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)459 attn_output = torch.matmul(attn_weights, value_states)460 461 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):462 raise ValueError(463 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"464 f" {attn_output.size()}"465 )466 467 attn_output = attn_output.transpose(1, 2).contiguous()468 469 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)470 471 if self.config.pretraining_tp > 1:472 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)473 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)474 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])475 else:476 attn_output = self.o_proj(attn_output)477 478 if not output_attentions:479 attn_weights = None480 481 return attn_output, attn_weights, past_key_value482 483 484class MiniCPMFlashAttention2(MiniCPMAttention):485 """486 MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays487 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of488 flash attention and deal with padding tokens in case the input contains any of them.489 """490 491 def __init__(self, *args, **kwargs):492 super().__init__(*args, **kwargs)493 494 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.495 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.496 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).497 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()498 499 def forward(500 self,501 hidden_states: torch.Tensor,502 attention_mask: Optional[torch.LongTensor] = None,503 position_ids: Optional[torch.LongTensor] = None,504 past_key_value: Optional[Cache] = None,505 output_attentions: bool = False,506 use_cache: bool = False,507 **kwargs,508 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:509 # MiniCPMFlashAttention2 attention does not support output_attentions510 if "padding_mask" in kwargs:511 warnings.warn(512 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"513 )514 515 # overwrite attention_mask with padding_mask516 attention_mask = kwargs.pop("padding_mask")517 518 output_attentions = False519 520 bsz, q_len, _ = hidden_states.size()521 522 query_states = self.q_proj(hidden_states)523 key_states = self.k_proj(hidden_states)524 value_states = self.v_proj(hidden_states)525 526 # Flash attention requires the input to have the shape527 # batch_size x seq_length x head_dim x hidden_dim528 # therefore we just need to keep the original shape529 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)530 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)531 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)532 533 kv_seq_len = key_states.shape[-2]534 if past_key_value is not None:535 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)536 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)537 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)538 539 if past_key_value is not None:540 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models541 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)542 543 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache544 # to be able to avoid many of these transpose/reshape/view.545 query_states = query_states.transpose(1, 2)546 key_states = key_states.transpose(1, 2)547 value_states = value_states.transpose(1, 2)548 549 dropout_rate = self.attention_dropout if self.training else 0.0550 551 # In PEFT, usually we cast the layer norms in float32 for training stability reasons552 # therefore the input hidden states gets silently casted in float32. Hence, we need553 # cast them back in the correct dtype just to be sure everything works as expected.554 # This might slowdown training & inference so it is recommended to not cast the LayerNorms555 # in fp32. (MiniCPMRMSNorm handles it correctly)556 557 input_dtype = query_states.dtype558 if input_dtype == torch.float32:559 # Handle the case where the model is quantized560 if hasattr(self.config, "_pre_quantization_dtype"):561 target_dtype = self.config._pre_quantization_dtype562 else:563 target_dtype = self.q_proj.weight.dtype564 565 logger.warning_once(566 f"The input hidden states seems to be silently casted in float32, this might be related to"567 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"568 f" {target_dtype}."569 )570 571 query_states = query_states.to(target_dtype)572 key_states = key_states.to(target_dtype)573 value_states = value_states.to(target_dtype)574 575 attn_output = self._flash_attention_forward(576 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate577 )578 579 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()580 attn_output = self.o_proj(attn_output)581 582 if not output_attentions:583 attn_weights = None584 585 return attn_output, attn_weights, past_key_value586 587 def _flash_attention_forward(588 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None589 ):590 """591 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token592 first unpad the input, then computes the attention scores and pad the final attention scores.593 594 Args:595 query_states (`torch.Tensor`):596 Input query states to be passed to Flash Attention API597 key_states (`torch.Tensor`):598 Input key states to be passed to Flash Attention API599 value_states (`torch.Tensor`):600 Input value states to be passed to Flash Attention API601 attention_mask (`torch.Tensor`):602 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the603 position of padding tokens and 1 for the position of non-padding tokens.604 dropout (`int`, *optional*):605 Attention dropout606 softmax_scale (`float`, *optional*):607 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)608 """609 if not self._flash_attn_uses_top_left_mask:610 causal = self.is_causal611 else:612 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.613 causal = self.is_causal and query_length != 1614 # Contains at least one padding token in the sequence615 if attention_mask is not None:616 batch_size = query_states.shape[0]617 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(618 query_states, key_states, value_states, attention_mask, query_length619 )620 621 cu_seqlens_q, cu_seqlens_k = cu_seq_lens622 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens623 attn_output_unpad = flash_attn_varlen_func(624 query_states,625 key_states,626 value_states,627 cu_seqlens_q=cu_seqlens_q,628 cu_seqlens_k=cu_seqlens_k,629 max_seqlen_q=max_seqlen_in_batch_q,630 max_seqlen_k=max_seqlen_in_batch_k,631 dropout_p=dropout,632 softmax_scale=softmax_scale,633 causal=causal,634 )635 636 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)637 else:638 attn_output = flash_attn_func(639 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal640 )641 642 return attn_output643 644 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):645 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)646 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape647 648 key_layer = index_first_axis(649 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k650 )651 value_layer = index_first_axis(652 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k653 )654 if query_length == kv_seq_len:655 query_layer = index_first_axis(656 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k657 )658 cu_seqlens_q = cu_seqlens_k659 max_seqlen_in_batch_q = max_seqlen_in_batch_k660 indices_q = indices_k661 elif query_length == 1:662 max_seqlen_in_batch_q = 1663 cu_seqlens_q = torch.arange(664 batch_size + 1, dtype=torch.int32, device=query_layer.device665 ) # There is a memcpy here, that is very bad.666 indices_q = cu_seqlens_q[:-1]667 query_layer = query_layer.squeeze(1)668 else:669 # The -q_len: slice assumes left padding.670 attention_mask = attention_mask[:, -query_length:]671 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)672 673 return (674 query_layer,675 key_layer,676 value_layer,677 indices_q,678 (cu_seqlens_q, cu_seqlens_k),679 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),680 )681 682 683class MiniCPMSdpaAttention(MiniCPMAttention):684 """685 MiniCPM attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from686 `MiniCPMAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to687 SDPA API.688 """689 690 # Adapted from MiniCPMAttention.forward691 def forward(692 self,693 hidden_states: torch.Tensor,694 attention_mask: Optional[torch.Tensor] = None,695 position_ids: Optional[torch.LongTensor] = None,696 past_key_value: Optional[Cache] = None,697 output_attentions: bool = False,698 use_cache: bool = False,699 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:700 if output_attentions:701 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.702 logger.warning_once(703 "MiniCPMModel is using MiniCPMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "704 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'705 )706 return super().forward(707 hidden_states=hidden_states,708 attention_mask=attention_mask,709 position_ids=position_ids,710 past_key_value=past_key_value,711 output_attentions=output_attentions,712 use_cache=use_cache,713 )714 715 bsz, q_len, _ = hidden_states.size()716 717 query_states = self.q_proj(hidden_states)718 key_states = self.k_proj(hidden_states)719 value_states = self.v_proj(hidden_states)720 721 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)722 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)723 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)724 725 kv_seq_len = key_states.shape[-2]726 if past_key_value is not None:727 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)728 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)729 730 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)731 732 if past_key_value is not None:733 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models734 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)735 736 key_states = repeat_kv(key_states, self.num_key_value_groups)737 value_states = repeat_kv(value_states, self.num_key_value_groups)738 739 if attention_mask is not None:740 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):741 raise ValueError(742 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"743 )744 745 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,746 # Reference: https://github.com/pytorch/pytorch/issues/112577.747 if query_states.device.type == "cuda" and attention_mask is not None:748 query_states = query_states.contiguous()749 key_states = key_states.contiguous()750 value_states = value_states.contiguous()751 752 attn_output = torch.nn.functional.scaled_dot_product_attention(753 query_states,754 key_states,755 value_states,756 attn_mask=attention_mask,757 dropout_p=self.attention_dropout if self.training else 0.0,758 # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.759 is_causal=self.is_causal and attention_mask is None and q_len > 1,760 )761 762 attn_output = attn_output.transpose(1, 2).contiguous()763 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)764 765 attn_output = self.o_proj(attn_output)766 767 return attn_output, None, past_key_value768 769 770MINICPM_ATTENTION_CLASSES = {771 "eager": MiniCPMAttention,772 "flash_attention_2": MiniCPMFlashAttention2,773 "sdpa": MiniCPMSdpaAttention,774}775 776 777class MiniCPMDecoderLayer(nn.Module):778 def __init__(self, config: MiniCPMConfig, layer_idx: int):779 super().__init__()780 self.hidden_size = config.hidden_size781 self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)782 783 self.mlp = MiniCPMMLP(config)784 self.input_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)785 self.post_attention_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)786 787 self.scale_depth = config.scale_depth788 self.num_hidden_layers = config.num_hidden_layers789 790 def forward(791 self,792 hidden_states: torch.Tensor,793 attention_mask: Optional[torch.Tensor] = None,794 position_ids: Optional[torch.LongTensor] = None,795 past_key_value: Optional[Tuple[torch.Tensor]] = None,796 output_attentions: Optional[bool] = False,797 use_cache: Optional[bool] = False,798 **kwargs,799 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:800 """801 Args:802 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`803 attention_mask (`torch.FloatTensor`, *optional*):804 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,805 query_sequence_length, key_sequence_length)` if default attention is used.806 output_attentions (`bool`, *optional*):807 Whether or not to return the attentions tensors of all attention layers. See `attentions` under808 returned tensors for more detail.809 use_cache (`bool`, *optional*):810 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding811 (see `past_key_values`).812 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states813 """814 if "padding_mask" in kwargs:815 warnings.warn(816 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"817 )818 819 residual = hidden_states820 hidden_states = self.input_layernorm(hidden_states)821 # Self Attention822 hidden_states, self_attn_weights, present_key_value = self.self_attn(823 hidden_states=hidden_states,824 attention_mask=attention_mask,825 position_ids=position_ids,826 past_key_value=past_key_value,827 output_attentions=output_attentions,828 use_cache=use_cache,829 **kwargs,830 )831 832 hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))833 834 # Fully Connected835 residual = hidden_states836 hidden_states = self.post_attention_layernorm(hidden_states)837 838 hidden_states = self.mlp(hidden_states)839 hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))840 841 outputs = (hidden_states,)842 843 if output_attentions:844 outputs += (self_attn_weights,)845 846 if use_cache:847 outputs += (present_key_value,)848 849 return outputs850 851 852MINICPM_START_DOCSTRING = r"""853 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the854 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads855 etc.)856 857 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.858 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage859 and behavior.860 861 Parameters:862 config ([`MiniCPMConfig`]):863 Model configuration class with all the parameters of the model. Initializing with a config file does not864 load the weights associated with the model, only the configuration. Check out the865 [`~PreTrainedModel.from_pretrained`] method to load the model weights.866"""867 868 869@add_start_docstrings(870 "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",871 MINICPM_START_DOCSTRING,872)873class MiniCPMPreTrainedModel(PreTrainedModel):874 config_class = MiniCPMConfig875 base_model_prefix = "model"876 supports_gradient_checkpointing = True877 _no_split_modules = ["MiniCPMDecoderLayer"]878 _skip_keys_device_placement = "past_key_values"879 _supports_flash_attn_2 = True880 _supports_sdpa = True881 _supports_cache_class = True882 883 def _init_weights(self, module):884 std = self.config.initializer_range885 if isinstance(module, nn.Linear):886 module.weight.data.normal_(mean=0.0, std=std)887 if module.bias is not None:888 module.bias.data.zero_()889 elif isinstance(module, nn.Embedding):890 module.weight.data.normal_(mean=0.0, std=std)891 if module.padding_idx is not None:892 module.weight.data[module.padding_idx].zero_()893 894 895MINICPM_INPUTS_DOCSTRING = r"""896 Args:897 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):898 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide899 it.900 901 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and902 [`PreTrainedTokenizer.__call__`] for details.903 904 [What are input IDs?](../glossary#input-ids)905 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):906 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:907 908 - 1 for tokens that are **not masked**,909 - 0 for tokens that are **masked**.910 911 [What are attention masks?](../glossary#attention-mask)912 913 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and914 [`PreTrainedTokenizer.__call__`] for details.915 916 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see917 `past_key_values`).918 919 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]920 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more921 information on the default strategy.922 923 - 1 indicates the head is **not masked**,924 - 0 indicates the head is **masked**.925 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):926 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,927 config.n_positions - 1]`.928 929 [What are position IDs?](../glossary#position-ids)930 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):931 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention932 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`933 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.934 935 Two formats are allowed:936 - a [`~cache_utils.Cache`] instance;937 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of938 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy939 cache format.940 941 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the942 legacy cache format will be returned.943 944 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't945 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`946 of shape `(batch_size, sequence_length)`.947 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):948 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This949 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the950 model's internal embedding lookup matrix.951 use_cache (`bool`, *optional*):952 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see953 `past_key_values`).954 output_attentions (`bool`, *optional*):955 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned956 tensors for more detail.957 output_hidden_states (`bool`, *optional*):958 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for959 more detail.960 return_dict (`bool`, *optional*):961 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.962"""963 964 965@add_start_docstrings(966 "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",967 MINICPM_START_DOCSTRING,968)969class MiniCPMModel(MiniCPMPreTrainedModel):970 """971 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]972 973 Args:974 config: MiniCPMConfig975 """976 977 def __init__(self, config: MiniCPMConfig):978 super().__init__(config)979 self.padding_idx = config.pad_token_id980 self.vocab_size = config.vocab_size981 982 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)983 self.layers = nn.ModuleList(984 [MiniCPMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]985 )986 self._use_sdpa = config._attn_implementation == "sdpa"987 self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"988 989 self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)990 991 self.gradient_checkpointing = False992 # Initialize weights and apply final processing993 self.post_init()994 995 def get_input_embeddings(self):996 return self.embed_tokens997 998 def set_input_embeddings(self, value):999 self.embed_tokens = value1000 1001 @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)1002 def forward(1003 self,1004 input_ids: torch.LongTensor = None,1005 attention_mask: Optional[torch.Tensor] = None,1006 position_ids: Optional[torch.LongTensor] = None,1007 past_key_values: Optional[List[torch.FloatTensor]] = None,1008 inputs_embeds: Optional[torch.FloatTensor] = None,1009 use_cache: Optional[bool] = None,1010 output_attentions: Optional[bool] = None,1011 output_hidden_states: Optional[bool] = None,1012 return_dict: Optional[bool] = None,1013 ) -> Union[Tuple, BaseModelOutputWithPast]:1014 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1015 output_hidden_states = (1016 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1017 )1018 use_cache = use_cache if use_cache is not None else self.config.use_cache1019 1020 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1021 1022 # retrieve input_ids and inputs_embeds1023 if input_ids is not None and inputs_embeds is not None:1024 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")1025 elif input_ids is not None:1026 batch_size, seq_length = input_ids.shape[:2]1027 elif inputs_embeds is not None:1028 batch_size, seq_length = inputs_embeds.shape[:2]1029 else:1030 raise ValueError("You have to specify either input_ids or inputs_embeds")1031 1032 if self.gradient_checkpointing and self.training:1033 if use_cache:1034 logger.warning_once(1035 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."1036 )1037 use_cache = False1038 1039 past_key_values_length = 01040 if use_cache:1041 use_legacy_cache = not isinstance(past_key_values, Cache)1042 if use_legacy_cache:1043 past_key_values = DynamicCache.from_legacy_cache(past_key_values)1044 past_key_values_length = past_key_values.get_usable_length(seq_length)1045 1046 if position_ids is None:1047 device = input_ids.device if input_ids is not None else inputs_embeds.device1048 position_ids = torch.arange(1049 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device1050 )1051 position_ids = position_ids.unsqueeze(0)1052 1053 if inputs_embeds is None:1054 inputs_embeds = self.embed_tokens(input_ids) * self.config.scale_emb1055 1056 if self._use_flash_attention_2:1057 # 2d mask is passed through the layers1058 attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None1059 elif self._use_sdpa and not output_attentions:1060 # output_attentions=True can not be supported when using SDPA, and we fall back on1061 # the manual implementation that requires a 4D causal mask in all cases.1062 attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(1063 attention_mask,1064 (batch_size, seq_length),1065 inputs_embeds,1066 past_key_values_length,1067 )1068 else:1069 # 4d mask is passed through the layers1070 attention_mask = _prepare_4d_causal_attention_mask(1071 attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length1072 )1073 1074 # embed positions1075 hidden_states = inputs_embeds1076 1077 # decoder layers1078 all_hidden_states = () if output_hidden_states else None1079 all_self_attns = () if output_attentions else None1080 next_decoder_cache = None1081 1082 for decoder_layer in self.layers:1083 if output_hidden_states:1084 all_hidden_states += (hidden_states,)1085 1086 if self.gradient_checkpointing and self.training:1087 layer_outputs = self._gradient_checkpointing_func(1088 decoder_layer.__call__,1089 hidden_states,1090 attention_mask,1091 position_ids,1092 past_key_values,1093 output_attentions,1094 use_cache,1095 )1096 else:1097 layer_outputs = decoder_layer(1098 hidden_states,1099 attention_mask=attention_mask,1100 position_ids=position_ids,1101 past_key_value=past_key_values,1102 output_attentions=output_attentions,1103 use_cache=use_cache,1104 )1105 1106 hidden_states = layer_outputs[0]1107 1108 if use_cache:1109 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1110 1111 if output_attentions:1112 all_self_attns += (layer_outputs[1],)1113 1114 hidden_states = self.norm(hidden_states)1115 1116 # add hidden states from the last decoder layer1117 if output_hidden_states:1118 all_hidden_states += (hidden_states,)1119 1120 next_cache = None1121 if use_cache:1122 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache1123 if not return_dict:1124 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1125 return BaseModelOutputWithPast(1126 last_hidden_state=hidden_states,1127 past_key_values=next_cache,1128 hidden_states=all_hidden_states,1129 attentions=all_self_attns,1130 )1131 1132 1133class MiniCPMForCausalLM(MiniCPMPreTrainedModel):1134 _tied_weights_keys = ["lm_head.weight"]1135 1136 def __init__(self, config):1137 super().__init__(config)1138 self.model = MiniCPMModel(config)1139 self.vocab_size = config.vocab_size1140 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1141 1142 # Initialize weights and apply final processing1143 self.post_init()1144 1145 def get_input_embeddings(self):1146 return self.model.embed_tokens1147 1148 def set_input_embeddings(self, value):1149 self.model.embed_tokens = value1150 1151 def get_output_embeddings(self):1152 return self.lm_head1153 1154 def set_output_embeddings(self, new_embeddings):1155 self.lm_head = new_embeddings1156 1157 def set_decoder(self, decoder):1158 self.model = decoder1159 1160 def get_decoder(self):1161 return self.model1162 1163 @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)1164 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)1165 def forward(1166 self,1167 input_ids: torch.LongTensor = None,1168 attention_mask: Optional[torch.Tensor] = None,1169 position_ids: Optional[torch.LongTensor] = None,1170 past_key_values: Optional[List[torch.FloatTensor]] = None,1171 inputs_embeds: Optional[torch.FloatTensor] = None,1172 labels: Optional[torch.LongTensor] = None,1173 use_cache: Optional[bool] = None,1174 output_attentions: Optional[bool] = None,1175 output_hidden_states: Optional[bool] = None,1176 return_dict: Optional[bool] = None,1177 ) -> Union[Tuple, CausalLMOutputWithPast]:1178 r"""1179 Args:1180 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1181 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1182 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1183 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1184 1185 Returns:1186 1187 Example:1188 1189 ```python1190 >>> from transformers import AutoTokenizer, MiniCPMForCausalLM1191 1192 >>> model = MiniCPMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)1193 >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)1194 1195 >>> prompt = "Hey, are you conscious? Can you talk to me?"1196 >>> inputs = tokenizer(prompt, return_tensors="pt")1197 1198 >>> # Generate1199 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)1200 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]