| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Adapted from |
| 5 | # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py |
| 6 | # Copyright 2024 The ModelBest team. |
| 7 | # Copyright 2023 The vLLM team. |
| 8 | # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. |
| 9 | # |
| 10 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 11 | # and OPT implementations in this library. It has been modified from its |
| 12 | # original forms to accommodate minor architectural differences compared |
| 13 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 14 | # |
| 15 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 16 | # you may not use this file except in compliance with the License. |
| 17 | # You may obtain a copy of the License at |
| 18 | # |
| 19 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 20 | # |
| 21 | # Unless required by applicable law or agreed to in writing, software |
| 22 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 23 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 24 | # See the License for the specific language governing permissions and |
| 25 | # limitations under the License. |
| 26 | """Inference-only MiniCPM3 model compatible with HuggingFace weights.""" |
| 27 | |
| 28 | import torch |
| 29 | from torch import nn |
| 30 | from transformers import PreTrainedConfig |
| 31 | |
| 32 | from vllm.config import CacheConfig, VllmConfig |
| 33 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 34 | from vllm.model_executor.layers.attention import Attention |
| 35 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 36 | from vllm.model_executor.layers.linear import ( |
| 37 | ColumnParallelLinear, |
| 38 | ReplicatedLinear, |
| 39 | RowParallelLinear, |
| 40 | ) |
| 41 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 42 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 43 | from vllm.model_executor.models.minicpm import ( |
| 44 | MiniCPMDecoderLayer, |
| 45 | MiniCPMForCausalLM, |
| 46 | MiniCPMModel, |
| 47 | ) |
| 48 | |
| 49 | from .utils import make_layers |
| 50 | |
| 51 | |
| 52 | class MiniCPM3Attention(nn.Module): |
| 53 | def __init__( |
| 54 | self, |
| 55 | config: PreTrainedConfig, |
| 56 | hidden_size: int, |
| 57 | num_heads: int, |
| 58 | qk_nope_head_dim: int, |
| 59 | qk_rope_head_dim: int, |
| 60 | v_head_dim: int, |
| 61 | q_lora_rank: int, |
| 62 | kv_lora_rank: int, |
| 63 | max_position_embeddings: int = 8192, |
| 64 | cache_config: CacheConfig | None = None, |
| 65 | quant_config: QuantizationConfig | None = None, |
| 66 | prefix: str = "", |
| 67 | ) -> None: |
| 68 | super().__init__() |
| 69 | self.hidden_size = hidden_size |
| 70 | self.qk_nope_head_dim = qk_nope_head_dim |
| 71 | self.qk_rope_head_dim = qk_rope_head_dim |
| 72 | self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim |
| 73 | self.v_head_dim = v_head_dim |
| 74 | self.q_lora_rank = q_lora_rank |
| 75 | self.kv_lora_rank = kv_lora_rank |
| 76 | self.num_heads = num_heads |
| 77 | |
| 78 | tp_size = get_tensor_model_parallel_world_size() |
| 79 | assert self.num_heads % tp_size == 0 |
| 80 | self.num_local_heads = num_heads // tp_size |
| 81 | |
| 82 | self.scaling = self.qk_head_dim**-0.5 |
| 83 | self.max_position_embeddings = max_position_embeddings |
| 84 | |
| 85 | self.q_a_proj = ReplicatedLinear( |
| 86 | self.hidden_size, self.q_lora_rank, bias=False, quant_config=quant_config |
| 87 | ) |
| 88 | self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) |
| 89 | self.q_b_proj = ColumnParallelLinear( |
| 90 | q_lora_rank, |
| 91 | self.num_heads * self.qk_head_dim, |
| 92 | bias=False, |
| 93 | quant_config=quant_config, |
| 94 | prefix=f"{prefix}.q_b_proj", |
| 95 | ) |
| 96 | |
| 97 | self.kv_a_proj_with_mqa = ReplicatedLinear( |
| 98 | self.hidden_size, |
| 99 | self.kv_lora_rank + self.qk_rope_head_dim, |
| 100 | bias=False, |
| 101 | quant_config=quant_config, |
| 102 | prefix=f"{prefix}.kv_a_proj_with_mqa", |
| 103 | ) |
| 104 | self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) |
| 105 | self.kv_b_proj = ColumnParallelLinear( |
| 106 | self.kv_lora_rank, |
| 107 | self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), |
| 108 | bias=False, |
| 109 | quant_config=quant_config, |
| 110 | prefix=f"{prefix}.kv_b_proj", |
| 111 | ) |
| 112 | # O projection. |
| 113 | self.o_proj = RowParallelLinear( |
| 114 | self.num_heads * self.v_head_dim, |
| 115 | self.hidden_size, |
| 116 | bias=False, |
| 117 | quant_config=quant_config, |
| 118 | prefix=f"{prefix}.o_proj", |
| 119 | ) |
| 120 | |
| 121 | self.rotary_emb = get_rope( |
| 122 | self.qk_rope_head_dim, |
| 123 | max_position=max_position_embeddings, |
| 124 | rope_parameters=config.rope_parameters, |
| 125 | ) |
| 126 | self.attn = Attention( |
| 127 | self.num_local_heads, |
| 128 | self.qk_head_dim, |
| 129 | self.scaling, |
| 130 | num_kv_heads=self.num_local_heads, |
| 131 | cache_config=cache_config, |
| 132 | quant_config=quant_config, |
| 133 | prefix=f"{prefix}.attn", |
| 134 | ) |
| 135 | |
| 136 | def forward( |
| 137 | self, |
| 138 | positions: torch.Tensor, |
| 139 | hidden_states: torch.Tensor, |
| 140 | ) -> torch.Tensor: |
| 141 | q, _ = self.q_a_proj(hidden_states) |
| 142 | q = self.q_a_layernorm(q) |
| 143 | q, _ = self.q_b_proj(q) |
| 144 | q = q.view(-1, self.num_local_heads, self.qk_head_dim) |
| 145 | _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) |
| 146 | latent_cache, _ = self.kv_a_proj_with_mqa(hidden_states) |
| 147 | kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) |
| 148 | latent_cache = latent_cache.unsqueeze(1) |
| 149 | kv_a = self.kv_a_layernorm(kv_a) |
| 150 | kv, _ = self.kv_b_proj(kv_a) |
| 151 | kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) |
| 152 | k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) |
| 153 | |
| 154 | k_pe = latent_cache[:, :, self.kv_lora_rank :] |
| 155 | |
| 156 | q_pe, k_pe = self.rotary_emb( |
| 157 | positions, |
| 158 | q_pe.reshape(-1, self.num_local_heads * self.qk_rope_head_dim), |
| 159 | k_pe.reshape(-1, self.qk_rope_head_dim), |
| 160 | ) |
| 161 | q_pe = q_pe.view(-1, self.num_local_heads, self.qk_rope_head_dim) |
| 162 | k_pe = k_pe.view(-1, 1, self.qk_rope_head_dim) |
| 163 | |
| 164 | q[..., self.qk_nope_head_dim :] = q_pe |
| 165 | |
| 166 | k = torch.empty_like(q) |
| 167 | |
| 168 | k[..., : self.qk_nope_head_dim] = k_nope |
| 169 | k[..., self.qk_nope_head_dim :] = k_pe |
| 170 | |
| 171 | q = q.reshape(-1, self.num_local_heads * self.qk_head_dim) |
| 172 | k = k.view(-1, self.num_local_heads * self.qk_head_dim) |
| 173 | v = torch.nn.functional.pad( |
| 174 | v, [0, self.qk_head_dim - self.v_head_dim], value=0 |
| 175 | ).view(-1, self.num_local_heads * self.qk_head_dim) |
| 176 | |
| 177 | attn_output = self.attn(q, k, v) |
| 178 | attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[ |
| 179 | ..., : self.v_head_dim |
| 180 | ].reshape(-1, self.num_local_heads * self.v_head_dim) |
| 181 | |
| 182 | output, _ = self.o_proj(attn_output) |
| 183 | return output |
| 184 | |
| 185 | |
| 186 | class MiniCPM3DecoderLayer(MiniCPMDecoderLayer): |
| 187 | def _init_attn_block(self): |
| 188 | self.input_layernorm = RMSNorm( |
| 189 | self.config.hidden_size, eps=self.config.rms_norm_eps |
| 190 | ) |
| 191 | self.self_attn = MiniCPM3Attention( |
| 192 | config=self.config, |
| 193 | hidden_size=self.hidden_size, |
| 194 | num_heads=self.config.num_attention_heads, |
| 195 | qk_nope_head_dim=self.config.qk_nope_head_dim, |
| 196 | qk_rope_head_dim=self.config.qk_rope_head_dim, |
| 197 | v_head_dim=self.config.v_head_dim, |
| 198 | q_lora_rank=self.config.q_lora_rank, |
| 199 | kv_lora_rank=self.config.kv_lora_rank, |
| 200 | max_position_embeddings=self.max_position_embeddings, |
| 201 | cache_config=self.cache_config, |
| 202 | quant_config=self.quant_config, |
| 203 | prefix=f"{self.prefix}.self_attn", |
| 204 | ) |
| 205 | |
| 206 | |
| 207 | class MiniCPM3Model(MiniCPMModel): |
| 208 | def _init_layers( |
| 209 | self, |
| 210 | prefix: str, |
| 211 | config: PreTrainedConfig, |
| 212 | cache_config: CacheConfig | None, |
| 213 | quant_config: QuantizationConfig | None, |
| 214 | ): |
| 215 | self.start_layer, self.end_layer, self.layers = make_layers( |
| 216 | config.num_hidden_layers, |
| 217 | lambda prefix: MiniCPM3DecoderLayer( |
| 218 | config, cache_config, quant_config, prefix=prefix |
| 219 | ), |
| 220 | prefix=f"{prefix}.layers", |
| 221 | ) |
| 222 | |
| 223 | |
| 224 | class MiniCPM3ForCausalLM(MiniCPMForCausalLM): |
| 225 | packed_modules_mapping = { |
| 226 | "gate_up_proj": [ |
| 227 | "gate_proj", |
| 228 | "up_proj", |
| 229 | ], |
| 230 | } |
| 231 | |
| 232 | def _init_model(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 233 | return MiniCPM3Model(vllm_config=vllm_config, prefix=prefix) |