A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[Feature] minicpm eagle support (#18943)
Signed-off-by: huangyuxiang03 <huangyx0321@gmail.com> Co-authored-by: huangyuxiang03 <huangyx0321@gmail.com>
Shawn Huang committed
May 30, 2025 at 21:45 UTC
e1fadf11976ecc87ed60c9fc1bf8f81271c68c18
4 files changed
+399
-2
tests/models/registry.py
+5
@@ -434,6 +434,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
434
trust_remote_code=True,
435
speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
436
tokenizer="meta-llama/Llama-3.1-8B-Instruct"),
437
+ "EagleMiniCPMForCausalLM": _HfExamplesInfo("openbmb/MiniCPM-1B-sft-bf16",
438
+ trust_remote_code=True,
439
+ is_available_online=False,
440
+ speculative_model="openbmb/MiniCPM-2B-sft-bf16",
441
+ tokenizer="openbmb/MiniCPM-2B-sft-bf16"),
442
"MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",
443
trust_remote_code=True,
444
speculative_model="XiaomiMiMo/MiMo-7B-RL")
vllm/model_executor/models/minicpm.py
+3
-2
@@ -242,6 +242,7 @@ class MiniCPMAttention(nn.Module):
242
base=rope_theta,
243
rope_scaling=rope_scaling,
244
)
245
+
246
self.attn = Attention(self.num_heads,
247
self.head_dim,
248
self.scaling,
@@ -444,6 +445,7 @@ class MiniCPMModel(nn.Module):
445
for weight_name in ["w1", "w2", "w3"]
446
]
447
params_dict = dict(self.named_parameters())
448
+
449
loaded_params: set[str] = set()
450
for name, loaded_weight in weights:
451
if "rotary_emb.inv_freq" in name:
@@ -567,7 +569,7 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
569
inputs_embeds: Optional[torch.Tensor] = None,
570
) -> Union[torch.Tensor, IntermediateTensors]:
571
hidden_states = self.model(input_ids, positions, intermediate_tensors,
570
- inputs_embeds)
572
+ inputs_embeds) / self.scale_width
573
return hidden_states
574
575
def compute_logits(
@@ -575,7 +577,6 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
577
hidden_states: torch.Tensor,
578
sampling_metadata: SamplingMetadata,
579
) -> Optional[torch.Tensor]:
578
- hidden_states = hidden_states / self.scale_width
580
logits = self.logits_processor(self.lm_head, hidden_states,
581
sampling_metadata)
582
return logits
vllm/model_executor/models/minicpm_eagle.py
new
+390
@@ -0,0 +1,390 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+# Adapted from
4
+# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
5
+# Copyright 2023 The vLLM team.
6
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
7
+#
8
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
9
+# and OPT implementations in this library. It has been modified from its
10
+# original forms to accommodate minor architectural differences compared
11
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
12
+#
13
+# Licensed under the Apache License, Version 2.0 (the "License");
14
+# you may not use this file except in compliance with the License.
15
+# You may obtain a copy of the License at
16
+#
17
+# http://www.apache.org/licenses/LICENSE-2.0
18
+#
19
+# Unless required by applicable law or agreed to in writing, software
20
+# distributed under the License is distributed on an "AS IS" BASIS,
21
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22
+# See the License for the specific language governing permissions and
23
+# limitations under the License.
24
+"""Inference-only EagleMiniCPM model compatible with HuggingFace weights."""
25
+import math
26
+from collections.abc import Iterable
27
+from typing import Optional, Union
28
+
29
+import torch
30
+from torch import nn
31
+from transformers import PretrainedConfig
32
+
33
+from vllm.compilation.decorators import support_torch_compile
34
+from vllm.config import CacheConfig, VllmConfig
35
+from vllm.model_executor.layers.layernorm import RMSNorm
36
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
+from vllm.model_executor.layers.quantization import QuantizationConfig
38
+from vllm.model_executor.layers.vocab_parallel_embedding import (
39
+ DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
40
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
41
+from vllm.model_executor.sampling_metadata import SamplingMetadata
42
+from vllm.sequence import IntermediateTensors
43
+
44
+from .interfaces import SupportsLoRA, SupportsPP
45
+from .minicpm import MiniCPMAttention as EagleMiniCPMAttention
46
+from .minicpm import MiniCPMMLP as EagleMiniCPMMLP
47
+from .minicpm import MiniCPMMoE as EagleMiniCPMMoE
48
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
49
+ make_empty_intermediate_tensors_factory, maybe_prefix)
50
+
51
+
52
+class EagleMiniCPMDecoderLayer(nn.Module):
53
+
54
+ def __init__(
55
+ self,
56
+ config: PretrainedConfig,
57
+ cache_config: Optional[CacheConfig] = None,
58
+ quant_config: Optional[QuantizationConfig] = None,
59
+ prefix: str = "",
60
+ ) -> None:
61
+ super().__init__()
62
+ self.config = config
63
+ self.cache_config = cache_config
64
+ self.quant_config = quant_config
65
+ self.hidden_size = config.hidden_size
66
+ self.rope_theta = getattr(config, "rope_theta", 10000)
67
+ self.rope_scaling = getattr(config, "rope_scaling", None)
68
+ self.max_position_embeddings = getattr(config,
69
+ "max_position_embeddings", 8192)
70
+ self.prefix = prefix
71
+ self._init_attn_block()
72
+ self._init_ffn_block()
73
+
74
+ def _init_attn_block(self):
75
+ self.input_layernorm = RMSNorm(self.config.hidden_size,
76
+ eps=self.config.rms_norm_eps)
77
+ self.self_attn = EagleMiniCPMAttention(
78
+ hidden_size=self.hidden_size,
79
+ num_heads=self.config.num_attention_heads,
80
+ num_kv_heads=self.config.num_key_value_heads,
81
+ rope_theta=self.rope_theta,
82
+ rope_scaling=self.rope_scaling,
83
+ max_position_embeddings=self.max_position_embeddings,
84
+ cache_config=self.cache_config,
85
+ quant_config=self.quant_config,
86
+ prefix=f"{self.prefix}.self_attn",
87
+ )
88
+
89
+ def _init_ffn_block(self):
90
+ self.post_attention_layernorm = RMSNorm(self.config.hidden_size,
91
+ eps=self.config.rms_norm_eps)
92
+ self.num_experts = getattr(self.config, "num_experts", 0)
93
+ if self.num_experts == 0:
94
+ self.mlp = EagleMiniCPMMLP(
95
+ hidden_size=self.hidden_size,
96
+ intermediate_size=self.config.intermediate_size,
97
+ hidden_act=self.config.hidden_act,
98
+ hidden_act_param=getattr(self.config, "hidden_act_param", 0.),
99
+ quant_config=self.quant_config,
100
+ )
101
+ else:
102
+ self.mlp = EagleMiniCPMMoE(
103
+ num_experts=self.config.num_experts,
104
+ top_k=self.config.num_experts_per_tok,
105
+ hidden_size=self.config.hidden_size,
106
+ intermediate_size=self.config.intermediate_size)
107
+
108
+ def forward(
109
+ self,
110
+ positions: torch.Tensor,
111
+ hidden_states: torch.Tensor,
112
+ residual: Optional[torch.Tensor],
113
+ ) -> tuple[torch.Tensor, torch.Tensor]:
114
+ # Self Attention
115
+ residual = hidden_states
116
+ hidden_states = self.input_layernorm(hidden_states)
117
+ hidden_states = self.self_attn(
118
+ positions=positions,
119
+ hidden_states=hidden_states,
120
+ )
121
+ hidden_states = residual + hidden_states * \
122
+ (self.config.scale_depth / math.sqrt(self.config.mup_denominator))
123
+
124
+ # Fully Connected
125
+ residual = hidden_states
126
+ hidden_states = self.post_attention_layernorm(hidden_states)
127
+ hidden_states = self.mlp(hidden_states)
128
+ hidden_states = residual + hidden_states * \
129
+ (self.config.scale_depth / math.sqrt(self.config.mup_denominator))
130
+
131
+ return hidden_states, None
132
+
133
+
134
+@support_torch_compile
135
+class EagleMiniCPMModel(nn.Module):
136
+
137
+ def __init__(self,
138
+ *,
139
+ vllm_config: VllmConfig,
140
+ prefix: str = "",
141
+ start_layer: int = 0):
142
+ super().__init__()
143
+
144
+ config = vllm_config.speculative_config.draft_model_config.hf_config
145
+ cache_config = vllm_config.cache_config
146
+ quant_config = vllm_config.quant_config
147
+ lora_config = vllm_config.lora_config
148
+
149
+ self.config = config
150
+ self.cache_config = cache_config
151
+ self.quant_config = quant_config
152
+ lora_vocab = (lora_config.lora_extra_vocab_size *
153
+ (lora_config.max_loras or 1)) if lora_config else 0
154
+ self.vocab_size = config.vocab_size + lora_vocab
155
+ self.org_vocab_size = config.vocab_size
156
+ self.fc = torch.nn.Linear(self.config.hidden_size * 2,
157
+ self.config.hidden_size,
158
+ bias=False)
159
+ self.input_norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
160
+ self.input_norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
161
+ self.embed_tokens = VocabParallelEmbedding(
162
+ self.vocab_size,
163
+ config.hidden_size,
164
+ org_num_embeddings=config.vocab_size,
165
+ )
166
+ self.num_experts = getattr(self.config, "num_experts", 0)
167
+ self._init_layers(prefix, config, cache_config, quant_config,
168
+ start_layer)
169
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
170
+ self.make_empty_intermediate_tensors = (
171
+ make_empty_intermediate_tensors_factory(
172
+ ["hidden_states", "residual"], self.config.hidden_size))
173
+
174
+ def _init_layers(
175
+ self,
176
+ prefix: str,
177
+ config: PretrainedConfig,
178
+ cache_config: Optional[CacheConfig],
179
+ quant_config: Optional[QuantizationConfig],
180
+ start_layer: int,
181
+ ):
182
+ self.eagle_layers = nn.ModuleList([
183
+ EagleMiniCPMDecoderLayer(
184
+ config,
185
+ cache_config,
186
+ quant_config,
187
+ f"{prefix}.eagle_layers.{i + start_layer}",
188
+ ) for i in range(self.config.num_hidden_layers)
189
+ ])
190
+
191
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
192
+ embedding = self.embed_tokens(input_ids)
193
+ return embedding * self.config.scale_emb
194
+
195
+ def forward(
196
+ self,
197
+ input_ids: torch.Tensor,
198
+ positions: torch.Tensor,
199
+ hidden_states: torch.Tensor,
200
+ ) -> Union[torch.Tensor, IntermediateTensors]:
201
+ input_embeds = self.get_input_embeddings(input_ids)
202
+ input_embeds = self.input_norm1(input_embeds)
203
+ hidden_states = self.input_norm2(hidden_states)
204
+
205
+ hidden_states = self.fc(
206
+ torch.cat((input_embeds, hidden_states), dim=-1))
207
+ residual = None
208
+ for layer in self.eagle_layers:
209
+ hidden_states, residual = layer(
210
+ positions,
211
+ hidden_states,
212
+ residual,
213
+ )
214
+
215
+ return hidden_states, hidden_states
216
+
217
+ def load_weights(self, weights: Iterable[tuple[str,
218
+ torch.Tensor]]) -> set[str]:
219
+ stacked_params_mapping = [
220
+ # (param_name, shard_name, shard_id)
221
+ ("qkv_proj", "q_proj", "q"),
222
+ ("qkv_proj", "k_proj", "k"),
223
+ ("qkv_proj", "v_proj", "v"),
224
+ ("gate_up_proj", "gate_proj", 0),
225
+ ("gate_up_proj", "up_proj", 1),
226
+ ]
227
+ expert_params_mapping = [
228
+ # (param_name, weight_name, expert_id)
229
+ ("ws" if weight_name in ["w1", "w3"] else "w2s",
230
+ f"experts.{expert_id}.{weight_name}.weight", expert_id)
231
+ for expert_id in range(self.num_experts)
232
+ for weight_name in ["w1", "w2", "w3"]
233
+ ]
234
+ params_dict = dict(self.named_parameters())
235
+
236
+ loaded_params: set[str] = set()
237
+ for name, loaded_weight in weights:
238
+ if "rotary_emb.inv_freq" in name:
239
+ continue
240
+ if ("rotary_emb.cos_cached" in name
241
+ or "rotary_emb.sin_cached" in name):
242
+ # Models trained using ColossalAI may include these tensors in
243
+ # the checkpoint. Skip them.
244
+ continue
245
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
246
+ if weight_name not in name:
247
+ continue
248
+ name = name.replace(weight_name, param_name)
249
+ # Skip loading extra bias for GPTQ models.
250
+ if name.endswith(".bias") and name not in params_dict:
251
+ continue
252
+ if is_pp_missing_parameter(name, self):
253
+ continue
254
+ param = params_dict[name]
255
+ weight_loader = param.weight_loader
256
+ weight_loader(param, loaded_weight, shard_id)
257
+ break
258
+ else:
259
+ for param_name, weight_name, expert_id in expert_params_mapping:
260
+ if weight_name not in name:
261
+ continue
262
+ name = name.replace(weight_name, param_name)
263
+ if is_pp_missing_parameter(name, self):
264
+ continue
265
+ param = params_dict[name]
266
+ weight_loader = param.weight_loader
267
+ weight_loader(param,
268
+ loaded_weight,
269
+ weight_name,
270
+ expert_id=expert_id)
271
+ break
272
+ else:
273
+ # Skip loading extra bias for GPTQ models.
274
+ if name.endswith(".bias") and name not in params_dict:
275
+ continue
276
+ if is_pp_missing_parameter(name, self):
277
+ continue
278
+ param = params_dict[name]
279
+ weight_loader = getattr(param, "weight_loader",
280
+ default_weight_loader)
281
+
282
+ weight_loader(param, loaded_weight)
283
+ loaded_params.add(name)
284
+ return loaded_params
285
+
286
+
287
+class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
288
+ packed_modules_mapping = {
289
+ "qkv_proj": [
290
+ "q_proj",
291
+ "k_proj",
292
+ "v_proj",
293
+ ],
294
+ "gate_up_proj": [
295
+ "gate_proj",
296
+ "up_proj",
297
+ ],
298
+ }
299
+
300
+ # LoRA specific attributes
301
+ embedding_modules = {
302
+ "embed_tokens": "input_embeddings",
303
+ "lm_head": "output_embeddings",
304
+ }
305
+ embedding_padding_modules = ["lm_head"]
306
+
307
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
308
+ super().__init__()
309
+ config = vllm_config.speculative_config.draft_model_config.hf_config
310
+ cache_config = vllm_config.cache_config
311
+ quant_config = vllm_config.quant_config
312
+ lora_config = vllm_config.lora_config
313
+
314
+ self.prefix = prefix
315
+ self.vllm_config = vllm_config
316
+ self.config = config
317
+ self.lora_config = lora_config
318
+ self.cache_config = cache_config
319
+ self.quant_config = quant_config
320
+
321
+ target_layer_num = vllm_config.model_config.get_num_layers(
322
+ vllm_config.parallel_config)
323
+
324
+ self.model = self._init_model(vllm_config=vllm_config,
325
+ prefix=maybe_prefix(prefix, "model"),
326
+ start_layer=target_layer_num)
327
+
328
+ unpadded_vocab_size = config.vocab_size
329
+ if lora_config:
330
+ unpadded_vocab_size += lora_config.lora_extra_vocab_size
331
+ self.lm_head = ParallelLMHead(
332
+ unpadded_vocab_size,
333
+ config.hidden_size,
334
+ org_num_embeddings=config.vocab_size,
335
+ padding_size=DEFAULT_VOCAB_PADDING_SIZE
336
+ # We need bigger padding if using lora for kernel
337
+ # compatibility
338
+ if not lora_config else lora_config.lora_vocab_padding_size,
339
+ quant_config=quant_config,
340
+ )
341
+ if config.tie_word_embeddings:
342
+ self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
343
+ self.scale_width = self.config.hidden_size / self.config.dim_model_base
344
+
345
+ self.logits_processor = LogitsProcessor(unpadded_vocab_size,
346
+ config.vocab_size)
347
+ self.make_empty_intermediate_tensors = (
348
+ self.model.make_empty_intermediate_tensors)
349
+
350
+ def _init_model(self,
351
+ *,
352
+ vllm_config: VllmConfig,
353
+ prefix: str = "",
354
+ start_layer: int = 0):
355
+ return EagleMiniCPMModel(vllm_config=vllm_config,
356
+ prefix=prefix,
357
+ start_layer=start_layer)
358
+
359
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
360
+ return self.model.get_input_embeddings(input_ids)
361
+
362
+ def forward(
363
+ self,
364
+ input_ids: torch.Tensor,
365
+ positions: torch.Tensor,
366
+ hidden_states: torch.Tensor,
367
+ ) -> tuple[torch.Tensor, torch.Tensor]:
368
+ hidden_states, hidden_states2 = self.model(input_ids, positions,
369
+ hidden_states)
370
+ hidden_states = hidden_states / self.scale_width
371
+ hidden_states2 = hidden_states2 / self.scale_width
372
+ return hidden_states, hidden_states2
373
+
374
+ def compute_logits(
375
+ self,
376
+ hidden_states: torch.Tensor,
377
+ sampling_metadata: SamplingMetadata,
378
+ ) -> Optional[torch.Tensor]:
379
+ logits = self.logits_processor(self.lm_head, hidden_states,
380
+ sampling_metadata)
381
+ return logits
382
+
383
+ def load_weights(self, weights: Iterable[tuple[str,
384
+ torch.Tensor]]) -> set[str]:
385
+ loader = AutoWeightsLoader(
386
+ self,
387
+ skip_prefixes=(["lm_head."]
388
+ if self.config.tie_word_embeddings else None),
389
+ )
390
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -223,6 +223,7 @@ _SPECULATIVE_DECODING_MODELS = {
223
"MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
224
"EAGLEModel": ("eagle", "EAGLE"),
225
"EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
226
+ "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
227
"Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
228
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
229
"MedusaModel": ("medusa", "Medusa"),