openbmb/AgentCPM-Report
305285
1# coding=utf-82# Copyright 2025 The OpenBMB Team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" PyTorch MiniCPM model."""16import math17import re18import warnings19from typing import Any, Dict, List, Optional, Tuple, Union20 21import torch22import torch.nn.functional as F23import torch.utils.checkpoint24from torch import nn25from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss26from transformers.activations import ACT2FN27from transformers.cache_utils import Cache, DynamicCache#, CacheLayerMixin, DynamicLayer28from transformers.modeling_attn_mask_utils import (29 AttentionMaskConverter,30 _prepare_4d_attention_mask,31 _prepare_4d_causal_attention_mask,32 _prepare_4d_causal_attention_mask_for_sdpa,33)34from transformers.modeling_outputs import (35 BaseModelOutputWithPast,36 CausalLMOutputWithPast,37 SequenceClassifierOutputWithPast,38)39from transformers.modeling_utils import PreTrainedModel40from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_1341from transformers.utils import (42 add_start_docstrings,43 add_start_docstrings_to_model_forward,44 is_flash_attn_greater_or_equal_2_10,45 logging,46 replace_return_docstrings,47)48from transformers.utils.import_utils import is_torch_fx_available49from transformers.modeling_flash_attention_utils import _flash_attention_forward50from .configuration_minicpm import MiniCPMConfig51 52try:53 from flash_attn import flash_attn_func, flash_attn_varlen_func54 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa55 from infllm_v2 import (56 infllmv2_attn_stage1,57 infllmv2_attn_varlen_func,58 infllmv2_attn_with_kvcache,59 max_pooling_1d,60 max_pooling_1d_varlen61 )62except:63 pass64 65from functools import lru_cache66 67 68def compressed_attention(69 q: torch.Tensor,70 k: torch.Tensor,71 v: torch.Tensor,72 kernel_size: int,73 kernel_stride: int,74 block_size: int,75 topk: int,76 cu_seqlens_q: torch.Tensor,77 cu_seqlens_k: torch.Tensor,78 max_seqlen_q: int,79 max_seqlen_k: int,80 sm_scale: float = None,81 init_blocks: int = 1,82 local_blocks: int = 2,83 cache_lens: torch.Tensor = None,84) -> Tuple[torch.Tensor, torch.Tensor]:85 """Attention between query and compressed key and value. Compute attention output and topk block idx used in topk_sparse_attention.86 87 Args:88 q (torch.Tensor): shape [total_q_len, num_q_heads, head_dim]89 k (torch.Tensor): shape [total_kv_len, num_kv_heads, head_dim]90 v (torch.Tensor): shape [total_kv_len, num_kv_heads, head_dim]91 kernel_size (int): kernel size in compress_key_value92 kernel_stride (int): stride of compress_key_value93 block_size (int): key value block size for topk sparse attention.94 topk (int): number of blocks for each query.95 cu_seqlens_q (torch.Tensor): shape [batch_size + 1], similar to cu_seqlens_q in flash_attn_func_varlen.96 cu_seqlens_k (torch.Tensor): shape [batch_size + 1], similar to cu_seqlens_k in flash_attn_func_varlen.97 max_seqlen_q (int): max q len of the batch.98 max_seqlen_k (int): max k len of the batch.99 sm_scale (float, optional): softmax scale. Defaults to None, means 1/sqrt(head_dim).100 init_blocks (int, optional): Number of init blocks for each query. Defaults to 1.101 local_blocks (int, optional): Number of local blocks for each query. Defaults to 2.102 cache_lens (torch.Tensor, optional): shape [batch_size], used to record the cache length of each query. Defaults to None.103 104 Returns:105 Tuple[torch.Tensor, torch.Tensor]: attention output and topk_idx used in topk_sparse_attention106 """107 with torch.no_grad():108 batch_size = cu_seqlens_q.shape[0] - 1109 110 # Check if it's prefilling stage111 is_prefilling = cache_lens is None or (cache_lens == 0).all().item()112 113 # prefilling stage114 if is_prefilling:115 # Calculate q_idx for each query position in each batch116 cache_lens = torch.zeros(batch_size, dtype=torch.int32, device=q.device) 117 q_idx = torch.cat([118 (torch.arange(cu_seqlens_q[i + 1] - cu_seqlens_q[i], device=q.device) + 119 max_seqlen_q - (cu_seqlens_q[i + 1] - cu_seqlens_q[i])) // block_size120 for i in range(batch_size)121 ], dim=0) # shape: [total_q_len]122 # decoding stage123 else:124 # Each batch has only one query (last position). Shape: [batch_size] = [total_q_len] in decoding125 q_idx = cache_lens // block_size126 127 # compute attention score128 score = infllmv2_attn_stage1(129 q.contiguous(),130 k.contiguous(),131 v.contiguous(),132 cu_seqlens_q=cu_seqlens_q,133 cu_seqlens_k=cu_seqlens_k,134 max_seqlen_q=max_seqlen_q,135 max_seqlen_k=max_seqlen_k,136 causal=is_prefilling)137 # Shape: [num_heads, total_q_len, num_blocks]138 score = score[:, :q_idx.shape[0], :]139 140 # Shape: [num_heads, total_q_len, num_blocks]141 block_score = max_pooling_1d_varlen(142 score.contiguous(),143 cu_seqlens_q,144 cu_seqlens_k,145 cache_lens,146 max_seqlen_q,147 max_seqlen_k,148 local_blocks=local_blocks,149 init_blocks=init_blocks,150 block_size=block_size,151 stride=kernel_stride)152 153 # get topk154 topk = min(topk, block_score.shape[-1])155 topk_idx = block_score.topk(topk, dim=-1).indices.sort(-1).values156 topk_idx[topk_idx > q_idx[None, :, None]] = -1157 topk_idx = topk_idx.to(torch.int32)158 159 return topk_idx160 161 162@lru_cache(maxsize=16)163def calc_chunks_with_stride(cu_seqlen, chunk_size, kernel_stride):164 """165 Compute the chunks that require Sparse attention, with stride support.166 167 Args:168 cu_seqlen (torch.Tensor): Cumulative sequence lengths for each sample.169 chunk_size (int): Chunk size used for Sparse attention.170 kernel_stride (int): Stride size when sliding over the sequence.171 172 Returns:173 filtered_indices (torch.Tensor): Indices used to directly index into the key/value tensors.174 cu_seqlens_compressed (torch.Tensor): Cumulative sequence lengths after compression.175 """176 # 1. Compute the length of each sequence177 batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1]178 179 # 2. Compute the start positions of chunks for each sequence (with stride)180 max_seq_len = torch.max(batch_sizes)181 max_num_chunks_per_seq = (max_seq_len - chunk_size) // kernel_stride + 1182 chunk_start_offsets = torch.arange(0, max_num_chunks_per_seq * kernel_stride, kernel_stride, device=cu_seqlen.device)183 seq_starts = cu_seqlen[:-1]184 chunk_start_in_seq = seq_starts[:, None] + chunk_start_offsets[None, :] # [batch_size, max_num_chunks_per_seq]185 186 # 3. Filter out chunks that exceed sequence length or are smaller than the full chunk size187 chunk_end_in_seq = chunk_start_in_seq + chunk_size188 valid_chunk_mask = (chunk_end_in_seq <= (seq_starts[:, None] + batch_sizes[:, None]))189 190 # 4. Filter valid chunk start positions using the valid_chunk_mask191 valid_chunk_starts = chunk_start_in_seq[valid_chunk_mask] # [num_valid_chunks]192 del chunk_start_in_seq193 # 5. Generate filtered_indices194 chunk_indices = torch.arange(195 0, chunk_size, device=cu_seqlen.device196 )[None, :] # [1, chunk_size]197 filtered_indices = valid_chunk_starts[:, None] + chunk_indices # [num_valid_chunks, chunk_size]198 filtered_indices = filtered_indices.view(-1) # Flatten to 1D indices199 200 # 6. Compute compressed cumulative sequence lengths201 num_filtered_chunks_per_batch = valid_chunk_mask.sum(dim=1) # Number of valid chunks per batch202 cu_seqlens_compressed = torch.zeros(203 len(cu_seqlen), dtype=torch.int32, device=cu_seqlen.device204 )205 cu_seqlens_compressed[1:] = num_filtered_chunks_per_batch.cumsum(dim=0)206 del num_filtered_chunks_per_batch, chunk_start_offsets, seq_starts, chunk_end_in_seq, valid_chunk_mask, chunk_indices207 return filtered_indices, cu_seqlens_compressed208 209 210class CompressK(torch.nn.Module):211 def __init__(self, head_num_k, head_dim, kernel_size, kernel_stride=16):212 """213 Module for compressing key (K) representations.214 215 Args:216 head_num_k (int): Number of key attention heads.217 head_dim (int): Dimension of each attention head.218 kernel_size (int): Size of each chunk used for compression.219 kernel_stride (int, optional): Stride used when dividing input into chunks. Default is 16.220 """221 super().__init__()222 self.kernel_size = kernel_size223 self.head_num_k = head_num_k224 self.head_dim = head_dim225 self.kernel_stride = kernel_stride226 227 def forward(self, k: torch.Tensor, cu_seqlens):228 """229 Forward pass for compressing the key (K) tensor.230 231 Args:232 k (torch.Tensor): Input key tensor of shape (total_seq_len, num_heads, head_dim).233 cu_seqlens (torch.Tensor): Cumulative sequence lengths for each sample in the batch, typically used for handling variable-length sequences.234 235 Returns:236 compress_k (torch.Tensor): Compressed key tensor.237 cu_seqlens_compressed (torch.Tensor): Updated cumulative sequence lengths after compression.238 239 """240 # Compute chunk-related metadata, with stride support241 filtered_k_indices, cu_seqlens_compressed = calc_chunks_with_stride(242 cu_seqlens, self.kernel_size, self.kernel_stride243 )244 245 # Extract filtered key vectors246 filtered_k = k.index_select(0, filtered_k_indices.view(-1))247 248 # split249 filtered_k = filtered_k.view(filtered_k.shape[0] // self.kernel_size, self.kernel_size, self.head_num_k, self.head_dim) # [l, block_size,h,d]250 251 compressed_k = filtered_k.mean(dim=1)252 return compressed_k, cu_seqlens_compressed253 254 255 256# class InfLLMv2CacheLayer(DynamicLayer):257# def __init__(self):258# super().__init__()259# # Initialize any additional attributes specific to InfLLMv2CacheLayer260# self.no_rope_keys = torch.tensor([], dtype=torch.float32)261# self.compress_k_cache = []262# self.no_compress_k_cache = []263# self.cached_compressed_cu_seqlens = torch.tensor([], dtype=torch.int32)264# self.compress_k_cache_varlen = torch.tensor([], dtype=torch.float32)265 266# def update_no_rope_key(self, key_states):267# if self.no_rope_keys.numel() == 0:268# self.no_rope_keys = key_states269# else:270# self.no_rope_keys = torch.cat([self.no_rope_keys, key_states], dim=1)271# return self.no_rope_keys272 273# def update_compress_k(self, key_states, cu_seqlens=None):274# if len(self.compress_k_cache) == 0:275# if cu_seqlens is not None:276# self.cached_compressed_cu_seqlens = cu_seqlens.clone()277# self.compress_k_cache_varlen = key_states278# split_sizes = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()279# self.compress_k_cache = list(torch.split(key_states, split_sizes))280# else:281# for index, k in enumerate(key_states):282# if k is not None:283# self.compress_k_cache[index] = torch.cat([self.compress_k_cache[index], k], dim=0)284# new_seq_lens = torch.tensor([tensor.shape[0] for tensor in self.compress_k_cache], dtype=torch.int32)285# new_cumsum = torch.cumsum(new_seq_lens, dim=0, dtype=torch.int32)286 287# self.compress_k_cache_varlen = torch.cat(self.compress_k_cache, dim=0)288# self.cached_compressed_cu_seqlens = torch.cat([torch.tensor([0], dtype=torch.int32), new_cumsum]).to(self.compress_k_cache_varlen.device)289# return self.compress_k_cache_varlen, self.cached_compressed_cu_seqlens290 291# def update_no_compress_k(self, key_states, kernel_size=32, kernel_stride=16):292# k_chunk_list = []293# for index, k in enumerate(key_states):294# if len(self.no_compress_k_cache) <= index:295# self.no_compress_k_cache.append(k)296# else:297# self.no_compress_k_cache[index] = torch.cat([self.no_compress_k_cache[index], k], dim=0)298# current_len = self.no_compress_k_cache[index].shape[0]299# if current_len >= kernel_size:300# k_chunk_list.append(self.no_compress_k_cache[index][:kernel_size])301# self.no_compress_k_cache[index] = self.no_compress_k_cache[index][kernel_stride:]302# else:303# k_chunk_list.append(None)304# return k_chunk_list305 306# class InfLLMv2Cache(DynamicCache):307# def __init__(self,308# config,num_hidden_layers: Optional[int] = None) -> None:309# super().__init__(config=config)310# self.layers = [InfLLMv2CacheLayer() for _ in range(num_hidden_layers)] if num_hidden_layers else []311# self._seen_tokens = 0312 313# def update(self, key_states, value_states, layer_idx, cache_kwargs=None):314# if layer_idx == 0:315# self._seen_tokens += key_states.shape[-2]316# return self.layers[layer_idx].update(key_states, value_states, cache_kwargs)317 318# def update_no_rope_key(self, key_states, layer_idx, cache_kwargs=None):319# return self.layers[layer_idx].update_no_rope_key(key_states)320 321# def update_compress_k(self, key_states, layer_idx, cu_seqlens=None, cache_kwargs=None):322# return self.layers[layer_idx].update_compress_k(key_states, cu_seqlens)323 324# def update_no_compress_k(self, key_states, layer_idx, kernel_size=32, kernel_stride=16, cache_kwargs=None):325# return self.layers[layer_idx].update_no_compress_k(key_states, kernel_size, kernel_stride)326 327# def crop(self, max_length):328# for layer in self.layers:329# layer.crop(max_length)330 331# def batch_repeat_interleave(self, repeats):332# for layer in self.layers:333# layer.batch_repeat_interleave(repeats)334 335# def batch_select_indices(self, indices):336# for layer in self.layers:337# layer.batch_select_indices(indices)338 339 340# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.341# It means that the function will not be traced through and simply appear as a node in the graph.342if is_torch_fx_available():343 if not is_torch_greater_or_equal_than_1_13:344 import torch.fx345 346 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)347 348 349logger = logging.get_logger(__name__)350 351_CONFIG_FOR_DOC = 'MiniCPMConfig'352 353 354def _get_unpad_data(attention_mask):355 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)356 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()357 max_seqlen_in_batch = seqlens_in_batch.max().item()358 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))359 return (360 indices,361 cu_seqlens,362 max_seqlen_in_batch,363 )364 365 366 367 368# @torch.jit.script # type: ignore369def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):370 old_dtype = hidden.dtype371 variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)372 hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)373 return hidden * weight374 375 376class MiniCPMRMSNorm(nn.Module):377 def __init__(self, hidden_size, eps=1e-6):378 """379 MiniCPMRMSNorm is equivalent to T5LayerNorm380 """381 super().__init__()382 self.weight = nn.Parameter(torch.ones(hidden_size))383 self.variance_epsilon = eps384 385 def forward(self, hidden_states):386 return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)387 388 389ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)390 391 392class MiniCPMRotaryEmbedding(nn.Module):393 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):394 super().__init__()395 396 self.dim = dim397 self.max_position_embeddings = max_position_embeddings398 self.base = base399 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))400 self.register_buffer('inv_freq', inv_freq, persistent=False)401 402 # Build here to make `torch.jit.trace` work.403 self._set_cos_sin_cache(404 # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()405 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32406 )407 408 def _set_cos_sin_cache(self, seq_len, device, dtype):409 self.max_seq_len_cached = seq_len410 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)411 freqs = torch.outer(t, self.inv_freq)412 # Different from paper, but it uses a different permutation in order to obtain the same calculation413 emb = torch.cat((freqs, freqs), dim=-1)414 415 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)416 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)417 418 def forward(self, x, seq_len=None):419 # x: [bs, num_attention_heads, seq_len, head_size]420 if seq_len > self.max_seq_len_cached:421 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)422 423 return (424 self.cos_cached[:seq_len].to(dtype=x.dtype),425 self.sin_cached[:seq_len].to(dtype=x.dtype),426 )427 428 429class MiniCPMLongRoPE(MiniCPMRotaryEmbedding):430 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""431 432 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, short_factor=None, long_factor=None, original_max_position_embeddings=None):433 self.short_factor = short_factor434 self.long_factor = long_factor435 self.original_max_position_embeddings = original_max_position_embeddings436 scale = (max_position_embeddings / self.original_max_position_embeddings)437 self.scaling_factor = math.sqrt(1 + math.log(scale) / math.log(self.original_max_position_embeddings))438 super().__init__(dim, max_position_embeddings, base, device)439 440 def _set_cos_sin_cache(self, seq_len, device, dtype):441 self.max_seq_len_cached = seq_len442 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)443 if seq_len > self.original_max_position_embeddings:444 ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=device)445 else:446 ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=device)447 448 freqs = torch.mul(449 torch.outer(t, 1.0 / ext_factors).to(device=device),450 self.inv_freq.to(device=device).to(dtype)451 )452 # Different from paper, but it uses a different permutation in order to obtain the same calculation453 emb = torch.cat((freqs, freqs), dim=-1)454 self.register_buffer('cos_cached', emb.cos().to(dtype) * self.scaling_factor, persistent=False)455 self.register_buffer('sin_cached', emb.sin().to(dtype) * self.scaling_factor, persistent=False)456 457 458class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):459 """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""460 461 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):462 self.scaling_factor = scaling_factor463 super().__init__(dim, max_position_embeddings, base, device)464 465 def _set_cos_sin_cache(self, seq_len, device, dtype):466 self.max_seq_len_cached = seq_len467 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)468 t = t / self.scaling_factor469 470 freqs = torch.outer(t, self.inv_freq)471 # Different from paper, but it uses a different permutation in order to obtain the same calculation472 emb = torch.cat((freqs, freqs), dim=-1)473 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)474 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)475 476 477class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):478 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""479 480 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):481 self.scaling_factor = scaling_factor482 super().__init__(dim, max_position_embeddings, base, device)483 484 def _set_cos_sin_cache(self, seq_len, device, dtype):485 self.max_seq_len_cached = seq_len486 487 if seq_len > self.max_position_embeddings:488 base = self.base * (489 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)490 ) ** (self.dim / (self.dim - 2))491 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))492 self.register_buffer('inv_freq', inv_freq, persistent=False)493 494 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)495 496 freqs = torch.outer(t, self.inv_freq)497 # Different from paper, but it uses a different permutation in order to obtain the same calculation498 emb = torch.cat((freqs, freqs), dim=-1)499 500 self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)501 self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)502 503 504def rotate_half(x):505 """Rotates half the hidden dims of the input."""506 x1 = x[..., : x.shape[-1] // 2]507 x2 = x[..., x.shape[-1] // 2:]508 return torch.cat((-x2, x1), dim=-1)509 510 511def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):512 """Applies Rotary Position Embedding to the query and key tensors.513 514 Args:515 q (`torch.Tensor`): The query tensor.516 k (`torch.Tensor`): The key tensor.517 cos (`torch.Tensor`): The cosine part of the rotary embedding.518 sin (`torch.Tensor`): The sine part of the rotary embedding.519 position_ids (`torch.Tensor`):520 The position indices of the tokens corresponding to the query and key tensors. For example, this can be521 used to pass offsetted position ids when working with a KV-cache.522 unsqueeze_dim (`int`, *optional*, defaults to 1):523 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and524 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note525 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and526 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes527 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have528 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.529 Returns:530 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.531 """532 # cos = cos[position_ids].unsqueeze(unsqueeze_dim)533 # sin = sin[position_ids].unsqueeze(unsqueeze_dim)534 # q_embed = (q * cos) + (rotate_half(q) * sin)535 # k_embed = (k * cos) + (rotate_half(k) * sin)536 orig_dtype = k.dtype537 cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]538 sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]539 q_fp32 = q.to(dtype=torch.float32, device=q.device)540 k_fp32 = k.to(dtype=torch.float32, device=k.device)541 q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)542 k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)543 return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)544 545 546class MiniCPMMLP(nn.Module):547 def __init__(self, config):548 super().__init__()549 self.config = config550 self.hidden_size = config.hidden_size551 self.intermediate_size = config.intermediate_size552 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)553 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)554 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)555 self.act_fn = ACT2FN[config.hidden_act]556 557 def forward(self, x):558 if self.config.pretraining_tp > 1:559 slice = self.intermediate_size // self.config.pretraining_tp560 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)561 up_proj_slices = self.up_proj.weight.split(slice, dim=0)562 down_proj_slices = self.down_proj.weight.split(slice, dim=1)563 564 gate_proj = torch.cat(565 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1566 )567 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)568 569 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)570 down_proj = [571 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)572 ]573 down_proj = sum(down_proj)574 else:575 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))576 577 return down_proj578 579def _unpad_one_tensor(hidden_states, attention_mask):580 # Unpad the hidden states using the indices581 indices, cu_seqlens, max_seqlen_in_batch = _get_unpad_data(attention_mask)582 batch_size, seq_len = hidden_states.shape[:2]583 584 # Get the remaining dimensions585 remaining_dims = hidden_states.shape[2:]586 587 # Reshape to (batch_size * seq_len, *remaining_dims)588 reshaped_states = hidden_states.reshape(batch_size * seq_len, *remaining_dims)589 590 # Apply unpadding using indices591 unpadded_states = index_first_axis(reshaped_states, indices)592 593 return unpadded_states, indices, cu_seqlens, max_seqlen_in_batch594 595def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:596 """597 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,598 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)599 """600 batch, num_key_value_heads, slen, head_dim = hidden_states.shape601 if n_rep == 1:602 return hidden_states603 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)604 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)605 606 607class MiniCPMAttention(nn.Module):608 """Multi-headed attention from 'Attention Is All You Need' paper"""609 610 def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):611 super().__init__()612 self.config = config613 self.layer_idx = layer_idx614 if layer_idx is None:615 logger.warning_once(616 f'Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will '617 'to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` '618 'when creating this class.'619 )620 621 self.attention_dropout = config.attention_dropout622 self.hidden_size = config.hidden_size623 self.num_heads = config.num_attention_heads624 self.head_dim = self.hidden_size // self.num_heads625 self.num_key_value_heads = config.num_key_value_heads626 self.num_key_value_groups = self.num_heads // self.num_key_value_heads627 self.max_position_embeddings = config.max_position_embeddings628 self.rope_theta = config.rope_theta629 self.is_causal = True630 631 if (self.head_dim * self.num_heads) != self.hidden_size:632 raise ValueError(633 f'hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}'634 f' and `num_heads`: {self.num_heads}).'635 )636 637 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)638 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)639 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)640 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)641 self._init_rope()642 self.softmax_scale = self.head_dim ** (-0.5)643 644 def _init_rope(self):645 if self.config.rope_scaling is None:646 self.rotary_emb = MiniCPMRotaryEmbedding(647 self.head_dim,648 max_position_embeddings=self.max_position_embeddings,649 base=self.rope_theta,650 )651 else:652 scaling_type = self.config.rope_scaling['rope_type']653 scaling_factor = self.config.rope_scaling.get('factor', None)654 if scaling_type == 'linear':655 self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(656 self.head_dim,657 max_position_embeddings=self.max_position_embeddings,658 scaling_factor=scaling_factor,659 base=self.rope_theta,660 )661 elif scaling_type == 'dynamic':662 self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(663 self.head_dim,664 max_position_embeddings=self.max_position_embeddings,665 scaling_factor=scaling_factor,666 base=self.rope_theta,667 )668 elif scaling_type == 'longrope':669 self.rotary_emb = MiniCPMLongRoPE(670 self.head_dim,671 max_position_embeddings=self.max_position_embeddings,672 short_factor=self.config.rope_scaling['short_factor'],673 long_factor=self.config.rope_scaling['long_factor'],674 base=self.rope_theta,675 original_max_position_embeddings=self.config.rope_scaling['original_max_position_embeddings']676 )677 else:678 raise ValueError(f'Unknown RoPE scaling type {scaling_type}')679 680 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):681 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()682 683 def forward(684 self,685 hidden_states: torch.Tensor,686 attention_mask: Optional[torch.Tensor] = None,687 position_ids: Optional[torch.LongTensor] = None,688 past_key_value: Optional[Cache] = None,689 output_attentions: bool = False,690 use_cache: bool = False,691 **kwargs,692 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:693 if 'padding_mask' in kwargs:694 warnings.warn(695 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'696 )697 698 bsz, q_len, _ = hidden_states.size()699 700 if self.config.pretraining_tp > 1:701 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp702 query_slices = self.q_proj.weight.split(703 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0704 )705 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)706 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)707 708 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]709 query_states = torch.cat(query_states, dim=-1)710 711 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]712 key_states = torch.cat(key_states, dim=-1)713 714 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]715 value_states = torch.cat(value_states, dim=-1)716 717 else:718 query_states = self.q_proj(hidden_states)719 key_states = self.k_proj(hidden_states)720 value_states = self.v_proj(hidden_states)721 722 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)723 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)724 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)725 726 kv_seq_len = position_ids.max().item() + 1727 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)728 729 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)730 731 if past_key_value is not None:732 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models733 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)734 735 key_states = repeat_kv(key_states, self.num_key_value_groups)736 value_states = repeat_kv(value_states, self.num_key_value_groups)737 738 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)739 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):740 raise ValueError(741 f'Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is'742 f' {attn_weights.size()}'743 )744 745 if attention_mask is not None:746 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):747 raise ValueError(748 f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'749 )750 attn_weights = attn_weights + attention_mask751 752 # upcast attention to fp32753 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)754 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)755 attn_output = torch.matmul(attn_weights, value_states)756 757 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):758 raise ValueError(759 f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'760 f' {attn_output.size()}'761 )762 763 attn_output = attn_output.transpose(1, 2).contiguous()764 765 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)766 767 if self.config.pretraining_tp > 1:768 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)769 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)770 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])771 else:772 attn_output = self.o_proj(attn_output)773 774 if not output_attentions:775 attn_weights = None776 777 return attn_output, attn_weights, past_key_value778 779 780class MiniCPMFlashAttention2(MiniCPMAttention):781 """782 MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays783 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of784 flash attention and deal with padding tokens in case the input contains any of them.785 """786 787 def __init__(self, *args, **kwargs):788 super().__init__(*args, **kwargs)789 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.790 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, 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.791 # 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).792 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()793 794 def forward(795 self,796 hidden_states: torch.Tensor,797 attention_mask: Optional[torch.LongTensor] = None,798 position_ids: Optional[torch.LongTensor] = None,799 past_key_value: Optional[Cache] = None,800 output_attentions: bool = False,801 use_cache: bool = False,802 **kwargs,803 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:804 # MiniCPMFlashAttention2 attention does not support output_attentions805 if 'padding_mask' in kwargs:806 warnings.warn(807 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'808 )809 810 # overwrite attention_mask with padding_mask811 attention_mask = kwargs.pop('padding_mask')812 813 output_attentions = False814 815 bsz, q_len, _ = hidden_states.size()816 817 query_states = self.q_proj(hidden_states)818 key_states = self.k_proj(hidden_states)819 value_states = self.v_proj(hidden_states)820 821 # Flash attention requires the input to have the shape822 # batch_size x seq_length x head_dim x hidden_dim823 # therefore we just need to keep the original shape824 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)825 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)826 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)827 828 kv_seq_len = position_ids.max().item() + 1829 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)830 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)831 832 if past_key_value is not None:833 cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models834 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)835 836 # 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 cache837 # to be able to avoid many of these transpose/reshape/view.838 query_states = query_states.transpose(1, 2)839 key_states = key_states.transpose(1, 2)840 value_states = value_states.transpose(1, 2)841 842 dropout_rate = self.attention_dropout if self.training else 0.0843 844 # In PEFT, usually we cast the layer norms in float32 for training stability reasons845 # therefore the input hidden states gets silently casted in float32. Hence, we need846 # cast them back in the correct dtype just to be sure everything works as expected.847 # This might slowdown training & inference so it is recommended to not cast the LayerNorms848 # in fp32. (MiniCPMRMSNorm handles it correctly)849 850 input_dtype = query_states.dtype851 if input_dtype == torch.float32:852 # Handle the case where the model is quantized853 if hasattr(self.config, '_pre_quantization_dtype'):854 target_dtype = self.config._pre_quantization_dtype855 else:856 target_dtype = self.q_proj.weight.dtype857 858 logger.warning_once(859 f'The input hidden states seems to be silently casted in float32, this might be related to'860 f' the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in'861 f' {target_dtype}.'862 )863 864 query_states = query_states.to(target_dtype)865 key_states = key_states.to(target_dtype)866 value_states = value_states.to(target_dtype)867 868 if attention_mask is not None:869 attn_output = self._flash_attention_forward(870 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate871 )872 else:873 attn_output = _flash_attention_forward(874 query_states,875 key_states,876 value_states,877 attention_mask=attention_mask,878 query_length=q_len,879 is_causal=True,880 dropout=dropout_rate,881 softmax_scale=self.softmax_scale,882 position_ids=position_ids,883 )884 885 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()886 attn_output = self.o_proj(attn_output)887 888 if not output_attentions:889 attn_weights = None890 891 return attn_output, attn_weights, past_key_value892 893 def _flash_attention_forward(894 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None895 ):896 """897 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token898 first unpad the input, then computes the attention scores and pad the final attention scores.899 900 Args:901 query_states (`torch.Tensor`):902 Input query states to be passed to Flash Attention API903 key_states (`torch.Tensor`):904 Input key states to be passed to Flash Attention API905 value_states (`torch.Tensor`):906 Input value states to be passed to Flash Attention API907 attention_mask (`torch.Tensor`):908 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the909 position of padding tokens and 1 for the position of non-padding tokens.910 dropout (`int`, *optional*):911 Attention dropout912 softmax_scale (`float`, *optional*):913 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)914 """915 if not self._flash_attn_uses_top_left_mask:916 causal = self.is_causal917 else:918 # 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__.919 causal = self.is_causal and query_length != 1920 # Contains at least one padding token in the sequence921 if attention_mask is not None:922 batch_size = query_states.shape[0]923 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(924 query_states, key_states, value_states, attention_mask, query_length925 )926 927 cu_seqlens_q, cu_seqlens_k = cu_seq_lens928 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens929 attn_output_unpad = flash_attn_varlen_func(930 query_states,931 key_states,932 value_states,933 cu_seqlens_q=cu_seqlens_q,934 cu_seqlens_k=cu_seqlens_k,935 max_seqlen_q=max_seqlen_in_batch_q,936 max_seqlen_k=max_seqlen_in_batch_k,937 dropout_p=dropout,938 softmax_scale=softmax_scale,939 causal=causal,940 )941 942 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)943 else:944 attn_output = flash_attn_func(945 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal946 )947 948 return attn_output949 950 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):951 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)952 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape953 954 key_layer = index_first_axis(955 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k956 )957 value_layer = index_first_axis(958 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k959 )960 if query_length == kv_seq_len:961 query_layer = index_first_axis(962 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k963 )964 cu_seqlens_q = cu_seqlens_k965 max_seqlen_in_batch_q = max_seqlen_in_batch_k966 indices_q = indices_k967 elif query_length == 1:968 max_seqlen_in_batch_q = 1969 cu_seqlens_q = torch.arange(970 batch_size + 1, dtype=torch.int32, device=query_layer.device971 ) # There is a memcpy here, that is very bad.972 indices_q = cu_seqlens_q[:-1]973 query_layer = query_layer.squeeze(1)974 else:975 # The -q_len: slice assumes left padding.976 attention_mask = attention_mask[:, -query_length:]977 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)978 979 return (980 query_layer,981 key_layer,982 value_layer,983 indices_q,984 (cu_seqlens_q, cu_seqlens_k),985 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),986 )987 988 989# class MiniCPMInfLLMv2Attention(MiniCPMAttention):990# """991# MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays992# untouched. The only required change would be on the forward pass where it needs to correctly call the public API of993# flash attention and deal with padding tokens in case the input contains any of them.994# """995 996# def __init__(self, *args, **kwargs):997# super().__init__(*args, **kwargs)998# assert self.config._attn_implementation == 'flash_attention_2', 'Only flash_attention_2 is supported for sparse attention'999# # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.1000# # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, 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.1001# # 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).1002# self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()1003 1004# # -------sparse-------1005# self.kernel_size = self.config.sparse_config.get('kernel_size', 32)1006# self.kernel_stride = self.config.sparse_config.get('kernel_stride', 16)1007# self.init_blocks = self.config.sparse_config.get('init_blocks', 1)1008# self.block_size = self.config.sparse_config.get('block_size', 64)1009# self.window_size = self.config.sparse_config.get('window_size', 2048)1010# self.dense_len = self.config.sparse_config.get('dense_len', 8192)1011 1012# self.local_blocks = self.window_size // self.block_size # local_blocks1013# self.topk = self.config.sparse_config.get('topk', 64) + (self.window_size//self.block_size)1014# self.use_nope = self.config.sparse_config.get('use_nope', False)1015# self.compress_k = CompressK(self.num_key_value_heads, self.head_dim, kernel_size=self.kernel_size, kernel_stride=self.kernel_stride)1016 1017# def forward(1018# self,1019# hidden_states: torch.Tensor,1020# attention_mask: Optional[torch.LongTensor] = None,1021# position_ids: Optional[torch.LongTensor] = None,1022# past_key_value: Optional[Cache] = None,1023# output_attentions: bool = False,1024# use_cache: bool = False,1025# **kwargs,1026# ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:1027# # MiniCPMFlashAttention2 attention does not support output_attentions1028# if 'padding_mask' in kwargs:1029# warnings.warn(1030# 'Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`'1031# )1032 1033# # overwrite attention_mask with padding_mask1034# attention_mask = kwargs.pop('padding_mask')1035 1036# output_attentions = False1037 1038# bsz, q_len, _ = hidden_states.size()1039 1040# query_states = self.q_proj(hidden_states)1041# key_states = self.k_proj(hidden_states)1042# value_states = self.v_proj(hidden_states)1043 1044# # !save no rope1045# if self.use_nope:1046# query_states_no_rope = query_states.view(bsz, q_len, self.num_heads, self.head_dim)1047# key_states_no_rope = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)1048 1049# # Flash attention requires the input to have the shape1050# # batch_size x seq_length x head_dim x hidden_dim1051# # therefore we just need to keep the original shape1052# query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)1053# key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)1054# value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)1055 1056# kv_seq_len = position_ids.max().item() + 11057# cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)1058# query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)1059 1060# if past_key_value is not None:1061# cache_kwargs = {'sin': sin, 'cos': cos} # Specific to RoPE models1062# key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)1063 1064# # 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 cache1065# # to be able to avoid many of these transpose/reshape/view.1066# query_states = query_states.transpose(1, 2)1067# key_states = key_states.transpose(1, 2)1068# value_states = value_states.transpose(1, 2)1069# if self.use_nope:1070# key_states_no_rope = past_key_value.update_no_rope_key(key_states_no_rope, self.layer_idx)1071# no_rope_param = {1072# 'key_states_no_rope': key_states_no_rope,1073# 'query_states_no_rope': query_states_no_rope,1074# }1075# else:1076# no_rope_param = None1077 1078# dropout_rate = self.attention_dropout if self.training else 0.01079 1080# # In PEFT, usually we cast the layer norms in float32 for training stability reasons1081# # therefore the input hidden states gets silently casted in float32. Hence, we need1082# # cast them back in the correct dtype just to be sure everything works as expected.1083# # This might slowdown training & inference so it is recommended to not cast the LayerNorms1084# # in fp32. (MiniCPMRMSNorm handles it correctly)1085 1086# input_dtype = query_states.dtype1087# if input_dtype == torch.float32:1088# # Handle the case where the model is quantized1089# if hasattr(self.config, '_pre_quantization_dtype'):1090# target_dtype = self.config._pre_quantization_dtype1091# else:1092# target_dtype = self.q_proj.weight.dtype1093 1094# logger.warning_once(1095# f'The input hidden states seems to be silently casted in float32, this might be related to'1096# f' the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in'1097# f' {target_dtype}.'1098# )1099 1100# query_states = query_states.to(target_dtype)1101# key_states = key_states.to(target_dtype)1102# value_states = value_states.to(target_dtype)1103# if kv_seq_len < self.dense_len:1104# attn_output = self._flash_attention_forward_dense(1105# query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate)1106# else:1107# attn_output = self._sparse_attention_forward(1108# query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate,1109# no_rope_param=no_rope_param, # if past_key_value is not None else None,1110# past_key_value=past_key_value)1111# attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()1112# attn_output = self.o_proj(attn_output)1113 1114# if not output_attentions:1115# attn_weights = None1116 1117# return attn_output, attn_weights, past_key_value1118 1119# def _sparse_attention_forward(1120# self,1121# query_states,1122# key_states,1123# value_states,1124# attention_mask,1125# query_length,1126# dropout=0.0,1127# softmax_scale=None,1128# no_rope_param=None,1129# past_key_value=None):1130# """1131# Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1132# first unpad the input, then computes the attention scores and pad the final attention scores.1133 1134# Args:1135# query_states (`torch.Tensor`):1136# Input query states to be passed to Flash Attention API1137# key_states (`torch.Tensor`):1138# Input key states to be passed to Flash Attention API1139# value_states (`torch.Tensor`):1140# Input value states to be passed to Flash Attention API1141# attention_mask (`torch.Tensor`):1142# The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1143# position of padding tokens and 1 for the position of non-padding tokens.1144# dropout (`int`, *optional*):1145# Attention dropout1146# softmax_scale (`float`, *optional*):1147# The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1148# """1149# if not self._flash_attn_uses_top_left_mask:1150# causal = self.is_causal1151# else:1152# # 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__.1153# causal = self.is_causal and query_length != 11154# # Contains at least one padding token in the sequence1155# if attention_mask is not None:1156# batch_size = query_states.shape[0]1157# # assert batch_size == 1, 'Only batch_size=1 is supported at the moment.'1158# if past_key_value!=None:1159# compressed_k, compressed_cu_seqlens = self.get_compress_k(1160# key_states=key_states if self.use_nope ==False else no_rope_param['key_states_no_rope'], # This can be optimized a bit;1161# attention_mask=attention_mask,1162# past_key_value=past_key_value)1163 1164# query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(1165# query_states, key_states, value_states, attention_mask, query_length1166# )1167 1168# cu_seqlens_q, cu_seqlens_k = cu_seq_lens1169# max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1170# if no_rope_param != None:1171# if max_seqlen_in_batch_q == 1:1172# no_rope_param['query_states_no_rope'] = no_rope_param['query_states_no_rope'].squeeze(1)1173# else:1174# no_rope_param['query_states_no_rope'],_, _, _ = _unpad_one_tensor(no_rope_param['query_states_no_rope'],attention_mask=attention_mask)1175# if past_key_value==None:1176# # compress_k use varlen form1177# compressed_k, compressed_cu_seqlens = self.compress_k(key_states,cu_seqlens_k)1178 1179# attn_output_unpad = self.sparse_forward(1180# query_states,1181# key_states,1182# value_states,1183# cu_seqlens_q,1184# cu_seqlens_k,1185# max_seqlen_in_batch_q,1186# max_seqlen_in_batch_k,1187# no_rope_param=no_rope_param,1188# compressed_k=compressed_k,1189# compressed_cu_seqlens=compressed_cu_seqlens)1190 1191# attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)1192# else:1193# raise ValueError('Need attention mask')1194 1195# return attn_output1196 1197# def get_compress_k(self, key_states, attention_mask, past_key_value):1198# """1199# Get compressed key states and corresponding cumulative sequence lengths.1200