@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c5e3454e5a
Fork 0

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 5, 12:15 UTC

[Model] Add support for BharatGen's Param2MoE model (#38000)

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

bhargav-patel-29 committed Apr 6, 2026 at 13:49 UTC c5e3454e5adf063d5af75140c36a7f900a9e4c4c
4 files changed +906
docs/models/supported_models.md
+1
@@ -457,6 +457,7 @@ th {
457 | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
458 | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
459 | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
460 +| `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ |
461 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
462 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
463 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |
tests/models/registry.py
+4
@@ -461,6 +461,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
461 trust_remote_code=True,
462 is_available_online=False,
463 ),
464 + "Param2MoEForCausalLM": _HfExamplesInfo(
465 + "bharatgenai/Param2-17B-A2.4B-Thinking",
466 + trust_remote_code=True,
467 + ),
468 "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
469 "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
470 "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
vllm/model_executor/models/param2moe.py new
+900
@@ -0,0 +1,900 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Copyright 2026 BharatGen AI team. All rights reserved.
5 +#
6 +# This code has been modified to accommodate Param2MoE's GQA-based MoE architecture.
7 +#
8 +# Licensed under the Apache License, Version 2.0 (the "License");
9 +# you may not use this file except in compliance with the License.
10 +# You may obtain a copy of the License at
11 +#
12 +# http://www.apache.org/licenses/LICENSE-2.0
13 +#
14 +# Unless required by applicable law or agreed to in writing, software
15 +# distributed under the License is distributed on an "AS IS" BASIS,
16 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 +# See the License for the specific language governing permissions and
18 +# limitations under the License.
19 +# limitations under the License.
20 +from __future__ import annotations
21 +
22 +from collections.abc import Iterable, Iterator
23 +from itertools import islice
24 +
25 +import torch
26 +import torch.nn.functional as F
27 +from torch import nn
28 +
29 +from vllm.config import CacheConfig, VllmConfig
30 +from vllm.distributed import (
31 + get_pp_group,
32 + get_tensor_model_parallel_world_size,
33 +)
34 +from vllm.model_executor.layers.activation import SiluAndMul
35 +from vllm.model_executor.layers.attention import Attention
36 +from vllm.model_executor.layers.fused_moe import SharedFusedMoE
37 +from vllm.model_executor.layers.layernorm import RMSNorm
38 +from vllm.model_executor.layers.linear import (
39 + MergedColumnParallelLinear,
40 + QKVParallelLinear,
41 + RowParallelLinear,
42 +)
43 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
44 +from vllm.model_executor.layers.quantization import QuantizationConfig
45 +from vllm.model_executor.layers.rotary_embedding import get_rope
46 +from vllm.model_executor.layers.vocab_parallel_embedding import (
47 + ParallelLMHead,
48 + VocabParallelEmbedding,
49 +)
50 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51 +from vllm.sequence import IntermediateTensors
52 +
53 +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
54 +from .utils import (
55 + AutoWeightsLoader,
56 + PPMissingLayer,
57 + is_pp_missing_parameter,
58 + make_empty_intermediate_tensors_factory,
59 + make_layers,
60 + maybe_prefix,
61 +)
62 +
63 +
64 +def _is_expert_bias_name(name: str) -> bool:
65 + """True when the weight is the MoE router's per-expert score bias."""
66 + return name.endswith(".mlp.gate.expert_bias")
67 +
68 +
69 +def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
70 + if t.numel() == 0:
71 + return t
72 + return t - t.mean()
73 +
74 +
75 +def _rename_and_normalize_weights(
76 + weights: Iterable[tuple[str, torch.Tensor]],
77 +) -> Iterator[tuple[str, torch.Tensor]]:
78 + """
79 + Translate HuggingFace Param2MoE weight names to vLLM internal names
80 + and zero-mean the expert-bias tensor so the router stays balanced.
81 +
82 + Mapping table (HF → vLLM):
83 + model.word_embeddings.* → model.embed_tokens.*
84 + *.attention.query_key_value.* → *.self_attn.qkv_proj.*
85 + *.attention.dense.* → *.self_attn.o_proj.*
86 + *.attention.query_layernorm.* → *.self_attn.q_layernorm.*
87 + *.attention.key_layernorm.* → *.self_attn.k_layernorm.*
88 + *.mlp.gate.expert_bias → *.mlp.gate.e_score_correction_bias
89 + (also zero-meant for load balance)
90 + """
91 + for name, w in weights:
92 + # Embedding table
93 + name = name.replace("model.word_embeddings.", "model.embed_tokens.")
94 + # Fused QKV projection (HF: query_key_value → vLLM: qkv_proj)
95 + name = name.replace(".attention.query_key_value.", ".self_attn.qkv_proj.")
96 + # Output projection (HF: dense → vLLM: o_proj)
97 + name = name.replace(".attention.dense.", ".self_attn.o_proj.")
98 + # Per-head query norm
99 + name = name.replace(".attention.query_layernorm.", ".self_attn.q_layernorm.")
100 + # Per-head key norm
101 + name = name.replace(".attention.key_layernorm.", ".self_attn.k_layernorm.")
102 + # Catch any remaining .attention. → .self_attn. prefixes
103 + # (e.g. future bias params on the projection layers)
104 + name = name.replace(".attention.", ".self_attn.")
105 +
106 + # Expert-score bias: rename + zero-mean
107 + if name.endswith(".mlp.gate.expert_bias"):
108 + name = name.replace(
109 + ".mlp.gate.expert_bias",
110 + ".mlp.gate.e_score_correction_bias",
111 + )
112 + w = _zero_mean_tensor(w)
113 +
114 + yield name, w
115 +
116 +
117 +class Param2MoEAttention(nn.Module):
118 + """
119 + Grouped-Query Attention (GQA) for Param2MoE.
120 +
121 + Notable differences from a vanilla GQA layer:
122 + * The checkpoint fuses Q, K, V into a single ``query_key_value`` weight.
123 + vLLM receives it already renamed to ``qkv_proj`` by the weight-name
124 + translator and splits it during ``load_weights``.
125 + * Optional per-head RMS norms on Q and K (``use_qk_norm=True``).
126 + """
127 +
128 + def __init__(
129 + self,
130 + config,
131 + cache_config: CacheConfig | None = None,
132 + quant_config: QuantizationConfig | None = None,
133 + prefix: str = "",
134 + ) -> None:
135 + super().__init__()
136 +
137 + self.hidden_size = config.hidden_size
138 + self.num_heads = config.num_attention_heads
139 + self.num_kv_heads = config.num_key_value_heads
140 + self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)
141 + self.use_qk_norm: bool = getattr(config, "use_qk_norm", False)
142 +
143 + tp_size = get_tensor_model_parallel_world_size()
144 + assert self.num_heads % tp_size == 0, (
145 + f"num_attention_heads ({self.num_heads}) must be divisible "
146 + f"by tensor-parallel world size ({tp_size})."
147 + )
148 + assert self.num_kv_heads % tp_size == 0, (
149 + f"num_key_value_heads ({self.num_kv_heads}) must be divisible "
150 + f"by tensor-parallel world size ({tp_size})."
151 + )
152 + self.num_local_heads = self.num_heads // tp_size
153 + self.num_local_kv_heads = self.num_kv_heads // tp_size
154 +
155 + # Sizes after TP split (used in forward to split qkv output)
156 + self.q_size_local = self.num_local_heads * self.head_dim
157 + self.kv_size_local = self.num_local_kv_heads * self.head_dim
158 +
159 + self.scaling = self.head_dim**-0.5
160 +
161 + self.qkv_proj = QKVParallelLinear(
162 + hidden_size=self.hidden_size,
163 + head_size=self.head_dim,
164 + total_num_heads=self.num_heads,
165 + total_num_kv_heads=self.num_kv_heads,
166 + bias=getattr(config, "use_qkv_bias", False),
167 + quant_config=quant_config,
168 + prefix=f"{prefix}.qkv_proj",
169 + )
170 +
171 + self.o_proj = RowParallelLinear(
172 + input_size=self.num_heads * self.head_dim,
173 + output_size=self.hidden_size,
174 + bias=getattr(config, "use_bias", False),
175 + quant_config=quant_config,
176 + prefix=f"{prefix}.o_proj",
177 + )
178 +
179 + if self.use_qk_norm:
180 + self.q_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
181 + self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
182 +
183 + # `partial_rotary_factor` defaults to 1.0 (full RoPE) if not in config
184 + partial_rotary_factor: float = getattr(config, "partial_rotary_factor", 1.0)
185 + rope_dim = int(self.head_dim * partial_rotary_factor)
186 +
187 + rope_parameters: dict = {
188 + "rope_type": "default",
189 + "base": config.rope_theta,
190 + }
191 + if config.rope_scaling is not None:
192 + rope_parameters.update(config.rope_scaling)
193 + # Normalise key: some checkpoints use "type", vLLM wants "rope_type"
194 + if "type" in rope_parameters and "rope_type" not in rope_parameters:
195 + rope_parameters["rope_type"] = rope_parameters.pop("type")
196 +
197 + self.rotary_emb = get_rope(
198 + rope_dim,
199 + max_position=config.max_position_embeddings,
200 + rope_parameters=rope_parameters,
201 + is_neox_style=True,
202 + )
203 +
204 + self.attn = Attention(
205 + num_heads=self.num_heads,
206 + head_size=self.head_dim,
207 + scale=self.scaling,
208 + num_kv_heads=self.num_kv_heads,
209 + cache_config=cache_config,
210 + quant_config=quant_config,
211 + prefix=f"{prefix}.attn",
212 + )
213 +
214 + def forward(
215 + self,
216 + positions: torch.Tensor,
217 + hidden_states: torch.Tensor,
218 + ) -> torch.Tensor:
219 + # 1. Fused QKV projection → split into local Q / K / V
220 + qkv, _ = self.qkv_proj(hidden_states)
221 + q, k, v = qkv.split(
222 + [self.q_size_local, self.kv_size_local, self.kv_size_local],
223 + dim=-1,
224 + )
225 +
226 + # 2. Optional per-head QK norms
227 + # Reshape to (T, num_local_heads, head_dim), norm, reshape back.
228 + if self.use_qk_norm:
229 + T = q.shape[0]
230 + q = self.q_layernorm(q.view(T, self.num_local_heads, self.head_dim)).view(
231 + T, self.q_size_local
232 + )
233 + k = self.k_layernorm(
234 + k.view(T, self.num_local_kv_heads, self.head_dim)
235 + ).view(T, self.kv_size_local)
236 +
237 + # 3. Rotary position embeddings
238 + q, k = self.rotary_emb(positions, q, k)
239 +
240 + # 4. Paged attention
241 + attn_output = self.attn(q, k, v)
242 +
243 + # 5. Output projection
244 + output, _ = self.o_proj(attn_output)
245 + return output
246 +
247 +
248 +class Param2MoEMLP(nn.Module):
249 + """SwiGLU feed-forward block used for dense layers."""
250 +
251 + def __init__(
252 + self,
253 + intermediate_size: int,
254 + config,
255 + quant_config: QuantizationConfig | None = None,
256 + reduce_results: bool = True,
257 + prefix: str = "",
258 + ) -> None:
259 + super().__init__()
260 +
261 + self.gate_up_proj = MergedColumnParallelLinear(
262 + input_size=config.hidden_size,
263 + output_sizes=[intermediate_size, intermediate_size],
264 + bias=False,
265 + quant_config=quant_config,
266 + prefix=f"{prefix}.gate_up_proj",
267 + )
268 + self.down_proj = RowParallelLinear(
269 + input_size=intermediate_size,
270 + output_size=config.hidden_size,
271 + bias=False,
272 + quant_config=quant_config,
273 + reduce_results=reduce_results,
274 + prefix=f"{prefix}.down_proj",
275 + )
276 + self.act_fn = SiluAndMul()
277 +
278 + def forward(self, x: torch.Tensor) -> torch.Tensor:
279 + gate_up, _ = self.gate_up_proj(x)
280 + x = self.act_fn(gate_up)
281 + x, _ = self.down_proj(x)
282 + return x
283 +
284 +
285 +class Param2MoEMoEBlock(nn.Module):
286 + """
287 + Mixture-of-Experts block for Param2MoE.
288 +
289 + Routing:
290 + * Sigmoid scoring (config.score_function = "sigmoid")
291 + * Grouped top-k (n_group, topk_group)
292 + * Per-expert bias (gate.expert_bias → e_score_correction_bias)
293 + * routed_scaling_factor normalisation
294 +
295 + One set of shared (always-active) experts is added on top.
296 + """
297 +
298 + def __init__(
299 + self,
300 + config,
301 + quant_config: QuantizationConfig | None = None,
302 + prefix: str = "",
303 + ) -> None:
304 + super().__init__()
305 +
306 + self.config = config
307 + self.tp_size = get_tensor_model_parallel_world_size()
308 + self.hidden_size = config.hidden_size
309 +
310 + self.num_experts: int = config.num_experts
311 + self.top_k: int = config.num_experts_per_tok
312 + self.routed_scaling_factor: float = getattr(
313 + config, "routed_scaling_factor", 1.0
314 + )
315 +
316 + self.n_group: int | None = getattr(config, "n_group", None)
317 + self.topk_group: int | None = getattr(config, "topk_group", None)
318 + self.use_grouped_topk: bool = (
319 + self.n_group is not None and self.topk_group is not None
320 + )
321 +
322 + self.norm_expert_prob: bool = getattr(config, "norm_topk_prob", True)
323 + self.score_function: str = getattr(config, "score_function", "sigmoid")
324 +
325 + self.gate = nn.Linear(
326 + self.hidden_size,
327 + self.num_experts,
328 + bias=False,
329 + )
330 +
331 + if getattr(config, "moe_router_enable_expert_bias", True):
332 + self.gate.e_score_correction_bias = nn.Parameter(
333 + torch.zeros(self.num_experts, dtype=torch.float32)
334 + )
335 + else:
336 + self.gate.e_score_correction_bias = None # type: ignore[assignment]
337 +
338 + self.num_shared_experts: int = getattr(config, "num_shared_experts", 1)
339 + if self.num_shared_experts > 0:
340 + # If moe_shared_expert_intermediate_size is present in the config
341 + # it already encodes the TOTAL intermediate size across all shared
342 + # experts (i.e. it equals moe_intermediate_size * num_shared_experts).
343 + # Do NOT multiply again. Fall back to computing the product only
344 + # when the dedicated field is absent.
345 + if (
346 + hasattr(config, "moe_shared_expert_intermediate_size")
347 + and config.moe_shared_expert_intermediate_size is not None
348 + ):
349 + shared_int: int = config.moe_shared_expert_intermediate_size
350 + else:
351 + shared_int = config.moe_intermediate_size * self.num_shared_experts
352 + self.shared_experts = Param2MoEMLP(
353 + intermediate_size=shared_int,
354 + config=config,
355 + quant_config=quant_config,
356 + reduce_results=False,
357 + prefix=f"{prefix}.shared_experts",
358 + )
359 + else:
360 + self.shared_experts = None # type: ignore[assignment]
361 +
362 + self.experts = SharedFusedMoE(
363 + shared_experts=self.shared_experts,
364 + num_experts=self.num_experts,
365 + top_k=self.top_k,
366 + hidden_size=self.hidden_size,
367 + intermediate_size=config.moe_intermediate_size,
368 + reduce_results=False,
369 + renormalize=self.norm_expert_prob,
370 + quant_config=quant_config,
371 + prefix=f"{prefix}.experts",
372 + scoring_func=self.score_function,
373 + e_score_correction_bias=self.gate.e_score_correction_bias,
374 + num_expert_group=self.n_group,
375 + topk_group=self.topk_group,
376 + use_grouped_topk=self.use_grouped_topk,
377 + routed_scaling_factor=self.routed_scaling_factor,
378 + )
379 +
380 + def maybe_get_fused_moe(self) -> SharedFusedMoE:
381 + return self.experts
382 +
383 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
384 + num_tokens, hidden_dim = hidden_states.shape
385 + hidden_states = hidden_states.view(-1, hidden_dim)
386 +
387 + # Router: both input and weight must be float32 for numerical
388 + # stability (mirrors the original Param2MoEGate behaviour).
389 + # The gate nn.Linear weight lives in the model dtype (bfloat16),
390 + # so we must cast both explicitly via F.linear instead of calling
391 + # self.gate() which would hit a dtype mismatch.
392 + router_logits = F.linear(
393 + hidden_states.float(),
394 + self.gate.weight.float(),
395 + ).to(hidden_states.dtype)
396 +
397 + final_hidden = self.experts(
398 + hidden_states=hidden_states,
399 + router_logits=router_logits,
400 + )
401 +
402 + if self.shared_experts is not None:
403 + shared_output, expert_output = final_hidden
404 + else:
405 + shared_output, expert_output = None, final_hidden
406 +
407 + if shared_output is not None:
408 + expert_output = expert_output + shared_output
409 +
410 + if self.tp_size > 1:
411 + expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
412 + expert_output
413 + )
414 +
415 + return expert_output.view(num_tokens, hidden_dim)
416 +
417 +
418 +class Param2MoEDecoderLayer(nn.Module):
419 + """
420 + Single transformer decoder block.
421 +
422 + Dense for the first ``first_k_dense_replace`` layers; MoE thereafter.
423 + """
424 +
425 + def __init__(
426 + self,
427 + vllm_config: VllmConfig,
428 + prefix: str = "",
429 + ) -> None:
430 + super().__init__()
431 +
432 + config = vllm_config.model_config.hf_config
433 + cache_config = vllm_config.cache_config
434 + quant_config = vllm_config.quant_config
435 +
436 + hidden_size = config.hidden_size
437 + # Derive the layer index from the prefix (e.g. "model.layers.3")
438 + layer_idx = int(prefix.split(".")[-1])
439 +
440 + self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
441 + self.self_attn = Param2MoEAttention(
442 + config=config,
443 + cache_config=cache_config,
444 + quant_config=quant_config,
445 + prefix=f"{prefix}.self_attn",
446 + )
447 + self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
448 +
449 + first_k_dense: int = getattr(config, "first_k_dense_replace", 1)
450 + is_moe_layer = config.num_experts is not None and layer_idx >= first_k_dense
451 +
452 + if is_moe_layer:
453 + self.mlp = Param2MoEMoEBlock(
454 + config=config,
455 + quant_config=quant_config,
456 + prefix=f"{prefix}.mlp",
457 + )
458 + else:
459 + self.mlp = Param2MoEMLP( # type: ignore[assignment]
460 + intermediate_size=config.intermediate_size,
461 + config=config,
462 + quant_config=quant_config,
463 + reduce_results=True,
464 + prefix=f"{prefix}.mlp",
465 + )
466 +
467 + def forward(
468 + self,
469 + hidden_states: torch.Tensor,
470 + positions: torch.Tensor,
471 + residual: torch.Tensor | None,
472 + ) -> tuple[torch.Tensor, torch.Tensor]:
473 + # Pre-norm + attention
474 + if residual is None:
475 + residual = hidden_states
476 + hidden_states = self.input_layernorm(hidden_states)
477 + else:
478 + hidden_states, residual = self.input_layernorm(hidden_states, residual)
479 +
480 + hidden_states = self.self_attn(
481 + positions=positions,
482 + hidden_states=hidden_states,
483 + )
484 +
485 + # Pre-norm + MLP
486 + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
487 + hidden_states = self.mlp(hidden_states)
488 + return hidden_states, residual
489 +
490 +
491 +class Param2MoEModel(nn.Module):
492 + def __init__(
493 + self,
494 + *,
495 + vllm_config: VllmConfig,
496 + prefix: str = "",
497 + ) -> None:
498 + super().__init__()
499 +
500 + config = vllm_config.model_config.hf_config
501 + quant_config = vllm_config.quant_config
502 +
503 + self.config = config
504 + self.vocab_size = config.vocab_size
505 + self.embed_dim = config.hidden_size
506 + self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False)
507 +
508 + # Embedding (HF name: word_embeddings → vLLM name: embed_tokens)
509 + if get_pp_group().is_first_rank or (
510 + self.tie_word_embeddings and get_pp_group().is_last_rank
511 + ):
512 + self.embed_tokens = VocabParallelEmbedding(
513 + self.vocab_size,
514 + self.embed_dim,
515 + quant_config=quant_config,
516 + prefix=f"{prefix}.embed_tokens",
517 + )
518 + else:
519 + self.embed_tokens = PPMissingLayer()
520 +
521 + self.start_layer, self.end_layer, self.layers = make_layers(
522 + config.num_hidden_layers,
523 + lambda prefix: Param2MoEDecoderLayer(
524 + vllm_config=vllm_config,
525 + prefix=prefix,
526 + ),
527 + prefix=f"{prefix}.layers",
528 + )
529 +
530 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
531 + ["hidden_states", "residual"], config.hidden_size
532 + )
533 +
534 + if get_pp_group().is_last_rank:
535 + self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
536 + else:
537 + self.norm = PPMissingLayer()
538 +
539 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
540 + return self.embed_tokens(input_ids)
541 +
542 + def forward(
543 + self,
544 + input_ids: torch.Tensor,
545 + positions: torch.Tensor,
546 + intermediate_tensors: IntermediateTensors | None,
547 + inputs_embeds: torch.Tensor | None = None,
548 + ) -> torch.Tensor | IntermediateTensors:
549 + if get_pp_group().is_first_rank:
550 + if inputs_embeds is not None:
551 + hidden_states = inputs_embeds
552 + else:
553 + hidden_states = self.embed_input_ids(input_ids)
554 + residual = None
555 + else:
556 + assert intermediate_tensors is not None
557 + hidden_states = intermediate_tensors["hidden_states"]
558 + residual = intermediate_tensors["residual"]
559 +
560 + for layer in islice(self.layers, self.start_layer, self.end_layer):
561 + hidden_states, residual = layer(hidden_states, positions, residual)
562 +
563 + if not get_pp_group().is_last_rank:
564 + return IntermediateTensors(
565 + {"hidden_states": hidden_states, "residual": residual}
566 + )
567 +
568 + if residual is None:
569 + hidden_states = self.norm(hidden_states)
570 + else:
571 + hidden_states, _ = self.norm(hidden_states, residual)
572 + return hidden_states
573 +
574 + def load_weights(
575 + self,
576 + weights: Iterable[tuple[str, torch.Tensor]],
577 + ) -> set[str]:
578 + """
579 + Custom weight loader for the inner Param2MoEModel.
580 +
581 + Receives weights that have already been renamed/normalised by the
582 + outer model and whose ``model.`` prefix has been stripped by
583 + ``AutoWeightsLoader``. Handles:
584 + 1. Fused QKV split (query_key_value → qkv_proj q/k/v shards).
585 + 2. gate_proj + up_proj → gate_up_proj stacking (dense + shared-exp).
586 + 3. Routed-expert weights via the fused-MoE mapping.
587 + 4. All remaining weights via their default loader.
588 + """
589 + config = self.config
590 + num_heads: int = config.num_attention_heads
591 + num_kv_heads: int = config.num_key_value_heads
592 + head_dim: int = config.head_dim or (config.hidden_size // num_heads)
593 + q_split = num_heads * head_dim
594 + kv_split = num_kv_heads * head_dim
595 +
596 + stacked_params_mapping = [
597 + # (vllm_param_name, ckpt_weight_name, shard_id)
598 + ("gate_up_proj", "gate_proj", 0),
599 + ("gate_up_proj", "up_proj", 1),
600 + ]
601 +
602 + params_dict = dict(self.named_parameters(remove_duplicate=False))
603 + loaded_params: set[str] = set()
604 + expert_params_mapping = self.get_expert_mapping()
605 +
606 + for name, loaded_weight in weights:
607 + # ------------------------------------------------------------------
608 + # 1. Fused QKV: split into q / k / v shards for QKVParallelLinear
609 + # ------------------------------------------------------------------
610 + if name.endswith(".self_attn.qkv_proj.weight"):
611 + if name not in params_dict:
612 + continue
613 + if is_pp_missing_parameter(name, self):
614 + continue
615 +
616 + param = params_dict[name]
617 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
618 + q_w = loaded_weight[:q_split, :]
619 + k_w = loaded_weight[q_split : q_split + kv_split, :]
620 + v_w = loaded_weight[q_split + kv_split :, :]
621 + weight_loader(param, q_w, "q")
622 + weight_loader(param, k_w, "k")
623 + weight_loader(param, v_w, "v")
624 + loaded_params.add(name)
625 + continue
626 +
627 + # ------------------------------------------------------------------
628 + # 2. gate_proj / up_proj → gate_up_proj (dense MLP + shared-exp.)
629 + # ------------------------------------------------------------------
630 + matched_stacked = False
631 + for param_name, weight_name, shard_id in stacked_params_mapping:
632 + if weight_name not in name:
633 + continue
634 + if "mlp.experts" in name: # routed experts handled below
635 + continue
636 + new_name = name.replace(weight_name, param_name)
637 + if new_name.endswith(".bias") and new_name not in params_dict:
638 + continue
639 + if new_name not in params_dict:
640 + continue
641 + if is_pp_missing_parameter(new_name, self):
642 + continue
643 +
644 + param = params_dict[new_name]
645 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
646 + weight_loader(param, loaded_weight, shard_id)
647 + loaded_params.add(new_name)
648 + matched_stacked = True
649 + break
650 +
651 + if matched_stacked:
652 + continue
653 +
654 + # ------------------------------------------------------------------
655 + # 3. Routed expert weights → fused-MoE kernel layout
656 + # ------------------------------------------------------------------
657 + matched_expert = False
658 + for (
659 + param_name,
660 + weight_name,
661 + expert_id,
662 + shard_id,
663 + ) in expert_params_mapping:
664 + if weight_name not in name:
665 + continue
666 + new_name = name.replace(weight_name, param_name)
667 + if is_pp_missing_parameter(new_name, self):
668 + continue
669 + if new_name not in params_dict:
670 + continue
671 +
672 + param = params_dict[new_name]
673 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
674 + weight_loader(
675 + param,
676 + loaded_weight,
677 + name,
678 + shard_id=shard_id,
679 + expert_id=expert_id,
680 + )
681 + loaded_params.add(new_name)
682 + matched_expert = True
683 + break
684 +
685 + if matched_expert:
686 + continue
687 +
688 + # ------------------------------------------------------------------
689 + # 4. All other weights: direct load (layernorms, embed_tokens, …)
690 + # ------------------------------------------------------------------
691 + if name.endswith(".bias") and name not in params_dict:
692 + continue
693 + if name not in params_dict:
694 + continue
695 + if is_pp_missing_parameter(name, self):
696 + continue
697 +
698 + param = params_dict[name]
699 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
700 + try:
701 + weight_loader(param, loaded_weight)
702 + except Exception as e:
703 + raise RuntimeError(
704 + f"[param2moe] Failed to load weight '{name}' "
705 + f"with shape {tuple(loaded_weight.shape)} "
706 + f"into param type {type(param).__name__}: {e}"
707 + ) from e
708 + loaded_params.add(name)
709 +
710 + return loaded_params
711 +
712 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
713 + return SharedFusedMoE.make_expert_params_mapping(
714 + self,
715 + ckpt_gate_proj_name="gate_proj",
716 + ckpt_down_proj_name="down_proj",
717 + ckpt_up_proj_name="up_proj",
718 + num_experts=self.config.num_experts,
719 + )
720 +
721 +
722 +class Param2MoEMixtureOfExperts(MixtureOfExperts):
723 + """Implements the vLLM MixtureOfExperts protocol for Param2MoE."""
724 +
725 + expert_weights: list[torch.Tensor]
726 +
727 + def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None:
728 + if example_moe is None:
729 + raise RuntimeError(
730 + "No Param2MoEMoEBlock found in model.layers. "
731 + "Check first_k_dense_replace and num_experts in config."
732 + )
733 + self.num_logical_experts = example_moe.num_experts
734 + self.num_routed_experts = example_moe.num_experts
735 + self.num_shared_experts = example_moe.num_shared_experts
736 +
737 + self.num_physical_experts = self.num_logical_experts
738 + self.num_local_physical_experts = self.num_logical_experts
739 + self.num_redundant_experts = 0
740 +
741 + def update_physical_experts_metadata(
742 + self,
743 + num_physical_experts: int,
744 + num_local_physical_experts: int,
745 + ) -> None:
746 + self.num_physical_experts = num_physical_experts
747 + self.num_local_physical_experts = num_local_physical_experts
748 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
749 +
750 + for moe in self.moe_mlp_layers:
751 + moe.n_physical_experts = num_physical_experts
752 + moe.n_local_physical_experts = num_local_physical_experts
753 + moe.n_redundant_experts = self.num_redundant_experts
754 +
755 + fused = moe.experts
756 + if hasattr(fused, "n_local_physical_experts"):
757 + fused.n_local_physical_experts = num_local_physical_experts
758 + if hasattr(fused, "n_physical_experts"):
759 + fused.n_physical_experts = num_physical_experts
760 + if hasattr(fused, "n_redundant_experts"):
761 + fused.n_redundant_experts = self.num_redundant_experts
762 + if hasattr(fused, "update_expert_map"):
763 + fused.update_expert_map()
764 +
765 + def set_eplb_state(
766 + self,
767 + expert_load_view: torch.Tensor,
768 + logical_to_physical_map: torch.Tensor,
769 + logical_replica_count: torch.Tensor,
770 + ) -> None:
771 + self.expert_weights.clear()
772 + for layer_idx, layer in enumerate(self.moe_layers):
773 + if hasattr(layer, "get_expert_weights"):
774 + self.expert_weights.append(layer.get_expert_weights())
775 + if hasattr(layer, "set_eplb_state"):
776 + layer.set_eplb_state(
777 + moe_layer_idx=layer_idx,
778 + expert_load_view=expert_load_view,
779 + logical_to_physical_map=logical_to_physical_map,
780 + logical_replica_count=logical_replica_count,
781 + )
782 +
783 +
784 +class Param2MoEForCausalLM(
785 + nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts
786 +):
787 + """
788 + vLLM-native Param2MoE CausalLM.
789 +
790 + Uses Grouped-Query Attention (GQA) with a Sigmoid-scored,
791 + grouped-topk Mixture-of-Experts MLP.
792 + """
793 +
794 + # LoRA packed-module mapping. The fused gate_up_proj handles
795 + # gate_proj and up_proj from the checkpoint.
796 + packed_modules_mapping = {
797 + "qkv_proj": ["query_key_value"],
798 + "gate_up_proj": ["gate_proj", "up_proj"],
799 + }
800 +
801 + # Modules eligible for LoRA adaptation.
802 + supported_lora_modules = [
803 + "qkv_proj",
804 + "o_proj",
805 + "gate_up_proj",
806 + "down_proj",
807 + ]
808 +
809 + # Embedding layers and their weight-tying counterparts.
810 + embedding_modules = {
811 + "embed_tokens": "input_embeddings",
812 + "lm_head": "output_embeddings",
813 + }
814 +
815 + # Modules that need vocab-size padding for LoRA.
816 + embedding_padding_modules = ["lm_head"]
817 +
818 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
819 + super().__init__()
820 +
821 + config = vllm_config.model_config.hf_config
822 + quant_config = vllm_config.quant_config
823 +
824 + self.config = config
825 + self.quant_config = quant_config
826 +
827 + self.model = Param2MoEModel(
828 + vllm_config=vllm_config,
829 + prefix=maybe_prefix(prefix, "model"),
830 + )
831 +
832 + self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False)
833 + if get_pp_group().is_last_rank:
834 + if self.tie_word_embeddings:
835 + self.lm_head = self.model.embed_tokens
836 + else:
837 + self.lm_head = ParallelLMHead(
838 + config.vocab_size,
839 + config.hidden_size,
840 + quant_config=quant_config,
841 + prefix=maybe_prefix(prefix, "lm_head"),
842 + )
843 + self.logits_processor = LogitsProcessor(config.vocab_size)
844 + else:
845 + self.lm_head = PPMissingLayer()
846 + self.logits_processor = None # type: ignore[assignment]
847 +
848 + self.make_empty_intermediate_tensors = (
849 + self.model.make_empty_intermediate_tensors
850 + )
851 +
852 + self.expert_weights: list[torch.Tensor] = []
853 + self.num_moe_layers: int = 0
854 + self.moe_layers: list = []
855 + self.moe_mlp_layers: list = []
856 +
857 + example_moe: Param2MoEMoEBlock | None = None
858 + for layer in self.model.layers:
859 + if isinstance(layer, PPMissingLayer):
860 + continue
861 + if isinstance(layer.mlp, Param2MoEMoEBlock):
862 + example_moe = layer.mlp
863 + self.moe_mlp_layers.append(layer.mlp)
864 + self.moe_layers.append(layer.mlp.experts)
865 + self.num_moe_layers += 1
866 +
867 + if self.config.num_experts is not None:
868 + self.extract_moe_parameters(example_moe)
869 +
870 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
871 + return self.model.embed_input_ids(input_ids)
872 +
873 + def forward(
874 + self,
875 + input_ids: torch.Tensor,
876 + positions: torch.Tensor,
877 + intermediate_tensors: IntermediateTensors | None = None,
878 + inputs_embeds: torch.Tensor | None = None,
879 + ) -> torch.Tensor | IntermediateTensors:
880 + return self.model(
881 + input_ids=input_ids,
882 + positions=positions,
883 + intermediate_tensors=intermediate_tensors,
884 + inputs_embeds=inputs_embeds,
885 + )
886 +
887 + def compute_logits(
888 + self,
889 + hidden_states: torch.Tensor,
890 + ) -> torch.Tensor | None:
891 + if not get_pp_group().is_last_rank:
892 + return None
893 + return self.logits_processor(self.lm_head, hidden_states)
894 +
895 + def load_weights(
896 + self,
897 + weights: Iterable[tuple[str, torch.Tensor]],
898 + ) -> set[str]:
899 + loader = AutoWeightsLoader(self)
900 + return loader.load_weights(_rename_and_normalize_weights(weights))
vllm/model_executor/models/registry.py
+1
@@ -182,6 +182,7 @@ _TEXT_GENERATION_MODELS = {
182 "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
183 "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
184 "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
185 + "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"),
186 "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
187 "PhiForCausalLM": ("phi", "PhiForCausalLM"),
188 "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),