Remove even more unnecessary `load_weights` methods (#48496)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 18, 2026 at 09:40 UTC
c233d90aa826df072872df47b201450059be8e71
44 files changed
+639
-2076
vllm/model_executor/models/AXK1.py
+1
-13
@@ -79,6 +79,7 @@ from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsP
79
from .utils import (
80
AutoWeightsLoader,
81
PPMissingLayer,
82
+ get_spec_layer_idx_from_weight_name,
83
is_pp_missing_parameter,
84
make_empty_intermediate_tensors_factory,
85
make_layers,
@@ -1143,16 +1144,3 @@ class AXK1ForCausalLM(
1144
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1145
loader = AutoWeightsLoader(self)
1146
return loader.load_weights(weights)
1146
-
1147
-
1148
-def get_spec_layer_idx_from_weight_name(
1149
- config: AXK1Config, weight_name: str
1150
-) -> int | None:
1151
- if config.num_nextn_predict_layers and config.num_nextn_predict_layers > 0:
1152
- layer_idx = config.num_hidden_layers
1153
- for i in range(config.num_nextn_predict_layers):
1154
- if weight_name.startswith(
1155
- f"model.layers.{layer_idx + i}."
1156
- ) or weight_name.startswith(f"layers.{layer_idx + i}."):
1157
- return layer_idx + i
1158
- return None
vllm/model_executor/models/aria.py
+22
-128
@@ -11,21 +11,13 @@ from transformers.models.aria.processing_aria import AriaProcessor
11
12
from vllm.config import VllmConfig
13
from vllm.config.multimodal import BaseDummyOptions
14
-from vllm.distributed import get_tensor_model_parallel_rank
14
from vllm.inputs import MultiModalDataDict
15
from vllm.model_executor.layers.activation import get_act_fn
17
-from vllm.model_executor.layers.fused_moe import (
18
- FusedMoE,
19
- RoutedExperts,
20
-)
16
+from vllm.model_executor.layers.fused_moe import FusedMoE
17
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.quantization import QuantizationConfig
20
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
25
-from vllm.model_executor.model_loader.weight_utils import (
26
- default_weight_loader,
27
- maybe_remap_kv_scale_name,
28
-)
21
from vllm.multimodal import MULTIMODAL_REGISTRY
22
from vllm.multimodal.inputs import (
23
MultiModalFieldConfig,
@@ -51,7 +43,6 @@ from .llama import LlamaDecoderLayer, LlamaMLP, LlamaModel
43
from .utils import (
44
AutoWeightsLoader,
45
WeightsMapper,
54
- is_pp_missing_parameter,
46
maybe_prefix,
47
)
48
@@ -94,34 +85,18 @@ class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant):
85
# Identity layer
86
self.post_layernorm = nn.Identity()
87
88
+ hf_to_vllm_mapper = WeightsMapper(
89
+ orig_to_new_stacked={
90
+ ".q_proj": (".qkv_proj", "q"),
91
+ ".k_proj": (".qkv_proj", "k"),
92
+ ".v_proj": (".qkv_proj", "v"),
93
+ }
94
+ )
95
+
96
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
98
- stacked_params_mapping = [
99
- # (param_name, shard_name, shard_id)
100
- ("qkv_proj", "q_proj", "q"),
101
- ("qkv_proj", "k_proj", "k"),
102
- ("qkv_proj", "v_proj", "v"),
103
- ]
104
- params_dict = dict(self.named_parameters())
105
- loaded_params: set[str] = set()
106
- for name, loaded_weight in weights:
107
- # NOTE: post_layernorm is not used in Aria
108
- if "post_layernorm" in name:
109
- continue
110
-
111
- for param_name, weight_name, shard_id in stacked_params_mapping:
112
- if weight_name not in name:
113
- continue
114
- name = name.replace(weight_name, param_name)
115
- param = params_dict[name]
116
- weight_loader = param.weight_loader
117
- weight_loader(param, loaded_weight, shard_id)
118
- break
119
- else:
120
- param = params_dict[name]
121
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
122
- weight_loader(param, loaded_weight)
123
- loaded_params.add(name)
124
- return loaded_params
97
+ # NOTE: post_layernorm is not used in Aria.
98
+ loader = AutoWeightsLoader(self, skip_substrs=["post_layernorm"])
99
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
100
101
102
class AriaProjectorMLP(nn.Module):
@@ -217,39 +192,6 @@ class AriaProjector(nn.Module):
192
return out
193
194
220
-class AriaRoutedExperts(RoutedExperts):
221
- def weight_loader(
222
- self, param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str
223
- ) -> None:
224
- # Override the weight_loader to handle the expert weights in the Aria
225
- # model, which are already packed with experts, and merge the gate and
226
- # up weights for each expert.
227
- # Note: Loading expert weights with quantization is not supported
228
- tp_rank = get_tensor_model_parallel_rank()
229
- tp_size = self.moe_config.tp_size
230
- if shard_id == "w13":
231
- # the shape of loaded_weight is
232
- # (num_experts, hidden_size, 2 * moe_intermediate_size)
233
- if tp_size > 1:
234
- up, gate = loaded_weight.chunk(2, dim=-1)
235
- up_current_rank = up.chunk(tp_size, dim=-1)[tp_rank]
236
- gate_current_rank = gate.chunk(tp_size, dim=-1)[tp_rank]
237
- up_and_gate = torch.cat(
238
- [up_current_rank, gate_current_rank], dim=-1
239
- ).transpose(1, 2)
240
- param.data.copy_(up_and_gate)
241
- else:
242
- param.data.copy_(loaded_weight.transpose(1, 2))
243
- elif shard_id == "w2":
244
- # the shape of loaded_weight is
245
- # (num_experts, moe_intermediate_size, hidden_size)
246
- if tp_size > 1:
247
- down_current_rank = loaded_weight.chunk(tp_size, dim=1)[tp_rank]
248
- param.data.copy_(down_current_rank.transpose(1, 2))
249
- else:
250
- param.data.copy_(loaded_weight.transpose(1, 2))
251
-
252
-
195
class AriaTextMoELayer(nn.Module):
196
"""
197
Mixture of Experts (MoE) Layer for the AriaMoE model.
@@ -288,7 +230,6 @@ class AriaTextMoELayer(nn.Module):
230
intermediate_size=config.intermediate_size,
231
quant_config=quant_config,
232
prefix=f"{prefix}.experts",
291
- routed_experts_cls=AriaRoutedExperts,
233
)
234
235
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -335,70 +276,23 @@ class AriaTextModel(LlamaModel, SupportsQuant):
276
packed_modules_mapping = {
277
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
278
"gate_up_proj": ["gate_proj", "up_proj"],
338
- "experts.routed_experts.w13_weight": ["experts.fc1.weight"],
339
- "experts.routed_experts.w2_weight": ["experts.fc2.weight"],
279
}
280
281
+ # Aria packs all experts into single (transposed) fc1/fc2 tensors, which is
282
+ # exactly the pre-fused checkpoint layout FusedMoE self-loads once fc1/fc2
283
+ # are renamed to the fused gate_up_proj/down_proj names.
284
+ hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper | WeightsMapper(
285
+ orig_to_new_substr={
286
+ "experts.fc1": "experts.gate_up_proj",
287
+ "experts.fc2": "experts.down_proj",
288
+ }
289
+ )
290
+
291
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
292
super().__init__(
293
vllm_config=vllm_config, prefix=prefix, layer_type=AriaTextDecoderLayer
294
)
295
347
- # Adapted from LlamaModel.load_weights with the modification of adding
348
- # the expert weights mapping to `stacked_params_mapping`
349
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
350
- stacked_params_mapping = [
351
- # (param_name, shard_name, shard_id)
352
- (".qkv_proj", ".q_proj", "q"),
353
- (".qkv_proj", ".k_proj", "k"),
354
- (".qkv_proj", ".v_proj", "v"),
355
- (".gate_up_proj", ".gate_proj", 0),
356
- (".gate_up_proj", ".up_proj", 1),
357
- ("experts.routed_experts.w13_weight", "experts.fc1.weight", "w13"),
358
- ("experts.routed_experts.w2_weight", "experts.fc2.weight", "w2"),
359
- ]
360
- params_dict = dict(self.named_parameters())
361
- loaded_params: set[str] = set()
362
- for name, loaded_weight in weights:
363
- if "rotary_emb.inv_freq" in name:
364
- continue
365
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
366
- # Models trained using ColossalAI may include these tensors in
367
- # the checkpoint. Skip them.
368
- continue
369
- for param_name, weight_name, shard_id in stacked_params_mapping:
370
- if weight_name not in name:
371
- continue
372
- name = name.replace(weight_name, param_name)
373
- # Skip loading extra bias for GPTQ models.
374
- if name.endswith(".bias") and name not in params_dict:
375
- continue
376
-
377
- if is_pp_missing_parameter(name, self):
378
- continue
379
-
380
- param = params_dict[name]
381
- weight_loader = param.weight_loader
382
- weight_loader(param, loaded_weight, shard_id)
383
- break
384
- else:
385
- # Skip loading extra bias for GPTQ models.
386
- if name.endswith(".bias") and name not in params_dict:
387
- continue
388
- # Remapping the name of FP8 kv-scale.
389
- name = maybe_remap_kv_scale_name(name, params_dict)
390
- if name is None:
391
- continue
392
-
393
- if is_pp_missing_parameter(name, self):
394
- continue
395
-
396
- param = params_dict[name]
397
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
398
- weight_loader(param, loaded_weight)
399
- loaded_params.add(name)
400
- return loaded_params
401
-
296
297
class AriaProcessingInfo(BaseProcessingInfo):
298
def get_hf_config(self):
vllm/model_executor/models/bailing_moe.py
+24
-92
@@ -41,10 +41,7 @@ from vllm.distributed import (
41
)
42
from vllm.model_executor.layers.activation import SiluAndMul
43
from vllm.model_executor.layers.attention import Attention
44
-from vllm.model_executor.layers.fused_moe import (
45
- FusedMoE,
46
- fused_moe_make_expert_params_mapping,
47
-)
44
+from vllm.model_executor.layers.fused_moe import FusedMoE
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
47
MergedColumnParallelLinear,
@@ -58,14 +55,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
55
ParallelLMHead,
56
VocabParallelEmbedding,
57
)
61
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
58
from vllm.sequence import IntermediateTensors
59
60
from .interfaces import SupportsLoRA, SupportsPP
61
from .utils import (
62
AutoWeightsLoader,
63
PPMissingLayer,
68
- is_pp_missing_parameter,
64
+ WeightsMapper,
65
make_empty_intermediate_tensors_factory,
66
make_layers,
67
maybe_prefix,
@@ -379,6 +375,16 @@ class BailingMoeBlock(nn.Module):
375
376
@support_torch_compile
377
class BailingMoeModel(nn.Module):
378
+ hf_to_vllm_mapper = WeightsMapper(
379
+ orig_to_new_stacked={
380
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
381
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
382
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
383
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
384
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
385
+ }
386
+ )
387
+
388
def __init__(
389
self,
390
*,
@@ -468,89 +474,9 @@ class BailingMoeModel(nn.Module):
474
hidden_states, _ = self.norm(hidden_states, residual)
475
return hidden_states
476
471
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
472
- return fused_moe_make_expert_params_mapping(
473
- self,
474
- ckpt_gate_proj_name="gate_proj",
475
- ckpt_down_proj_name="down_proj",
476
- ckpt_up_proj_name="up_proj",
477
- num_experts=self.config.num_experts,
478
- )
479
-
477
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
481
- stacked_params_mapping = [
482
- # (param_name, shard_name, shard_id)
483
- ("gate_up_proj", "gate_proj", 0),
484
- ("gate_up_proj", "up_proj", 1),
485
- ]
486
-
487
- params_dict = dict(self.named_parameters(remove_duplicate=False))
488
- loaded_params: set[str] = set()
489
- expert_params_mapping = self.get_expert_mapping()
490
- for name, loaded_weight in weights:
491
- if (
492
- hasattr(self.config, "norm_head")
493
- and self.config.norm_head
494
- and "lm_head.weight" in name
495
- ):
496
- loaded_weight = F.normalize(loaded_weight, dim=0, p=2, eps=1e-7)
497
-
498
- for param_name, weight_name, shard_id in stacked_params_mapping:
499
- if weight_name not in name:
500
- continue
501
- if "mlp.experts" in name:
502
- continue
503
- name = name.replace(weight_name, param_name)
504
- # Skip loading extra bias for GPTQ models.
505
- if name.endswith(".bias") and name not in params_dict:
506
- continue
507
- if name not in params_dict:
508
- continue
509
-
510
- if is_pp_missing_parameter(name, self):
511
- continue
512
-
513
- param = params_dict[name]
514
- weight_loader = param.weight_loader
515
- weight_loader(param, loaded_weight, shard_id)
516
- break
517
- else:
518
- for mapping in expert_params_mapping:
519
- param_name, weight_name, expert_id, shard_id = mapping
520
- if weight_name not in name:
521
- continue
522
- name = name.replace(weight_name, param_name)
523
-
524
- if is_pp_missing_parameter(name, self):
525
- continue
526
- if name not in params_dict:
527
- continue
528
- param = params_dict[name]
529
- weight_loader = param.weight_loader
530
- weight_loader(
531
- param,
532
- loaded_weight,
533
- name,
534
- shard_id=shard_id,
535
- expert_id=expert_id,
536
- )
537
- break
538
- else:
539
- if name.endswith(".bias") and name not in params_dict:
540
- continue
541
- if name not in params_dict:
542
- continue
543
-
544
- if is_pp_missing_parameter(name, self):
545
- continue
546
-
547
- param = params_dict[name]
548
- weight_loader = getattr(
549
- param, "weight_loader", default_weight_loader
550
- )
551
- weight_loader(param, loaded_weight)
552
- loaded_params.add(name)
553
- return loaded_params
478
+ loader = AutoWeightsLoader(self)
479
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
480
481
482
class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
@@ -622,15 +548,21 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
548
logits = self.logits_processor(self.lm_head, hidden_states)
549
return logits
550
551
+ def _normalize_lm_head(
552
+ self, weights: Iterable[tuple[str, torch.Tensor]]
553
+ ) -> Iterable[tuple[str, torch.Tensor]]:
554
+ norm_head = getattr(self.config, "norm_head", False)
555
+ for name, loaded_weight in weights:
556
+ if norm_head and "lm_head.weight" in name:
557
+ loaded_weight = F.normalize(loaded_weight, dim=0, p=2, eps=1e-7)
558
+ yield name, loaded_weight
559
+
560
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
561
loader = AutoWeightsLoader(
562
self,
563
skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
564
)
630
- return loader.load_weights(weights)
631
-
632
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
633
- return self.model.get_expert_mapping()
565
+ return loader.load_weights(self._normalize_lm_head(weights))
566
567
568
class BailingMoeV2ForCausalLM(BailingMoeForCausalLM):
vllm/model_executor/models/bloom.py
+15
-30
@@ -47,13 +47,11 @@ 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
50
from vllm.sequence import IntermediateTensors
51
52
from .interfaces import SupportsPP, SupportsQuant
53
from .utils import (
54
AutoWeightsLoader,
56
- is_pp_missing_parameter,
55
make_empty_intermediate_tensors_factory,
56
make_layers,
57
maybe_prefix,
@@ -297,36 +295,23 @@ class BloomModel(nn.Module):
295
hidden_states = self.ln_f(hidden_states)
296
return hidden_states
297
300
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
301
- params_dict = dict(self.named_parameters(remove_duplicate=False))
302
- loaded_params: set[str] = set()
298
+ def _repack_qkv(
299
+ self, weights: Iterable[tuple[str, torch.Tensor]]
300
+ ) -> Iterable[tuple[str, torch.Tensor]]:
301
+ # BLOOM's fused QKV is laid out as (num_heads * 3 * head_size) on its
302
+ # output dim (0), while vLLM expects (3 * num_heads * head_size).
303
+ num_heads = self.config.num_attention_heads
304
for name, loaded_weight in weights:
304
- if is_pp_missing_parameter(name, self):
305
- continue
306
- param = params_dict[name]
307
-
305
if "query_key_value" in name:
309
- # NOTE: BLOOM's fused QKV's output_dim has the shape of
310
- # (num_heads * 3 * head_size), while the
311
- # required shape is (3 * num_heads * head_size).
312
- # Thus, we need weight conversion.
313
- output_dim = getattr(param, "output_dim", None)
314
- num_heads = self.config.num_attention_heads
315
- if output_dim is not None:
316
- loaded_weight_shape = loaded_weight.shape
317
- loaded_weight = loaded_weight.view(
318
- loaded_weight_shape[:output_dim]
319
- + (num_heads, 3, -1)
320
- + loaded_weight_shape[output_dim + 1 :]
321
- )
322
- loaded_weight = loaded_weight.transpose(output_dim, output_dim + 1)
323
- loaded_weight = loaded_weight.reshape(loaded_weight_shape)
324
-
325
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
326
- weight_loader(param, loaded_weight)
327
- loaded_params.add(name)
328
-
329
- return loaded_params
306
+ shape = loaded_weight.shape
307
+ loaded_weight = loaded_weight.view((num_heads, 3, -1) + shape[1:])
308
+ loaded_weight = loaded_weight.transpose(0, 1)
309
+ loaded_weight = loaded_weight.reshape(shape)
310
+ yield name, loaded_weight
311
+
312
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
313
+ loader = AutoWeightsLoader(self)
314
+ return loader.load_weights(self._repack_qkv(weights))
315
316
317
class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant):
vllm/model_executor/models/cohere2_moe.py
+1
-3
@@ -16,9 +16,7 @@ from vllm.distributed import (
16
)
17
from vllm.model_executor.layers.activation import SiluAndMul
18
from vllm.model_executor.layers.attention import Attention
19
-from vllm.model_executor.layers.fused_moe import (
20
- FusedMoE,
21
-)
19
+from vllm.model_executor.layers.fused_moe import FusedMoE
20
from vllm.model_executor.layers.linear import (
21
MergedColumnParallelLinear,
22
QKVParallelLinear,
vllm/model_executor/models/deepencoder.py
+3
-10
@@ -22,9 +22,9 @@ from vllm.model_executor.custom_op import PluggableLayer
22
from vllm.model_executor.layers.attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
26
from .clip import CLIPEncoder, CLIPVisionEmbeddings
27
+from .utils import AutoWeightsLoader
28
29
30
class MLPBlock(nn.Module):
@@ -671,12 +671,5 @@ class DeepCLIPVisionTransformer(nn.Module):
671
return encoder_outputs
672
673
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
674
- params_dict = dict(self.named_parameters())
675
- loaded_params: set[str] = set()
676
-
677
- for name, loaded_weight in weights:
678
- param = params_dict[name]
679
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
680
- weight_loader(param, loaded_weight)
681
- loaded_params.add(name)
682
- return loaded_params
674
+ loader = AutoWeightsLoader(self)
675
+ return loader.load_weights(weights)
vllm/model_executor/models/deepseek_mtp.py
+4
-1
@@ -33,9 +33,12 @@ from .deepseek_v2 import (
33
DeepseekV2MixtureOfExperts,
34
DeepseekV2MoE,
35
_try_load_fp8_indexer_wk,
36
+)
37
+from .utils import (
38
+ get_pp_missing_layer_names,
39
get_spec_layer_idx_from_weight_name,
40
+ maybe_prefix,
41
)
38
-from .utils import get_pp_missing_layer_names, maybe_prefix
42
43
44
def _restore_full_token_layout_if_needed(
vllm/model_executor/models/deepseek_v2.py
+1
-18
@@ -108,6 +108,7 @@ from .interfaces import (
108
from .utils import (
109
PPMissingLayer,
110
get_pp_missing_layer_names,
111
+ get_spec_layer_idx_from_weight_name,
112
is_pp_missing_parameter,
113
make_empty_intermediate_tensors_factory,
114
make_layers,
@@ -1918,21 +1919,3 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
1919
1920
class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM):
1921
pass
1921
-
1922
-
1923
-# Compatibility with
1924
-# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
1925
-def get_spec_layer_idx_from_weight_name(
1926
- config: DeepseekV2Config | DeepseekV3Config, weight_name: str
1927
-) -> int | None:
1928
- if (
1929
- hasattr(config, "num_nextn_predict_layers")
1930
- and config.num_nextn_predict_layers > 0
1931
- ):
1932
- layer_idx = config.num_hidden_layers
1933
- for i in range(config.num_nextn_predict_layers):
1934
- if weight_name.startswith(
1935
- f"model.layers.{layer_idx + i}."
1936
- ) or weight_name.startswith(f"layers.{layer_idx + i}."):
1937
- return layer_idx + i
1938
- return None
vllm/model_executor/models/ernie45_moe.py
+31
-136
@@ -23,8 +23,7 @@
23
# limitations under the License.
24
"""Inference-only ErineMoE model compatible with HuggingFace weights."""
25
26
-import typing
27
-from collections.abc import Callable, Iterable
26
+from collections.abc import Iterable
27
from itertools import islice
28
from typing import Any
29
@@ -42,11 +41,7 @@ from vllm.distributed import (
41
from vllm.logger import init_logger
42
from vllm.model_executor.layers.activation import SiluAndMul
43
from vllm.model_executor.layers.attention import Attention
45
-from vllm.model_executor.layers.fused_moe import (
46
- FusedMoE,
47
- MoERunner,
48
- fused_moe_make_expert_params_mapping,
49
-)
44
+from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
47
MergedColumnParallelLinear,
@@ -61,10 +56,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead,
57
VocabParallelEmbedding,
58
)
64
-from vllm.model_executor.model_loader.weight_utils import (
65
- default_weight_loader,
66
- maybe_remap_kv_scale_name,
67
-)
59
from vllm.sequence import IntermediateTensors
60
from vllm.transformers_utils.config import set_default_rope_theta
61
@@ -72,8 +63,8 @@ from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
63
from .utils import (
64
AutoWeightsLoader,
65
PPMissingLayer,
66
+ WeightsMapper,
67
extract_layer_index,
76
- is_pp_missing_parameter,
68
make_empty_intermediate_tensors_factory,
69
make_layers,
70
maybe_prefix,
@@ -407,6 +398,19 @@ class Ernie4_5_MoeDecoderLayer(nn.Module):
398
399
@support_torch_compile
400
class Ernie4_5_MoeModel(nn.Module):
401
+ hf_to_vllm_mapper = WeightsMapper(
402
+ orig_to_new_stacked={
403
+ ".q_proj": (".qkv_proj", "q"),
404
+ ".k_proj": (".qkv_proj", "k"),
405
+ ".v_proj": (".qkv_proj", "v"),
406
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
407
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
408
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
409
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
410
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
411
+ }
412
+ )
413
+
414
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
415
super().__init__()
416
@@ -486,132 +490,26 @@ class Ernie4_5_MoeModel(nn.Module):
490
491
return hidden_states
492
489
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
490
- # Params for weights, fp8 weight scales, fp8 activation scales
491
- # (param_name, weight_name, expert_id, shard_id)
492
- return fused_moe_make_expert_params_mapping(
493
- self,
494
- ckpt_gate_proj_name="gate_proj",
495
- ckpt_down_proj_name="down_proj",
496
- ckpt_up_proj_name="up_proj",
497
- num_experts=self.config.moe_num_experts,
498
- num_redundant_experts=self.num_redundant_experts,
499
- )
500
-
501
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
502
- stacked_params_mapping = [
503
- # (param_name, shard_name, shard_id)
504
- ("qkv_proj", "q_proj", "q"),
505
- ("qkv_proj", "k_proj", "k"),
506
- ("qkv_proj", "v_proj", "v"),
507
- ("gate_up_proj", "gate_proj", 0),
508
- ("gate_up_proj", "up_proj", 1),
509
- ]
510
-
511
- params_dict = dict(self.named_parameters())
512
- loaded_params: set[str] = set()
513
- expert_params_mapping = self.get_expert_mapping()
493
+ def _preprocess(
494
+ self, weights: Iterable[tuple[str, torch.Tensor]]
495
+ ) -> Iterable[tuple[str, torch.Tensor]]:
496
for name, loaded_weight in weights:
515
- if self.config.tie_word_embeddings and name.endswith("lm_head.weight"):
516
- continue
517
- # MTP will be supported soon.
518
- if "mtp" in name:
519
- continue
520
-
497
+ # moe_statics.e_score_correction_bias is stored with a leading
498
+ # singleton dim and under a different module name.
499
if "e_score_correction_bias" in name:
500
name = name.replace("moe_statics", "gate")
501
loaded_weight = loaded_weight.squeeze(0)
502
+ yield name, loaded_weight
503
525
- for param_name, weight_name, shard_id in stacked_params_mapping:
526
- # Skip non-stacked layers and experts (experts handled below).
527
- if weight_name not in name:
528
- continue
529
-
530
- if ("mlp.experts." in name) and name not in params_dict:
531
- continue
532
- name = name.replace(weight_name, param_name)
533
- # Skip loading extra bias for GPTQ models.
534
- if (
535
- name.endswith(".bias") or name.endswith("_bias")
536
- ) and name not in params_dict:
537
- continue
538
- # Skip layers on other devices.
539
- if is_pp_missing_parameter(name, self):
540
- continue
541
-
542
- param = params_dict[name]
543
- weight_loader = param.weight_loader
544
- weight_loader(param, loaded_weight, shard_id)
545
- break
546
- else:
547
- is_expert_weight = False
548
- for mapping in expert_params_mapping:
549
- param_name, weight_name, expert_id, shard_id = mapping
550
-
551
- if weight_name not in name:
552
- continue
553
-
554
- # Anyway, this is an expert weight and should not be
555
- # attempted to load as other weights later
556
- is_expert_weight = True
557
-
558
- # Do not modify `name` since the loop may continue here
559
- # Instead, create a new variable
560
- name_mapped = name.replace(weight_name, param_name)
561
- # Skip layers on other devices.
562
- if is_pp_missing_parameter(name_mapped, self):
563
- continue
564
-
565
- # Skip loading extra bias for GPTQ models.
566
- if (
567
- name_mapped.endswith(".bias") or name_mapped.endswith("_bias")
568
- ) and name_mapped not in params_dict:
569
- continue
570
- param = params_dict[name_mapped]
571
- # We should ask the weight loader to return success or not
572
- # here since otherwise we may skip experts with other
573
- # available replicas.
574
- weight_loader = typing.cast(
575
- Callable[..., bool], param.weight_loader
576
- )
577
- success = weight_loader(
578
- param,
579
- loaded_weight,
580
- name_mapped,
581
- shard_id=shard_id,
582
- expert_id=expert_id,
583
- return_success=True,
584
- )
585
- if success:
586
- name = name_mapped
587
- break
588
- else:
589
- if is_expert_weight:
590
- # We've checked that this is an expert weight
591
- # However it's not mapped locally to this rank
592
- # So we simply skip it
593
- continue
594
-
595
- # Skip loading extra bias for GPTQ models.
596
- if (
597
- name.endswith(".bias") or name.endswith("_bias")
598
- ) and name not in params_dict:
599
- continue
600
- # Skip layers on other devices.
601
- if is_pp_missing_parameter(name, self):
602
- continue
603
- # Remapping the name of FP8 kv-scale.
604
- name = maybe_remap_kv_scale_name(name, params_dict)
605
- if name is None:
606
- continue
607
-
608
- param = params_dict[name]
609
- weight_loader = getattr(
610
- param, "weight_loader", default_weight_loader
611
- )
612
- weight_loader(param, loaded_weight)
613
- loaded_params.add(name)
614
- return loaded_params
504
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
505
+ loader = AutoWeightsLoader(
506
+ self,
507
+ skip_substrs=["mtp"],
508
+ ignore_unexpected_suffixes=[".bias", "_bias"],
509
+ )
510
+ return loader.load_weights(
511
+ self._preprocess(weights), mapper=self.hf_to_vllm_mapper
512
+ )
513
514
515
class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExperts):
@@ -741,6 +639,3 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe
639
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
640
)
641
return loader.load_weights(weights)
744
-
745
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
746
- return self.model.get_expert_mapping()
vllm/model_executor/models/funaudiochat.py
+17
-32
@@ -30,7 +30,6 @@ from vllm.config.multimodal import BaseDummyOptions
30
from vllm.inputs import MultiModalDataDict
31
from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
32
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
33
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
from vllm.multimodal.inputs import (
35
MultiModalFieldConfig,
@@ -53,7 +52,12 @@ from vllm.sequence import IntermediateTensors
52
from vllm.utils.import_utils import _has_module
53
54
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
56
-from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
55
+from .utils import (
56
+ AutoWeightsLoader,
57
+ WeightsMapper,
58
+ init_vllm_registered_model,
59
+ maybe_prefix,
60
+)
61
62
63
class _SinusoidsPositionEmbedding(nn.Module):
@@ -79,6 +83,14 @@ class _SinusoidsPositionEmbedding(nn.Module):
83
class FunAudioChatAudioAttention(nn.Module):
84
"""Multi-headed attention used inside the continuous audio tower."""
85
86
+ hf_to_vllm_mapper = WeightsMapper(
87
+ orig_to_new_stacked={
88
+ ".q_proj": (".qkv_proj", "q"),
89
+ ".k_proj": (".qkv_proj", "k"),
90
+ ".v_proj": (".qkv_proj", "v"),
91
+ }
92
+ )
93
+
94
def __init__(self, config: Any):
95
super().__init__()
96
self.embed_dim = int(config.d_model)
@@ -123,42 +135,15 @@ class FunAudioChatAudioAttention(nn.Module):
135
bias=True,
136
)
137
126
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
127
- stacked_params_mapping = [
128
- ("qkv_proj", "q_proj", "q"),
129
- ("qkv_proj", "k_proj", "k"),
130
- ("qkv_proj", "v_proj", "v"),
131
- ]
132
-
133
- params_dict = dict(self.named_parameters())
138
with torch.no_grad():
139
if self.qkv_proj.bias is not None:
140
# HF FunAudioChat uses bias=False for k_proj. Ensure the missing
141
# shard starts as zeros, while allowing q/v shards to load.
142
self.qkv_proj.bias.zero_()
143
140
- loaded_params: set[str] = set()
141
- for name, loaded_weight in weights:
142
- for param_name, shard_name, shard_id in stacked_params_mapping:
143
- if shard_name not in name:
144
- continue
145
- name = name.replace(shard_name, param_name)
146
- param = params_dict[name]
147
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
148
- weight_loader(param, loaded_weight, shard_id)
149
- break
150
- else:
151
- # Skip loading extra bias for GPTQ models.
152
- if name.endswith(".bias") and name not in params_dict:
153
- continue
154
-
155
- param = params_dict[name]
156
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
157
- weight_loader(param, loaded_weight)
158
-
159
- loaded_params.add(name)
160
-
161
- return loaded_params
144
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
145
+ loader = AutoWeightsLoader(self)
146
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
147
148
def forward(
149
self,
vllm/model_executor/models/gemma3n.py
+22
-57
@@ -44,18 +44,14 @@ 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 VocabParallelEmbedding
47
-from vllm.model_executor.model_loader.weight_utils import (
48
- default_weight_loader,
49
- maybe_remap_kv_scale_name,
50
-)
47
from vllm.sequence import IntermediateTensors
48
from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata
49
50
from .interfaces import SupportsQuant
51
from .utils import (
52
AutoWeightsLoader,
53
+ WeightsMapper,
54
extract_layer_index,
58
- is_pp_missing_parameter,
55
make_layers,
56
maybe_prefix,
57
)
@@ -788,6 +784,25 @@ class Gemma3nCrossDecoder(nn.Module):
784
enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill
785
)
786
class Gemma3nTextModel(nn.Module, SupportsQuant):
787
+ # Decoder layers, altup_unembed_projections and norm live on the text
788
+ # model; every other submodule lives under self_decoder.
789
+ hf_to_vllm_mapper = WeightsMapper(
790
+ orig_to_new_prefix={
791
+ "embed_tokens.": "self_decoder.embed_tokens.",
792
+ "embed_tokens_per_layer.": "self_decoder.embed_tokens_per_layer.",
793
+ "per_layer_model_projection.": "self_decoder.per_layer_model_projection.",
794
+ "per_layer_projection_norm.": "self_decoder.per_layer_projection_norm.",
795
+ "altup_projections.": "self_decoder.altup_projections.",
796
+ },
797
+ orig_to_new_stacked={
798
+ ".q_proj": (".qkv_proj", "q"),
799
+ ".k_proj": (".qkv_proj", "k"),
800
+ ".v_proj": (".qkv_proj", "v"),
801
+ ".gate_proj": (".gate_up_proj", 0),
802
+ ".up_proj": (".gate_up_proj", 1),
803
+ },
804
+ )
805
+
806
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
807
super().__init__()
808
config = vllm_config.model_config.hf_config
@@ -1036,58 +1051,8 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):
1051
return self.norm(hidden_states)
1052
1053
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1039
- stacked_params_mapping = [
1040
- # (param_name, shard_name, shard_id)
1041
- ("qkv_proj", "q_proj", "q"),
1042
- ("qkv_proj", "k_proj", "k"),
1043
- ("qkv_proj", "v_proj", "v"),
1044
- ("gate_up_proj", "gate_proj", 0),
1045
- ("gate_up_proj", "up_proj", 1),
1046
- ]
1047
- params_dict = dict(self.named_parameters())
1048
- loaded_params: set[str] = set()
1049
- for name, loaded_weight in weights:
1050
- # decoder layer weights, altup_unembed_projections and rmsnorm
1051
- # are initialized in text model, others are in self decoder
1052
- if (
1053
- not name.startswith("layers")
1054
- and not name.startswith("altup_unembed_projections")
1055
- and not name.startswith("norm")
1056
- ):
1057
- name = f"self_decoder.{name}"
1058
-
1059
- for param_name, shard_name, shard_id in stacked_params_mapping:
1060
- if shard_name not in name:
1061
- continue
1062
- # Avoid spurious match with ".up_proj".
1063
- if "altup_projections" in name:
1064
- continue
1065
- name = name.replace(shard_name, param_name)
1066
- # Skip loading extra bias for GPTQ models.
1067
- if name.endswith(".bias") and name not in params_dict:
1068
- continue
1069
- if is_pp_missing_parameter(name, self):
1070
- continue
1071
- param = params_dict[name]
1072
- weight_loader = param.weight_loader
1073
- weight_loader(param, loaded_weight, shard_id)
1074
- break
1075
- else:
1076
- # Skip loading extra bias for GPTQ models.
1077
- if name.endswith(".bias") and name not in params_dict:
1078
- continue
1079
- # Remapping the name of FP8 kv-scale.
1080
- name = maybe_remap_kv_scale_name(name, params_dict)
1081
- if name is None:
1082
- continue
1083
- if is_pp_missing_parameter(name, self):
1084
- continue
1085
- param = params_dict[name]
1086
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1087
- weight_loader(param, loaded_weight)
1088
- loaded_params.add(name)
1089
-
1090
- return loaded_params
1054
+ loader = AutoWeightsLoader(self)
1055
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1056
1057
1058
class Gemma3nForCausalLM(nn.Module):
vllm/model_executor/models/glm4.py
-13
@@ -303,16 +303,3 @@ class Glm4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
303
]
304
loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
305
return loader.load_weights(weights)
306
-
307
-
308
-def get_spec_layer_idx_from_weight_name(
309
- config: Glm4Config, weight_name: str
310
-) -> int | None:
311
- if hasattr(config, "num_nextn_predict_layers") and (
312
- config.num_nextn_predict_layers > 0
313
- ):
314
- layer_idx = config.num_hidden_layers
315
- for i in range(config.num_nextn_predict_layers):
316
- if f"layers.{layer_idx + i}." in weight_name:
317
- return layer_idx + i
318
- return None
vllm/model_executor/models/glm4_moe.py
+24
-208
@@ -24,8 +24,7 @@
24
"""Inference-only GLM-4.5, GLM-4.6, GLM-4.7 model
25
compatible with HuggingFace weights."""
26
27
-import typing
28
-from collections.abc import Callable, Iterable
27
+from collections.abc import Iterable
28
from itertools import islice
29
30
import torch
@@ -43,10 +42,7 @@ from vllm.distributed import (
42
from vllm.logger import init_logger
43
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.attention import Attention
46
-from vllm.model_executor.layers.fused_moe import (
47
- FusedMoE,
48
- fused_moe_make_expert_params_mapping,
49
-)
45
+from vllm.model_executor.layers.fused_moe import FusedMoE
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
48
MergedColumnParallelLinear,
@@ -60,20 +56,18 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead,
57
VocabParallelEmbedding,
58
)
63
-from vllm.model_executor.model_loader.weight_utils import (
64
- default_weight_loader,
65
- maybe_remap_kv_scale_name,
66
-)
59
from vllm.sequence import IntermediateTensors
60
61
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
62
from .utils import (
63
AutoWeightsLoader,
64
PPMissingLayer,
73
- is_pp_missing_parameter,
65
+ WeightsMapper,
66
make_empty_intermediate_tensors_factory,
67
make_layers,
68
+ maybe_fuse_shared_experts,
69
maybe_prefix,
70
+ skip_spec_layers,
71
)
72
73
logger = init_logger(__name__)
@@ -410,6 +404,19 @@ class Glm4MoeDecoderLayer(nn.Module):
404
}
405
)
406
class Glm4MoeModel(nn.Module):
407
+ hf_to_vllm_mapper = WeightsMapper(
408
+ orig_to_new_stacked={
409
+ ".q_proj": (".qkv_proj", "q"),
410
+ ".k_proj": (".qkv_proj", "k"),
411
+ ".v_proj": (".qkv_proj", "v"),
412
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
413
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
414
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
415
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
416
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
417
+ }
418
+ )
419
+
420
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
421
super().__init__()
422
@@ -480,189 +487,14 @@ class Glm4MoeModel(nn.Module):
487
hidden_states, _ = self.norm(hidden_states, residual)
488
return hidden_states
489
483
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
484
- # Params for weights, fp8 weight scales, fp8 activation scales
485
- # (param_name, weight_name, expert_id, shard_id)
486
- # FSE widens the mapping by n_shared_experts slots; see deepseek_v2.py.
487
- num_experts = self.config.n_routed_experts
488
- if (
489
- rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
490
- and self.config.n_shared_experts
491
- ):
492
- num_experts += self.config.n_shared_experts
493
- return fused_moe_make_expert_params_mapping(
494
- self,
495
- ckpt_gate_proj_name="gate_proj",
496
- ckpt_down_proj_name="down_proj",
497
- ckpt_up_proj_name="up_proj",
498
- num_experts=num_experts,
499
- )
500
-
490
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
502
- rocm_aiter_moe_shared_expert_enabled = (
503
- rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
491
+ weights = maybe_fuse_shared_experts(
492
+ skip_spec_layers(weights, self.config),
493
+ n_routed_experts=self.config.n_routed_experts,
494
+ n_shared_experts=self.config.n_shared_experts or 1,
495
)
505
- stacked_params_mapping = [
506
- # (param_name, shard_name, shard_id)
507
- ("qkv_proj", "q_proj", "q"),
508
- ("qkv_proj", "k_proj", "k"),
509
- ("qkv_proj", "v_proj", "v"),
510
- ("gate_up_proj", "gate_proj", 0),
511
- ("gate_up_proj", "up_proj", 1),
512
- ]
513
-
514
- params_dict = dict(self.named_parameters())
515
- loaded_params: set[str] = set()
516
- expert_params_mapping = self.get_expert_mapping()
517
- for name, loaded_weight in weights:
518
- spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
519
- if spec_layer is not None:
520
- continue
521
-
522
- is_fusion_moe_shared_experts_layer = (
523
- rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
524
- )
525
-
526
- for param_name, weight_name, shard_id in stacked_params_mapping:
527
- # Skip non-stacked layers and experts (experts handled below).
528
- if weight_name not in name:
529
- continue
530
- # We have mlp.experts[0].gate_proj in the checkpoint.
531
- # Since we handle the experts below in expert_params_mapping,
532
- # we need to skip here BEFORE we update the name, otherwise
533
- # name will be updated to mlp.experts[0].gate_up_proj, which
534
- # will then be updated below in expert_params_mapping
535
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
536
- if ("mlp.experts." in name) and name not in params_dict:
537
- continue
538
- if is_fusion_moe_shared_experts_layer:
539
- continue
540
-
541
- name = name.replace(weight_name, param_name)
542
- # Skip loading extra bias for GPTQ models.
543
- if name.endswith(".bias") and name not in params_dict:
544
- continue
545
-
546
- name = maybe_remap_kv_scale_name(name, params_dict)
547
- if name is None:
548
- continue
549
- if is_pp_missing_parameter(name, self):
550
- continue
551
-
552
- param = params_dict[name]
553
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
554
- if weight_loader == default_weight_loader:
555
- weight_loader(param, loaded_weight)
556
- else:
557
- weight_loader(param, loaded_weight, shard_id)
558
- break
559
- else:
560
- is_expert_weight = False
561
-
562
- # FSE: split a widened mlp.shared_experts tensor into
563
- # n_shared_experts chunks; see deepseek_v2.py for details.
564
- num_chunks = 1
565
- split_dim = 0
566
- chunk_size = 0
567
- if is_fusion_moe_shared_experts_layer:
568
- num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
569
- split_dim = (
570
- 1
571
- if ("down_proj.weight" in name and loaded_weight.ndim > 1)
572
- else 0
573
- )
574
- total = loaded_weight.shape[split_dim]
575
- if total % num_chunks != 0:
576
- raise ValueError(
577
- f"FSE shared-expert weight {name} has dim "
578
- f"{total} along axis {split_dim} which is not "
579
- f"divisible by n_shared_experts={num_chunks}."
580
- )
581
- chunk_size = total // num_chunks
582
-
583
- for j in range(num_chunks):
584
- chunk_name = name
585
- weight_to_load = loaded_weight
586
-
587
- if is_fusion_moe_shared_experts_layer:
588
- chunk_slice = slice(j * chunk_size, (j + 1) * chunk_size)
589
- if loaded_weight.ndim == 1:
590
- weight_to_load = loaded_weight[chunk_slice]
591
- elif split_dim == 0:
592
- weight_to_load = loaded_weight[chunk_slice, :]
593
- else:
594
- weight_to_load = loaded_weight[:, chunk_slice]
595
- # Synthesize an expert-style name for expert mapping.
596
- chunk_name = name.replace(
597
- "mlp.shared_experts",
598
- f"mlp.experts.{self.config.n_routed_experts + j}",
599
- )
600
-
601
- for mapping in expert_params_mapping:
602
- param_name, weight_name, expert_id, shard_id = mapping
603
- if weight_name not in chunk_name:
604
- continue
605
-
606
- # Anyway, this is an expert weight and should not be
607
- # attempted to load as other weights later
608
- is_expert_weight = True
609
-
610
- # Do not modify `name` since the loop may continue here
611
- # Instead, create a new variable
612
- name_mapped = chunk_name.replace(weight_name, param_name)
613
-
614
- if is_pp_missing_parameter(name_mapped, self):
615
- continue
616
-
617
- param = params_dict[name_mapped]
618
- # We should ask the weight loader to return success
619
- # or not here since otherwise we may skip experts
620
- # with other available replicas.
621
- weight_loader = typing.cast(
622
- Callable[..., bool], param.weight_loader
623
- )
624
- success = weight_loader(
625
- param,
626
- weight_to_load,
627
- name_mapped,
628
- shard_id=shard_id,
629
- expert_id=expert_id,
630
- return_success=True,
631
- )
632
- if success:
633
- if not is_fusion_moe_shared_experts_layer:
634
- name = name_mapped
635
- else:
636
- loaded_params.add(name_mapped)
637
- break
638
- else:
639
- if is_expert_weight:
640
- # We've checked that this is an expert weight
641
- # However it's not mapped locally to this rank
642
- # So we simply skip it
643
- continue
644
-
645
- # Skip loading extra bias for GPTQ models.
646
- if name.endswith(".bias") and name not in params_dict:
647
- continue
648
-
649
- # Remapping the name of FP8 kv-scale.
650
- name = maybe_remap_kv_scale_name(name, params_dict)
651
- if name is None:
652
- continue
653
-
654
- if is_pp_missing_parameter(name, self):
655
- continue
656
-
657
- param = params_dict[name]
658
- weight_loader = getattr(
659
- param, "weight_loader", default_weight_loader
660
- )
661
- weight_loader(param, loaded_weight)
662
- if name is not None and not is_fusion_moe_shared_experts_layer:
663
- loaded_params.add(name)
664
-
665
- return loaded_params
496
+ loader = AutoWeightsLoader(self)
497
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
498
499
500
class Glm4MixtureOfExperts(MixtureOfExperts):
@@ -776,19 +608,3 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExper
608
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
609
loader = AutoWeightsLoader(self)
610
return loader.load_weights(weights)
779
-
780
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
781
- return self.model.get_expert_mapping()
782
-
783
-
784
-def get_spec_layer_idx_from_weight_name(
785
- config: Glm4MoeConfig, weight_name: str
786
-) -> int | None:
787
- if hasattr(config, "num_nextn_predict_layers") and (
788
- config.num_nextn_predict_layers > 0
789
- ):
790
- layer_idx = config.num_hidden_layers
791
- for i in range(config.num_nextn_predict_layers):
792
- if f"layers.{layer_idx + i}." in weight_name:
793
- return layer_idx + i
794
- return None
vllm/model_executor/models/glm4_moe_lite.py
+1
-13
@@ -70,6 +70,7 @@ from .interfaces import SupportsLoRA, SupportsPP
70
from .utils import (
71
AutoWeightsLoader,
72
PPMissingLayer,
73
+ get_spec_layer_idx_from_weight_name,
74
is_pp_missing_parameter,
75
make_empty_intermediate_tensors_factory,
76
make_layers,
@@ -628,16 +629,3 @@ class Glm4MoeLiteForCausalLM(
629
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
630
loader = AutoWeightsLoader(self)
631
return loader.load_weights(weights)
631
-
632
-
633
-def get_spec_layer_idx_from_weight_name(
634
- config: "Glm4MoeLiteConfig", weight_name: str
635
-) -> int | None:
636
- if hasattr(config, "num_nextn_predict_layers") and (
637
- config.num_nextn_predict_layers > 0
638
- ):
639
- layer_idx = config.num_hidden_layers
640
- for i in range(config.num_nextn_predict_layers):
641
- if f"layers.{layer_idx + i}." in weight_name:
642
- return layer_idx + i
643
- return None
vllm/model_executor/models/glm4_moe_lite_mtp.py
+1
-2
@@ -54,10 +54,9 @@ from .glm4_moe_lite import (
54
Glm4MixtureOfExperts,
55
Glm4MoeLite,
56
Glm4MoeLiteDecoderLayer,
57
- get_spec_layer_idx_from_weight_name,
57
)
58
from .interfaces import SupportsPP
60
-from .utils import maybe_prefix
59
+from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
60
61
62
class SharedHead(nn.Module):
vllm/model_executor/models/glm4_moe_mtp.py
+1
-2
@@ -51,9 +51,8 @@ from .glm4_moe import (
51
Glm4MixtureOfExperts,
52
Glm4MoE,
53
Glm4MoeDecoderLayer,
54
- get_spec_layer_idx_from_weight_name,
54
)
56
-from .utils import maybe_prefix
55
+from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
56
57
58
class SharedHead(nn.Module):
vllm/model_executor/models/glm_ocr_mtp.py
+2
-1
@@ -41,13 +41,14 @@ from vllm.model_executor.model_loader.weight_utils import (
41
from vllm.platforms import current_platform
42
from vllm.sequence import IntermediateTensors
43
44
-from .glm4 import Glm4DecoderLayer, get_spec_layer_idx_from_weight_name
44
+from .glm4 import Glm4DecoderLayer
45
from .glm4_moe_lite_mtp import (
46
Glm4MoeLiteMultiTokenPredictor,
47
SharedHead,
48
)
49
from .interfaces import SupportsPP
50
from .utils import (
51
+ get_spec_layer_idx_from_weight_name,
52
is_pp_missing_parameter,
53
maybe_prefix,
54
)
vllm/model_executor/models/glmasr.py
+16
-38
@@ -65,7 +65,12 @@ from .interfaces import (
65
SupportsPP,
66
SupportsTranscription,
67
)
68
-from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
68
+from .utils import (
69
+ AutoWeightsLoader,
70
+ WeightsMapper,
71
+ init_vllm_registered_model,
72
+ maybe_prefix,
73
+)
74
from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj
75
76
@@ -395,6 +400,14 @@ class GlmAsrEncoder(nn.Module):
400
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
401
}
402
403
+ hf_to_vllm_mapper = WeightsMapper(
404
+ orig_to_new_stacked={
405
+ ".q_proj": (".qkv_proj", "q"),
406
+ ".k_proj": (".qkv_proj", "k"),
407
+ ".v_proj": (".qkv_proj", "v"),
408
+ }
409
+ )
410
+
411
def __init__(
412
self,
413
config,
@@ -496,44 +509,9 @@ class GlmAsrEncoder(nn.Module):
509
return _GlmAsrEncoderOutput(last_hidden_state=hidden_states)
510
511
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
499
- """Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
500
- from vllm.model_executor.model_loader.weight_utils import default_weight_loader
501
-
512
weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
503
-
504
- stacked_params_mapping = [
505
- # (param_name, shard_name, shard_id)
506
- ("qkv_proj", "q_proj", "q"),
507
- ("qkv_proj", "k_proj", "k"),
508
- ("qkv_proj", "v_proj", "v"),
509
- ]
510
- params_dict = dict(self.named_parameters())
511
- loaded_params: set[str] = set()
512
-
513
- for name, loaded_weight in weights:
514
- for param_name, weight_name, shard_id in stacked_params_mapping:
515
- if weight_name not in name:
516
- continue
517
- name = name.replace(weight_name, param_name)
518
- # Skip loading extra bias for GPTQ models.
519
- if name.endswith(".bias") and name not in params_dict:
520
- continue
521
-
522
- param = params_dict[name]
523
- weight_loader = param.weight_loader
524
- weight_loader(param, loaded_weight, shard_id)
525
- break
526
- else:
527
- # Default weight loading for non-stacked params
528
- if name.endswith(".bias") and name not in params_dict:
529
- continue
530
- if name not in params_dict:
531
- continue
532
- param = params_dict[name]
533
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
534
- weight_loader(param, loaded_weight)
535
- loaded_params.add(name)
536
- return loaded_params
513
+ loader = AutoWeightsLoader(self)
514
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
515
516
517
class GlmAsrFeatureInputs(TensorSchema):
vllm/model_executor/models/gpt2.py
+17
-26
@@ -47,13 +47,11 @@ 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
50
from vllm.sequence import IntermediateTensors
51
52
from .interfaces import SupportsCrossEncoding, SupportsPP
53
from .utils import (
54
AutoWeightsLoader,
56
- is_pp_missing_parameter,
55
make_empty_intermediate_tensors_factory,
56
make_layers,
57
maybe_prefix,
@@ -241,32 +239,25 @@ class GPT2Model(nn.Module):
239
hidden_states = self.ln_f(hidden_states)
240
return hidden_states
241
244
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
245
- params_dict = dict(self.named_parameters(remove_duplicate=False))
246
- loaded_params: set[str] = set()
242
+ def _transpose_conv1d(
243
+ self, weights: Iterable[tuple[str, torch.Tensor]]
244
+ ) -> Iterable[tuple[str, torch.Tensor]]:
245
+ # HF's GPT-2 uses Conv1D instead of Linear, so its 2D weights are
246
+ # stored transposed relative to what vLLM expects.
247
+ # Note(zhuohan): the logic below might break quantized models.
248
for name, loaded_weight in weights:
248
- if ".attn.bias" in name or ".attn.masked_bias" in name:
249
- # Skip attention mask.
250
- # NOTE: "c_attn.bias" should not be skipped.
251
- continue
252
-
253
- if is_pp_missing_parameter(name, self):
254
- continue
255
-
256
- param = params_dict[name]
257
- # The HF's GPT-2 implementation uses Conv1D instead of Linear.
258
- # Because of this, we need to transpose the weights.
259
- # Note(zhuohan): the logic below might break quantized models.
260
- for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]:
261
- if conv1d_weight_name not in name:
262
- continue
263
- if not name.endswith(".weight"):
264
- continue
249
+ if name.endswith(".weight") and any(
250
+ proj in name for proj in ("c_attn", "c_proj", "c_fc")
251
+ ):
252
loaded_weight = loaded_weight.t()
266
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
267
- weight_loader(param, loaded_weight)
268
- loaded_params.add(name)
269
- return loaded_params
253
+ yield name, loaded_weight
254
+
255
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
256
+ # Skip attention mask buffers; NOTE: "c_attn.bias" must not be skipped.
257
+ loader = AutoWeightsLoader(
258
+ self, skip_substrs=[".attn.bias", ".attn.masked_bias"]
259
+ )
260
+ return loader.load_weights(self._transpose_conv1d(weights))
261
262
263
class GPT2LMHeadModel(nn.Module, SupportsPP):
vllm/model_executor/models/gpt_neox.py
+17
-39
@@ -43,13 +43,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
43
ParallelLMHead,
44
VocabParallelEmbedding,
45
)
46
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
46
from vllm.sequence import IntermediateTensors
47
48
from .interfaces import SupportsPP
49
from .utils import (
50
AutoWeightsLoader,
52
- is_pp_missing_parameter,
51
make_empty_intermediate_tensors_factory,
52
make_layers,
53
maybe_prefix,
@@ -249,45 +247,25 @@ class GPTNeoXModel(nn.Module):
247
hidden_states = self.final_layer_norm(hidden_states)
248
return hidden_states
249
252
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
253
- params_dict = dict(self.named_parameters())
254
- loaded_params: set[str] = set()
250
+ def _repack_qkv(
251
+ self, weights: Iterable[tuple[str, torch.Tensor]]
252
+ ) -> Iterable[tuple[str, torch.Tensor]]:
253
+ # GPT-NeoX's fused QKV is laid out as (num_heads * 3 * head_size) on
254
+ # its output dim (0), while vLLM expects (3 * num_heads * head_size).
255
+ num_heads = self.config.num_attention_heads
256
for name, loaded_weight in weights:
256
- if (
257
- "attention.bias" in name
258
- or "attention.masked_bias" in name
259
- or "rotary_emb.inv_freq" in name
260
- ):
261
- continue
262
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
263
- # Models trained using OpenRLHF may include
264
- # these tensors in the checkpoint. Skip them.
265
- continue
266
- if is_pp_missing_parameter(name, self):
267
- continue
268
- param = params_dict[name]
269
-
257
if "query_key_value" in name:
271
- # NOTE: GPT-NeoX's fused QKV's output_dim has the shape of
272
- # (num_heads * 3 * head_size), while the
273
- # required shape is (3 * num_heads * head_size).
274
- # Thus, we need weight conversion.
275
- output_dim = getattr(param, "output_dim", None)
276
- num_heads = self.config.num_attention_heads
277
- if output_dim is not None:
278
- loaded_weight_shape = loaded_weight.shape
279
- loaded_weight = loaded_weight.view(
280
- loaded_weight_shape[:output_dim]
281
- + (num_heads, 3, -1)
282
- + loaded_weight_shape[output_dim + 1 :]
283
- )
284
- loaded_weight = loaded_weight.transpose(output_dim, output_dim + 1)
285
- loaded_weight = loaded_weight.reshape(loaded_weight_shape)
286
-
287
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
288
- weight_loader(param, loaded_weight)
289
- loaded_params.add(name)
290
- return loaded_params
258
+ shape = loaded_weight.shape
259
+ loaded_weight = loaded_weight.view((num_heads, 3, -1) + shape[1:])
260
+ loaded_weight = loaded_weight.transpose(0, 1)
261
+ loaded_weight = loaded_weight.reshape(shape)
262
+ yield name, loaded_weight
263
+
264
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
265
+ loader = AutoWeightsLoader(
266
+ self, skip_substrs=["attention.bias", "attention.masked_bias"]
267
+ )
268
+ return loader.load_weights(self._repack_qkv(weights))
269
270
271
class GPTNeoXForCausalLM(nn.Module, SupportsPP):
vllm/model_executor/models/hy_v3.py
+1
-15
@@ -73,6 +73,7 @@ from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
73
from .utils import (
74
AutoWeightsLoader,
75
PPMissingLayer,
76
+ get_spec_layer_idx_from_weight_name,
77
is_pp_missing_parameter,
78
make_empty_intermediate_tensors_factory,
79
make_layers,
@@ -650,21 +651,6 @@ class HYV3Model(nn.Module, MixtureOfExperts):
651
return loaded_params
652
653
653
-def get_spec_layer_idx_from_weight_name(
654
- config: PretrainedConfig, weight_name: str
655
-) -> int | None:
656
- # HYV3MTP is enabled only when num_nextn_predict_layers is greater than 1
657
- if (
658
- hasattr(config, "num_nextn_predict_layers")
659
- and config.num_nextn_predict_layers > 0
660
- ):
661
- layer_idx = config.num_hidden_layers
662
- for i in range(config.num_nextn_predict_layers):
663
- if weight_name.startswith(f"model.layers.{layer_idx + i}."):
664
- return layer_idx + i
665
- return None
666
-
667
-
654
class HYV3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
655
packed_modules_mapping = {
656
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
vllm/model_executor/models/hy_v3_mtp.py
+6
-2
@@ -51,8 +51,12 @@ from vllm.v1.outputs import SamplerOutput
51
from vllm.v1.sample.metadata import SamplingMetadata
52
from vllm.v1.sample.sampler import Sampler
53
54
-from .hy_v3 import HYV3DecoderLayer, get_spec_layer_idx_from_weight_name
55
-from .utils import is_pp_missing_parameter, maybe_prefix
54
+from .hy_v3 import HYV3DecoderLayer
55
+from .utils import (
56
+ get_spec_layer_idx_from_weight_name,
57
+ is_pp_missing_parameter,
58
+ maybe_prefix,
59
+)
60
61
62
def _is_moe(config: PretrainedConfig) -> bool:
vllm/model_executor/models/intern_vit.py
+3
-9
@@ -36,8 +36,8 @@ from vllm.model_executor.layers.linear import (
36
RowParallelLinear,
37
)
38
from vllm.model_executor.layers.quantization import QuantizationConfig
39
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
39
40
+from .utils import AutoWeightsLoader
41
from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model
42
43
NORM2FN = {
@@ -445,11 +445,5 @@ class InternVisionModel(nn.Module):
445
return encoder_outputs
446
447
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
448
- params_dict = dict(self.named_parameters())
449
- loaded_params: set[str] = set()
450
- for name, loaded_weight in weights:
451
- param = params_dict[name]
452
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
453
- weight_loader(param, loaded_weight)
454
- loaded_params.add(name)
455
- return loaded_params
448
+ loader = AutoWeightsLoader(self)
449
+ return loader.load_weights(weights)
vllm/model_executor/models/interns1_vit.py
+4
-9
@@ -20,7 +20,8 @@ from vllm.model_executor.layers.conv import Conv2dLayer
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
22
from vllm.model_executor.layers.quantization import QuantizationConfig
23
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23
+
24
+from .utils import AutoWeightsLoader
25
26
NORM2FN = {
27
"rms_norm": RMSNorm,
@@ -433,11 +434,5 @@ class InternS1VisionModel(nn.Module):
434
return encoder_outputs
435
436
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
436
- params_dict = dict(self.named_parameters())
437
- loaded_params: set[str] = set()
438
- for name, loaded_weight in weights:
439
- param = params_dict[name]
440
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
441
- weight_loader(param, loaded_weight)
442
- loaded_params.add(name)
443
- return loaded_params
437
+ loader = AutoWeightsLoader(self)
438
+ return loader.load_weights(weights)
vllm/model_executor/models/jamba.py
+9
-90
@@ -14,10 +14,7 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.model_executor.layers.attention import Attention
17
-from vllm.model_executor.layers.fused_moe import (
18
- FusedMoE,
19
- fused_moe_make_expert_params_mapping,
20
-)
17
+from vllm.model_executor.layers.fused_moe import FusedMoE
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (
20
QKVParallelLinear,
@@ -38,7 +35,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35
ParallelLMHead,
36
VocabParallelEmbedding,
37
)
41
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.model_executor.models.llama import LlamaMLP as JambaMLP
39
from vllm.sequence import IntermediateTensors
40
@@ -52,7 +48,6 @@ from .interfaces import (
48
from .utils import (
49
AutoWeightsLoader,
50
WeightsMapper,
55
- is_pp_missing_parameter,
51
make_empty_intermediate_tensors_factory,
52
make_layers,
53
maybe_prefix,
@@ -378,87 +373,6 @@ class JambaModel(nn.Module):
373
hidden_states, _ = self.final_layernorm(hidden_states, residual)
374
return hidden_states
375
381
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
382
- # Params for weights, fp8 weight scales, fp8 activation scales
383
- # (param_name, weight_name, expert_id, shard_id)
384
- return fused_moe_make_expert_params_mapping(
385
- self,
386
- ckpt_gate_proj_name="gate_proj",
387
- ckpt_down_proj_name="down_proj",
388
- ckpt_up_proj_name="up_proj",
389
- num_experts=self.config.num_experts,
390
- )
391
-
392
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
393
- stacked_params_mapping = [
394
- # (param_name, shard_name, shard_id)
395
- ("qkv_proj", "q_proj", "q"),
396
- ("qkv_proj", "k_proj", "k"),
397
- ("qkv_proj", "v_proj", "v"),
398
- (".gate_up_proj", ".gate_proj", 0),
399
- (".gate_up_proj", ".up_proj", 1),
400
- ]
401
-
402
- params_dict = dict(self.named_parameters())
403
- loaded_params: set[str] = set()
404
- expert_params_mapping = self.get_expert_mapping()
405
- for name, loaded_weight in weights:
406
- if "rotary_emb.inv_freq" in name:
407
- continue
408
- for param_name, weight_name, shard_id in stacked_params_mapping:
409
- if weight_name not in name:
410
- continue
411
- if "experts" in name:
412
- continue
413
- name = name.replace(weight_name, param_name)
414
- # Skip loading extra bias for GPTQ models.
415
- if name.endswith(".bias") and name not in params_dict:
416
- continue
417
- # Skip layers on other devices.
418
- if is_pp_missing_parameter(name, self):
419
- continue
420
- param = params_dict[name]
421
- weight_loader = param.weight_loader
422
- weight_loader(param, loaded_weight, shard_id)
423
- break
424
- else:
425
- for (
426
- param_name,
427
- weight_name,
428
- expert_id,
429
- shard_id,
430
- ) in expert_params_mapping:
431
- if weight_name not in name:
432
- continue
433
-
434
- if is_pp_missing_parameter(name, self):
435
- continue
436
- name = name.replace(weight_name, param_name)
437
- param = params_dict[name]
438
- weight_loader = param.weight_loader
439
- weight_loader(
440
- param,
441
- loaded_weight,
442
- name,
443
- shard_id=shard_id,
444
- expert_id=expert_id,
445
- )
446
- break
447
- else:
448
- # Skip loading extra bias for GPTQ models.
449
- if name.endswith(".bias") and name not in params_dict:
450
- continue
451
- if is_pp_missing_parameter(name, self):
452
- continue
453
-
454
- param = params_dict[name]
455
- weight_loader = getattr(
456
- param, "weight_loader", default_weight_loader
457
- )
458
- weight_loader(param, loaded_weight)
459
- loaded_params.add(name)
460
- return loaded_params
461
-
376
377
class JambaForCausalLM(
378
nn.Module,
@@ -470,6 +384,14 @@ class JambaForCausalLM(
384
):
385
hf_to_vllm_mapper = WeightsMapper(
386
orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"},
387
+ orig_to_new_stacked={
388
+ ".q_proj": (".qkv_proj", "q"),
389
+ ".k_proj": (".qkv_proj", "k"),
390
+ ".v_proj": (".qkv_proj", "v"),
391
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
392
+ ".feed_forward.gate_proj": (".feed_forward.gate_up_proj", 0),
393
+ ".feed_forward.up_proj": (".feed_forward.gate_up_proj", 1),
394
+ },
395
)
396
packed_modules_mapping = {
397
"qkv_proj": [
@@ -577,9 +499,6 @@ class JambaForCausalLM(
499
loader = AutoWeightsLoader(self)
500
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
501
580
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
581
- return self.model.get_expert_mapping()
582
-
502
503
class JambaForSequenceClassification(JambaForCausalLM):
504
is_pooling_model = True
vllm/model_executor/models/kimi_linear.py
+1
-13
@@ -52,6 +52,7 @@ from .interfaces import HasInnerState, IsHybrid, MixtureOfExperts, SupportsPP
52
from .utils import (
53
AutoWeightsLoader,
54
PPMissingLayer,
55
+ get_spec_layer_idx_from_weight_name,
56
is_pp_missing_parameter,
57
make_layers,
58
maybe_prefix,
@@ -643,16 +644,3 @@ class KimiLinearForCausalLM(
644
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
645
)
646
return loader.load_weights(weights)
646
-
647
-
648
-def get_spec_layer_idx_from_weight_name(
649
- config: KimiLinearConfig, weight_name: str
650
-) -> int | None:
651
- if hasattr(config, "num_nextn_predict_layers") and (
652
- config.num_nextn_predict_layers > 0
653
- ):
654
- layer_idx = config.num_hidden_layers
655
- for i in range(config.num_nextn_predict_layers):
656
- if weight_name.startswith(f"model.layers.{layer_idx + i}."):
657
- return layer_idx + i
658
- return None
vllm/model_executor/models/laguna.py
+40
-169
@@ -2,8 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only Laguna model compatible with HuggingFace weights."""
4
5
-import typing
6
-from collections.abc import Callable, Iterable
5
+from collections.abc import Iterable
6
from itertools import islice
7
8
import torch
@@ -20,10 +19,7 @@ from vllm.distributed import (
19
)
20
from vllm.logger import init_logger
21
from vllm.model_executor.layers.attention import Attention
23
-from vllm.model_executor.layers.fused_moe import (
24
- FusedMoE,
25
- fused_moe_make_expert_params_mapping,
26
-)
22
+from vllm.model_executor.layers.fused_moe import FusedMoE
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
@@ -38,10 +34,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
34
ParallelLMHead,
35
VocabParallelEmbedding,
36
)
41
-from vllm.model_executor.model_loader.weight_utils import (
42
- default_weight_loader,
43
- maybe_remap_kv_scale_name,
44
-)
37
from vllm.model_executor.models.interfaces import (
38
EagleModelMixin,
39
SupportsEagle3,
@@ -51,8 +43,8 @@ from vllm.model_executor.models.interfaces import (
43
from vllm.model_executor.models.utils import (
44
AutoWeightsLoader,
45
PPMissingLayer,
46
+ WeightsMapper,
47
extract_layer_index,
55
- is_pp_missing_parameter,
48
make_empty_intermediate_tensors_factory,
49
make_layers,
50
maybe_prefix,
@@ -575,6 +567,14 @@ class LagunaDecoderLayer(nn.Module):
567
568
@support_torch_compile
569
class LagunaModel(nn.Module, EagleModelMixin):
570
+ hf_to_vllm_mapper = WeightsMapper(
571
+ orig_to_new_stacked={
572
+ ".q_proj": (".qkv_proj", "q"),
573
+ ".k_proj": (".qkv_proj", "k"),
574
+ ".v_proj": (".qkv_proj", "v"),
575
+ }
576
+ )
577
+
578
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
579
super().__init__()
580
@@ -675,167 +675,41 @@ class LagunaModel(nn.Module, EagleModelMixin):
675
return hidden_states, aux_hidden_states
676
return hidden_states
677
678
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
679
- """Get expert parameter mapping for weight loading.
680
-
681
- Returns mapping tuples of (param_name, weight_name, expert_id, shard_id)
682
- that handle both weights and quantization scales.
683
- """
684
- return fused_moe_make_expert_params_mapping(
685
- self,
686
- ckpt_gate_proj_name="gate_proj",
687
- ckpt_down_proj_name="down_proj",
688
- ckpt_up_proj_name="up_proj",
689
- num_experts=self.config.num_experts,
690
- num_redundant_experts=self.num_redundant_experts,
691
- )
692
-
693
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
694
- stacked_params_mapping = [
695
- # (param_name, shard_name, shard_id)
696
- ("qkv_proj", "q_proj", "q"),
697
- ("qkv_proj", "k_proj", "k"),
698
- ("qkv_proj", "v_proj", "v"),
699
- # gate_proj and up_proj are loaded as separate Linears (see
700
- # LagunaMLP) so no merge entry is needed here.
701
- ]
702
-
703
- # Suffixes to skip for GPTQ/modelopt models if param doesn't exist
704
- ignore_suffixes = (
705
- ".bias",
706
- "_bias",
707
- ".k_scale",
708
- "_k_scale",
709
- ".v_scale",
710
- "_v_scale",
711
- ".weight_scale",
712
- "_weight_scale",
713
- ".input_scale",
714
- "_input_scale",
715
- )
716
-
717
- params_dict = dict(self.named_parameters())
718
- loaded_params: set[str] = set()
719
- expert_params_mapping = self.get_expert_mapping()
720
-
678
+ def _slice_sink(
679
+ self, weights: Iterable[tuple[str, torch.Tensor]]
680
+ ) -> Iterable[tuple[str, torch.Tensor]]:
681
+ tp_size = get_tensor_model_parallel_world_size()
682
tp_rank = get_tensor_model_parallel_rank()
722
-
683
for name, loaded_weight in weights:
684
# Handle attention sinks (distributed across ranks). Derive the
685
# per-rank slice from the parameter's own shape so per-layer
686
# variations in head count are handled correctly.
687
if "sink" in name:
728
- param = params_dict.get(name)
729
- if param is not None:
730
- layer_heads_per_rank = param.shape[0]
731
- layer_head_start = tp_rank * layer_heads_per_rank
732
- narrow_weight = loaded_weight.narrow(
733
- 0, layer_head_start, layer_heads_per_rank
734
- )
735
- param.data.copy_(narrow_weight)
736
- loaded_params.add(name)
737
- continue
738
-
739
- # Handle stacked params (QKV, gate_up for
740
- # non-expert layers and shared_expert)
741
- for param_name, weight_name, shard_id in stacked_params_mapping:
742
- if weight_name not in name:
743
- continue
744
- # Skip expert weights - handled below via expert_params_mapping
745
- if "mlp.experts" in name and "shared_expert" not in name:
746
- continue
747
- name = name.replace(weight_name, param_name)
748
-
749
- if name.endswith(ignore_suffixes) and name not in params_dict:
750
- continue
751
- if is_pp_missing_parameter(name, self):
752
- continue
753
- # Remap FP8 kv_scale names for backwards compatibility
754
- if name.endswith("scale"):
755
- name = maybe_remap_kv_scale_name(name, params_dict)
756
- if name is None:
757
- continue
758
- if name not in params_dict:
759
- continue
760
-
761
- param = params_dict[name]
762
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
763
- if weight_loader == default_weight_loader:
764
- weight_loader(param, loaded_weight)
765
- else:
766
- weight_loader(param, loaded_weight, shard_id)
767
- loaded_params.add(name)
768
- break
769
- else:
770
- # Try expert params mapping (handles weights + quantization scales)
771
- is_expert_weight = False
772
- for mapping in expert_params_mapping:
773
- param_name, weight_name, expert_id, shard_id = mapping
774
- if weight_name not in name:
775
- continue
776
-
777
- # Mark as expert weight so we skip regular loading below
778
- is_expert_weight = True
779
-
780
- # Create mapped name without modifying original
781
- name_mapped = name.replace(weight_name, param_name)
782
-
783
- if is_pp_missing_parameter(name_mapped, self):
784
- continue
785
- if (
786
- name_mapped.endswith(ignore_suffixes)
787
- and name_mapped not in params_dict
788
- ):
789
- continue
790
- if name_mapped not in params_dict:
791
- continue
792
-
793
- param = params_dict[name_mapped]
794
- # Use return_success to handle expert parallelism correctly
795
- weight_loader = typing.cast(
796
- Callable[..., bool], param.weight_loader
797
- )
798
- success = weight_loader(
799
- param,
800
- loaded_weight,
801
- name_mapped,
802
- shard_id=shard_id,
803
- expert_id=expert_id,
804
- return_success=True,
805
- )
806
- if success:
807
- loaded_params.add(name_mapped)
808
- break
809
- else:
810
- # Expert weight not mapped to this rank - skip
811
- if is_expert_weight:
812
- continue
813
-
814
- # Remap kv_scale names before the ignore_suffixes filter:
815
- # the suffix list includes .k_scale/.v_scale, so filtering
816
- # first drops the checkpoint key before remap can rewrite
817
- # it to the .attn.* name that exists in params_dict.
818
- name = maybe_remap_kv_scale_name(name, params_dict)
819
- if name is None:
820
- continue
821
-
822
- if name.endswith(ignore_suffixes) and name not in params_dict:
823
- continue
824
-
825
- if is_pp_missing_parameter(name, self):
826
- continue
827
-
828
- if name not in params_dict:
829
- continue
830
-
831
- param = params_dict[name]
832
- weight_loader = getattr(
833
- param, "weight_loader", default_weight_loader
834
- )
835
- weight_loader(param, loaded_weight)
836
- loaded_params.add(name)
837
-
838
- return loaded_params
688
+ heads_per_rank = loaded_weight.shape[0] // tp_size
689
+ loaded_weight = loaded_weight.narrow(
690
+ 0, tp_rank * heads_per_rank, heads_per_rank
691
+ )
692
+ yield name, loaded_weight
693
+
694
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
695
+ loader = AutoWeightsLoader(
696
+ self,
697
+ ignore_unexpected_suffixes=[
698
+ ".bias",
699
+ "_bias",
700
+ ".k_scale",
701
+ "_k_scale",
702
+ ".v_scale",
703
+ "_v_scale",
704
+ ".weight_scale",
705
+ "_weight_scale",
706
+ ".input_scale",
707
+ "_input_scale",
708
+ ],
709
+ )
710
+ return loader.load_weights(
711
+ self._slice_sink(weights), mapper=self.hf_to_vllm_mapper
712
+ )
713
714
715
class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3):
@@ -892,9 +766,6 @@ class LagunaForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3):
766
logits = self.logits_processor(self.lm_head, hidden_states)
767
return logits
768
895
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
896
- return self.model.get_expert_mapping()
897
-
769
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
770
loader = AutoWeightsLoader(
771
self,
vllm/model_executor/models/lfm2_moe.py
+20
-107
@@ -15,10 +15,7 @@ from vllm.distributed import (
15
)
16
from vllm.model_executor.layers.activation import SiluAndMul
17
from vllm.model_executor.layers.attention import Attention
18
-from vllm.model_executor.layers.fused_moe import (
19
- FusedMoE,
20
- fused_moe_make_expert_params_mapping,
21
-)
18
+from vllm.model_executor.layers.fused_moe import FusedMoE
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (
21
MergedColumnParallelLinear,
@@ -40,10 +37,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
37
ParallelLMHead,
38
VocabParallelEmbedding,
39
)
43
-from vllm.model_executor.model_loader.weight_utils import (
44
- default_weight_loader,
45
- maybe_remap_moe_expert_param_name,
46
-)
40
from vllm.sequence import IntermediateTensors
41
from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig
42
@@ -60,7 +53,6 @@ from .utils import (
53
PPMissingLayer,
54
WeightsMapper,
55
extract_layer_index,
63
- is_pp_missing_parameter,
56
make_empty_intermediate_tensors_factory,
57
make_layers,
58
maybe_prefix,
@@ -167,6 +159,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module):
159
scoring_func="sigmoid",
160
e_score_correction_bias=self.gate.e_score_correction_bias,
161
routed_scaling_factor=self.routed_scaling_factor,
162
+ ckpt_names=("w1", "w2", "w3"),
163
)
164
165
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -402,6 +395,22 @@ class Lfm2MoeShortConvDecoderLayer(nn.Module):
395
396
@support_torch_compile
397
class Lfm2MoeModel(nn.Module):
398
+ hf_to_vllm_mapper = WeightsMapper(
399
+ orig_to_new_substr={
400
+ ".conv.": ".short_conv.",
401
+ "expert_bias": "gate.e_score_correction_bias",
402
+ },
403
+ orig_to_new_stacked={
404
+ ".q_proj": (".qkv_proj", "q"),
405
+ ".k_proj": (".qkv_proj", "k"),
406
+ ".v_proj": (".qkv_proj", "v"),
407
+ # Scoped (with trailing dots) to the dense MLP so routed experts.*.w1/w3
408
+ # (loaded by FusedMoE) are left untouched and .w1 does not match inside .w13
409
+ ".feed_forward.w1.": (".feed_forward.w13.", 0),
410
+ ".feed_forward.w3.": (".feed_forward.w13.", 1),
411
+ },
412
+ )
413
+
414
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
415
super().__init__()
416
@@ -487,102 +496,9 @@ class Lfm2MoeModel(nn.Module):
496
hidden_states, _ = self.embedding_norm(hidden_states, residual)
497
return hidden_states
498
490
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
491
- return fused_moe_make_expert_params_mapping(
492
- self,
493
- ckpt_gate_proj_name="w1",
494
- ckpt_down_proj_name="w2",
495
- ckpt_up_proj_name="w3",
496
- num_experts=self.config.num_experts,
497
- num_redundant_experts=self.num_redundant_experts,
498
- )
499
-
499
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
501
- stacked_params_mapping = [
502
- (".qkv_proj", ".q_proj", "q"),
503
- (".qkv_proj", ".k_proj", "k"),
504
- (".qkv_proj", ".v_proj", "v"),
505
- (".w13", ".w1", 0),
506
- (".w13", ".w3", 1),
507
- ]
508
- params_dict = dict(self.named_parameters())
509
- loaded_params: set[str] = set()
510
- expert_params_mapping = self.get_expert_mapping()
511
- for name, loaded_weight in weights:
512
- if "expert_bias" in name:
513
- name = name.replace("expert_bias", "gate.e_score_correction_bias")
514
-
515
- if ".conv." in name:
516
- name = name.replace(".conv.", ".short_conv.", 1)
517
-
518
- for param_name, weight_name, shard_id in stacked_params_mapping:
519
- # Skip non-stacked layers and experts (experts handled below).
520
- # Use segment-boundary matching (trailing dot) to prevent
521
- # e.g. ".w1" from matching inside ".w13" in pre-fused keys.
522
- if weight_name + "." not in name:
523
- continue
524
-
525
- if ("feed_forward.experts." in name) and name not in params_dict:
526
- continue
527
- name = name.replace(weight_name + ".", param_name + ".")
528
- # Skip loading extra bias for GPTQ models.
529
- if (
530
- name.endswith(".bias") or name.endswith("_bias")
531
- ) and name not in params_dict:
532
- continue
533
- # Skip layers on other devices.
534
- if is_pp_missing_parameter(name, self):
535
- continue
536
-
537
- param = params_dict[name]
538
- weight_loader = param.weight_loader
539
- weight_loader(param, loaded_weight, shard_id)
540
- break
541
- else:
542
- for mapping in expert_params_mapping:
543
- param_name, weight_name, expert_id, shard_id = mapping
544
-
545
- if weight_name not in name:
546
- continue
547
-
548
- name = name.replace(weight_name, param_name)
549
- # Skip layers on other devices.
550
- if is_pp_missing_parameter(name, self):
551
- continue
552
-
553
- # Skip loading extra bias for GPTQ models.
554
- if (
555
- name.endswith(".bias") or name.endswith("_bias")
556
- ) and name not in params_dict:
557
- continue
558
- param = params_dict[name]
559
-
560
- weight_loader = param.weight_loader
561
- weight_loader(
562
- param,
563
- loaded_weight,
564
- name,
565
- shard_id=shard_id,
566
- expert_id=expert_id,
567
- )
568
- break
569
- else:
570
- # Skip loading extra bias for GPTQ models.
571
- if (
572
- name.endswith(".bias") or name.endswith("_bias")
573
- ) and name not in params_dict:
574
- continue
575
- # Skip layers on other devices.
576
- if is_pp_missing_parameter(name, self):
577
- continue
578
- name = maybe_remap_moe_expert_param_name(name, params_dict)
579
- param = params_dict[name]
580
- weight_loader = getattr(
581
- param, "weight_loader", default_weight_loader
582
- )
583
- weight_loader(param, loaded_weight)
584
- loaded_params.add(name)
585
- return loaded_params
500
+ loader = AutoWeightsLoader(self, ignore_unexpected_suffixes=[".bias", "_bias"])
501
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
502
503
504
class Lfm2MoeForCausalLM(
@@ -763,6 +679,3 @@ class Lfm2MoeForCausalLM(
679
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
680
)
681
return loader.load_weights(weights)
766
-
767
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
768
- return self.model.get_expert_mapping()
vllm/model_executor/models/molmo.py
+1
-16
@@ -46,7 +46,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
46
ParallelLMHead,
47
VocabParallelEmbedding,
48
)
49
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
49
from vllm.model_executor.models.module_mapping import MultiModelKeys
50
from vllm.multimodal import MULTIMODAL_REGISTRY
51
from vllm.multimodal.inputs import (
@@ -76,7 +75,6 @@ from .interfaces import (
75
from .utils import (
76
AutoWeightsLoader,
77
WeightsMapper,
79
- is_pp_missing_parameter,
78
make_empty_intermediate_tensors_factory,
79
make_layers,
80
maybe_prefix,
@@ -888,20 +886,7 @@ class MolmoModel(nn.Module, SupportsQuant):
886
return hidden_states
887
888
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
891
- params_dict = dict(self.named_parameters())
892
- loaded_params: set[str] = set()
893
-
894
- for name, loaded_weight in weights:
895
- if name.endswith(".bias") and name not in params_dict:
896
- continue
897
- if is_pp_missing_parameter(name, self):
898
- continue
899
-
900
- param = params_dict[name]
901
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
902
- weight_loader(param, loaded_weight)
903
- loaded_params.add(name)
904
- return loaded_params
889
+ return AutoWeightsLoader(self).load_weights(weights)
890
891
892
def _lowest_multiple(x: int, k: int) -> int:
vllm/model_executor/models/molmo2.py
+1
-16
@@ -51,7 +51,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
51
ParallelLMHead,
52
VocabParallelEmbedding,
53
)
54
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
54
from vllm.model_executor.models.module_mapping import MultiModelKeys
55
from vllm.multimodal import MULTIMODAL_REGISTRY
56
from vllm.multimodal.inputs import (
@@ -89,7 +88,6 @@ from .utils import (
88
WeightsMapper,
89
_merge_multimodal_embeddings,
90
extract_layer_index,
92
- is_pp_missing_parameter,
91
make_empty_intermediate_tensors_factory,
92
make_layers,
93
maybe_prefix,
@@ -1221,20 +1219,7 @@ class Molmo2TextModel(nn.Module, SupportsQuant):
1219
return hidden_states
1220
1221
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1224
- params_dict = dict(self.named_parameters())
1225
- loaded_params: set[str] = set()
1226
-
1227
- for name, loaded_weight in weights:
1228
- if name.endswith(".bias") and name not in params_dict:
1229
- continue
1230
- if is_pp_missing_parameter(name, self):
1231
- continue
1232
-
1233
- param = params_dict[name]
1234
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1235
- weight_loader(param, loaded_weight)
1236
- loaded_params.add(name)
1237
- return loaded_params
1222
+ return AutoWeightsLoader(self).load_weights(weights)
1223
1224
1225
def get_patches_grid_size(
vllm/model_executor/models/nemotron_h.py
+11
-101
@@ -18,8 +18,7 @@
18
# limitations under the License.
19
"""Inference-only NemotronH model."""
20
21
-import typing
22
-from collections.abc import Callable, Iterable, Mapping
21
+from collections.abc import Iterable, Mapping
22
from itertools import islice
23
24
import torch
@@ -59,10 +58,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
58
ParallelLMHead,
59
VocabParallelEmbedding,
60
)
62
-from vllm.model_executor.model_loader.weight_utils import (
63
- default_weight_loader,
64
- maybe_remap_kv_scale_name,
65
-)
61
from vllm.model_executor.models.interfaces import (
62
EagleModelMixin,
63
HasInnerState,
@@ -78,7 +73,6 @@ from vllm.model_executor.models.interfaces import (
73
from vllm.model_executor.models.utils import (
74
AutoWeightsLoader,
75
WeightsMapper,
81
- is_pp_missing_parameter,
76
make_empty_intermediate_tensors_factory,
77
make_layers,
78
maybe_prefix,
@@ -222,6 +216,7 @@ class NemotronHMoE(nn.Module):
216
intermediate_size=config.moe_intermediate_size,
217
renormalize=config.norm_topk_prob,
218
quant_config=quant_config,
219
+ ckpt_names=("up_proj", "down_proj", ""),
220
use_grouped_topk=True,
221
num_expert_group=config.n_group,
222
topk_group=config.topk_group,
@@ -680,9 +675,12 @@ class NemotronHModel(nn.Module, EagleModelMixin):
675
return max_experts
676
677
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
678
+ # Consumed by `get_moe_expert_mapping` (bitsandbytes / LoRA); the main
679
+ # weight load is self-served by each `RoutedExperts` layer. Sized to the
680
+ # MAX expert count so heterogeneous puzzle models load every expert.
681
if self.has_moe:
682
# (param_name, weight_name, expert_id, shard_id)
685
- expert_params_mapping = fused_moe_make_expert_params_mapping(
683
+ return fused_moe_make_expert_params_mapping(
684
# - FusedMoe.w1 (aka gate_proj) should be up_proj since that's
685
# what the activation is applied to
686
# - FusedMoe.w3 (aka up_proj) should be ignored since we're
@@ -694,102 +692,9 @@ class NemotronHModel(nn.Module, EagleModelMixin):
692
num_experts=self._get_max_n_routed_experts(),
693
num_redundant_experts=getattr(self, "num_redundant_experts", 0),
694
)
697
- return expert_params_mapping
695
696
return []
697
701
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
702
- stacked_params_mapping = [
703
- # (param_name, shard_name, shard_id)
704
- ("qkv_proj", "q_proj", "q"),
705
- ("qkv_proj", "k_proj", "k"),
706
- ("qkv_proj", "v_proj", "v"),
707
- ]
708
-
709
- expert_params_mapping = self.get_expert_mapping()
710
-
711
- params_dict = dict(self.named_parameters())
712
- loaded_params: set[str] = set()
713
- for name, loaded_weight in weights:
714
- if "scale" in name or "zero_point" in name:
715
- # Remapping the name of FP8 kv-scale.
716
- name = maybe_remap_kv_scale_name(name, params_dict)
717
- if name is None:
718
- continue
719
-
720
- # Skip MTP/spec decode layers early (before stacked params mapping)
721
- if name.startswith("mtp."):
722
- continue
723
-
724
- # load stacked params
725
- for param_name, weight_name, shard_id in stacked_params_mapping:
726
- if weight_name not in name:
727
- continue
728
- name = name.replace(weight_name, param_name)
729
- # Skip loading extra bias for GPTQ models.
730
- if name.endswith(".bias") and name not in params_dict:
731
- continue
732
-
733
- if is_pp_missing_parameter(name, self):
734
- continue
735
-
736
- param = params_dict[name]
737
- weight_loader = param.weight_loader
738
- weight_loader(param, loaded_weight, shard_id)
739
- break
740
-
741
- # load other params
742
- else:
743
- is_expert_weight = False
744
- for mapping in expert_params_mapping:
745
- param_name, weight_name, expert_id, shard_id = mapping
746
- if weight_name not in name:
747
- continue
748
-
749
- # Anyway, this is an expert weight and should not be
750
- # attempted to load as other weights later
751
- is_expert_weight = True
752
-
753
- # Do not modify `name` since the loop may continue here
754
- # Instead, create a new variable
755
- name_mapped = name.replace(weight_name, param_name)
756
-
757
- if is_pp_missing_parameter(name_mapped, self):
758
- continue
759
- param = params_dict[name_mapped]
760
- # We should ask the weight loader to return success or not
761
- # here since otherwise we may skip experts with other
762
- # available replicas.
763
- weight_loader = typing.cast(
764
- Callable[..., bool], param.weight_loader
765
- )
766
- success = weight_loader(
767
- param,
768
- loaded_weight,
769
- name_mapped,
770
- shard_id=shard_id,
771
- expert_id=expert_id,
772
- return_success=True,
773
- )
774
- if success:
775
- name = name_mapped
776
- break
777
- else:
778
- if is_expert_weight:
779
- continue
780
-
781
- if is_pp_missing_parameter(name, self):
782
- continue
783
-
784
- param = params_dict[name]
785
- weight_loader = getattr(
786
- param, "weight_loader", default_weight_loader
787
- )
788
- weight_loader(param, loaded_weight)
789
-
790
- loaded_params.add(name)
791
- return loaded_params
792
-
698
699
class NemotronHForCausalLM(
700
nn.Module,
@@ -809,6 +714,11 @@ class NemotronHForCausalLM(
714
hf_to_vllm_mapper = WeightsMapper(
715
orig_to_new_prefix={"backbone": "model"},
716
orig_to_new_substr={"A_log": "A", "embeddings": "embed_tokens"},
717
+ orig_to_new_stacked={
718
+ ".q_proj": (".qkv_proj", "q"),
719
+ ".k_proj": (".qkv_proj", "k"),
720
+ ".v_proj": (".qkv_proj", "v"),
721
+ },
722
)
723
724
packed_modules_mapping = {
vllm/model_executor/models/olmo_hybrid.py
+20
-77
@@ -63,16 +63,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
63
ParallelLMHead,
64
VocabParallelEmbedding,
65
)
66
-from vllm.model_executor.model_loader.weight_utils import (
67
- default_weight_loader,
68
-)
66
from vllm.sequence import IntermediateTensors
67
68
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
69
from .utils import (
70
AutoWeightsLoader,
71
+ WeightsMapper,
72
extract_layer_index,
75
- is_pp_missing_parameter,
73
make_empty_intermediate_tensors_factory,
74
make_layers,
75
maybe_prefix,
@@ -302,6 +299,23 @@ class OlmoHybridDecoderLayer(nn.Module):
299
300
@support_torch_compile
301
class OlmoHybridModel(nn.Module):
302
+ hf_to_vllm_mapper = WeightsMapper(
303
+ orig_to_new_stacked={
304
+ ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
305
+ ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
306
+ ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
307
+ ".linear_attn.q_proj": (".linear_attn.in_proj_qkvg", 0),
308
+ ".linear_attn.k_proj": (".linear_attn.in_proj_qkvg", 1),
309
+ ".linear_attn.v_proj": (".linear_attn.in_proj_qkvg", 2),
310
+ ".linear_attn.g_proj": (".linear_attn.in_proj_qkvg", 3),
311
+ ".q_conv1d": (".conv1d", 0),
312
+ ".k_conv1d": (".conv1d", 1),
313
+ ".v_conv1d": (".conv1d", 2),
314
+ ".gate_proj": (".gate_up_proj", 0),
315
+ ".up_proj": (".gate_up_proj", 1),
316
+ }
317
+ )
318
+
319
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
320
super().__init__()
321
self.config = vllm_config.model_config.hf_config
@@ -359,79 +373,8 @@ class OlmoHybridModel(nn.Module):
373
return hidden_states
374
375
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
362
- stacked_params_mapping = [
363
- ("qkv_proj", "q_proj", "q"),
364
- ("qkv_proj", "k_proj", "k"),
365
- ("qkv_proj", "v_proj", "v"),
366
- ("gate_up_proj", "gate_proj", 0),
367
- ("gate_up_proj", "up_proj", 1),
368
- ]
369
-
370
- linear_attn_stacked_params_mapping = [
371
- ("in_proj_qkvg", "q_proj", 0),
372
- ("in_proj_qkvg", "k_proj", 1),
373
- ("in_proj_qkvg", "v_proj", 2),
374
- ("in_proj_qkvg", "g_proj", 3),
375
- ("conv1d", "q_conv1d", 0),
376
- ("conv1d", "k_conv1d", 1),
377
- ("conv1d", "v_conv1d", 2),
378
- ]
379
-
380
- params_dict = dict(self.named_parameters(remove_duplicate=False))
381
- loaded_params: set[str] = set()
382
-
383
- for name, loaded_weight in weights:
384
- if is_pp_missing_parameter(name, self):
385
- continue
386
-
387
- handled = False
388
-
389
- if "linear_attn" in name:
390
- for (
391
- param_name,
392
- weight_name,
393
- shard_id,
394
- ) in linear_attn_stacked_params_mapping:
395
- if weight_name not in name:
396
- continue
397
- mapped_name = name.replace(weight_name, param_name)
398
- if mapped_name.endswith(".bias") and (
399
- mapped_name not in params_dict
400
- ):
401
- continue
402
- if mapped_name not in params_dict:
403
- continue
404
- param = params_dict[mapped_name]
405
- weight_loader = param.weight_loader
406
- weight_loader(param, loaded_weight, shard_id)
407
- name = mapped_name
408
- handled = True
409
- break
410
- else:
411
- for param_name, weight_name, shard_id in stacked_params_mapping:
412
- if weight_name not in name:
413
- continue
414
- name = name.replace(weight_name, param_name)
415
- if name.endswith(".bias") and name not in params_dict:
416
- continue
417
- if name not in params_dict:
418
- continue
419
- param = params_dict[name]
420
- weight_loader = param.weight_loader
421
- weight_loader(param, loaded_weight, shard_id)
422
- handled = True
423
- break
424
-
425
- if not handled:
426
- if name.endswith(".bias") and name not in params_dict:
427
- continue
428
- if name not in params_dict:
429
- continue
430
- param = params_dict[name]
431
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
432
- weight_loader(param, loaded_weight)
433
- loaded_params.add(name)
434
- return loaded_params
376
+ loader = AutoWeightsLoader(self)
377
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
378
379
380
class OlmoHybridForCausalLM(
vllm/model_executor/models/openpangu.py
+39
-152
@@ -21,8 +21,7 @@
21
# See the License for the specific language governing permissions and
22
# limitations under the License.
23
24
-import typing
25
-from collections.abc import Callable, Iterable
24
+from collections.abc import Iterable
25
from typing import Any
26
27
import torch
@@ -44,10 +43,7 @@ from vllm.model_executor.layers.attention import (
43
Attention,
44
StaticSinkAttention,
45
)
47
-from vllm.model_executor.layers.fused_moe import (
48
- FusedMoE,
49
- fused_moe_make_expert_params_mapping,
50
-)
46
+from vllm.model_executor.layers.fused_moe import FusedMoE
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
ColumnParallelLinear,
@@ -64,10 +60,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
60
ParallelLMHead,
61
VocabParallelEmbedding,
62
)
67
-from vllm.model_executor.model_loader.weight_utils import (
68
- default_weight_loader,
69
- maybe_remap_kv_scale_name,
70
-)
63
from vllm.model_executor.models.interfaces import (
64
MixtureOfExperts,
65
SupportsLoRA,
@@ -76,8 +68,8 @@ from vllm.model_executor.models.interfaces import (
68
from vllm.model_executor.models.utils import (
69
AutoWeightsLoader,
70
PPMissingLayer,
71
+ WeightsMapper,
72
extract_layer_index,
80
- is_pp_missing_parameter,
73
make_empty_intermediate_tensors_factory,
74
make_layers,
75
maybe_prefix,
@@ -1056,153 +1048,48 @@ class OpenPanguModel(nn.Module):
1048
hidden_states, _ = self.norm(hidden_states, residual)
1049
return hidden_states
1050
1059
- def load_attn_mlp_weight(
1060
- self,
1061
- attn_mlp_replace_mapping: list[tuple[str, str, int]],
1062
- params_dict: dict[str, Any],
1063
- weight_name: str,
1064
- loaded_weight: torch.Tensor,
1065
- loaded_params: set[str],
1066
- ) -> bool:
1067
- for param_name, origin_name, shard_id in attn_mlp_replace_mapping:
1068
- if origin_name not in weight_name or (
1069
- ("mlp.experts." in weight_name) and weight_name not in params_dict
1070
- ):
1071
- continue
1072
- weight_name_mapped = weight_name.replace(origin_name, param_name)
1073
- if (
1074
- param_name == "fused_qkv_a_proj"
1075
- and weight_name_mapped not in params_dict
1076
- ):
1077
- continue
1078
- else:
1079
- weight_name = weight_name_mapped
1080
- if weight_name.endswith(".bias") and weight_name not in params_dict:
1081
- continue
1082
- if is_pp_missing_parameter(weight_name, self):
1083
- continue
1084
-
1085
- param = params_dict[weight_name]
1086
- weight_loader = param.weight_loader
1087
- weight_loader(param, loaded_weight, shard_id)
1088
- loaded_params.add(weight_name)
1089
- return True
1090
- return False
1091
-
1092
- def load_expert_weight(
1093
- self,
1094
- expert_merge_mapping: list[tuple[str, str, int, str]],
1095
- params_dict: dict[str, Any],
1096
- weight_name: str,
1097
- loaded_weight: torch.Tensor,
1098
- loaded_params: set[str],
1099
- flag_dict: dict[str, bool],
1100
- ) -> bool:
1101
- for mapping in expert_merge_mapping:
1102
- param_name, origin_name, expert_id, shard_id = mapping
1103
- if origin_name not in weight_name:
1104
- continue
1105
- flag_dict["is_expert_weight"] = True
1106
- weight_name_mapped = weight_name.replace(origin_name, param_name)
1107
- if is_pp_missing_parameter(weight_name_mapped, self):
1108
- continue
1109
- param = params_dict[weight_name_mapped]
1110
- weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
1111
- success = weight_loader(
1112
- param,
1113
- loaded_weight,
1114
- weight_name_mapped,
1115
- shard_id=shard_id,
1116
- expert_id=expert_id,
1117
- return_success=True,
1118
- )
1119
- if success:
1120
- weight_name = weight_name_mapped
1121
- loaded_params.add(weight_name_mapped)
1122
- return True
1123
- return False
1124
-
1125
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1126
- attn_mlp_replace_mapping = [
1127
- (".qkv_proj", ".q_proj", "q"),
1128
- (".qkv_proj", ".k_proj", "k"),
1129
- (".qkv_proj", ".v_proj", "v"),
1130
- (".fused_qkv_a_proj", ".q_a_proj", 0),
1131
- (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
1132
- (".gate_up_proj", ".gate_proj", 0),
1133
- (".gate_up_proj", ".up_proj", 1),
1134
- ]
1135
- has_experts = hasattr(self.config, "n_routed_experts")
1136
- if has_experts:
1137
- expert_merge_mapping = fused_moe_make_expert_params_mapping(
1138
- self,
1139
- ckpt_gate_proj_name="gate_proj",
1140
- ckpt_down_proj_name="down_proj",
1141
- ckpt_up_proj_name="up_proj",
1142
- num_experts=self.config.n_routed_experts,
1143
- num_redundant_experts=self.num_redundant_experts,
1144
- )
1145
-
1146
- params_dict = dict(self.named_parameters())
1147
- loaded_params: set[str] = set()
1051
+ def _filter_spec_layers(
1052
+ self, weights: Iterable[tuple[str, torch.Tensor]]
1053
+ ) -> Iterable[tuple[str, torch.Tensor]]:
1054
+ # Skip the MTP (spec-decode) layers that follow the main model's layers.
1055
+ num_mtp = getattr(self.config, "num_nextn_predict_layers", 0)
1056
for name, loaded_weight in weights:
1149
- if "rotary_emb.inv_freq" in name:
1150
- continue
1151
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
1152
- continue
1153
-
1154
- if (
1155
- "layers" in name
1156
- and hasattr(self.config, "num_nextn_predict_layers")
1157
- and (self.config.num_nextn_predict_layers > 0)
1158
- ):
1057
+ if "layers" in name and num_mtp > 0:
1058
layer_idx = int(name.split("layers.")[-1].split(".")[0])
1059
mtp_idx = layer_idx - self.config.num_hidden_layers
1161
- if mtp_idx >= 0 and mtp_idx < self.config.num_nextn_predict_layers:
1162
- continue # skip spec decode layers for main model
1163
-
1164
- flag_dict = {"is_expert_weight": False}
1165
- if (
1166
- self.load_attn_mlp_weight(
1167
- attn_mlp_replace_mapping,
1168
- params_dict,
1169
- name,
1170
- loaded_weight,
1171
- loaded_params,
1172
- )
1173
- or has_experts
1174
- and self.load_expert_weight(
1175
- expert_merge_mapping,
1176
- params_dict,
1177
- name,
1178
- loaded_weight,
1179
- loaded_params,
1180
- flag_dict,
1181
- )
1182
- ):
1183
- continue
1184
- else:
1185
- if flag_dict["is_expert_weight"]:
1186
- continue
1187
- if name.endswith(".bias") and name not in params_dict:
1188
- continue
1189
- name = maybe_remap_kv_scale_name(name, params_dict)
1190
- if name.endswith("e_score_correction_bias"):
1191
- name = name.replace(
1192
- "e_score_correction_bias", "gate.e_score_correction_bias"
1193
- )
1194
- if name is None:
1195
- continue
1196
- if is_pp_missing_parameter(name, self):
1060
+ if 0 <= mtp_idx < num_mtp:
1061
continue
1062
+ yield name, loaded_weight
1063
1199
- param = params_dict[name]
1200
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1201
- weight_loader(param, loaded_weight)
1202
- loaded_params.add(name)
1203
-
1064
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1065
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
1066
+ stacked: dict[str, tuple[str, int | str]] = {
1067
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
1068
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
1069
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
1070
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
1071
+ }
1072
+ # The attention layout is config-dependent: MLA with a fused low-rank
1073
+ # projection, MLA without it, or standard qkv. Only add the stacked
1074
+ # entries whose fused target actually exists as a parameter.
1075
+ param_names = list(dict(self.named_parameters()))
1076
+ if any(".fused_qkv_a_proj." in n for n in param_names):
1077
+ stacked[".q_a_proj"] = (".fused_qkv_a_proj", 0)
1078
+ stacked[".kv_a_proj_with_mqa"] = (".fused_qkv_a_proj", 1)
1079
+ if any(".qkv_proj." in n for n in param_names):
1080
+ stacked[".q_proj"] = (".qkv_proj", "q")
1081
+ stacked[".k_proj"] = (".qkv_proj", "k")
1082
+ stacked[".v_proj"] = (".qkv_proj", "v")
1083
+ mapper = WeightsMapper(
1084
+ orig_to_new_substr={
1085
+ "e_score_correction_bias": "gate.e_score_correction_bias",
1086
+ },
1087
+ orig_to_new_stacked=stacked,
1088
+ )
1089
+ loader = AutoWeightsLoader(self)
1090
+ loaded = loader.load_weights(self._filter_spec_layers(weights), mapper=mapper)
1091
self.post_weight_load()
1205
- return loaded_params
1092
+ return loaded
1093
1094
def post_weight_load(self) -> None:
1095
for name, module in self.named_modules():
vllm/model_executor/models/param2moe.py
+36
-183
@@ -32,10 +32,7 @@ from vllm.distributed import (
32
)
33
from vllm.model_executor.layers.activation import SiluAndMul
34
from vllm.model_executor.layers.attention import Attention
35
-from vllm.model_executor.layers.fused_moe import (
36
- FusedMoE,
37
- fused_moe_make_expert_params_mapping,
38
-)
35
+from vllm.model_executor.layers.fused_moe import FusedMoE
36
from vllm.model_executor.layers.layernorm import RMSNorm
37
from vllm.model_executor.layers.linear import (
38
MergedColumnParallelLinear,
@@ -49,14 +46,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
46
ParallelLMHead,
47
VocabParallelEmbedding,
48
)
52
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
49
from vllm.sequence import IntermediateTensors
50
51
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
52
from .utils import (
53
AutoWeightsLoader,
54
PPMissingLayer,
59
- is_pp_missing_parameter,
55
+ WeightsMapper,
56
make_empty_intermediate_tensors_factory,
57
make_layers,
58
maybe_prefix,
@@ -74,45 +70,17 @@ def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
70
return t - t.mean()
71
72
77
-def _rename_and_normalize_weights(
73
+def _normalize_expert_bias(
74
weights: Iterable[tuple[str, torch.Tensor]],
75
) -> Iterator[tuple[str, torch.Tensor]]:
80
- """
81
- Translate HuggingFace Param2MoE weight names to vLLM internal names
82
- and zero-mean the expert-bias tensor so the router stays balanced.
83
-
84
- Mapping table (HF → vLLM):
85
- model.word_embeddings.* → model.embed_tokens.*
86
- *.attention.query_key_value.* → *.self_attn.qkv_proj.*
87
- *.attention.dense.* → *.self_attn.o_proj.*
88
- *.attention.query_layernorm.* → *.self_attn.q_layernorm.*
89
- *.attention.key_layernorm.* → *.self_attn.k_layernorm.*
90
- *.mlp.gate.expert_bias → *.mlp.gate.e_score_correction_bias
91
- (also zero-meant for load balance)
76
+ """Zero-mean the MoE router's per-expert score bias for load balance.
77
+
78
+ The rename to ``e_score_correction_bias`` is done by the mapper; only the
79
+ tensor adjustment lives here, since a WeightsMapper cannot transform data.
80
"""
81
for name, w in weights:
94
- # Embedding table
95
- name = name.replace("model.word_embeddings.", "model.embed_tokens.")
96
- # Fused QKV projection (HF: query_key_value → vLLM: qkv_proj)
97
- name = name.replace(".attention.query_key_value.", ".self_attn.qkv_proj.")
98
- # Output projection (HF: dense → vLLM: o_proj)
99
- name = name.replace(".attention.dense.", ".self_attn.o_proj.")
100
- # Per-head query norm
101
- name = name.replace(".attention.query_layernorm.", ".self_attn.q_layernorm.")
102
- # Per-head key norm
103
- name = name.replace(".attention.key_layernorm.", ".self_attn.k_layernorm.")
104
- # Catch any remaining .attention. → .self_attn. prefixes
105
- # (e.g. future bias params on the projection layers)
106
- name = name.replace(".attention.", ".self_attn.")
107
-
108
- # Expert-score bias: rename + zero-mean
109
- if name.endswith(".mlp.gate.expert_bias"):
110
- name = name.replace(
111
- ".mlp.gate.expert_bias",
112
- ".mlp.gate.e_score_correction_bias",
113
- )
82
+ if _is_expert_bias_name(name):
83
w = _zero_mean_tensor(w)
115
-
84
yield name, w
85
86
@@ -123,7 +91,8 @@ class Param2MoEAttention(nn.Module):
91
Notable differences from a vanilla GQA layer:
92
* The checkpoint fuses Q, K, V into a single ``query_key_value`` weight.
93
vLLM receives it already renamed to ``qkv_proj`` by the weight-name
126
- translator and splits it during ``load_weights``.
94
+ translator and loads it directly; ``QKVParallelLinear`` splits the
95
+ fused ``[Q|K|V]`` tensor internally.
96
* Optional per-head RMS norms on Q and K (``use_qk_norm=True``).
97
"""
98
@@ -472,6 +441,28 @@ class Param2MoEDecoderLayer(nn.Module):
441
442
443
class Param2MoEModel(nn.Module):
444
+ hf_to_vllm_mapper = WeightsMapper(
445
+ orig_to_new_substr={
446
+ "word_embeddings.": "embed_tokens.",
447
+ # Fused query_key_value is renamed to qkv_proj and loaded directly;
448
+ # QKVParallelLinear splits the fused [Q|K|V] tensor internally.
449
+ ".attention.query_key_value.": ".self_attn.qkv_proj.",
450
+ ".attention.dense.": ".self_attn.o_proj.",
451
+ ".attention.query_layernorm.": ".self_attn.q_layernorm.",
452
+ ".attention.key_layernorm.": ".self_attn.k_layernorm.",
453
+ # Catch-all for any remaining .attention. names (must come last).
454
+ ".attention.": ".self_attn.",
455
+ ".mlp.gate.expert_bias": ".mlp.gate.e_score_correction_bias",
456
+ },
457
+ orig_to_new_stacked={
458
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
459
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
460
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
461
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
462
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
463
+ },
464
+ )
465
+
466
def __init__(
467
self,
468
*,
@@ -558,147 +549,9 @@ class Param2MoEModel(nn.Module):
549
self,
550
weights: Iterable[tuple[str, torch.Tensor]],
551
) -> set[str]:
561
- """
562
- Custom weight loader for the inner Param2MoEModel.
563
-
564
- Receives weights that have already been renamed/normalised by the
565
- outer model and whose ``model.`` prefix has been stripped by
566
- ``AutoWeightsLoader``. Handles:
567
- 1. Fused QKV split (query_key_value → qkv_proj q/k/v shards).
568
- 2. gate_proj + up_proj → gate_up_proj stacking (dense + shared-exp).
569
- 3. Routed-expert weights via the fused-MoE mapping.
570
- 4. All remaining weights via their default loader.
571
- """
572
- config = self.config
573
- num_heads: int = config.num_attention_heads
574
- num_kv_heads: int = config.num_key_value_heads
575
- head_dim: int = config.head_dim or (config.hidden_size // num_heads)
576
- q_split = num_heads * head_dim
577
- kv_split = num_kv_heads * head_dim
578
-
579
- stacked_params_mapping = [
580
- # (vllm_param_name, ckpt_weight_name, shard_id)
581
- ("gate_up_proj", "gate_proj", 0),
582
- ("gate_up_proj", "up_proj", 1),
583
- ]
584
-
585
- params_dict = dict(self.named_parameters(remove_duplicate=False))
586
- loaded_params: set[str] = set()
587
- expert_params_mapping = self.get_expert_mapping()
588
-
589
- for name, loaded_weight in weights:
590
- # ------------------------------------------------------------------
591
- # 1. Fused QKV: split into q / k / v shards for QKVParallelLinear
592
- # ------------------------------------------------------------------
593
- if name.endswith(".self_attn.qkv_proj.weight"):
594
- if name not in params_dict:
595
- continue
596
- if is_pp_missing_parameter(name, self):
597
- continue
598
-
599
- param = params_dict[name]
600
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
601
- q_w = loaded_weight[:q_split, :]
602
- k_w = loaded_weight[q_split : q_split + kv_split, :]
603
- v_w = loaded_weight[q_split + kv_split :, :]
604
- weight_loader(param, q_w, "q")
605
- weight_loader(param, k_w, "k")
606
- weight_loader(param, v_w, "v")
607
- loaded_params.add(name)
608
- continue
609
-
610
- # ------------------------------------------------------------------
611
- # 2. gate_proj / up_proj → gate_up_proj (dense MLP + shared-exp.)
612
- # ------------------------------------------------------------------
613
- matched_stacked = False
614
- for param_name, weight_name, shard_id in stacked_params_mapping:
615
- if weight_name not in name:
616
- continue
617
- if "mlp.experts" in name: # routed experts handled below
618
- continue
619
- new_name = name.replace(weight_name, param_name)
620
- if new_name.endswith(".bias") and new_name not in params_dict:
621
- continue
622
- if new_name not in params_dict:
623
- continue
624
- if is_pp_missing_parameter(new_name, self):
625
- continue
626
-
627
- param = params_dict[new_name]
628
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
629
- weight_loader(param, loaded_weight, shard_id)
630
- loaded_params.add(new_name)
631
- matched_stacked = True
632
- break
633
-
634
- if matched_stacked:
635
- continue
636
-
637
- # ------------------------------------------------------------------
638
- # 3. Routed expert weights → fused-MoE kernel layout
639
- # ------------------------------------------------------------------
640
- matched_expert = False
641
- for (
642
- param_name,
643
- weight_name,
644
- expert_id,
645
- shard_id,
646
- ) in expert_params_mapping:
647
- if weight_name not in name:
648
- continue
649
- new_name = name.replace(weight_name, param_name)
650
- if is_pp_missing_parameter(new_name, self):
651
- continue
652
- if new_name not in params_dict:
653
- continue
654
-
655
- param = params_dict[new_name]
656
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
657
- weight_loader(
658
- param,
659
- loaded_weight,
660
- name,
661
- shard_id=shard_id,
662
- expert_id=expert_id,
663
- )
664
- loaded_params.add(new_name)
665
- matched_expert = True
666
- break
667
-
668
- if matched_expert:
669
- continue
670
-
671
- # ------------------------------------------------------------------
672
- # 4. All other weights: direct load (layernorms, embed_tokens, …)
673
- # ------------------------------------------------------------------
674
- if name.endswith(".bias") and name not in params_dict:
675
- continue
676
- if name not in params_dict:
677
- continue
678
- if is_pp_missing_parameter(name, self):
679
- continue
680
-
681
- param = params_dict[name]
682
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
683
- try:
684
- weight_loader(param, loaded_weight)
685
- except Exception as e:
686
- raise RuntimeError(
687
- f"[param2moe] Failed to load weight '{name}' "
688
- f"with shape {tuple(loaded_weight.shape)} "
689
- f"into param type {type(param).__name__}: {e}"
690
- ) from e
691
- loaded_params.add(name)
692
-
693
- return loaded_params
694
-
695
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
696
- return fused_moe_make_expert_params_mapping(
697
- self,
698
- ckpt_gate_proj_name="gate_proj",
699
- ckpt_down_proj_name="down_proj",
700
- ckpt_up_proj_name="up_proj",
701
- num_experts=self.config.num_experts,
552
+ loader = AutoWeightsLoader(self)
553
+ return loader.load_weights(
554
+ _normalize_expert_bias(weights), mapper=self.hf_to_vllm_mapper
555
)
556
557
@@ -859,4 +712,4 @@ class Param2MoEForCausalLM(
712
weights: Iterable[tuple[str, torch.Tensor]],
713
) -> set[str]:
714
loader = AutoWeightsLoader(self)
862
- return loader.load_weights(_rename_and_normalize_weights(weights))
715
+ return loader.load_weights(weights)
vllm/model_executor/models/qwen3_5.py
+10
-9
@@ -93,6 +93,7 @@ from .utils import (
93
extract_layer_index,
94
make_empty_intermediate_tensors_factory,
95
make_layers,
96
+ maybe_fuse_shared_experts,
97
maybe_prefix,
98
)
99
@@ -261,20 +262,20 @@ class Qwen3_5Model(Qwen3NextModel):
262
self.aux_hidden_state_layers: tuple[int, ...] = ()
263
264
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
264
- mapper = self.hf_to_vllm_mapper
265
# FSE must match construction (Qwen3NextSparseMoeBlock): reroute the
266
# shared expert into the extra fused slot only when AITER FSE is both
267
# requested and compatible with the quant spec.
268
- is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and (
269
- _is_shared_expert_fse_compatible(self.quant_config)
270
- )
271
- if is_fse:
272
- num_routed = self.config.num_experts
273
- mapper = mapper | WeightsMapper(
274
- orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
268
+ if "moe" in self.config.model_type:
269
+ weights = maybe_fuse_shared_experts(
270
+ weights,
271
+ enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
272
+ and _is_shared_expert_fse_compatible(self.quant_config),
273
+ n_routed_experts=self.config.num_experts,
274
+ n_shared_experts=1,
275
+ ckpt_prefix="mlp.shared_expert",
276
)
277
loader = AutoWeightsLoader(self)
277
- return loader.load_weights(weights, mapper=mapper)
278
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
279
280
281
class Qwen3_5ForCausalLMBase(
vllm/model_executor/models/qwen3_5_mtp.py
+11
-10
@@ -41,9 +41,9 @@ from .interfaces import (
41
from .utils import (
42
AutoWeightsLoader,
43
PPMissingLayer,
44
- WeightsMapper,
44
_merge_multimodal_embeddings,
45
make_empty_intermediate_tensors_factory,
46
+ maybe_fuse_shared_experts,
47
maybe_prefix,
48
)
49
@@ -174,17 +174,18 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
174
return hidden_states
175
176
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
177
- mapper = self.hf_to_vllm_mapper
178
- is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and (
179
- _is_shared_expert_fse_compatible(get_current_vllm_config().quant_config)
177
+ weights = maybe_fuse_shared_experts(
178
+ weights,
179
+ enabled=rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
180
+ and _is_shared_expert_fse_compatible(
181
+ get_current_vllm_config().quant_config
182
+ ),
183
+ n_routed_experts=getattr(self.config, "num_experts", 0),
184
+ n_shared_experts=1,
185
+ ckpt_prefix="mlp.shared_expert",
186
)
181
- if is_fse:
182
- num_routed = getattr(self.config, "num_experts", 0)
183
- mapper = mapper | WeightsMapper(
184
- orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
185
- )
187
loader = AutoWeightsLoader(self)
187
- return loader.load_weights(weights, mapper=mapper)
188
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
189
190
191
@support_torch_compile(
vllm/model_executor/models/qwen3_dflash.py
+26
-43
@@ -31,10 +31,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
31
ParallelLMHead,
32
VocabParallelEmbedding,
33
)
34
-from vllm.model_executor.model_loader.weight_utils import (
35
- default_weight_loader,
36
- maybe_remap_kv_scale_name,
37
-)
34
from vllm.multimodal.inputs import NestedTensors
35
from vllm.transformers_utils.config import set_default_rope_theta
36
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
@@ -44,6 +40,7 @@ from .qwen2 import Qwen2MLP as Qwen3MLP
40
from .qwen3 import Qwen3ForCausalLM
41
from .utils import (
42
AutoWeightsLoader,
43
+ WeightsMapper,
44
get_draft_quant_config,
45
maybe_prefix,
46
process_eagle_weight,
@@ -333,6 +330,17 @@ class DFlashQwen3DecoderLayer(nn.Module):
330
331
@support_torch_compile
332
class DFlashQwen3Model(nn.Module):
333
+ hf_to_vllm_mapper = WeightsMapper(
334
+ orig_to_new_substr={"midlayer.": "layers.0."},
335
+ orig_to_new_stacked={
336
+ ".q_proj": (".qkv_proj", "q"),
337
+ ".k_proj": (".qkv_proj", "k"),
338
+ ".v_proj": (".qkv_proj", "v"),
339
+ ".gate_proj": (".gate_up_proj", 0),
340
+ ".up_proj": (".gate_up_proj", 1),
341
+ },
342
+ )
343
+
344
def __init__(
345
self,
346
*,
@@ -624,51 +632,26 @@ class DFlashQwen3Model(nn.Module):
632
hidden_states, _ = self.norm(hidden_states, residual)
633
return hidden_states
634
627
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
628
- stacked_params_mapping = [
629
- (".qkv_proj", ".q_proj", "q"),
630
- (".qkv_proj", ".k_proj", "k"),
631
- (".qkv_proj", ".v_proj", "v"),
632
- (".gate_up_proj", ".gate_proj", 0),
633
- (".gate_up_proj", ".up_proj", 1),
634
- ]
635
- params_dict = dict(self.named_parameters())
636
- loaded_params: set[str] = set()
637
- tp_rank = get_tensor_model_parallel_rank()
635
+ def _preprocess(
636
+ self, weights: Iterable[tuple[str, torch.Tensor]]
637
+ ) -> Iterable[tuple[str, torch.Tensor]]:
638
tp_size = get_tensor_model_parallel_world_size()
639
+ tp_rank = get_tensor_model_parallel_rank()
640
for name, loaded_weight in weights:
640
- if "midlayer." in name:
641
- name = name.replace("midlayer.", "layers.0.")
642
- if "scale" in name:
643
- name = maybe_remap_kv_scale_name(name, params_dict)
644
- if name is None:
645
- continue
641
if "attention_sink_bias" in name:
647
- if name not in params_dict:
648
- continue
642
# Sink bias is per-head; shard it across TP ranks like the
643
# attention heads themselves.
651
- param = params_dict[name]
644
heads_per_rank = loaded_weight.shape[0] // tp_size
653
- head_start = tp_rank * heads_per_rank
654
- narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank)
655
- param.data.copy_(narrow_weight)
656
- loaded_params.add(name)
657
- continue
658
- for param_name, weight_name, shard_id in stacked_params_mapping:
659
- if weight_name not in name:
660
- continue
661
- name = name.replace(weight_name, param_name)
662
- param = params_dict[name]
663
- weight_loader = param.weight_loader
664
- weight_loader(param, loaded_weight, shard_id)
665
- break
666
- else:
667
- param = params_dict[name]
668
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
669
- weight_loader(param, loaded_weight)
670
- loaded_params.add(name)
671
- return loaded_params
645
+ loaded_weight = loaded_weight.narrow(
646
+ 0, tp_rank * heads_per_rank, heads_per_rank
647
+ )
648
+ yield name, loaded_weight
649
+
650
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
651
+ loader = AutoWeightsLoader(self)
652
+ return loader.load_weights(
653
+ self._preprocess(weights), mapper=self.hf_to_vllm_mapper
654
+ )
655
656
657
class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
vllm/model_executor/models/qwen3_next.py
+8
-10
@@ -71,6 +71,7 @@ from .utils import (
71
extract_layer_index,
72
make_empty_intermediate_tensors_factory,
73
make_layers,
74
+ maybe_fuse_shared_experts,
75
maybe_prefix,
76
)
77
@@ -713,17 +714,14 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
714
return hidden_states
715
716
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
716
- mapper = self.hf_to_vllm_mapper
717
- if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
718
- # AITER fused-shared-experts: route the shared_expert checkpoint
719
- # weights into the extra fused expert slot. Merge (not mutate) so the
720
- # shared class mapper isn't permanently altered.
721
- num_routed = getattr(self.config, "num_experts", 0)
722
- mapper = mapper | WeightsMapper(
723
- orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
724
- )
717
+ weights = maybe_fuse_shared_experts(
718
+ weights,
719
+ n_routed_experts=getattr(self.config, "num_experts", 0),
720
+ n_shared_experts=1,
721
+ ckpt_prefix="mlp.shared_expert",
722
+ )
723
loader = AutoWeightsLoader(self)
726
- return loader.load_weights(weights, mapper=mapper)
724
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
725
726
727
class QwenNextMixtureOfExperts(MixtureOfExperts):
vllm/model_executor/models/qwen3_next_mtp.py
+8
-11
@@ -7,7 +7,6 @@ from collections.abc import Iterable
7
import torch
8
from torch import nn
9
10
-from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed.parallel_state import get_pp_group
@@ -30,8 +29,8 @@ from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
29
30
from .utils import (
31
AutoWeightsLoader,
33
- WeightsMapper,
32
make_empty_intermediate_tensors_factory,
33
+ maybe_fuse_shared_experts,
34
maybe_prefix,
35
)
36
@@ -154,16 +153,14 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
153
return hidden_states
154
155
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
157
- mapper = self.hf_to_vllm_mapper
158
- if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
159
- # AITER fused-shared-experts: route the shared_expert checkpoint
160
- # weights into the extra fused expert slot.
161
- num_routed = self.config.num_experts
162
- mapper = mapper | WeightsMapper(
163
- orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
164
- )
156
+ weights = maybe_fuse_shared_experts(
157
+ weights,
158
+ n_routed_experts=self.config.num_experts,
159
+ n_shared_experts=1,
160
+ ckpt_prefix="mlp.shared_expert",
161
+ )
162
loader = AutoWeightsLoader(self)
166
- return loader.load_weights(weights, mapper=mapper)
163
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
164
165
166
@support_torch_compile
vllm/model_executor/models/sarvam.py
+16
-97
@@ -35,11 +35,7 @@ from vllm.distributed import (
35
get_tensor_model_parallel_world_size,
36
)
37
from vllm.model_executor.layers.activation import SiluAndMul
38
-from vllm.model_executor.layers.fused_moe import (
39
- FusedMoE,
40
- MoERunner,
41
- fused_moe_make_expert_params_mapping,
42
-)
38
+from vllm.model_executor.layers.fused_moe import FusedMoE, MoERunner
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
ColumnParallelLinear,
@@ -55,7 +51,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
51
ParallelLMHead,
52
VocabParallelEmbedding,
53
)
58
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
54
from vllm.sequence import IntermediateTensors
55
56
from .bailing_moe import BailingMoeForCausalLM
@@ -63,7 +58,7 @@ from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
58
from .utils import (
59
AutoWeightsLoader,
60
PPMissingLayer,
66
- is_pp_missing_parameter,
61
+ WeightsMapper,
62
make_empty_intermediate_tensors_factory,
63
make_layers,
64
maybe_prefix,
@@ -448,6 +443,16 @@ class SarvamMLABlock(nn.Module):
443
444
445
class SarvamMLAModel(nn.Module):
446
+ hf_to_vllm_mapper = WeightsMapper(
447
+ orig_to_new_stacked={
448
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
449
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
450
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
451
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
452
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
453
+ }
454
+ )
455
+
456
def __init__(
457
self,
458
*,
@@ -532,97 +537,14 @@ class SarvamMLAModel(nn.Module):
537
hidden_states, _ = self.norm(hidden_states, residual)
538
return hidden_states
539
535
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
536
- return fused_moe_make_expert_params_mapping(
537
- self,
538
- ckpt_gate_proj_name="gate_proj",
539
- ckpt_down_proj_name="down_proj",
540
- ckpt_up_proj_name="up_proj",
541
- num_experts=self.config.num_experts,
542
- )
543
-
540
def load_weights(
541
self,
542
weights: Iterable[tuple[str, torch.Tensor]],
543
) -> set[str]:
548
- """Load weights with stacked gate+up and MoE expert remapping."""
549
- weights = _normalized_weights(weights)
550
- stacked_params_mapping = [
551
- ("gate_up_proj", "gate_proj", 0),
552
- ("gate_up_proj", "up_proj", 1),
553
- ]
554
-
555
- params_dict = dict(self.named_parameters(remove_duplicate=False))
556
- loaded_params: set[str] = set()
557
- expert_params_mapping = self.get_expert_mapping()
558
-
559
- for name, loaded_weight in weights:
560
- for param_name, weight_name, shard_id in stacked_params_mapping:
561
- if weight_name not in name:
562
- continue
563
- if "mlp.experts" in name:
564
- continue
565
- new_name = name.replace(weight_name, param_name)
566
- if new_name.endswith(".bias") and new_name not in params_dict:
567
- continue
568
- if new_name not in params_dict:
569
- continue
570
- if is_pp_missing_parameter(new_name, self):
571
- continue
572
-
573
- param = params_dict[new_name]
574
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
575
- weight_loader(param, loaded_weight, shard_id)
576
- loaded_params.add(new_name)
577
- break
578
- else:
579
- mapped = False
580
- for (
581
- param_name,
582
- weight_name,
583
- expert_id,
584
- shard_id,
585
- ) in expert_params_mapping:
586
- if weight_name not in name:
587
- continue
588
-
589
- new_name = name.replace(weight_name, param_name)
590
- if is_pp_missing_parameter(new_name, self):
591
- continue
592
- if new_name not in params_dict:
593
- continue
594
-
595
- param = params_dict[new_name]
596
- weight_loader = getattr(
597
- param, "weight_loader", default_weight_loader
598
- )
599
- weight_loader(
600
- param,
601
- loaded_weight,
602
- name,
603
- shard_id=shard_id,
604
- expert_id=expert_id,
605
- )
606
- loaded_params.add(new_name)
607
- mapped = True
608
- break
609
-
610
- if mapped:
611
- continue
612
-
613
- if name.endswith(".bias") and name not in params_dict:
614
- continue
615
- if name not in params_dict:
616
- continue
617
- if is_pp_missing_parameter(name, self):
618
- continue
619
-
620
- param = params_dict[name]
621
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
622
- weight_loader(param, loaded_weight)
623
- loaded_params.add(name)
624
-
625
- return loaded_params
544
+ loader = AutoWeightsLoader(self)
545
+ return loader.load_weights(
546
+ _normalized_weights(weights), mapper=self.hf_to_vllm_mapper
547
+ )
548
549
550
class SarvamMixtureOfExperts(MixtureOfExperts):
@@ -764,9 +686,6 @@ class SarvamMLAForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SarvamMixtureOfE
686
)
687
return loader.load_weights(weights)
688
767
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
768
- return self.model.get_expert_mapping()
769
-
689
690
class SarvamMoEForCausalLM(BailingMoeForCausalLM):
691
"""Same as BailingMoeForCausalLM, but normalizes gate expert_bias pre-load."""
vllm/model_executor/models/step3p5.py
+1
-15
@@ -53,6 +53,7 @@ from .utils import (
53
PPMissingLayer,
54
WeightsMapper,
55
extract_layer_index,
56
+ get_spec_layer_idx_from_weight_name,
57
is_pp_missing_parameter,
58
make_empty_intermediate_tensors_factory,
59
make_layers,
@@ -978,18 +979,3 @@ class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
979
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
980
loader = AutoWeightsLoader(self)
981
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
981
-
982
-
983
-def get_spec_layer_idx_from_weight_name(
984
- config: ModelConfig, weight_name: str
985
-) -> int | None:
986
- if hasattr(config, "num_nextn_predict_layers") and (
987
- config.num_nextn_predict_layers > 0
988
- ):
989
- layer_idx = config.num_hidden_layers
990
- for i in range(config.num_nextn_predict_layers):
991
- if weight_name.startswith(
992
- f"layers.{layer_idx + i}." # Step3p5Model
993
- ) or weight_name.startswith(f"model.layers.{layer_idx + i}."): # Step3p5MTP
994
- return layer_idx + i
995
- return None
vllm/model_executor/models/step3p5_mtp.py
+2
-2
@@ -18,8 +18,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
18
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
19
from vllm.sequence import IntermediateTensors
20
21
-from .step3p5 import Step3p5DecoderLayer, get_spec_layer_idx_from_weight_name
22
-from .utils import maybe_prefix
21
+from .step3p5 import Step3p5DecoderLayer
22
+from .utils import get_spec_layer_idx_from_weight_name, maybe_prefix
23
24
logger = init_logger(__name__)
25
vllm/model_executor/models/telechat2.py
+34
-58
@@ -26,19 +26,23 @@ import torch
26
import torch.nn as nn
27
28
from vllm.config import VllmConfig
29
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29
from vllm.model_executor.models.llama import LlamaForCausalLM, LlamaModel
30
31
from .llama import LlamaDecoderLayer
33
-from .utils import (
34
- AutoWeightsLoader,
35
- PPMissingLayer,
36
- WeightsMapper,
37
- is_pp_missing_parameter,
38
-)
32
+from .utils import AutoWeightsLoader, PPMissingLayer, WeightsMapper
33
34
35
class TeleChat2Model(LlamaModel):
36
+ hf_to_vllm_mapper = WeightsMapper(
37
+ orig_to_new_stacked={
38
+ ".query": (".qkv_proj", "q"),
39
+ ".key": (".qkv_proj", "k"),
40
+ ".value": (".qkv_proj", "v"),
41
+ ".gate_proj": (".gate_up_proj", 0),
42
+ ".up_proj": (".gate_up_proj", 1),
43
+ }
44
+ )
45
+
46
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
47
hf_config = vllm_config.model_config.hf_config
48
@@ -65,62 +69,34 @@ class TeleChat2Model(LlamaModel):
69
layer.mlp.gate_up_proj.bias = None
70
layer.mlp.gate_up_proj.skip_bias_add = True
71
68
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
69
- stacked_params_mapping = [
70
- ("gate_up_proj", "gate_proj", 0),
71
- ("gate_up_proj", "up_proj", 1),
72
- ]
73
- params_dict = dict(self.named_parameters())
74
- loaded_params: set[str] = set()
72
+ def _split_key_value(
73
+ self, weights: Iterable[tuple[str, torch.Tensor]]
74
+ ) -> Iterable[tuple[str, torch.Tensor]]:
75
+ # TeleChat2 stores k/v as a single per-head-interleaved `key_value`
76
+ # tensor. De-interleave it into separate k/v so the qkv_proj mapper
77
+ # can stack them.
78
total_num_heads = self.config.n_head
79
head_dim = self.config.hidden_size // total_num_heads
80
for name, loaded_weight in weights:
81
if "self_attn.key_value" in name:
79
- k_weight = []
80
- v_weight = []
81
- for i in range(total_num_heads):
82
- start = i * head_dim * 2
83
- k_weight.append(loaded_weight[start : start + head_dim, :])
84
- v_weight.append(
85
- loaded_weight[start + head_dim : start + 2 * head_dim :]
86
- )
87
- k_weight = torch.cat(k_weight, dim=0)
88
- v_weight = torch.cat(v_weight, dim=0)
89
- name = name.replace("key_value", "qkv_proj")
90
- if is_pp_missing_parameter(name, self):
91
- continue
92
- param = params_dict[name]
93
- weight_loader = param.weight_loader
94
- weight_loader(param, k_weight, "k")
95
- weight_loader(param, v_weight, "v")
96
- elif "query" in name:
97
- name = name.replace("query", "qkv_proj")
98
- if is_pp_missing_parameter(name, self):
99
- continue
100
- param = params_dict[name]
101
- weight_loader = param.weight_loader
102
- weight_loader(param, loaded_weight, "q")
82
+ starts = [i * head_dim * 2 for i in range(total_num_heads)]
83
+ k_weight = torch.cat(
84
+ [loaded_weight[s : s + head_dim, :] for s in starts], dim=0
85
+ )
86
+ v_weight = torch.cat(
87
+ [loaded_weight[s + head_dim : s + 2 * head_dim, :] for s in starts],
88
+ dim=0,
89
+ )
90
+ yield name.replace("key_value", "key"), k_weight
91
+ yield name.replace("key_value", "value"), v_weight
92
else:
104
- for param_name, weight_name, shard_id in stacked_params_mapping:
105
- if weight_name not in name:
106
- continue
107
- name = name.replace(weight_name, param_name)
108
- if is_pp_missing_parameter(name, self):
109
- continue
110
- param = params_dict[name]
111
- weight_loader = param.weight_loader
112
- weight_loader(param, loaded_weight, shard_id)
113
- break
114
- else:
115
- if is_pp_missing_parameter(name, self):
116
- continue
117
- param = params_dict[name]
118
- weight_loader = getattr(
119
- param, "weight_loader", default_weight_loader
120
- )
121
- weight_loader(param, loaded_weight)
122
- loaded_params.add(name)
123
- return loaded_params
93
+ yield name, loaded_weight
94
+
95
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
96
+ loader = AutoWeightsLoader(self)
97
+ return loader.load_weights(
98
+ self._split_key_value(weights), mapper=self.hf_to_vllm_mapper
99
+ )
100
101
102
class TeleChat2ForCausalLM(LlamaForCausalLM):
vllm/model_executor/models/utils.py
+111
@@ -35,6 +35,7 @@ if TYPE_CHECKING:
35
from transformers import PretrainedConfig
36
from transformers.conversion_mapping import WeightRenaming
37
38
+ from vllm.config.model import ModelConfig
39
from vllm.model_executor.layers.quantization import QuantizationConfig
40
41
logger = init_logger(__name__)
@@ -424,6 +425,116 @@ class AutoWeightsLoader:
425
return autoloaded_weights
426
427
428
+def maybe_fuse_shared_experts(
429
+ weights: Iterable[tuple[str, torch.Tensor]],
430
+ *,
431
+ n_routed_experts: int,
432
+ n_shared_experts: int,
433
+ ckpt_prefix: str = "mlp.shared_experts",
434
+ enabled: bool | None = None,
435
+) -> Iterable[tuple[str, torch.Tensor]]:
436
+ """Route AITER fused-shared-expert checkpoint weights into fused slots.
437
+
438
+ When AITER fused-shared-experts is active, shared experts are packed into
439
+ the routed expert tensor. The checkpoint stores them under `ckpt_prefix`
440
+ as a single (possibly widened) tensor; this splits it into
441
+ `n_shared_experts` chunks named `mlp.experts.{n_routed_experts + j}` so
442
+ the `RoutedExperts` loader treats them as extra experts. Yields the input
443
+ unchanged when the fusion is inactive, so callers can wrap unconditionally.
444
+
445
+ Args:
446
+ weights: Iterable of `(name, tensor)` checkpoint pairs.
447
+ n_routed_experts: Number of routed experts; offsets the fused slots.
448
+ n_shared_experts: Number of shared experts packed into the tensor.
449
+ ckpt_prefix: Checkpoint module name of the shared experts.
450
+ enabled: Whether AITER fused-shared-experts is active. Defaults to
451
+ `rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()`; pass an
452
+ explicit value only when the model gates on something more (e.g.
453
+ quant-spec compatibility) and it must match its construction-time
454
+ decision.
455
+
456
+ Yields:
457
+ `(name, tensor)` pairs with shared experts routed to fused slots.
458
+ """
459
+ if enabled is None:
460
+ from vllm._aiter_ops import rocm_aiter_ops
461
+
462
+ enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
463
+ if not enabled:
464
+ yield from weights
465
+ return
466
+
467
+ # Match on the dotted boundary so e.g. "mlp.shared_expert." does not also
468
+ # catch a sibling "mlp.shared_expert_gate".
469
+ prefix = f"{ckpt_prefix}."
470
+ for name, loaded_weight in weights:
471
+ if prefix not in name:
472
+ yield name, loaded_weight
473
+ continue
474
+ # gate/up split on the output dim; down_proj on the input dim.
475
+ split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0
476
+ total = loaded_weight.shape[split_dim]
477
+ if total % n_shared_experts != 0:
478
+ raise ValueError(
479
+ f"FSE shared-expert weight {name!r} has size {total} along axis "
480
+ f"{split_dim}, not divisible by n_shared_experts={n_shared_experts}."
481
+ )
482
+ chunk = total // n_shared_experts
483
+ for j in range(n_shared_experts):
484
+ sl = slice(j * chunk, (j + 1) * chunk)
485
+ if loaded_weight.ndim == 1:
486
+ chunk_weight = loaded_weight[sl]
487
+ elif split_dim == 0:
488
+ chunk_weight = loaded_weight[sl, :]
489
+ else:
490
+ chunk_weight = loaded_weight[:, sl]
491
+ yield (
492
+ name.replace(prefix, f"mlp.experts.{n_routed_experts + j}."),
493
+ chunk_weight,
494
+ )
495
+
496
+
497
+def get_spec_layer_idx_from_weight_name(
498
+ config: "ModelConfig", weight_name: str
499
+) -> int | None:
500
+ """Return the MTP layer index a weight belongs to, or None.
501
+
502
+ Args:
503
+ config: The model config; must expose `num_hidden_layers` and, for MTP
504
+ checkpoints, `num_nextn_predict_layers`.
505
+ weight_name: Checkpoint weight name to classify.
506
+
507
+ Returns:
508
+ The absolute layer index for an MTP-layer weight, else None.
509
+ """
510
+ if not (n := getattr(config, "num_nextn_predict_layers", 0)):
511
+ return None
512
+ base = config.num_hidden_layers
513
+ for i in range(n):
514
+ if weight_name.startswith((f"model.layers.{base + i}.", f"layers.{base + i}.")):
515
+ return base + i
516
+ return None
517
+
518
+
519
+def skip_spec_layers(
520
+ weights: Iterable[tuple[str, torch.Tensor]], config: "ModelConfig"
521
+) -> Iterable[tuple[str, torch.Tensor]]:
522
+ """Drop MTP spec-layer weights (loaded by the MTP head, not the base model).
523
+
524
+ Args:
525
+ weights: Iterable of `(name, tensor)` checkpoint pairs.
526
+ config: The model config, passed to `get_spec_layer_idx_from_weight_name`.
527
+
528
+ Yields:
529
+ `(name, tensor)` pairs whose weight is not an MTP-layer weight.
530
+ """
531
+ return (
532
+ (name, w)
533
+ for name, w in weights
534
+ if get_spec_layer_idx_from_weight_name(config, name) is None
535
+ )
536
+
537
+
538
def init_vllm_registered_model(
539
vllm_config: VllmConfig,
540
*,