Remove more unnecessary `load_weights` methods (#47058)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jun 30, 2026 at 15:22 UTC
1ab952293587a53dd80c339324596f4168482290
63 files changed
+1028
-3778
tests/models/multimodal/processing/test_moss_audio.py
-48
@@ -17,7 +17,6 @@ from vllm.model_executor.models.moss_audio import (
17
MOSS_AUDIO_PLACEHOLDER,
18
MOSS_AUDIO_TOKEN,
19
MOSS_AUDIO_TOKEN_ID,
20
- GatedMLP,
20
MossAudioConfig,
21
MossAudioDummyInputsBuilder,
22
MossAudioEncoder,
@@ -594,53 +593,6 @@ def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
593
assert set(forward_tensors.tensors) == set(tensors.tensors)
594
595
597
-@pytest.mark.parametrize("tp_size", [1, 2])
598
-def test_moss_audio_gated_mlp_tp_shapes_and_loading(monkeypatch, tp_size):
599
- from vllm.config import VllmConfig, set_current_vllm_config
600
- from vllm.config.device import DeviceConfig
601
-
602
- _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=tp_size)
603
- with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
604
- mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
605
-
606
- params = dict(mlp.named_parameters())
607
- assert params["gate_up_proj.weight"].shape == torch.Size([16 // tp_size, 4])
608
- assert params["down_proj.weight"].shape == torch.Size([6, 8 // tp_size])
609
-
610
- gate_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
611
- up_weight = torch.arange(100, 132, dtype=torch.float32).reshape(8, 4)
612
- down_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8)
613
- loaded = mlp.load_weights(
614
- [
615
- ("gate_proj.weight", gate_weight),
616
- ("up_proj.weight", up_weight),
617
- ("down_proj.weight", down_weight),
618
- ]
619
- )
620
-
621
- assert loaded == {"gate_up_proj.weight", "down_proj.weight"}
622
- shard = 8 // tp_size
623
- assert torch.equal(params["gate_up_proj.weight"][:shard], gate_weight[:shard])
624
- assert torch.equal(params["gate_up_proj.weight"][shard:], up_weight[:shard])
625
- assert torch.equal(params["down_proj.weight"], down_weight[:, : 8 // tp_size])
626
-
627
- with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
628
- packed_mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
629
- packed_params = dict(packed_mlp.named_parameters())
630
- loaded = packed_mlp.load_weights(
631
- [("gate_up_proj.weight", torch.cat([gate_weight, up_weight], dim=0))]
632
- )
633
- assert loaded == {"gate_up_proj.weight"}
634
- assert torch.equal(
635
- packed_params["gate_up_proj.weight"][:shard],
636
- gate_weight[:shard],
637
- )
638
- assert torch.equal(
639
- packed_params["gate_up_proj.weight"][shard:],
640
- up_weight[:shard],
641
- )
642
-
643
-
596
def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
597
from vllm.config import VllmConfig, set_current_vllm_config
598
from vllm.config.device import DeviceConfig
vllm/lora/layers/fused_moe.py
+1
@@ -39,6 +39,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
39
self._ep_check()
40
41
routed_experts = self.base_layer.routed_experts
42
+ routed_experts.lora_base_layer_prefix = "base_layer."
43
assert not routed_experts.quant_method.is_monolithic, (
44
"Monolithic kernels are not supported for Fused MoE LoRA."
45
)
vllm/lora/utils.py
+21
-27
@@ -372,32 +372,26 @@ def process_packed_modules_mapping(
372
model: nn.Module, force_2d_moe: bool = False
373
) -> dict[str, list[str]]:
374
if is_moe_model(model):
375
- if moe_packed_mapping := get_moe_expert_mapping(model):
376
- # This method generates and returns a dictionary mapping packed module
377
- # names to lists of their corresponding submodule names. It includes
378
- # both static mappings and dynamic mappings for expert layers, where
379
- # the expert indices are expanded based on the configured number
380
- # of routed experts.
381
- packed_modules_mapping = get_packed_modules_mapping(model)
382
- # The 2D mapping is needed when the model itself is 2D, or when
383
- # the engine forces the universal 2D wrapper via
384
- # enable_mixed_moe_lora_format (so 3D models can also load 2D
385
- # adapters through FusedMoEWithLoRA).
386
- if (not model.is_3d_moe_weight) or force_2d_moe:
387
- # Filter out malformed entries: non-gated MoE has empty
388
- # ckpt_up_proj_name which results in weight_name containing ".."
389
- # (e.g., "experts.0.." instead of "experts.0.layer_name.")
390
- packed_modules_mapping["experts"] = [
391
- weight_name.rstrip(".")
392
- for _, weight_name, _, _ in moe_packed_mapping
393
- if ".." not in weight_name
394
- ]
395
-
396
- return packed_modules_mapping
397
- else:
398
- raise AttributeError(
399
- "To support LoRA for MoE model, "
400
- "'get_expert_mapping' must be implemented"
401
- )
375
+ # This method generates and returns a dictionary mapping packed module
376
+ # names to lists of their corresponding submodule names. It includes
377
+ # both static mappings and dynamic mappings for expert layers, where
378
+ # the expert indices are expanded based on the configured number
379
+ # of routed experts.
380
+ packed_modules_mapping = get_packed_modules_mapping(model)
381
+ # The 2D mapping is needed when the model itself is 2D, or when
382
+ # the engine forces the universal 2D wrapper via
383
+ # enable_mixed_moe_lora_format (so 3D models can also load 2D
384
+ # adapters through FusedMoEWithLoRA).
385
+ if (not model.is_3d_moe_weight) or force_2d_moe:
386
+ # Filter out malformed entries: non-gated MoE has empty
387
+ # ckpt_up_proj_name which results in weight_name containing ".."
388
+ # (e.g., "experts.0.." instead of "experts.0.layer_name.")
389
+ packed_modules_mapping["experts"] = [
390
+ weight_name.rstrip(".")
391
+ for _, weight_name, _, _ in get_moe_expert_mapping(model)
392
+ if ".." not in weight_name
393
+ ]
394
+
395
+ return packed_modules_mapping
396
else:
397
return get_packed_modules_mapping(model)
vllm/model_executor/layers/fused_moe/layer.py
+4
-2
@@ -127,7 +127,7 @@ def FusedMoE(
127
num_redundant_experts: int = 0,
128
has_bias: bool = False,
129
is_sequence_parallel: bool = False,
130
- expert_mapping: list[tuple[str, str, int, str]] | None = None,
130
+ ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
131
n_shared_experts: int | None = None,
132
router_logits_dtype: torch.dtype | None = None,
133
gate: torch.nn.Module | None = None,
@@ -356,7 +356,9 @@ def FusedMoE(
356
moe_config,
357
quant_config,
358
expert_map_manager=expert_map_manager,
359
- expert_mapping=expert_mapping,
359
+ ckpt_gate_proj_name=ckpt_names[0],
360
+ ckpt_down_proj_name=ckpt_names[1],
361
+ ckpt_up_proj_name=ckpt_names[2],
362
# Extra params that are needed by quant_methods, pass along for now
363
# Prefer getting these from other sources, e.g. moe_config or
364
# router object
vllm/model_executor/layers/fused_moe/routed_experts.py
+108
-24
@@ -58,7 +58,9 @@ class RoutedExperts(PluggableLayer):
58
moe_config: FusedMoEConfig,
59
quant_config: QuantizationConfig | None,
60
expert_map_manager: ExpertMapManager,
61
- expert_mapping: list[tuple[str, str, int, str]] | None = None,
61
+ ckpt_gate_proj_name: str = "gate_proj",
62
+ ckpt_down_proj_name: str = "down_proj",
63
+ ckpt_up_proj_name: str = "up_proj",
64
#
65
# Extra params that are needed by quant_methods, pass along for now
66
# Prefer getting these from other sources, e.g. moe_config or
@@ -81,7 +83,9 @@ class RoutedExperts(PluggableLayer):
83
self.layer_name = layer_name
84
self.moe_config = moe_config
85
self.quant_config = quant_config
84
- self.expert_mapping = expert_mapping
86
+ self.ckpt_gate_proj_name = ckpt_gate_proj_name
87
+ self.ckpt_down_proj_name = ckpt_down_proj_name
88
+ self.ckpt_up_proj_name = ckpt_up_proj_name
89
self.expert_map_manager = expert_map_manager
90
self.hidden_size = moe_config.hidden_dim
91
self.global_num_experts = moe_config.num_experts
@@ -167,6 +171,8 @@ class RoutedExperts(PluggableLayer):
171
172
self.quant_method.create_weights(layer=self, **moe_quant_params)
173
174
+ self.lora_base_layer_prefix = ""
175
+
176
# TODO(bnell): Temporary hack. Get rid of this.
177
def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
178
self.quant_method = quant_method
@@ -863,26 +869,33 @@ class RoutedExperts(PluggableLayer):
869
def load_weights(
870
self, weights: Iterable[tuple[str, torch.Tensor]]
871
) -> Iterable[str]:
866
- if (expert_mapping := self.expert_mapping) is None:
867
- raise ValueError(
868
- "`self.expert_mapping` must be provided to "
869
- "load weights using `self.load_weights`."
870
- )
872
+ expert_mapping = self.get_expert_mapping(include_fused=True)
873
+ unpadded_hidden = self.moe_config.hidden_dim_unpadded
874
for expert_name, loaded_weight in weights:
875
qual_name = f"{self.layer_name}.{expert_name}"
876
+ # Fused expert weights can be identified by their 3D tensors
877
+ is_fused = loaded_weight.dim() == 3
878
+ matched = False
879
for param_name, weight_name, expert_id, shard_id in expert_mapping:
880
if weight_name not in qual_name:
881
+ if matched and is_fused:
882
+ break
883
continue
884
+ matched = True
885
weight_name = qual_name.replace(weight_name, param_name)
886
param_name = weight_name.removeprefix(f"{self.layer_name}.")
887
param = getattr(self, param_name)
879
- # Fused expert weights can be identified by their 3D tensors
880
- if loaded_weight.dim() == 3:
881
- # Repurpose expert_id as shard_idx for deconcatenating w1 and w3
888
+ if is_fused:
889
if shard_id in {"w1", "w3"}:
883
- shard_idx = expert_id
884
- experts_shard = loaded_weight.chunk(2, dim=1)[shard_idx]
890
+ if loaded_weight.shape[-1] != unpadded_hidden:
891
+ # [..., hidden, intermediate] -> [..., intermediate, hidden]
892
+ loaded_weight = loaded_weight.transpose(-1, -2)
893
+ # Repurpose expert_id for deconcatenating w1 and w3
894
+ experts_shard = loaded_weight.chunk(2, dim=1)[expert_id]
895
else:
896
+ if loaded_weight.shape[-2] != unpadded_hidden:
897
+ # [..., intermediate, hidden] -> [..., hidden, intermediate]
898
+ loaded_weight = loaded_weight.transpose(-1, -2)
899
experts_shard = loaded_weight
900
start = 0
901
else:
@@ -912,6 +925,27 @@ class RoutedExperts(PluggableLayer):
925
)
926
yield param_name
927
928
+ def get_expert_mapping(
929
+ self,
930
+ ckpt_gate_proj_name: str | None = None,
931
+ ckpt_down_proj_name: str | None = None,
932
+ ckpt_up_proj_name: str | None = None,
933
+ include_fused: bool = False,
934
+ ) -> list[tuple[str, str, int, str]]:
935
+ moe_config = self.moe_config
936
+ num_fused_shared_experts = self.expert_map_manager.num_fused_shared_experts
937
+ num_redundant_experts = moe_config.num_experts - moe_config.num_logical_experts
938
+ return self.build_expert_params_mapping(
939
+ ckpt_gate_proj_name or self.ckpt_gate_proj_name,
940
+ ckpt_down_proj_name or self.ckpt_down_proj_name,
941
+ ckpt_up_proj_name or self.ckpt_up_proj_name,
942
+ num_experts=moe_config.num_logical_experts + num_fused_shared_experts,
943
+ num_redundant_experts=num_redundant_experts,
944
+ routed_experts_prefix="",
945
+ lora_base_layer_prefix=self.lora_base_layer_prefix,
946
+ include_fused=include_fused,
947
+ )
948
+
949
@staticmethod
950
def make_expert_params_mapping(
951
model: torch.nn.Module,
@@ -921,6 +955,36 @@ class RoutedExperts(PluggableLayer):
955
num_experts: int,
956
num_redundant_experts: int = 0,
957
routed_experts_prefix: str = "routed_experts",
958
+ ) -> list[tuple[str, str, int, str]]:
959
+ """Build the expert mapping, detecting the LoRA `base_layer.` prefix by
960
+ scanning `model`'s parameters.
961
+
962
+ Legacy entry point for models that still hand-roll `load_weights`; the
963
+ `RoutedExperts` weight loader uses `get_expert_mapping` /
964
+ `build_expert_params_mapping` instead (which take the prefix directly).
965
+ See `build_expert_params_mapping` for the returned tuple format.
966
+ """
967
+ has_base_layer = any(".base_layer." in n for n, _ in model.named_parameters())
968
+ return RoutedExperts.build_expert_params_mapping(
969
+ ckpt_gate_proj_name,
970
+ ckpt_down_proj_name,
971
+ ckpt_up_proj_name,
972
+ num_experts,
973
+ num_redundant_experts,
974
+ routed_experts_prefix,
975
+ "base_layer." if has_base_layer else "",
976
+ )
977
+
978
+ @staticmethod
979
+ def build_expert_params_mapping(
980
+ ckpt_gate_proj_name: str,
981
+ ckpt_down_proj_name: str,
982
+ ckpt_up_proj_name: str,
983
+ num_experts: int,
984
+ num_redundant_experts: int = 0,
985
+ routed_experts_prefix: str = "routed_experts",
986
+ lora_base_layer_prefix: str = "",
987
+ include_fused: bool = False,
988
) -> list[tuple[str, str, int, str]]:
989
"""
990
Create expert parameter mapping for weight loading with redundant experts.
@@ -929,12 +993,13 @@ class RoutedExperts(PluggableLayer):
993
when loading weights with EPLB redundant experts.
994
995
Args:
932
- model: The model containing the MoE layer
996
ckpt_gate_proj_name: Name of gate projection in checkpoint
997
ckpt_down_proj_name: Name of down projection in checkpoint
998
ckpt_up_proj_name: Name of up projection in checkpoint
999
num_experts: Number of logical (non-redundant) experts
1000
num_redundant_experts: Number of redundant experts
1001
+ lora_base_layer_prefix: Prefix to add if this layer is a LoRA base layer
1002
+ include_fused: Prepend the fused pre-fused-checkpoint entries
1003
1004
Returns:
1005
List of tuples (param_name, weight_name, expert_id, shard_id)
@@ -956,22 +1021,39 @@ class RoutedExperts(PluggableLayer):
1021
)
1022
)
1023
959
- base_layer = (
960
- "base_layer."
961
- if any(".base_layer." in name for name, _ in model.named_parameters())
962
- else ""
963
- )
964
-
1024
if routed_experts_prefix != "":
1025
routed_experts_prefix = f"{routed_experts_prefix}."
1026
968
- return [
1027
+ w13 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w13_"
1028
+ w2 = f"experts.{routed_experts_prefix}{lora_base_layer_prefix}w2_"
1029
+
1030
+ fused_mapping = []
1031
+ if include_fused:
1032
+ gate_up = None
1033
+ if ckpt_gate_proj_name == "gate_proj" and ckpt_up_proj_name == "up_proj":
1034
+ gate_up = "gate_up_proj"
1035
+ elif ckpt_gate_proj_name == "w1" and ckpt_up_proj_name == "w3":
1036
+ gate_up = "w13"
1037
+ else:
1038
+ logger.warning(
1039
+ "Unexpected gate/up projection names: %s, %s. "
1040
+ "Fused gate/up mapping will be skipped.",
1041
+ ckpt_gate_proj_name,
1042
+ ckpt_up_proj_name,
1043
+ )
1044
+ if gate_up is not None:
1045
+ fused_mapping = [
1046
+ # (param_name, weight_name, expert_id, shard_id)
1047
+ (f"{w13}weight", f"experts.{gate_up}", 0, "w1"),
1048
+ (f"{w13}weight", f"experts.{gate_up}", 1, "w3"),
1049
+ (f"{w2}weight", f"experts.{ckpt_down_proj_name}", 0, "w2"),
1050
+ ]
1051
+
1052
+ per_expert_mapping = [
1053
# (param_name, weight_name, expert_id, shard_id)
1054
(
971
- f"experts.{routed_experts_prefix}{base_layer}w13_"
972
- if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name]
973
- else f"experts.{routed_experts_prefix}{base_layer}w2_",
974
- f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.{base_layer}",
1055
+ w13 if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name] else w2,
1056
+ f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.{lora_base_layer_prefix}",
1057
expert_id,
1058
shard_id,
1059
)
@@ -983,6 +1065,8 @@ class RoutedExperts(PluggableLayer):
1065
]
1066
]
1067
1068
+ return fused_mapping + per_expert_mapping
1069
+
1070
def get_expert_weights(self) -> Iterable[torch.Tensor]:
1071
def _maybe_make_contiguous(
1072
name: str, p: torch.nn.Parameter
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+5
@@ -292,6 +292,11 @@ class MoERunner(MoERunnerInterface):
292
# For smuggling this layer into the fused moe custom op
293
register_layer_for_moe_forward_op(get_current_vllm_config(), self)
294
295
+ def load_weights(
296
+ self, weights: Iterable[tuple[str, torch.Tensor]]
297
+ ) -> Iterable[str]:
298
+ return self.routed_experts.load_weights(weights)
299
+
300
def _select_forward(self) -> Callable:
301
if current_platform.is_tpu() or current_platform.is_cpu():
302
# TODO: Once the OOM issue for the TPU backend is resolved, we
vllm/model_executor/model_loader/bitsandbytes_loader.py
-6
@@ -567,12 +567,6 @@ class BitsAndBytesModelLoader(BaseModelLoader):
567
568
if is_moe_model(model):
569
self.expert_params_mapping = get_moe_expert_mapping(model)
570
- if not self.expert_params_mapping:
571
- raise AttributeError(
572
- f"MoE Model {type(model).__name__} does not support "
573
- "BitsAndBytes quantization yet. Ensure this model has "
574
- "'get_expert_mapping' method."
575
- )
570
# For some models like Molmo, we need to use hf_to_vllm_mapper
571
# to ensure correct loading of weights.
572
if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
vllm/model_executor/models/afmoe.py
+12
-125
@@ -2,8 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only AfMoE 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
@@ -21,7 +20,6 @@ from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.fused_moe import (
21
FusedMoE,
22
MoERunner,
24
- fused_moe_make_expert_params_mapping,
23
)
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
@@ -36,10 +34,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
34
ParallelLMHead,
35
VocabParallelEmbedding,
36
)
39
-from vllm.model_executor.model_loader.weight_utils import (
40
- default_weight_loader,
41
- maybe_remap_kv_scale_name,
42
-)
37
from vllm.model_executor.models.interfaces import (
38
EagleModelMixin,
39
MixtureOfExperts,
@@ -53,7 +47,6 @@ from vllm.model_executor.models.utils import (
47
PPMissingLayer,
48
WeightsMapper,
49
extract_layer_index,
56
- is_pp_missing_parameter,
50
make_empty_intermediate_tensors_factory,
51
make_layers,
52
maybe_prefix,
@@ -481,120 +474,6 @@ class AfmoeModel(nn.Module, EagleModelMixin):
474
}
475
)
476
484
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
485
- # Params for weights, fp8 weight scales, fp8 activation scales
486
- # (param_name, weight_name, expert_id, shard_id)
487
- return fused_moe_make_expert_params_mapping(
488
- self,
489
- ckpt_gate_proj_name="gate_proj",
490
- ckpt_down_proj_name="down_proj",
491
- ckpt_up_proj_name="up_proj",
492
- num_experts=self.config.num_experts,
493
- )
494
-
495
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
496
- stacked_params_mapping = [
497
- # (param_name, shard_name, shard_id)
498
- ("qkv_proj", "q_proj", "q"),
499
- ("qkv_proj", "k_proj", "k"),
500
- ("qkv_proj", "v_proj", "v"),
501
- ("gate_up_proj", "gate_proj", 0),
502
- ("gate_up_proj", "up_proj", 1),
503
- ]
504
-
505
- params_dict = dict(self.named_parameters())
506
- loaded_params: set[str] = set()
507
- expert_params_mapping = self.get_expert_mapping()
508
-
509
- for name, loaded_weight in weights:
510
- for param_name, weight_name, shard_id in stacked_params_mapping:
511
- # Skip non-stacked layers and experts (experts handled below).
512
- if (weight_name not in name) or ("self_attn.gate_proj" in name):
513
- continue
514
- # We have mlp.experts[0].gate_proj in the checkpoint.
515
- # Since we handle the experts below in expert_params_mapping,
516
- # we need to skip here BEFORE we update the name, otherwise
517
- # name will be updated to mlp.experts[0].gate_up_proj, which
518
- # will then be updated below in expert_params_mapping
519
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
520
- if ("mlp.experts." in name) and name not in params_dict:
521
- continue
522
-
523
- name = name.replace(weight_name, param_name)
524
- # Skip loading extra bias for GPTQ models.
525
- if name.endswith(".bias") and name not in params_dict:
526
- continue
527
- if is_pp_missing_parameter(name, self):
528
- continue
529
-
530
- param = params_dict[name]
531
- weight_loader = param.weight_loader
532
- weight_loader(param, loaded_weight, shard_id)
533
- break
534
- else:
535
- is_expert_weight = False
536
- for mapping in expert_params_mapping:
537
- param_name, weight_name, expert_id, shard_id = mapping
538
- if weight_name not in name:
539
- continue
540
-
541
- # Anyway, this is an expert weight and should not be
542
- # attempted to load as other weights later
543
- is_expert_weight = True
544
-
545
- # Do not modify `name` since the loop may continue here
546
- # Instead, create a new variable
547
- name_mapped = name.replace(weight_name, param_name)
548
-
549
- if is_pp_missing_parameter(name_mapped, self):
550
- continue
551
-
552
- param = params_dict[name_mapped]
553
- # We should ask the weight loader to return success or not
554
- # here since otherwise we may skip experts with other
555
- # available replicas.
556
- weight_loader = typing.cast(
557
- Callable[..., bool], param.weight_loader
558
- )
559
- success = weight_loader(
560
- param,
561
- loaded_weight,
562
- name_mapped,
563
- shard_id=shard_id,
564
- expert_id=expert_id,
565
- return_success=True,
566
- )
567
- if success:
568
- name = name_mapped
569
- break
570
- else:
571
- if is_expert_weight:
572
- # We've checked that this is an expert weight
573
- # However it's not mapped locally to this rank
574
- # So we simply skip it
575
- continue
576
-
577
- # Skip loading extra bias for GPTQ models.
578
- if name.endswith(".bias") and name not in params_dict:
579
- continue
580
-
581
- # Remapping the name of FP8 kv-scale.
582
- name = maybe_remap_kv_scale_name(name, params_dict)
583
- if name is None:
584
- continue
585
-
586
- if is_pp_missing_parameter(name, self):
587
- continue
588
-
589
- param = params_dict[name]
590
- weight_loader = getattr(
591
- param, "weight_loader", default_weight_loader
592
- )
593
- weight_loader(param, loaded_weight)
594
- loaded_params.add(name)
595
-
596
- return loaded_params
597
-
477
478
class AfmoeForCausalLM(
479
nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA, MixtureOfExperts
@@ -615,6 +494,17 @@ class AfmoeForCausalLM(
494
orig_to_new_suffix={
495
".router.gate.weight": ".gate.weight",
496
},
497
+ orig_to_new_stacked={
498
+ # weight_name: (param_name, shard_id)
499
+ ".q_proj": (".qkv_proj", "q"),
500
+ ".k_proj": (".qkv_proj", "k"),
501
+ ".v_proj": (".qkv_proj", "v"),
502
+ # `.self_attn.gate_proj` is a gated-attention projection (not fused).
503
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
504
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
505
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
506
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
507
+ },
508
)
509
510
fall_back_to_pt_during_load = False
@@ -705,6 +595,3 @@ class AfmoeForCausalLM(
595
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
596
loader = AutoWeightsLoader(self)
597
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
708
-
709
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
710
- return self.model.get_expert_mapping()
vllm/model_executor/models/aimv2.py
+18
-32
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.linear import (
20
RowParallelLinear,
21
)
22
from vllm.model_executor.layers.quantization import QuantizationConfig
23
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
24
from vllm.transformers_utils.configs.ovis import AIMv2Config
25
26
@@ -194,6 +194,13 @@ class AIMv2Transformer(nn.Module):
194
195
196
class AIMv2Model(torch.nn.Module):
197
+ hf_to_vllm_mapper = WeightsMapper(
198
+ orig_to_new_stacked={
199
+ ".fc1": (".fc13", 0),
200
+ ".fc3": (".fc13", 1),
201
+ }
202
+ )
203
+
204
def __init__(
205
self,
206
config: AIMv2Config,
@@ -218,34 +225,13 @@ class AIMv2Model(torch.nn.Module):
225
return x
226
227
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
221
- stacked_params_mapping = [
222
- # (param_name, shard_name, shard_id)
223
- (".fc13", ".fc1", 0),
224
- (".fc13", ".fc3", 1),
225
- ]
226
- params_dict = dict(self.named_parameters())
227
- loaded_params: set[str] = set()
228
-
229
- for name, loaded_weight in weights:
230
- # post_layernorm is optional in SiglipVisionModel
231
- if (
232
- name.startswith("trunk.post_trunk_norm")
233
- and self.trunk.post_trunk_norm is None
234
- ):
235
- continue
236
-
237
- for param_name, weight_name, shard_id in stacked_params_mapping:
238
- if weight_name not in name:
239
- continue
240
- name = name.replace(weight_name, param_name)
241
-
242
- param = params_dict[name]
243
- weight_loader = param.weight_loader
244
- weight_loader(param, loaded_weight, shard_id)
245
- break
246
- else:
247
- param = params_dict[name]
248
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
249
- weight_loader(param, loaded_weight)
250
- loaded_params.add(name)
251
- return loaded_params
228
+ loader = AutoWeightsLoader(
229
+ self,
230
+ # post_trunk_norm is optional (absent for clip-skip backbones).
231
+ skip_prefixes=(
232
+ ["trunk.post_trunk_norm."]
233
+ if self.trunk.post_trunk_norm is None
234
+ else None
235
+ ),
236
+ )
237
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/apertus.py
+9
-62
@@ -53,10 +53,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
53
ParallelLMHead,
54
VocabParallelEmbedding,
55
)
56
-from vllm.model_executor.model_loader.weight_utils import (
57
- default_weight_loader,
58
- maybe_remap_kv_scale_name,
59
-)
56
from vllm.sequence import IntermediateTensors
57
from vllm.v1.attention.backend import AttentionType
58
@@ -70,8 +66,8 @@ from .interfaces import (
66
from .utils import (
67
AutoWeightsLoader,
68
PPMissingLayer,
69
+ WeightsMapper,
70
extract_layer_index,
74
- is_pp_missing_parameter,
71
make_empty_intermediate_tensors_factory,
72
make_layers,
73
maybe_prefix,
@@ -404,66 +400,17 @@ class ApertusModel(nn.Module, EagleModelMixin):
400
return hidden_states, aux_hidden_states
401
return hidden_states
402
407
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
408
- stacked_params_mapping = [
409
- # (param_name, shard_name, shard_id)
410
- (".qkv_proj", ".q_proj", "q"),
411
- (".qkv_proj", ".k_proj", "k"),
412
- (".qkv_proj", ".v_proj", "v"),
413
- ]
414
- params_dict = dict(self.named_parameters())
415
-
416
- # we need to load the buffers for beta and eps (XIELU)
417
- for name, buffer in self.named_buffers():
418
- if name.endswith(".beta") or name.endswith(".eps"):
419
- params_dict[name] = buffer
420
-
421
- loaded_params: set[str] = set()
422
- for name, loaded_weight in weights:
423
- if "rotary_emb.inv_freq" in name:
424
- continue
425
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
426
- # Models trained using ColossalAI may include these tensors in
427
- # the checkpoint. Skip them.
428
- continue
429
- if "scale" in name or "zero_point" in name:
430
- # Remapping the name of FP8 kv-scale.
431
- name = maybe_remap_kv_scale_name(name, params_dict)
432
- if name is None:
433
- continue
434
- for param_name, weight_name, shard_id in stacked_params_mapping:
435
- if weight_name not in name:
436
- continue
437
- name = name.replace(weight_name, param_name)
438
- # Skip loading extra bias for GPTQ models.
439
- if name.endswith(".bias") and name not in params_dict:
440
- continue
441
-
442
- if is_pp_missing_parameter(name, self):
443
- continue
444
-
445
- param = params_dict[name]
446
- weight_loader = param.weight_loader
447
- weight_loader(param, loaded_weight, shard_id)
448
- break
449
- else:
450
- # Skip loading extra bias for GPTQ models.
451
- if name.endswith(".bias") and name not in params_dict:
452
- continue
453
-
454
- if is_pp_missing_parameter(name, self):
455
- continue
456
-
457
- param = params_dict[name]
458
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
459
- weight_loader(param, loaded_weight)
460
- loaded_params.add(name)
461
- return loaded_params
462
-
403
404
class ApertusForCausalLM(
405
nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3
406
):
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
+ }
413
+ )
414
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
415
416
# LoRA specific attributes
@@ -548,4 +495,4 @@ class ApertusForCausalLM(
495
self,
496
skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
497
)
551
- return loader.load_weights(weights)
498
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/bert.py
+10
-40
@@ -369,6 +369,14 @@ class BertModel(nn.Module, SupportsQuant):
369
370
packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}
371
372
+ hf_to_vllm_mapper = WeightsMapper(
373
+ orig_to_new_stacked={
374
+ ".self.query": (".self.qkv_proj", "q"),
375
+ ".self.key": (".self.qkv_proj", "k"),
376
+ ".self.value": (".self.qkv_proj", "v"),
377
+ }
378
+ )
379
+
380
def __init__(
381
self,
382
*,
@@ -400,43 +408,9 @@ class BertModel(nn.Module, SupportsQuant):
408
409
return self.encoder(hidden_states)
410
403
- def _load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
404
- stacked_params_mapping = [
405
- # (param_name, shard_name, shard_id)
406
- ("qkv_proj", "query", "q"),
407
- ("qkv_proj", "key", "k"),
408
- ("qkv_proj", "value", "v"),
409
- ]
410
-
411
- loaded_stacked_params = []
412
- other_weights = []
413
- params_dict = dict(self.named_parameters())
414
- for name, loaded_weight in weights:
415
- for param_name, weight_name, shard_id in stacked_params_mapping:
416
- if weight_name not in name:
417
- continue
418
-
419
- name = name.replace(weight_name, param_name)
420
- if name not in params_dict:
421
- continue
422
- param = params_dict[name]
423
- weight_loader = param.weight_loader
424
- weight_loader(param, loaded_weight, shard_id)
425
- loaded_stacked_params.append(name)
426
- break
427
- else:
428
- if name in params_dict:
429
- other_weights.append((name, loaded_weight))
430
-
431
- return other_weights, loaded_stacked_params
432
-
411
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
434
- other_weights, loaded_stacked_params = self._load_weights(weights)
435
-
412
loader = AutoWeightsLoader(self, skip_prefixes=["pooler."])
437
- loaded_params = loader.load_weights(other_weights)
438
- loaded_params.update(loaded_stacked_params)
439
- return loaded_params
413
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
414
415
416
class BertPoolingModel(BertModel):
@@ -458,12 +432,8 @@ class BertPoolingModel(BertModel):
432
self.pooler = BertPooler(vllm_config.model_config)
433
434
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
461
- other_weights, loaded_stacked_params = self._load_weights(weights)
462
-
435
loader = AutoWeightsLoader(self)
464
- loaded_params = loader.load_weights(other_weights)
465
- loaded_params.update(loaded_stacked_params)
466
- return loaded_params
436
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
437
438
439
@default_pooling_type(seq_pooling_type="CLS")
vllm/model_executor/models/blip.py
+22
-34
@@ -19,9 +19,9 @@ from vllm.model_executor.layers.linear import (
19
RowParallelLinear,
20
)
21
from vllm.model_executor.layers.quantization import QuantizationConfig
22
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
23
from .interfaces import SupportsQuant
24
+from .utils import AutoWeightsLoader, WeightsMapper
25
26
27
def get_blip_patch_grid_length(*, image_size: int, patch_size: int) -> int:
@@ -268,6 +268,14 @@ class BlipVisionModel(nn.Module, SupportsQuant):
268
main_input_name = "pixel_values"
269
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
270
271
+ hf_to_vllm_mapper = WeightsMapper(
272
+ orig_to_new_stacked={
273
+ ".q_proj": (".qkv", "q"),
274
+ ".k_proj": (".qkv", "k"),
275
+ ".v_proj": (".qkv", "v"),
276
+ }
277
+ )
278
+
279
def __init__(
280
self,
281
config: BlipVisionConfig,
@@ -316,38 +324,18 @@ class BlipVisionModel(nn.Module, SupportsQuant):
324
return self.post_layernorm(hidden_states)
325
326
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
319
- stacked_params_mapping = [
320
- # (param_name, shard_name, shard_id)
321
- ("qkv_proj", "q_proj", "q"),
322
- ("qkv_proj", "k_proj", "k"),
323
- ("qkv_proj", "v_proj", "v"),
324
- ]
325
- params_dict = dict(self.named_parameters())
326
- loaded_params: set[str] = set()
327
- layer_count = len(self.encoder.layers)
328
-
329
- for name, loaded_weight in weights:
330
- # post_layernorm is not needed in BlipVisionModel
331
- if name.startswith("post_layernorm") and self.post_layernorm is None:
332
- continue
333
-
334
- # omit layers when num_hidden_layers_override is set
335
- if name.startswith("encoder.layers"):
336
- layer_idx = int(name.split(".")[2])
337
- if layer_idx >= layer_count:
327
+ skip_prefixes: list[str] = []
328
+ if self.post_layernorm is None:
329
+ skip_prefixes.append("post_layernorm.")
330
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
331
+
332
+ # omit layers when num_hidden_layers_override is set
333
+ def _filter(ws):
334
+ for name, weight in ws:
335
+ if name.startswith("encoder.layers.") and int(
336
+ name.split(".")[2]
337
+ ) >= len(self.encoder.layers):
338
continue
339
+ yield name, weight
340
340
- for param_name, weight_name, shard_id in stacked_params_mapping:
341
- if weight_name not in name:
342
- continue
343
- name = name.replace(weight_name, param_name)
344
- param = params_dict[name]
345
- weight_loader = param.weight_loader
346
- weight_loader(param, loaded_weight, shard_id)
347
- break
348
- else:
349
- param = params_dict[name]
350
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
351
- weight_loader(param, loaded_weight)
352
- loaded_params.add(name)
353
- return loaded_params
341
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/chameleon.py
+16
-83
@@ -37,10 +37,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
37
ParallelLMHead,
38
VocabParallelEmbedding,
39
)
40
-from vllm.model_executor.model_loader.weight_utils import (
41
- default_weight_loader,
42
- row_parallel_weight_loader,
43
-)
40
+from vllm.model_executor.model_loader.weight_utils import row_parallel_weight_loader
41
from vllm.model_executor.utils import set_weight_attrs
42
from vllm.multimodal import MULTIMODAL_REGISTRY
43
from vllm.multimodal.inputs import (
@@ -66,7 +63,8 @@ from .interfaces import (
63
SupportsQuant,
64
)
65
from .utils import (
69
- is_pp_missing_parameter,
66
+ AutoWeightsLoader,
67
+ WeightsMapper,
68
make_empty_intermediate_tensors_factory,
69
make_layers,
70
maybe_prefix,
@@ -922,6 +920,16 @@ class ChameleonForConditionalGeneration(
920
"gate_up_proj": ["gate_proj", "up_proj"],
921
}
922
923
+ hf_to_vllm_mapper = WeightsMapper(
924
+ orig_to_new_stacked={
925
+ ".q_proj": (".qkv_proj", "q"),
926
+ ".k_proj": (".qkv_proj", "k"),
927
+ ".v_proj": (".qkv_proj", "v"),
928
+ ".gate_proj": (".gate_up_proj", 0),
929
+ ".up_proj": (".gate_up_proj", 1),
930
+ },
931
+ )
932
+
933
@classmethod
934
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
935
if modality.startswith("image"):
@@ -1023,81 +1031,6 @@ class ChameleonForConditionalGeneration(
1031
return logits
1032
1033
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1026
- stacked_params_mapping = [
1027
- # (param_name, shard_name, shard_id)
1028
- (".qkv_proj", ".q_proj", "q"),
1029
- (".qkv_proj", ".k_proj", "k"),
1030
- (".qkv_proj", ".v_proj", "v"),
1031
- (".gate_up_proj", ".gate_proj", 0),
1032
- (".gate_up_proj", ".up_proj", 1),
1033
- ]
1034
- params_dict = dict(self.named_parameters())
1035
- loaded_params: set[str] = set()
1036
- for name, loaded_weight in weights:
1037
- if "rotary_emb.inv_freq" in name:
1038
- continue
1039
-
1040
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
1041
- # Models trained using ColossalAI may include these tensors in
1042
- # the checkpoint. Skip them.
1043
- continue
1044
-
1045
- # With tie_word_embeddings, we can skip lm_head.weight
1046
- # The weight might appear unnecessarily in the files if the model is
1047
- # processed with quantization, LoRA, fine-tuning, etc.
1048
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
1049
- continue
1050
-
1051
- use_default_weight_loading = False
1052
- if "vqmodel" in name:
1053
- if self.model.vqmodel is not None:
1054
- # We only do sharding for language model and
1055
- # not vqvae for now.
1056
- use_default_weight_loading = True
1057
- else:
1058
- for param_name, weight_name, shard_id in stacked_params_mapping:
1059
- if weight_name not in name:
1060
- continue
1061
- name = name.replace(weight_name, param_name)
1062
- # Skip loading extra bias for GPTQ models.
1063
- if name.endswith(".bias") and name not in params_dict:
1064
- continue
1065
- if is_pp_missing_parameter(name, self):
1066
- continue
1067
- param = params_dict[name]
1068
- weight_loader = param.weight_loader
1069
- weight_loader(param, loaded_weight, shard_id)
1070
- break
1071
- else:
1072
- # Skip loading extra bias for GPTQ models.
1073
- if name.endswith(".bias") and name not in params_dict:
1074
- continue
1075
- # Remapping the name of FP8 kv-scale.
1076
- if name.endswith("kv_scale"):
1077
- remapped_kv_scale_name = name.replace(
1078
- ".kv_scale", ".attn.kv_scale"
1079
- )
1080
- if remapped_kv_scale_name not in params_dict:
1081
- logger.warning_once(
1082
- "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501
1083
- name,
1084
- remapped_kv_scale_name,
1085
- )
1086
- continue
1087
- else:
1088
- name = remapped_kv_scale_name
1089
- if is_pp_missing_parameter(name, self):
1090
- continue
1091
- param = params_dict[name]
1092
- weight_loader = getattr(
1093
- param, "weight_loader", default_weight_loader
1094
- )
1095
- weight_loader(param, loaded_weight)
1096
- if use_default_weight_loading and name in params_dict:
1097
- if is_pp_missing_parameter(name, self):
1098
- continue
1099
- param = params_dict[name]
1100
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1101
- weight_loader(param, loaded_weight)
1102
- loaded_params.add(name)
1103
- return loaded_params
1034
+ skip_prefixes = ["lm_head."] if self.config.tie_word_embeddings else None
1035
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
1036
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/clip.py
+32
-61
@@ -29,7 +29,6 @@ from vllm.model_executor.layers.linear import (
29
from vllm.model_executor.layers.pooler import DispatchPooler
30
from vllm.model_executor.layers.quantization import QuantizationConfig
31
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
32
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32
from vllm.model_executor.models.interfaces import SupportsQuant
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
from vllm.multimodal.inputs import (
@@ -56,7 +55,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
56
from .interfaces import MultiModalEmbeddings, SupportsMultiModal
57
from .interfaces_base import default_pooling_type
59
-from .utils import AutoWeightsLoader, maybe_prefix
58
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
59
from .vision import (
60
VisionEncoderInfo,
61
VisionFeatureSelectStrategy,
@@ -555,6 +554,14 @@ class CLIPEncoder(nn.Module):
554
555
556
class CLIPTextTransformer(nn.Module):
557
+ hf_to_vllm_mapper = WeightsMapper(
558
+ orig_to_new_stacked={
559
+ ".q_proj": (".qkv_proj", "q"),
560
+ ".k_proj": (".qkv_proj", "k"),
561
+ ".v_proj": (".qkv_proj", "v"),
562
+ }
563
+ )
564
+
565
def __init__(
566
self,
567
config: CLIPTextConfig,
@@ -605,34 +612,19 @@ class CLIPTextTransformer(nn.Module):
612
return last_hidden_state
613
614
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
608
- stacked_params_mapping = [
609
- # (param_name, shard_name, shard_id)
610
- ("qkv_proj", "q_proj", "q"),
611
- ("qkv_proj", "k_proj", "k"),
612
- ("qkv_proj", "v_proj", "v"),
613
- ]
614
- params_dict = dict(self.named_parameters())
615
- loaded_params: set[str] = set()
616
-
617
- for name, loaded_weight in weights:
618
- for param_name, weight_name, shard_id in stacked_params_mapping:
619
- if weight_name not in name:
620
- continue
621
- name = name.replace(weight_name, param_name)
622
-
623
- param = params_dict[name]
624
- weight_loader = param.weight_loader
625
- weight_loader(param, loaded_weight, shard_id)
626
- break
627
- else:
628
- param = params_dict[name]
629
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
630
- weight_loader(param, loaded_weight)
631
- loaded_params.add(name)
632
- return loaded_params
615
+ loader = AutoWeightsLoader(self)
616
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
617
618
619
class CLIPVisionTransformer(nn.Module):
620
+ hf_to_vllm_mapper = WeightsMapper(
621
+ orig_to_new_stacked={
622
+ ".q_proj": (".qkv_proj", "q"),
623
+ ".k_proj": (".qkv_proj", "k"),
624
+ ".v_proj": (".qkv_proj", "v"),
625
+ }
626
+ )
627
+
628
def __init__(
629
self,
630
config: CLIPVisionConfig,
@@ -714,42 +706,21 @@ class CLIPVisionTransformer(nn.Module):
706
return encoder_outputs
707
708
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
717
- stacked_params_mapping = [
718
- # (param_name, shard_name, shard_id)
719
- ("qkv_proj", "q_proj", "q"),
720
- ("qkv_proj", "k_proj", "k"),
721
- ("qkv_proj", "v_proj", "v"),
722
- ]
723
- params_dict = dict(self.named_parameters())
724
- loaded_params: set[str] = set()
725
- layer_count = len(self.encoder.layers)
726
-
727
- for name, loaded_weight in weights:
728
- # post_layernorm is not needed in CLIPVisionModel
729
- if name.startswith("post_layernorm") and self.post_layernorm is None:
730
- continue
731
-
732
- # omit layers when num_hidden_layers_override is set
733
- if name.startswith("encoder.layers"):
734
- layer_idx = int(name.split(".")[2])
735
- if layer_idx >= layer_count:
709
+ skip_prefixes: list[str] = []
710
+ if self.post_layernorm is None:
711
+ skip_prefixes.append("post_layernorm.")
712
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
713
+
714
+ # Drop layers beyond num_hidden_layers_override.
715
+ def _filter(ws):
716
+ for name, w in ws:
717
+ if name.startswith("encoder.layers.") and int(
718
+ name.split(".")[2]
719
+ ) >= len(self.encoder.layers):
720
continue
721
+ yield name, w
722
738
- for param_name, weight_name, shard_id in stacked_params_mapping:
739
- if weight_name not in name:
740
- continue
741
- name = name.replace(weight_name, param_name)
742
-
743
- param = params_dict[name]
744
- weight_loader = param.weight_loader
745
- weight_loader(param, loaded_weight, shard_id)
746
- break
747
- else:
748
- param = params_dict[name]
749
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
750
- weight_loader(param, loaded_weight)
751
- loaded_params.add(name)
752
- return loaded_params
723
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
724
725
726
class CLIPVisionModel(nn.Module):
vllm/model_executor/models/cohere2_moe.py
+15
-83
@@ -18,7 +18,6 @@ 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
- fused_moe_make_expert_params_mapping,
21
)
22
from vllm.model_executor.layers.linear import (
23
MergedColumnParallelLinear,
@@ -31,8 +30,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.layers.rotary_embedding import get_rope
31
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
32
from vllm.model_executor.model_loader.weight_utils import (
34
- default_weight_loader,
35
- maybe_remap_kv_scale_name,
33
row_parallel_weight_loader,
34
)
35
from vllm.model_executor.utils import set_weight_attrs
@@ -43,8 +40,8 @@ from .commandr import LayerNorm
40
from .interfaces import SupportsPP, SupportsQuant
41
from .utils import (
42
AutoWeightsLoader,
43
+ WeightsMapper,
44
extract_layer_index,
47
- is_pp_missing_parameter,
45
make_empty_intermediate_tensors_factory,
46
make_layers,
47
maybe_prefix,
@@ -462,88 +459,23 @@ class Cohere2MoeModel(nn.Module):
459
hidden_states, _ = self.norm(hidden_states, residual)
460
return hidden_states
461
465
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
466
- stacked_params_mapping = [
467
- ("qkv_proj", "q_proj", "q"),
468
- ("qkv_proj", "k_proj", "k"),
469
- ("qkv_proj", "v_proj", "v"),
470
- ("gate_up_proj", "gate_proj", 0),
471
- ("gate_up_proj", "up_proj", 1),
472
- ]
473
-
474
- expert_params_mapping = fused_moe_make_expert_params_mapping(
475
- self,
476
- ckpt_gate_proj_name="gate_proj",
477
- ckpt_down_proj_name="down_proj",
478
- ckpt_up_proj_name="up_proj",
479
- num_experts=self.config.num_experts,
480
- )
481
-
482
- params_dict = dict(self.named_parameters())
483
- loaded_params: set[str] = set()
484
- for name, loaded_weight in weights:
485
- if "rotary_emb.inv_freq" in name:
486
- continue
487
-
488
- for param_name, shard_name, shard_id in stacked_params_mapping:
489
- if shard_name not in name:
490
- continue
491
- if "mlp.experts" in name:
492
- continue
493
- name = name.replace(shard_name, param_name)
494
- if name.endswith(".bias") and name not in params_dict:
495
- continue
496
- if is_pp_missing_parameter(name, self):
497
- continue
498
- param = params_dict[name]
499
- weight_loader = param.weight_loader
500
- weight_loader(param, loaded_weight, shard_id)
501
- break
502
- else:
503
- for mapping in expert_params_mapping:
504
- param_name, weight_name, expert_id, shard_id = mapping
505
- if weight_name not in name:
506
- continue
507
- name = name.replace(weight_name, param_name)
508
- if is_pp_missing_parameter(name, self):
509
- continue
510
- if (
511
- name.endswith(".bias") or name.endswith("_bias")
512
- ) and name not in params_dict:
513
- continue
514
- param = params_dict[name]
515
- weight_loader = param.weight_loader
516
- weight_loader(
517
- param,
518
- loaded_weight,
519
- name,
520
- shard_id=shard_id,
521
- expert_id=expert_id,
522
- )
523
- break
524
- else:
525
- if (
526
- name.endswith(".bias") or name.endswith("_bias")
527
- ) and name not in params_dict:
528
- continue
529
- if is_pp_missing_parameter(name, self):
530
- continue
531
- name = maybe_remap_kv_scale_name(name, params_dict)
532
- if name is None:
533
- continue
534
- param = params_dict[name]
535
- weight_loader = getattr(
536
- param, "weight_loader", default_weight_loader
537
- )
538
- weight_loader(param, loaded_weight)
539
- loaded_params.add(name)
540
-
541
- return loaded_params
542
-
462
463
class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
464
is_text_generation_model = True
465
466
+ hf_to_vllm_mapper = WeightsMapper(
467
+ orig_to_new_stacked={
468
+ # weight_name: (param_name, shard_id)
469
+ ".q_proj": (".qkv_proj", "q"),
470
+ ".k_proj": (".qkv_proj", "k"),
471
+ ".v_proj": (".qkv_proj", "v"),
472
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
473
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
474
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
475
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
476
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
477
+ }
478
+ )
479
packed_modules_mapping = {
480
"qkv_proj": [
481
"q_proj",
@@ -599,4 +531,4 @@ class Cohere2MoeForCausalLM(nn.Module, SupportsPP, SupportsQuant):
531
532
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
533
loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
602
- return loader.load_weights(weights)
534
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/deepseek_eagle3.py
+14
-41
@@ -20,10 +20,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
20
ParallelLMHead,
21
VocabParallelEmbedding,
22
)
23
-from vllm.model_executor.model_loader.weight_utils import (
24
- default_weight_loader,
25
- maybe_remap_kv_scale_name,
26
-)
23
from vllm.model_executor.models.deepseek_v2 import (
24
DeepseekV2ForCausalLM,
25
DeepseekV2MLAAttention,
@@ -34,6 +30,7 @@ from vllm.multimodal.inputs import NestedTensors
30
from .interfaces import LocalArgmaxMixin
31
from .utils import (
32
AutoWeightsLoader,
33
+ WeightsMapper,
34
get_draft_quant_config,
35
maybe_prefix,
36
process_eagle_weight,
@@ -270,44 +267,20 @@ class DeepseekV2Eagle3Model(nn.Module):
267
268
return hidden_states, aux_output
269
273
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
274
- stacked_params_mapping = [
275
- # (param_name, shard_name, shard_id)
276
- (".gate_up_proj", ".gate_proj", 0),
277
- (".gate_up_proj", ".up_proj", 1),
278
- (".fused_qkv_a_proj", ".q_a_proj", 0),
279
- (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
280
- ]
281
- params_dict = dict(self.named_parameters())
282
- loaded_params: set[str] = set()
270
+ # midlayer rename + gate_up / MLA fused_qkv_a merges
271
+ hf_to_vllm_mapper = WeightsMapper(
272
+ orig_to_new_substr={"midlayer.": "layers.0."},
273
+ orig_to_new_stacked={
274
+ ".gate_proj": (".gate_up_proj", 0),
275
+ ".up_proj": (".gate_up_proj", 1),
276
+ ".q_a_proj": (".fused_qkv_a_proj", 0),
277
+ ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1),
278
+ },
279
+ )
280
284
- for name, loaded_weight in weights:
285
- if "midlayer." in name:
286
- name = name.replace("midlayer.", "layers.0.")
287
-
288
- # Remapping the name FP8 kv-scale
289
- if "scale" in name:
290
- name = maybe_remap_kv_scale_name(name, params_dict)
291
- if name is None:
292
- continue
293
-
294
- for param_name, weight_name, shard_id in stacked_params_mapping:
295
- if weight_name not in name:
296
- continue
297
- name = name.replace(weight_name, param_name)
298
- param = params_dict[name]
299
- weight_loader = param.weight_loader
300
- weight_loader(param, loaded_weight, shard_id)
301
- break
302
- else:
303
- if name not in params_dict:
304
- continue
305
- param = params_dict[name]
306
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
307
- weight_loader(param, loaded_weight)
308
- loaded_params.add(name)
309
-
310
- return loaded_params
281
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
282
+ loader = AutoWeightsLoader(self)
283
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
284
285
286
class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):
vllm/model_executor/models/dots_ocr.py
+4
-31
@@ -32,7 +32,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
32
from vllm.model_executor.layers.rotary_embedding.common import (
33
ApplyRotaryEmb,
34
)
35
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35
from vllm.model_executor.models.interfaces import (
36
MultiModalEmbeddings,
37
SupportsLoRA,
@@ -355,36 +354,6 @@ class DotsSwiGLUFFN(nn.Module):
354
x, _ = self.fc2(x)
355
return x
356
358
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
359
- stacked_params_mapping = [
360
- ("fc13", "fc1", 0),
361
- ("fc13", "fc3", 1),
362
- ]
363
- params_dict = dict(self.named_parameters())
364
- loaded_params: set[str] = set()
365
- for name, loaded_weight in weights:
366
- for param_name, weight_name, shard_id in stacked_params_mapping:
367
- if weight_name not in name:
368
- continue
369
- name = name.replace(weight_name, param_name)
370
- # Skip loading extra bias for GPTQ models.
371
- if name.endswith(".bias") and name not in params_dict:
372
- continue
373
- param = params_dict[name]
374
- weight_loader = param.weight_loader
375
- weight_loader(param, loaded_weight, shard_id)
376
- break
377
- else:
378
- # Skip loading extra bias for GPTQ models.
379
- if name.endswith(".bias") and name not in params_dict:
380
- continue
381
-
382
- param = params_dict[name]
383
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
384
- weight_loader(param, loaded_weight)
385
- loaded_params.add(name)
386
- return loaded_params
387
-
357
358
class DotsPatchEmbed(nn.Module):
359
def __init__(self, config):
@@ -622,6 +591,10 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
591
".attn.qkv_proj.": ".attn.qkv.",
592
".attn.out_proj.": ".attn.proj.",
593
},
594
+ orig_to_new_stacked={
595
+ ".fc1.": (".fc13.", 0),
596
+ ".fc3.": (".fc13.", 1),
597
+ },
598
orig_to_new_prefix={
599
"lm_head.": "language_model.lm_head.",
600
"model.": "language_model.model.",
vllm/model_executor/models/ernie_mtp.py
+35
-94
@@ -27,7 +27,6 @@ from collections.abc import Iterable
27
28
import torch
29
import torch.nn as nn
30
-from transformers import PretrainedConfig
30
31
from vllm.config import VllmConfig
32
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -36,11 +35,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35
ParallelLMHead,
36
VocabParallelEmbedding,
37
)
39
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.sequence import IntermediateTensors
39
40
from .llama import LlamaDecoderLayer
43
-from .utils import is_pp_missing_parameter, maybe_prefix
41
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
42
43
44
class ErnieMultiTokenPredictorLayer(nn.Module):
@@ -147,6 +145,28 @@ class ErnieMTP(nn.Module):
145
super().__init__()
146
147
self.config = vllm_config.model_config.hf_config
148
+ # MTP weights are stored under a flat `mtp_*.0.` block in the
149
+ # checkpoint; rewrite them into `model.layers.{spec_layer}.*`.
150
+ spec_layer = self.config.num_hidden_layers
151
+ self.hf_to_vllm_mapper = WeightsMapper(
152
+ orig_to_new_substr={
153
+ "model.mtp_emb_norm.0.": f"model.layers.{spec_layer}.mtp_emb_norm.",
154
+ "model.mtp_hidden_norm.0.": (
155
+ f"model.layers.{spec_layer}.mtp_hidden_norm."
156
+ ),
157
+ "model.mtp_linear_proj.0.": (
158
+ f"model.layers.{spec_layer}.mtp_linear_proj."
159
+ ),
160
+ "model.mtp_block.0.": f"model.layers.{spec_layer}.mtp_block.",
161
+ },
162
+ orig_to_new_stacked={
163
+ ".q_proj": (".qkv_proj", "q"),
164
+ ".k_proj": (".qkv_proj", "k"),
165
+ ".v_proj": (".qkv_proj", "v"),
166
+ ".gate_proj": (".gate_up_proj", 0),
167
+ ".up_proj": (".gate_up_proj", 1),
168
+ },
169
+ )
170
self.model = ErnieMultiTokenPredictor(
171
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
172
)
@@ -185,94 +205,15 @@ class ErnieMTP(nn.Module):
205
return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)
206
207
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
188
- stacked_params_mapping = [
189
- ("qkv_proj", "q_proj", "q"),
190
- ("qkv_proj", "k_proj", "k"),
191
- ("qkv_proj", "v_proj", "v"),
192
- ("gate_up_proj", "gate_proj", 0),
193
- ("gate_up_proj", "up_proj", 1),
194
- ]
195
-
196
- params_dict = dict(self.named_parameters())
197
- loaded_params: set[str] = set()
198
- for name, loaded_weight in weights:
199
- if self.config.tie_word_embeddings and name.endswith("lm_head.weight"):
200
- continue
201
- if "rotary_emb.inv_freq" in name:
202
- continue
203
- if "mtp" in name:
204
- name = self._rewrite_spec_layer_name(self.config, name)
205
-
206
- for param_name, weight_name, shard_id in stacked_params_mapping:
207
- # Skip non-stacked layers and experts (experts handled below).
208
- if weight_name not in name:
209
- continue
210
- if "mtp" not in name:
211
- continue
212
- # We have mlp.experts[0].gate_proj in the checkpoint.
213
- # Since we handle the experts below in expert_params_mapping,
214
- # we need to skip here BEFORE we update the name, otherwise
215
- # name will be updated to mlp.experts[0].gate_up_proj, which
216
- # will then be updated below in expert_params_mapping
217
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
218
- if ("mlp.experts." in name) and name not in params_dict:
219
- continue
220
- name = name.replace(weight_name, param_name)
221
- # Skip loading extra bias for GPTQ models.
222
- if (
223
- name.endswith(".bias") or name.endswith("_bias")
224
- ) and name not in params_dict:
225
- continue
226
- # Skip layers on other devices.
227
- if is_pp_missing_parameter(name, self):
228
- continue
229
-
230
- param = params_dict[name]
231
- weight_loader = param.weight_loader
232
- weight_loader(param, loaded_weight, shard_id)
233
- break
234
- else:
235
- # Skip loading extra bias for GPTQ models.
236
- if (
237
- name.endswith(".bias") or name.endswith("_bias")
238
- ) and name not in params_dict:
239
- continue
240
- # Skip layers on other devices.
241
- if is_pp_missing_parameter(name, self):
242
- continue
243
-
244
- # According to DeepSeek-V3 Technical Report, MTP modules
245
- # shares embedding layer. We only load the first weights.
246
- if "mtp_" not in name and (
247
- "embed_tokens" not in name and "lm_head" not in name
248
- ):
249
- continue
250
-
251
- param = params_dict[name]
252
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
253
- weight_loader(param, loaded_weight)
254
- loaded_params.add(name)
255
- return loaded_params
256
-
257
- def _rewrite_spec_layer_name(self, config: PretrainedConfig, name: str) -> str:
258
- """
259
- Rewrite the weight name to match the format of the original model.
260
- """
261
- spec_layer_weight_names = [
262
- "embed_tokens",
263
- "mtp_emb_norm",
264
- "mtp_hidden_norm",
265
- "mtp_linear_proj",
266
- ]
267
- layer_idx = config.num_hidden_layers
268
- for weight_name in spec_layer_weight_names:
269
- if weight_name in name:
270
- name = name.replace(
271
- f"model.{weight_name}.0.",
272
- f"model.layers.{layer_idx}.{weight_name}.",
273
- )
274
- return name
275
- name = name.replace(
276
- "model.mtp_block.0.", f"model.layers.{layer_idx}.mtp_block."
277
- )
278
- return name
208
+ # Checkpoint bundles the full base model; only MTP, embed_tokens and
209
+ # lm_head weights belong to this module.
210
+ def _filter(
211
+ weights: Iterable[tuple[str, torch.Tensor]],
212
+ ) -> Iterable[tuple[str, torch.Tensor]]:
213
+ for name, weight in weights:
214
+ if any(k in name for k in ("mtp", "embed_tokens", "lm_head")):
215
+ yield name, weight
216
+
217
+ skip_prefixes = ["lm_head"] if self.config.tie_word_embeddings else []
218
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
219
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/exaone_moe.py
+29
-143
@@ -15,8 +15,7 @@
15
# limitations under the License.
16
"""Inference-only K-EXAONE-236B-A22B model compatible with HuggingFace weights."""
17
18
-import typing
19
-from collections.abc import Callable, Iterable
18
+from collections.abc import Iterable
19
from itertools import islice
20
21
import torch
@@ -32,7 +31,6 @@ from vllm.distributed import (
31
)
32
from vllm.model_executor.layers.fused_moe import (
33
FusedMoE,
35
- fused_moe_make_expert_params_mapping,
34
)
35
from vllm.model_executor.layers.layernorm import RMSNorm
36
from vllm.model_executor.layers.linear import ReplicatedLinear
@@ -43,10 +41,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
41
ParallelLMHead,
42
VocabParallelEmbedding,
43
)
46
-from vllm.model_executor.model_loader.weight_utils import (
47
- default_weight_loader,
48
- maybe_remap_kv_scale_name,
49
-)
44
from vllm.sequence import IntermediateTensors
45
46
from .exaone4 import Exaone4Attention as ExaoneMoeAttention
@@ -55,8 +49,8 @@ from .interfaces import SupportsLoRA, SupportsPP
49
from .utils import (
50
AutoWeightsLoader,
51
PPMissingLayer,
52
+ WeightsMapper,
53
extract_layer_index,
59
- is_pp_missing_parameter,
54
make_empty_intermediate_tensors_factory,
55
make_layers,
56
maybe_prefix,
@@ -325,142 +319,21 @@ class ExaoneMoeModel(nn.Module):
319
hidden_states, _ = self.norm(hidden_states, residual)
320
return hidden_states
321
328
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
329
- # Params for weights, fp8 weight scales, fp8 activation scales
330
- # (param_name, weight_name, expert_id, shard_id)
331
- return fused_moe_make_expert_params_mapping(
332
- self,
333
- ckpt_gate_proj_name="gate_proj",
334
- ckpt_down_proj_name="down_proj",
335
- ckpt_up_proj_name="up_proj",
336
- num_experts=self.config.num_experts,
337
- num_redundant_experts=self.num_redundant_experts,
338
- )
339
-
340
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
341
- stacked_params_mapping = [
342
- # (param_name, shard_name, shard_id)
343
- (".qkv_proj", ".q_proj", "q"),
344
- (".qkv_proj", ".k_proj", "k"),
345
- (".qkv_proj", ".v_proj", "v"),
346
- (".gate_up_proj", ".gate_proj", 0),
347
- (".gate_up_proj", ".up_proj", 1),
348
- ]
349
-
350
- # Skip loading extra parameters for GPTQ/modelopt models.
351
- ignore_suffixes = (
352
- ".bias",
353
- "_bias",
354
- ".k_scale",
355
- "_k_scale",
356
- ".v_scale",
357
- "_v_scale",
358
- ".weight_scale",
359
- "_weight_scale",
360
- ".input_scale",
361
- "_input_scale",
362
- )
363
-
364
- params_dict = dict(self.named_parameters())
365
- loaded_params: set[str] = set()
366
- expert_params_mapping = self.get_expert_mapping()
367
- for name, loaded_weight in weights:
368
- if name.startswith("mtp."):
369
- continue
370
- if "rotary_emb.inv_freq" in name:
371
- continue
372
- if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
373
- # Models trained using ColossalAI may include these tensors in
374
- # the checkpoint. Skip them.
375
- continue
376
- for param_name, weight_name, shard_id in stacked_params_mapping:
377
- if weight_name not in name:
378
- continue
379
- if "mlp.experts" in name:
380
- continue
381
- name = name.replace(weight_name, param_name)
382
- # Skip loading extra bias for GPTQ models.
383
- if name.endswith(".bias") and name not in params_dict:
384
- continue
385
- if is_pp_missing_parameter(name, self):
386
- continue
387
- if name not in params_dict:
388
- continue
389
-
390
- param = params_dict[name]
391
- weight_loader = param.weight_loader
392
- weight_loader(param, loaded_weight, shard_id)
393
- break
394
- else:
395
- is_expert_weight = False
396
- for mapping in expert_params_mapping:
397
- param_name, weight_name, expert_id, shard_id = mapping
398
- if weight_name not in name:
399
- continue
400
-
401
- # Anyway, this is an expert weight and should not be
402
- # attempted to load as other weights later
403
- is_expert_weight = True
404
-
405
- # Do not modify `name` since the loop may continue here
406
- # Instead, create a new variable
407
- name_mapped = name.replace(weight_name, param_name)
408
-
409
- if is_pp_missing_parameter(name_mapped, self):
410
- continue
411
-
412
- # Skip loading extra parameters for GPTQ/modelopt models.
413
- if (
414
- name_mapped.endswith(ignore_suffixes)
415
- and name_mapped not in params_dict
416
- ):
417
- continue
418
-
419
- param = params_dict[name_mapped]
420
- # We should ask the weight loader to return success or not
421
- # here since otherwise we may skip experts with other
422
- # available replicas.
423
- weight_loader = typing.cast(
424
- Callable[..., bool], param.weight_loader
425
- )
426
- success = weight_loader(
427
- param,
428
- loaded_weight,
429
- name_mapped,
430
- shard_id=shard_id,
431
- expert_id=expert_id,
432
- return_success=True,
433
- )
434
- if success:
435
- name = name_mapped
436
- break
437
- else:
438
- if is_expert_weight:
439
- continue
440
- # Skip loading extra bias for GPTQ models.
441
- if name.endswith(".bias") and name not in params_dict:
442
- continue
443
- # Skip loading extra parameters for GPTQ/modelopt models.
444
- if name.endswith(ignore_suffixes) and name not in params_dict:
445
- continue
446
- # Remapping the name of FP8 kv-scale.
447
- name = maybe_remap_kv_scale_name(name, params_dict)
448
- if name is None:
449
- continue
450
-
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
-
322
323
class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
324
+ hf_to_vllm_mapper = WeightsMapper(
325
+ orig_to_new_stacked={
326
+ # weight_name: (param_name, shard_id)
327
+ ".q_proj": (".qkv_proj", "q"),
328
+ ".k_proj": (".qkv_proj", "k"),
329
+ ".v_proj": (".qkv_proj", "v"),
330
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
331
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
332
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
333
+ ".shared_experts.gate_proj": (".shared_experts.gate_up_proj", 0),
334
+ ".shared_experts.up_proj": (".shared_experts.gate_up_proj", 1),
335
+ }
336
+ )
337
packed_modules_mapping = {
338
"qkv_proj": [
339
"q_proj",
@@ -554,5 +427,18 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
427
skip_prefixes=(
428
["lm_head.", "mtp."] if self.config.tie_word_embeddings else ["mtp."]
429
),
430
+ # Skip loading extra parameters for GPTQ/modelopt models.
431
+ ignore_unexpected_suffixes=[
432
+ ".bias",
433
+ "_bias",
434
+ ".k_scale",
435
+ "_k_scale",
436
+ ".v_scale",
437
+ "_v_scale",
438
+ ".weight_scale",
439
+ "_weight_scale",
440
+ ".input_scale",
441
+ "_input_scale",
442
+ ],
443
)
558
- return loader.load_weights(weights)
444
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/exaone_moe_mtp.py
+14
-55
@@ -18,15 +18,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
18
ParallelLMHead,
19
VocabParallelEmbedding,
20
)
21
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.model_executor.models.exaone_moe import ExaoneMoeDecoderLayer
22
from vllm.sequence import IntermediateTensors
23
25
-from .utils import (
26
- AutoWeightsLoader,
27
- is_pp_missing_parameter,
28
- maybe_prefix,
29
-)
24
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
25
26
logger = init_logger(__name__)
27
@@ -35,6 +30,17 @@ KVCache = tuple[torch.Tensor, torch.Tensor]
30
31
@support_torch_compile
32
class ExaoneMoeMultiTokenPredictor(nn.Module):
33
+ hf_to_vllm_mapper = WeightsMapper(
34
+ orig_to_new_stacked={
35
+ ".q_proj": (".qkv_proj", "q"),
36
+ ".k_proj": (".qkv_proj", "k"),
37
+ ".v_proj": (".qkv_proj", "v"),
38
+ # Scope to dense mlp; experts are handled separately.
39
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
40
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
41
+ }
42
+ )
43
+
44
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
45
super().__init__()
46
@@ -128,55 +134,8 @@ class ExaoneMoeMultiTokenPredictor(nn.Module):
134
return hidden_states
135
136
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
131
- stacked_params_mapping = [
132
- # (param_name, shard_name, shard_id)
133
- ("qkv_proj", "q_proj", "q"),
134
- ("qkv_proj", "k_proj", "k"),
135
- ("qkv_proj", "v_proj", "v"),
136
- ("gate_up_proj", "gate_proj", 0),
137
- ("gate_up_proj", "up_proj", 1),
138
- ]
139
-
140
- # Params for weights, fp8 weight scales, fp8 activation scales
141
- # (param_name, weight_name, expert_id, shard_id)
142
- params_dict = dict(self.named_parameters())
143
- loaded_params: set[str] = set()
144
- for name, loaded_weight in weights:
145
- if "rotary_emb.inv_freq" in name:
146
- continue
147
-
148
- for param_name, weight_name, shard_id in stacked_params_mapping:
149
- if weight_name not in name:
150
- continue
151
-
152
- if "mlp.experts" in name:
153
- continue
154
-
155
- name = name.replace(weight_name, param_name)
156
- # Skip loading extra bias for GPTQ models.
157
- if name.endswith(".bias") and name not in params_dict:
158
- continue
159
- # Skip layers on other devices.
160
- if is_pp_missing_parameter(name, self):
161
- continue
162
- if name not in params_dict:
163
- continue
164
- param = params_dict[name]
165
- weight_loader = param.weight_loader
166
- weight_loader(param, loaded_weight, shard_id)
167
- break
168
- else:
169
- # Skip loading extra bias for GPTQ models.
170
- if name.endswith(".bias") and name not in params_dict:
171
- continue
172
- if is_pp_missing_parameter(name, self):
173
- continue
174
-
175
- param = params_dict[name]
176
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
177
- weight_loader(param, loaded_weight)
178
- loaded_params.add(name)
179
- return loaded_params
137
+ loader = AutoWeightsLoader(self)
138
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
139
140
141
@support_torch_compile
vllm/model_executor/models/falcon_h1.py
+16
-63
@@ -35,10 +35,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35
ParallelLMHead,
36
VocabParallelEmbedding,
37
)
38
-from vllm.model_executor.model_loader.weight_utils import (
39
- default_weight_loader,
40
- maybe_remap_kv_scale_name,
41
-)
38
from vllm.sequence import IntermediateTensors
39
from vllm.transformers_utils.config import set_default_rope_theta
40
@@ -52,7 +48,7 @@ from .interfaces import (
48
from .utils import (
49
AutoWeightsLoader,
50
PPMissingLayer,
55
- is_pp_missing_parameter,
51
+ WeightsMapper,
52
make_empty_intermediate_tensors_factory,
53
make_layers,
54
maybe_prefix,
@@ -496,63 +492,6 @@ class FalconH1Model(nn.Module):
492
hidden_states = self.final_layernorm(hidden_states)
493
return hidden_states
494
499
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
500
- stacked_params_mapping = [
501
- # (param_name, shard_name, shard_id)
502
- ("qkv_proj", "q_proj", "q"),
503
- ("qkv_proj", "k_proj", "k"),
504
- ("qkv_proj", "v_proj", "v"),
505
- ("gate_up_proj", "gate_proj", 0),
506
- ("gate_up_proj", "up_proj", 1),
507
- ]
508
-
509
- params_dict = dict(self.named_parameters())
510
- loaded_params: set[str] = set()
511
- for name, loaded_weight in weights:
512
- if "rotary_emb.inv_freq" in name:
513
- continue
514
-
515
- if "A_log" in name:
516
- name = name.replace("A_log", "A")
517
-
518
- if "mamba" in name:
519
- name = name.replace("mamba", "mamba.mamba")
520
-
521
- if "scale" in name:
522
- # Remapping the name of kv-scale.
523
- name = maybe_remap_kv_scale_name(name, params_dict)
524
- if name is None:
525
- continue
526
-
527
- for param_name, weight_name, shard_id in stacked_params_mapping:
528
- if weight_name not in name:
529
- continue
530
-
531
- name = name.replace(weight_name, param_name)
532
- # Skip loading extra bias for GPTQ models.
533
- if name.endswith(".bias") and name not in params_dict:
534
- continue
535
- # Skip layers on other devices.
536
- if is_pp_missing_parameter(name, self):
537
- continue
538
- param = params_dict[name]
539
- weight_loader = param.weight_loader
540
- weight_loader(param, loaded_weight, shard_id)
541
- break
542
- else:
543
- # Skip loading extra bias for GPTQ models.
544
- if name.endswith(".bias") and name not in params_dict:
545
- continue
546
- if is_pp_missing_parameter(name, self):
547
- continue
548
-
549
- param = params_dict[name]
550
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
551
- weight_loader(param, loaded_weight)
552
- loaded_params.add(name)
553
-
554
- return loaded_params
555
-
495
496
class FalconH1ForCausalLM(
497
nn.Module,
@@ -567,6 +506,20 @@ class FalconH1ForCausalLM(
506
"gate_up_proj": ["gate_proj", "up_proj"],
507
}
508
509
+ hf_to_vllm_mapper = WeightsMapper(
510
+ orig_to_new_substr={
511
+ "A_log": "A",
512
+ "mamba": "mamba.mamba",
513
+ },
514
+ orig_to_new_stacked={
515
+ ".q_proj": (".qkv_proj", "q"),
516
+ ".k_proj": (".qkv_proj", "k"),
517
+ ".v_proj": (".qkv_proj", "v"),
518
+ ".gate_proj": (".gate_up_proj", 0),
519
+ ".up_proj": (".gate_up_proj", 1),
520
+ },
521
+ )
522
+
523
embedding_modules = {
524
"embed_tokens": "input_embeddings",
525
"lm_head": "output_embeddings",
@@ -694,4 +647,4 @@ class FalconH1ForCausalLM(
647
self,
648
skip_prefixes=(["lm_head."] if self.tie_word_embeddings else None),
649
)
697
- return loader.load_weights(weights)
650
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/funasr.py
+10
-30
@@ -31,7 +31,6 @@ from vllm.model_executor.layers.linear import (
31
)
32
from vllm.model_executor.layers.logits_processor import LogitsProcessor
33
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
34
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
34
from vllm.model_executor.models.whisper_utils import (
35
ISO639_1_SUPPORTED_LANGS,
36
)
@@ -618,6 +617,14 @@ class FunASRAudioInputs(TensorSchema):
617
618
619
class FunASREncoder(nn.Module):
620
+ hf_to_vllm_mapper = WeightsMapper(
621
+ orig_to_new_stacked={
622
+ ".self_attn.q_proj": (".self_attn.qkv", "q"),
623
+ ".self_attn.k_proj": (".self_attn.qkv", "k"),
624
+ ".self_attn.v_proj": (".self_attn.qkv", "v"),
625
+ }
626
+ )
627
+
628
def __init__(
629
self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
630
):
@@ -637,35 +644,8 @@ class FunASREncoder(nn.Module):
644
)
645
646
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
640
- """Load weights with mapping from HuggingFace format."""
641
- stacked_params_mapping = [
642
- # (param_name, shard_name, shard_id)
643
- ("self_attn.qkv.", "self_attn.q_proj.", "q"),
644
- ("self_attn.qkv.", "self_attn.k_proj.", "k"),
645
- ("self_attn.qkv.", "self_attn.v_proj.", "v"),
646
- ]
647
- params_dict = dict(self.named_parameters(remove_duplicate=False))
648
- loaded_params: set[str] = set()
649
-
650
- for name, loaded_weight in weights:
651
- for param_name, weight_name, shard_id in stacked_params_mapping:
652
- if weight_name not in name:
653
- continue
654
- name = name.replace(weight_name, param_name)
655
-
656
- param = params_dict[name]
657
- weight_loader = param.weight_loader
658
- weight_loader(param, loaded_weight, shard_id)
659
- break
660
- else:
661
- param = params_dict.get(name)
662
- if param is not None:
663
- weight_loader = getattr(
664
- param, "weight_loader", default_weight_loader
665
- )
666
- weight_loader(param, loaded_weight)
667
- loaded_params.add(name)
668
- return loaded_params
647
+ loader = AutoWeightsLoader(self)
648
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
649
650
651
class FunASRModel(nn.Module):
vllm/model_executor/models/gemma4_mtp.py
+4
-39
@@ -44,7 +44,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
44
ParallelLMHead,
45
VocabParallelEmbedding,
46
)
47
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
47
from vllm.sequence import IntermediateTensors
48
49
from .gemma4 import Gemma4MLP, _get_text_config
@@ -412,44 +411,6 @@ class Gemma4MultiTokenPredictor(nn.Module):
411
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
412
return self.embed_tokens(input_ids) * self.normalizer
413
415
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
416
- stacked_params_mapping = [
417
- ("gate_up_proj", "gate_proj", 0),
418
- ("gate_up_proj", "up_proj", 1),
419
- ]
420
-
421
- params_dict = dict(self.named_parameters())
422
- params_dict.update(dict(self.named_buffers()))
423
- loaded_params: set[str] = set()
424
-
425
- for name, loaded_weight in weights:
426
- if "rotary_emb.inv_freq" in name:
427
- continue
428
-
429
- for param_name, weight_name, shard_id in stacked_params_mapping:
430
- if weight_name not in name:
431
- continue
432
- name = name.replace(weight_name, param_name)
433
- if name.endswith(".bias") and name not in params_dict:
434
- continue
435
- if name not in params_dict:
436
- continue
437
- param = params_dict[name]
438
- weight_loader = param.weight_loader
439
- weight_loader(param, loaded_weight, shard_id)
440
- break
441
- else:
442
- if name.endswith(".bias") and name not in params_dict:
443
- continue
444
- if name not in params_dict:
445
- continue
446
- param = params_dict[name]
447
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
448
- weight_loader(param, loaded_weight)
449
- loaded_params.add(name)
450
-
451
- return loaded_params
452
-
414
def forward(
415
self,
416
input_ids: torch.Tensor,
@@ -502,6 +463,10 @@ class Gemma4MTP(nn.Module):
463
"pre_projection.": "model.pre_projection.",
464
"post_projection.": "model.post_projection.",
465
},
466
+ orig_to_new_stacked={
467
+ ".gate_proj": (".gate_up_proj", 0),
468
+ ".up_proj": (".gate_up_proj", 1),
469
+ },
470
)
471
472
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/glm4_1v.py
+12
-28
@@ -76,7 +76,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
76
from vllm.model_executor.layers.rotary_embedding.common import (
77
ApplyRotaryEmb,
78
)
79
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
79
from vllm.model_executor.models.module_mapping import MultiModelKeys
80
from vllm.multimodal import MULTIMODAL_REGISTRY
81
from vllm.multimodal.inputs import (
@@ -618,6 +617,16 @@ class Glm4vVisionEmbeddings(nn.Module):
617
618
619
class Glm4vVisionTransformer(nn.Module):
620
+ hf_to_vllm_mapper = WeightsMapper(
621
+ orig_to_new_stacked={
622
+ ".attn.q.": (".attn.qkv.", "q"),
623
+ ".attn.k.": (".attn.qkv.", "k"),
624
+ ".attn.v.": (".attn.qkv.", "v"),
625
+ ".gate_proj": (".gate_up_proj", 0),
626
+ ".up_proj": (".gate_up_proj", 1),
627
+ }
628
+ )
629
+
630
def __init__(
631
self,
632
text_config: Glm4vTextConfig,
@@ -958,33 +967,8 @@ class Glm4vVisionTransformer(nn.Module):
967
return x
968
969
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
961
- stacked_params_mapping = [
962
- # (param_name, shard_name, shard_id)
963
- ("attn.qkv.", "attn.q.", "q"),
964
- ("attn.qkv.", "attn.k.", "k"),
965
- ("attn.qkv.", "attn.v.", "v"),
966
- ("gate_up_proj", "gate_proj", 0),
967
- ("gate_up_proj", "up_proj", 1),
968
- ]
969
- params_dict = dict(self.named_parameters(remove_duplicate=False))
970
- loaded_params: set[str] = set()
971
-
972
- for name, loaded_weight in weights:
973
- for param_name, weight_name, shard_id in stacked_params_mapping:
974
- if weight_name not in name:
975
- continue
976
- name = name.replace(weight_name, param_name)
977
-
978
- param = params_dict[name]
979
- weight_loader = param.weight_loader
980
- weight_loader(param, loaded_weight, shard_id)
981
- break
982
- else:
983
- param = params_dict[name]
984
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
985
- weight_loader(param, loaded_weight)
986
- loaded_params.add(name)
987
- return loaded_params
970
+ loader = AutoWeightsLoader(self)
971
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
972
973
974
class Glm4vProcessingInfo(BaseProcessingInfo):
vllm/model_executor/models/hunyuan_vision.py
+10
-26
@@ -48,7 +48,6 @@ from vllm.model_executor.layers.linear import (
48
RowParallelLinear,
49
)
50
from vllm.model_executor.layers.quantization import QuantizationConfig
51
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51
from vllm.model_executor.models.module_mapping import MultiModelKeys
52
from vllm.multimodal import MULTIMODAL_REGISTRY
53
from vllm.multimodal.inputs import (
@@ -431,6 +430,14 @@ class HunYuanVisionPatchMerger(nn.Module):
430
431
432
class HunYuanVisionTransformer(nn.Module):
433
+ hf_to_vllm_mapper = WeightsMapper(
434
+ orig_to_new_stacked={
435
+ ".q_proj": (".qkv", "q"),
436
+ ".k_proj": (".qkv", "k"),
437
+ ".v_proj": (".qkv", "v"),
438
+ }
439
+ )
440
+
441
def __init__(
442
self,
443
vision_config: HunYuanVLVisionConfig,
@@ -529,31 +536,8 @@ class HunYuanVisionTransformer(nn.Module):
536
return image_embeds_list
537
538
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
532
- stacked_params_mapping = [
533
- # (param_name, shard_name, shard_id)
534
- (".qkv", ".q_proj", "q"),
535
- (".qkv", ".k_proj", "k"),
536
- (".qkv", ".v_proj", "v"),
537
- ]
538
- params_dict = dict(self.named_parameters(remove_duplicate=False))
539
- loaded_params: set[str] = set()
540
-
541
- for name, loaded_weight in weights:
542
- for param_name, weight_name, shard_id in stacked_params_mapping:
543
- if weight_name not in name:
544
- continue
545
- name = name.replace(weight_name, param_name)
546
-
547
- param = params_dict[name]
548
- weight_loader = param.weight_loader
549
- weight_loader(param, loaded_weight, shard_id)
550
- break
551
- else:
552
- param = params_dict[name]
553
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
554
- weight_loader(param, loaded_weight)
555
- loaded_params.add(name)
556
- return loaded_params
539
+ loader = AutoWeightsLoader(self)
540
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
541
542
543
def _hunyuan_vl_field_config(hf_inputs: Mapping[str, torch.Tensor]):
vllm/model_executor/models/idefics2_vision_model.py
+24
-36
@@ -38,8 +38,8 @@ from vllm.model_executor.layers.linear import (
38
RowParallelLinear,
39
)
40
from vllm.model_executor.layers.quantization import QuantizationConfig
41
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
41
42
+from .utils import AutoWeightsLoader, WeightsMapper
43
from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model
44
45
@@ -352,6 +352,14 @@ class Idefics2Encoder(nn.Module):
352
353
354
class Idefics2VisionTransformer(nn.Module):
355
+ hf_to_vllm_mapper = WeightsMapper(
356
+ orig_to_new_stacked={
357
+ ".q_proj": (".qkv_proj", "q"),
358
+ ".k_proj": (".qkv_proj", "k"),
359
+ ".v_proj": (".qkv_proj", "v"),
360
+ }
361
+ )
362
+
363
def __init__(
364
self,
365
config: Idefics2VisionConfig,
@@ -453,42 +461,22 @@ class Idefics2VisionTransformer(nn.Module):
461
return last_hidden_state
462
463
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
456
- stacked_params_mapping = [
457
- # (param_name, shard_name, shard_id)
458
- ("qkv_proj", "q_proj", "q"),
459
- ("qkv_proj", "k_proj", "k"),
460
- ("qkv_proj", "v_proj", "v"),
461
- ]
462
- params_dict = dict(self.named_parameters())
463
- loaded_params: set[str] = set()
464
- layer_count = len(self.encoder.layers)
465
-
466
- for name, loaded_weight in weights:
467
- # skip pooling header
468
- if name.startswith("head."):
469
- continue
464
+ # head is a pooling header absent from this model.
465
+ skip_prefixes = ["head."]
466
+ if not self.require_post_norm:
467
+ skip_prefixes.append("post_layernorm.")
468
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
469
471
- # post_layernorm is optional
472
- if name.startswith("post_layernorm.") and not self.require_post_norm:
473
- continue
470
+ layer_count = len(self.encoder.layers)
471
475
- # omit layers when num_hidden_layers_override is set
476
- if name.startswith("encoder.layers."):
477
- layer_idx = int(name.split(".")[2])
478
- if layer_idx >= layer_count:
472
+ def _filter(ws: Iterable[tuple[str, torch.Tensor]]):
473
+ # Drop layers beyond num_hidden_layers_override.
474
+ for name, w in ws:
475
+ if (
476
+ name.startswith("encoder.layers.")
477
+ and int(name.split(".")[2]) >= layer_count
478
+ ):
479
continue
480
+ yield name, w
481
481
- for param_name, weight_name, shard_id in stacked_params_mapping:
482
- if weight_name not in name or self.use_data_parallel:
483
- continue
484
- name = name.replace(weight_name, param_name)
485
- param = params_dict[name]
486
- weight_loader = param.weight_loader
487
- weight_loader(param, loaded_weight, shard_id)
488
- break
489
- else:
490
- param = params_dict[name]
491
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
492
- weight_loader(param, loaded_weight)
493
- loaded_params.add(name)
494
- return loaded_params
482
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/isaac.py
+10
-28
@@ -25,9 +25,6 @@ from vllm.model_executor.layers.linear import (
25
RowParallelLinear,
26
)
27
from vllm.model_executor.layers.quantization import QuantizationConfig
28
-from vllm.model_executor.model_loader.weight_utils import (
29
- default_weight_loader,
30
-)
28
from vllm.model_executor.models.interfaces import (
29
MultiModalEmbeddings,
30
SupportsLoRA,
@@ -652,6 +649,14 @@ class Siglip2Encoder(nn.Module):
649
650
651
class Siglip2VisionTransformer(nn.Module):
652
+ hf_to_vllm_mapper = WeightsMapper(
653
+ orig_to_new_stacked={
654
+ ".q_proj": (".qkv_proj", "q"),
655
+ ".k_proj": (".qkv_proj", "k"),
656
+ ".v_proj": (".qkv_proj", "v"),
657
+ }
658
+ )
659
+
660
def __init__(
661
self,
662
config: PixelShuffleSiglip2VisionConfig,
@@ -715,31 +720,8 @@ class Siglip2VisionTransformer(nn.Module):
720
return hidden_states
721
722
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
718
- stacked_params_mapping = [
719
- # (param_name, shard_name, shard_id)
720
- ("qkv_proj", "q_proj", "q"),
721
- ("qkv_proj", "k_proj", "k"),
722
- ("qkv_proj", "v_proj", "v"),
723
- ]
724
- params_dict = dict(self.named_parameters())
725
- loaded_params: set[str] = set()
726
-
727
- for name, loaded_weight in weights:
728
- for param_name, weight_name, shard_id in stacked_params_mapping:
729
- if weight_name not in name:
730
- continue
731
- name = name.replace(weight_name, param_name)
732
-
733
- param = params_dict[name]
734
- weight_loader = param.weight_loader
735
- weight_loader(param, loaded_weight, shard_id)
736
- break
737
- else:
738
- param = params_dict[name]
739
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
740
- weight_loader(param, loaded_weight)
741
- loaded_params.add(name)
742
- return loaded_params
723
+ loader = AutoWeightsLoader(self)
724
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
725
726
727
def _resolve_vision_token_id(model_config: ModelConfig, vision_token: str) -> int:
vllm/model_executor/models/keye.py
+10
-52
@@ -34,10 +34,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.layers.rotary_embedding.common import (
35
ApplyRotaryEmb,
36
)
37
-from vllm.model_executor.model_loader.weight_utils import (
38
- default_weight_loader,
39
- maybe_remap_kv_scale_name,
40
-)
37
from vllm.model_executor.models.module_mapping import MultiModelKeys
38
from vllm.multimodal import MULTIMODAL_REGISTRY
39
from vllm.multimodal.inputs import (
@@ -76,7 +72,6 @@ from .utils import (
72
AutoWeightsLoader,
73
WeightsMapper,
74
init_vllm_registered_model,
79
- is_pp_missing_parameter,
75
maybe_prefix,
76
)
77
from .vision import is_vit_use_data_parallel
@@ -718,6 +713,14 @@ class KeyeSiglipVisionModel(nn.Module):
713
config_class = PretrainedConfig
714
main_input_name = "pixel_values"
715
716
+ hf_to_vllm_mapper = WeightsMapper(
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
def __init__(
725
self,
726
config: PretrainedConfig,
@@ -776,53 +779,8 @@ class KeyeSiglipVisionModel(nn.Module):
779
)
780
781
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
779
- stacked_params_mapping = [
780
- ("qkv_proj", "q_proj", "q"),
781
- ("qkv_proj", "k_proj", "k"),
782
- ("qkv_proj", "v_proj", "v"),
783
- ]
784
- params_dict = dict(self.named_parameters(remove_duplicate=False))
785
- loaded_params: set[str] = set()
786
- for name, loaded_weight in weights:
787
- if "rotary_emb.inv_freq" in name:
788
- continue
789
- if "head.attention" in name or "head.layernorm" in name:
790
- continue
791
- if "head.mlp" in name or "head.probe" in name:
792
- continue
793
- for (
794
- param_name,
795
- weight_name,
796
- shard_id,
797
- ) in stacked_params_mapping:
798
- if weight_name not in name:
799
- continue
800
- name = name.replace(weight_name, param_name)
801
- if name.endswith(".bias") and name not in params_dict:
802
- continue
803
- if is_pp_missing_parameter(name, self):
804
- continue
805
- param = params_dict[name]
806
- weight_loader = param.weight_loader
807
- weight_loader(param, loaded_weight, shard_id)
808
- break
809
- else:
810
- if name.endswith(".bias") and name not in params_dict:
811
- continue
812
- name = maybe_remap_kv_scale_name(name, params_dict)
813
- if name is None:
814
- continue
815
- if is_pp_missing_parameter(name, self):
816
- continue
817
- param = params_dict[name]
818
- weight_loader = getattr(
819
- param,
820
- "weight_loader",
821
- default_weight_loader,
822
- )
823
- weight_loader(param, loaded_weight)
824
- loaded_params.add(name)
825
- return loaded_params
782
+ loader = AutoWeightsLoader(self, skip_prefixes=["vision_model.head."])
783
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
784
785
786
class Projector(nn.Module):
vllm/model_executor/models/kimi_audio.py
+10
-32
@@ -17,7 +17,6 @@ from vllm.config.multimodal import BaseDummyOptions
17
from vllm.config.speech_to_text import SpeechToTextParams
18
from vllm.inputs import PromptType, TokensPrompt
19
from vllm.model_executor.model_loader import DefaultModelLoader
20
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
20
from vllm.model_executor.models.interfaces import (
21
SupportsMultiModal,
22
SupportsPP,
@@ -81,6 +80,14 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
80
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
81
}
82
83
+ hf_to_vllm_mapper = WeightsMapper(
84
+ orig_to_new_stacked={
85
+ ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
86
+ ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
87
+ ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
88
+ },
89
+ )
90
+
91
def __init__(
92
self, *, vllm_config: VllmConfig, prefix: str = "", init_in_fp32: bool = False
93
):
@@ -102,37 +109,8 @@ class KimiAudioWhisperEncoder(WhisperEncoder):
109
)
110
111
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
105
- stacked_params_mapping = [
106
- # (param_name, shard_name, shard_id)
107
- ("qkv_proj", "q_proj", "q"),
108
- ("qkv_proj", "k_proj", "k"),
109
- ("qkv_proj", "v_proj", "v"),
110
- ]
111
- params_dict = dict(self.named_parameters())
112
- loaded_params: set[str] = set()
113
- for name, loaded_weight in weights:
114
- for param_name, weight_name, shard_id in stacked_params_mapping:
115
- if weight_name not in name:
116
- continue
117
- name = name.replace(weight_name, param_name)
118
- # Skip loading extra bias for GPTQ models.
119
- if name.endswith(".bias") and name not in params_dict:
120
- continue
121
-
122
- param = params_dict[name]
123
- weight_loader = param.weight_loader
124
- weight_loader(param, loaded_weight, shard_id)
125
- break
126
- else:
127
- # Skip loading extra bias for GPTQ models.
128
- if name.endswith(".bias") and name not in params_dict:
129
- continue
130
-
131
- param = params_dict[name]
132
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
133
- weight_loader(param, loaded_weight)
134
- loaded_params.add(name)
135
- return loaded_params
112
+ loader = AutoWeightsLoader(self)
113
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
114
115
116
# -----------------------------------------------------------------------------
vllm/model_executor/models/lfm2.py
+18
-42
@@ -32,7 +32,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
32
ParallelLMHead,
33
VocabParallelEmbedding,
34
)
35
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35
from vllm.sequence import IntermediateTensors
36
37
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP, SupportsQuant
@@ -41,7 +40,6 @@ from .utils import (
40
PPMissingLayer,
41
WeightsMapper,
42
extract_layer_index,
44
- is_pp_missing_parameter,
43
make_empty_intermediate_tensors_factory,
44
make_layers,
45
maybe_prefix,
@@ -297,6 +295,20 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
295
296
@support_torch_compile
297
class Lfm2Model(nn.Module):
298
+ # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
299
+ # with the inner .conv.conv child (ShortConv has a child self.conv, so
300
+ # naming the container .conv too makes _match_target_modules match both).
301
+ hf_to_vllm_mapper = WeightsMapper(
302
+ orig_to_new_substr={".conv.": ".short_conv."},
303
+ orig_to_new_stacked={
304
+ ".q_proj": (".qkv_proj", "q"),
305
+ ".k_proj": (".qkv_proj", "k"),
306
+ ".v_proj": (".qkv_proj", "v"),
307
+ ".w1": (".w13", 0),
308
+ ".w3": (".w13", 1),
309
+ },
310
+ )
311
+
312
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
313
super().__init__()
314
@@ -375,40 +387,8 @@ class Lfm2Model(nn.Module):
387
return hidden_states
388
389
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
378
- stacked_params_mapping = [
379
- (".qkv_proj", ".q_proj", "q"),
380
- (".qkv_proj", ".k_proj", "k"),
381
- (".qkv_proj", ".v_proj", "v"),
382
- (".w13", ".w1", 0),
383
- (".w13", ".w3", 1),
384
- ]
385
- params_dict = dict(self.named_parameters())
386
- loaded_params: set[str] = set()
387
- for name, loaded_weight in weights:
388
- if ".conv." in name:
389
- name = name.replace(".conv.", ".short_conv.", 1)
390
-
391
- for param_name, weight_name, shard_id in stacked_params_mapping:
392
- # Use segment-boundary matching (trailing dot) to prevent
393
- # e.g. ".w1" from matching inside ".w13" in pre-fused keys.
394
- if weight_name + "." not in name:
395
- continue
396
- name = name.replace(weight_name + ".", param_name + ".")
397
-
398
- if is_pp_missing_parameter(name, self):
399
- continue
400
- param = params_dict[name]
401
- weight_loader = param.weight_loader
402
- weight_loader(param, loaded_weight, shard_id)
403
- break
404
- else:
405
- if is_pp_missing_parameter(name, self):
406
- continue
407
- param = params_dict[name]
408
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
409
- weight_loader(param, loaded_weight)
410
- loaded_params.add(name)
411
- return loaded_params
390
+ loader = AutoWeightsLoader(self)
391
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
392
393
394
class Lfm2ForCausalLM(
@@ -427,12 +407,8 @@ class Lfm2ForCausalLM(
407
"in_proj": ["in_proj"],
408
}
409
430
- # HF uses .conv. but vLLM uses .short_conv. to avoid LoRA regex collision
431
- # with the inner .conv.conv child (ShortConv has a child self.conv, so
432
- # naming the container .conv too makes _match_target_modules match both)
433
- hf_to_vllm_mapper = WeightsMapper(
434
- orig_to_new_substr={".conv.": ".short_conv."},
435
- )
410
+ # Reuse the backbone mapper so LoRA/quantization see the same name mapping.
411
+ hf_to_vllm_mapper = Lfm2Model.hf_to_vllm_mapper
412
413
# LoRA specific attributes
414
embedding_modules = {
vllm/model_executor/models/lfm2_siglip2.py
+23
-37
@@ -23,9 +23,8 @@ from vllm.model_executor.layers.linear import (
23
RowParallelLinear,
24
)
25
from vllm.model_executor.layers.quantization import QuantizationConfig
26
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
26
28
-from .utils import maybe_prefix
27
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
28
from .vision import (
29
is_vit_use_data_parallel,
30
resolve_visual_encoder_outputs,
@@ -458,6 +457,14 @@ class Siglip2VisionTransformer(nn.Module):
457
458
459
class Siglip2Model(torch.nn.Module):
460
+ hf_to_vllm_mapper = WeightsMapper(
461
+ orig_to_new_stacked={
462
+ ".q_proj": (".qkv_proj", "q"),
463
+ ".k_proj": (".qkv_proj", "k"),
464
+ ".v_proj": (".qkv_proj", "v"),
465
+ }
466
+ )
467
+
468
def __init__(
469
self,
470
config: Siglip2VisionConfig,
@@ -501,42 +508,21 @@ class Siglip2Model(torch.nn.Module):
508
)
509
510
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
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
- layer_count = len(self.vision_model.encoder.layers)
511
+ skip_prefixes = []
512
+ if self.vision_model.post_layernorm is None:
513
+ skip_prefixes.append("vision_model.post_layernorm.")
514
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
515
514
- for name, loaded_weight in weights:
515
- # post_layernorm is optional in Siglip2Model
516
- if (
517
- name.startswith("vision_model.post_layernorm")
518
- and self.vision_model.post_layernorm is None
519
- ):
520
- continue
516
+ # Drop layers omitted by num_hidden_layers_override.
517
+ layer_count = len(self.vision_model.encoder.layers)
518
522
- # omit layers when num_hidden_layers_override is set
523
- if name.startswith("vision_model.encoder.layers"):
524
- layer_idx = int(name.split(".")[3])
525
- if layer_idx >= layer_count:
519
+ def _filter(ws):
520
+ for n, w in ws:
521
+ if (
522
+ n.startswith("vision_model.encoder.layers.")
523
+ and int(n.split(".")[3]) >= layer_count
524
+ ):
525
continue
526
+ yield n, w
527
528
- for param_name, weight_name, shard_id in stacked_params_mapping:
529
- if weight_name not in name:
530
- continue
531
- name = name.replace(weight_name, param_name)
532
-
533
- param = params_dict[name]
534
- weight_loader = param.weight_loader
535
- weight_loader(param, loaded_weight, shard_id)
536
- break
537
- else:
538
- param = params_dict[name]
539
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
540
- weight_loader(param, loaded_weight)
541
- loaded_params.add(name)
542
- return loaded_params
528
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/llama4_eagle.py
+19
-30
@@ -32,18 +32,33 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
32
ParallelLMHead,
33
VocabParallelEmbedding,
34
)
35
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35
from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausalLM
36
from vllm.model_executor.models.utils import extract_layer_index
37
38
from .interfaces import SupportsMultiModal
40
-from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight
39
+from .utils import (
40
+ AutoWeightsLoader,
41
+ WeightsMapper,
42
+ maybe_prefix,
43
+ process_eagle_weight,
44
+)
45
46
logger = init_logger(__name__)
47
48
49
@support_torch_compile
50
class LlamaModel(nn.Module):
51
+ hf_to_vllm_mapper = WeightsMapper(
52
+ orig_to_new_prefix={"model.": ""},
53
+ orig_to_new_stacked={
54
+ ".q_proj": (".qkv_proj", "q"),
55
+ ".k_proj": (".qkv_proj", "k"),
56
+ ".v_proj": (".qkv_proj", "v"),
57
+ ".gate_proj": (".gate_up_proj", 0),
58
+ ".up_proj": (".gate_up_proj", 1),
59
+ },
60
+ )
61
+
62
def __init__(
63
self,
64
*,
@@ -108,34 +123,8 @@ class LlamaModel(nn.Module):
123
return hidden_states, hidden_states
124
125
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
111
- stacked_params_mapping = [
112
- # (param_name, shard_name, shard_id)
113
- (".qkv_proj", ".q_proj", "q"),
114
- (".qkv_proj", ".k_proj", "k"),
115
- (".qkv_proj", ".v_proj", "v"),
116
- (".gate_up_proj", ".gate_proj", 0),
117
- (".gate_up_proj", ".up_proj", 1),
118
- ]
119
- params_dict = dict(self.named_parameters())
120
- loaded_params: set[str] = set()
121
- for name, loaded_weight in weights:
122
- name = name.removeprefix("model.")
123
- for param_name, weight_name, shard_id in stacked_params_mapping:
124
- if weight_name not in name:
125
- continue
126
- name = name.replace(weight_name, param_name)
127
- param = params_dict[name]
128
- weight_loader = param.weight_loader
129
- weight_loader(param, loaded_weight, shard_id)
130
- break
131
- else:
132
- param = params_dict[name]
133
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
134
- weight_loader(param, loaded_weight)
135
- loaded_params.add(name)
136
- for name in params_dict:
137
- assert name in loaded_params, f"{name} is not loaded!"
138
- return loaded_params
126
+ loader = AutoWeightsLoader(self)
127
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
128
129
def validate_and_update_config(
130
self, start_layer_id: int, quant_config: QuantizationConfig | None = None
vllm/model_executor/models/llama_eagle.py
+14
-34
@@ -14,14 +14,11 @@ from vllm.model_executor.layers.linear import ReplicatedLinear
14
from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
16
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
17
-from vllm.model_executor.model_loader.weight_utils import (
18
- default_weight_loader,
19
- maybe_remap_kv_scale_name,
20
-)
17
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
18
19
from .utils import (
20
AutoWeightsLoader,
21
+ WeightsMapper,
22
get_draft_quant_config,
23
maybe_prefix,
24
process_eagle_weight,
@@ -53,6 +50,17 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
50
51
@support_torch_compile
52
class LlamaModel(nn.Module):
53
+ hf_to_vllm_mapper = WeightsMapper(
54
+ orig_to_new_stacked={
55
+ # weight_name: (param_name, shard_id)
56
+ ".q_proj": (".qkv_proj", "q"),
57
+ ".k_proj": (".qkv_proj", "k"),
58
+ ".v_proj": (".qkv_proj", "v"),
59
+ ".gate_proj": (".gate_up_proj", 0),
60
+ ".up_proj": (".gate_up_proj", 1),
61
+ }
62
+ )
63
+
64
def __init__(
65
self,
66
*,
@@ -116,36 +124,8 @@ class LlamaModel(nn.Module):
124
return hidden_states, hidden_states
125
126
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
119
- stacked_params_mapping = [
120
- # (param_name, shard_name, shard_id)
121
- (".qkv_proj", ".q_proj", "q"),
122
- (".qkv_proj", ".k_proj", "k"),
123
- (".qkv_proj", ".v_proj", "v"),
124
- (".gate_up_proj", ".gate_proj", 0),
125
- (".gate_up_proj", ".up_proj", 1),
126
- ]
127
- params_dict = dict(self.named_parameters())
128
- loaded_params: set[str] = set()
129
- for name, loaded_weight in weights:
130
- # Remapping the name FP8 kv-scale or zero point.
131
- if "scale" in name or "zero_point" in name:
132
- name = maybe_remap_kv_scale_name(name, params_dict)
133
- if name is None:
134
- continue
135
- for param_name, weight_name, shard_id in stacked_params_mapping:
136
- if weight_name not in name:
137
- continue
138
- name = name.replace(weight_name, param_name)
139
- param = params_dict[name]
140
- weight_loader = param.weight_loader
141
- weight_loader(param, loaded_weight, shard_id)
142
- break
143
- else:
144
- param = params_dict[name]
145
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
146
- weight_loader(param, loaded_weight)
147
- loaded_params.add(name)
148
- return loaded_params
127
+ loader = AutoWeightsLoader(self)
128
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
129
130
131
class EagleLlamaForCausalLM(LlamaForCausalLM):
vllm/model_executor/models/llama_eagle3.py
+14
-36
@@ -18,15 +18,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
18
ParallelLMHead,
19
VocabParallelEmbedding,
20
)
21
-from vllm.model_executor.model_loader.weight_utils import (
22
- default_weight_loader,
23
- maybe_remap_kv_scale_name,
24
-)
21
from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM
22
from vllm.multimodal.inputs import NestedTensors
23
24
from .utils import (
25
AutoWeightsLoader,
26
+ WeightsMapper,
27
get_draft_quant_config,
28
maybe_prefix,
29
process_eagle_weight,
@@ -253,39 +250,20 @@ class LlamaModel(nn.Module):
250
251
return hidden_states, aux_output
252
253
+ hf_to_vllm_mapper = WeightsMapper(
254
+ orig_to_new_substr={"midlayer.": "layers.0."},
255
+ orig_to_new_stacked={
256
+ ".q_proj": (".qkv_proj", "q"),
257
+ ".k_proj": (".qkv_proj", "k"),
258
+ ".v_proj": (".qkv_proj", "v"),
259
+ ".gate_proj": (".gate_up_proj", 0),
260
+ ".up_proj": (".gate_up_proj", 1),
261
+ },
262
+ )
263
+
264
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
257
- stacked_params_mapping = [
258
- # (param_name, shard_name, shard_id)
259
- (".qkv_proj", ".q_proj", "q"),
260
- (".qkv_proj", ".k_proj", "k"),
261
- (".qkv_proj", ".v_proj", "v"),
262
- (".gate_up_proj", ".gate_proj", 0),
263
- (".gate_up_proj", ".up_proj", 1),
264
- ]
265
- params_dict = dict(self.named_parameters())
266
- loaded_params: set[str] = set()
267
- for name, loaded_weight in weights:
268
- if "midlayer." in name:
269
- name = name.replace("midlayer.", "layers.0.")
270
- # Remapping the name FP8 kv-scale or zero point.
271
- if "scale" in name or "zero_point" in name:
272
- name = maybe_remap_kv_scale_name(name, params_dict)
273
- if name is None:
274
- continue
275
- for param_name, weight_name, shard_id in stacked_params_mapping:
276
- if weight_name not in name:
277
- continue
278
- name = name.replace(weight_name, param_name)
279
- param = params_dict[name]
280
- weight_loader = param.weight_loader
281
- weight_loader(param, loaded_weight, shard_id)
282
- break
283
- else:
284
- param = params_dict[name]
285
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
286
- weight_loader(param, loaded_weight)
287
- loaded_params.add(name)
288
- return loaded_params
265
+ loader = AutoWeightsLoader(self)
266
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
267
268
269
class Eagle3LlamaForCausalLM(LlamaForCausalLM):
vllm/model_executor/models/mimo_mtp.py
+32
-104
@@ -34,11 +34,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
34
ParallelLMHead,
35
VocabParallelEmbedding,
36
)
37
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37
from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer
38
from vllm.sequence import IntermediateTensors
39
41
-from .utils import maybe_prefix
40
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
41
42
43
class MiMoMultiTokenPredictorLayer(nn.Module):
@@ -163,6 +162,28 @@ class MiMoMTP(nn.Module):
162
self.config.hidden_size,
163
prefix=maybe_prefix(prefix, "lm_head"),
164
)
165
+ # Checkpoint stores MTP layers 0-indexed and without the `mtp_block`
166
+ # wrapper around the transformer block; remap onto the offset index.
167
+ start = self.config.num_hidden_layers
168
+ self.hf_to_vllm_mapper = WeightsMapper(
169
+ orig_to_new_substr={
170
+ ".self_attn.": ".mtp_block.self_attn.",
171
+ ".mlp.": ".mtp_block.mlp.",
172
+ ".input_layernorm.": ".mtp_block.input_layernorm.",
173
+ ".post_attention_layernorm.": ".mtp_block.post_attention_layernorm.",
174
+ },
175
+ orig_to_new_stacked={
176
+ ".q_proj": (".qkv_proj", "q"),
177
+ ".k_proj": (".qkv_proj", "k"),
178
+ ".v_proj": (".qkv_proj", "v"),
179
+ ".gate_proj": (".gate_up_proj", 0),
180
+ ".up_proj": (".gate_up_proj", 1),
181
+ },
182
+ orig_to_new_prefix={
183
+ f"model.mtp_layers.{i}.": f"model.mtp_layers.{i + start}."
184
+ for i in range(self.config.num_nextn_predict_layers)
185
+ },
186
+ )
187
188
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
189
return self.model.embed_input_ids(input_ids)
@@ -190,105 +211,12 @@ class MiMoMTP(nn.Module):
211
return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)
212
213
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
193
- stacked_params_mapping = [
194
- ("qkv_proj", "q_proj", "q"),
195
- ("qkv_proj", "k_proj", "k"),
196
- ("qkv_proj", "v_proj", "v"),
197
- ("gate_up_proj", "gate_proj", 0),
198
- ("gate_up_proj", "up_proj", 1),
199
- ]
200
-
201
- params_dict = dict(self.named_parameters())
202
- loaded_params: set[str] = set()
203
- for name, loaded_weight in weights:
204
- if "rotary_emb.inv_freq" in name:
205
- continue
206
- name = self.map_model_name_to_mtp_param_name(name)
207
-
208
- for param_name, weight_name, shard_id in stacked_params_mapping:
209
- # Skip non-stacked layers and experts (experts handled below).
210
- if weight_name not in name:
211
- continue
212
- if "mtp_layers" not in name:
213
- break
214
- # We have mlp.experts[0].gate_proj in the checkpoint.
215
- # Since we handle the experts below in expert_params_mapping,
216
- # we need to skip here BEFORE we update the name, otherwise
217
- # name will be updated to mlp.experts[0].gate_up_proj, which
218
- # will then be updated below in expert_params_mapping
219
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
220
- if ("mlp.experts." in name) and name not in params_dict:
221
- continue
222
- name = name.replace(weight_name, param_name)
223
- # Skip loading extra bias for GPTQ models.
224
- if name.endswith(".bias") and name not in params_dict:
225
- continue
226
-
227
- param = params_dict[name]
228
- weight_loader = param.weight_loader
229
- weight_loader(param, loaded_weight, shard_id)
230
- break
231
- else:
232
- # Skip loading extra bias for GPTQ models.
233
- if name.endswith(".bias") and name not in params_dict:
234
- continue
235
- if "mtp_layers" not in name and (
236
- "embed_tokens" not in name and "lm_head" not in name
237
- ):
238
- continue
239
- param = params_dict[name]
240
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
241
- weight_loader(param, loaded_weight)
242
- loaded_params.add(name)
243
- return loaded_params
244
-
245
- def map_model_name_to_mtp_param_name(self, name: str) -> str:
246
- import regex as re
247
-
248
- # append mtp_start_layer_idx
249
- pattern = r"(model\.mtp_layers\.)(\d+)(\.)"
250
- match = re.match(pattern, name)
251
- if match:
252
- original_num = int(match.group(2))
253
- new_num = original_num + self.config.num_hidden_layers
254
- name = name.replace(match.group(), f"{match.group(1)}{new_num}.")
255
- # check for early turn
256
- name_without_prefix = [
257
- "token_layernorm",
258
- "hidden_layernorm",
259
- "input_proj",
260
- "final_layernorm",
261
- ]
262
- for sub_name in name_without_prefix:
263
- if sub_name in name:
264
- return name
265
- # add mtp_block
266
- pattern = r"(model\.mtp_layers\.\d+\.)"
267
- match = re.match(pattern, name)
268
- if match:
269
- name = name.replace(match.group(), match.group() + "mtp_block.")
270
- return name
271
-
272
- def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
273
- """
274
- Rewrite the weight name to match the format of the original model.
275
- Add .mtp_block for modules in transformer layer block for spec layer
276
- """
277
- spec_layer_weight_names = [
278
- "embed_tokens",
279
- "enorm",
280
- "hnorm",
281
- "eh_proj",
282
- "shared_head",
283
- ]
284
- spec_layer_weight = False
285
- for weight_name in spec_layer_weight_names:
286
- if weight_name in name:
287
- spec_layer_weight = True
288
- break
289
- if not spec_layer_weight:
290
- # treat rest weights as weights for transformer layer block
291
- name = name.replace(
292
- f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
293
- )
294
- return name
214
+ # The checkpoint carries the full model; keep only the MTP layers and
215
+ # the shared embedding/head.
216
+ def mtp_weights():
217
+ for name, weight in weights:
218
+ if "mtp_layers" in name or "embed_tokens" in name or "lm_head" in name:
219
+ yield name, weight
220
+
221
+ loader = AutoWeightsLoader(self)
222
+ return loader.load_weights(mtp_weights(), mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/mimo_v2_omni.py
+9
-23
@@ -29,7 +29,6 @@ from vllm.model_executor.layers.linear import (
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.model_executor.layers.rotary_embedding import get_rope
31
from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
32
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
32
from vllm.model_executor.models.vision import is_vit_use_data_parallel
33
from vllm.multimodal import MULTIMODAL_REGISTRY
34
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
@@ -378,6 +377,13 @@ class MiMoVisionBlock(nn.Module):
377
378
379
class MiMoVisionTransformer(nn.Module):
380
+ hf_to_vllm_mapper = WeightsMapper(
381
+ orig_to_new_stacked={
382
+ "mlp.gate_proj": ("mlp.gate_up_proj", 0),
383
+ "mlp.up_proj": ("mlp.gate_up_proj", 1),
384
+ }
385
+ )
386
+
387
def __init__(
388
self,
389
vision_cfg: PretrainedConfig,
@@ -627,28 +633,8 @@ class MiMoVisionTransformer(nn.Module):
633
return x
634
635
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
630
- stacked_params_mapping = [
631
- ("mlp.gate_up_proj", "mlp.gate_proj", 0),
632
- ("mlp.gate_up_proj", "mlp.up_proj", 1),
633
- ]
634
- params_dict = dict(self.named_parameters(remove_duplicate=False))
635
- loaded_params: set[str] = set()
636
-
637
- for name, loaded_weight in weights:
638
- for param_name, weight_name, shard_id in stacked_params_mapping:
639
- if weight_name not in name:
640
- continue
641
- name = name.replace(weight_name, param_name)
642
- param = params_dict[name]
643
- weight_loader = param.weight_loader
644
- weight_loader(param, loaded_weight, shard_id)
645
- break
646
- else:
647
- param = params_dict[name]
648
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
649
- weight_loader(param, loaded_weight)
650
- loaded_params.add(name)
651
- return loaded_params
636
+ loader = AutoWeightsLoader(self)
637
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
638
639
640
class MiMoV2OmniProcessingInfo(BaseProcessingInfo):
vllm/model_executor/models/minicpmv4_6.py
+10
-26
@@ -25,7 +25,6 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateShapeCalculator,
26
)
27
from vllm.model_executor.layers.quantization import QuantizationConfig
28
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
from vllm.multimodal import MULTIMODAL_REGISTRY
29
from vllm.multimodal.inputs import (
30
MultiModalFeatureSpec,
@@ -619,6 +618,14 @@ class MiniCPMV4_6ProcessingInfo(MiniCPMVProcessingInfo):
618
619
620
class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module):
621
+ hf_to_vllm_mapper = WeightsMapper(
622
+ orig_to_new_stacked={
623
+ "q_proj": ("qkv_proj", "q"),
624
+ "k_proj": ("qkv_proj", "k"),
625
+ "v_proj": ("qkv_proj", "v"),
626
+ }
627
+ )
628
+
629
def __init__(
630
self,
631
config,
@@ -667,31 +674,8 @@ class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module):
674
return out
675
676
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
670
- stacked_params_mapping = [
671
- ("qkv_proj", "q_proj", "q"),
672
- ("qkv_proj", "k_proj", "k"),
673
- ("qkv_proj", "v_proj", "v"),
674
- ]
675
- params_dict = dict(self.named_parameters())
676
- loaded_params: set[str] = set()
677
- for name, loaded_weight in weights:
678
- for param_name, weight_name, shard_id in stacked_params_mapping:
679
- if weight_name not in name:
680
- continue
681
- mapped_name = name.replace(weight_name, param_name, 1)
682
- if mapped_name not in params_dict:
683
- continue
684
- param = params_dict[mapped_name]
685
- param.weight_loader(param, loaded_weight, shard_id)
686
- break
687
- else:
688
- if name not in params_dict:
689
- continue
690
- param = params_dict[name]
691
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
692
- weight_loader(param, loaded_weight)
693
- loaded_params.add(name)
694
- return loaded_params
677
+ loader = AutoWeightsLoader(self)
678
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
679
680
681
class MiniCPMV4_6ViTWindowAttentionMerger(nn.Module):
vllm/model_executor/models/minimax_m2.py
+23
-124
@@ -41,7 +41,6 @@ from vllm.distributed import (
41
from vllm.model_executor.layers.attention import Attention
42
from vllm.model_executor.layers.fused_moe import (
43
FusedMoE,
44
- fused_moe_make_expert_params_mapping,
44
)
45
from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
46
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -57,17 +56,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
56
ParallelLMHead,
57
VocabParallelEmbedding,
58
)
60
-from vllm.model_executor.model_loader.weight_utils import (
61
- default_weight_loader,
62
- maybe_remap_kv_scale_name,
63
-)
59
from vllm.sequence import IntermediateTensors
60
61
from .interfaces import EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP
62
from .utils import (
63
AutoWeightsLoader,
64
PPMissingLayer,
70
- is_pp_missing_parameter,
65
+ WeightsMapper,
66
make_empty_intermediate_tensors_factory,
67
make_layers,
68
maybe_prefix,
@@ -111,6 +106,7 @@ class MiniMaxM2MoE(nn.Module):
106
quant_config=quant_config,
107
prefix=f"{prefix}.experts",
108
router_logits_dtype=torch.float32,
109
+ ckpt_names=("w1", "w2", "w3"),
110
)
111
112
self.gate = GateLinear(
@@ -327,6 +323,15 @@ class MiniMaxM2DecoderLayer(nn.Module):
323
class MiniMaxM2Model(nn.Module, EagleModelMixin):
324
fall_back_to_pt_during_load = False
325
326
+ hf_to_vllm_mapper = WeightsMapper(
327
+ orig_to_new_stacked={
328
+ # weight_name: (param_name, shard_id)
329
+ ".q_proj": (".qkv_proj", "q"),
330
+ ".k_proj": (".qkv_proj", "k"),
331
+ ".v_proj": (".qkv_proj", "v"),
332
+ }
333
+ )
334
+
335
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
336
super().__init__()
337
@@ -409,115 +414,23 @@ class MiniMaxM2Model(nn.Module, EagleModelMixin):
414
415
return hidden_states
416
412
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
413
- return fused_moe_make_expert_params_mapping(
417
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
418
+ # Skip spec-decode (MTP) layers; they are appended after the main
419
+ # decoder layers and have no destination in the main model.
420
+ skip_prefixes = None
421
+ num_mtp = getattr(self.config, "num_mtp_modules", 0)
422
+ if num_mtp:
423
+ base = self.config.num_hidden_layers
424
+ skip_prefixes = [f"layers.{base + i}." for i in range(num_mtp)]
425
+ loader = AutoWeightsLoader(
426
self,
415
- ckpt_gate_proj_name="w1",
416
- ckpt_down_proj_name="w2",
417
- ckpt_up_proj_name="w3",
418
- num_experts=self.config.num_local_experts,
427
+ skip_prefixes=skip_prefixes,
428
)
420
-
421
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
422
- stacked_params_mapping = [
423
- # (param_name, shard_name, shard_id)
424
- ("qkv_proj", "q_proj", "q"),
425
- ("qkv_proj", "k_proj", "k"),
426
- ("qkv_proj", "v_proj", "v"),
427
- ]
428
-
429
- # Params for weights, fp8 weight scales, fp8 activation scales
430
- # (param_name, weight_name, expert_id, shard_id)
431
- expert_params_mapping = self.get_expert_mapping()
432
-
433
- params_dict = dict(self.named_parameters())
434
- loaded_params: set[str] = set()
435
- for name, loaded_weight in weights:
436
- if "rotary_emb.inv_freq" in name:
437
- continue
438
-
439
- spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
440
- if spec_layer is not None:
441
- continue # skip spec decode layers for main model
442
-
443
- for param_name, weight_name, shard_id in stacked_params_mapping:
444
- # Skip non-stacked layers and experts (experts handled below).
445
- if weight_name not in name:
446
- continue
447
- # We have mlp.experts[0].gate_proj in the checkpoint.
448
- # Since we handle the experts below in expert_params_mapping,
449
- # we need to skip here BEFORE we update the name, otherwise
450
- # name will be updated to mlp.experts[0].gate_up_proj, which
451
- # will then be updated below in expert_params_mapping
452
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
453
- if ("mlp.experts." in name) and name not in params_dict:
454
- continue
455
- name = name.replace(weight_name, param_name)
456
- # Skip loading extra bias for GPTQ models.
457
- if name.endswith(".bias") and name not in params_dict:
458
- continue
459
-
460
- if is_pp_missing_parameter(name, self):
461
- continue
462
-
463
- # Remap qkv_proj.[kv]_scale to attn.[kv]_scale
464
- if name.endswith((".k_scale", ".v_scale")):
465
- remapped_name = maybe_remap_kv_scale_name(name, params_dict)
466
- if remapped_name is not None and remapped_name in params_dict:
467
- param = params_dict[remapped_name]
468
- weight_loader = getattr(
469
- param, "weight_loader", default_weight_loader
470
- )
471
- weight_loader(param, loaded_weight)
472
- break
473
-
474
- param = params_dict[name]
475
- weight_loader = param.weight_loader
476
- weight_loader(param, loaded_weight, shard_id)
477
- break
478
- else:
479
- for mapping in expert_params_mapping:
480
- param_name, weight_name, expert_id, shard_id = mapping
481
- if weight_name not in name:
482
- continue
483
- name = name.replace(weight_name, param_name)
484
-
485
- if is_pp_missing_parameter(name, self):
486
- continue
487
-
488
- param = params_dict[name]
489
- weight_loader = param.weight_loader
490
- weight_loader(
491
- param,
492
- loaded_weight,
493
- name,
494
- shard_id=shard_id,
495
- expert_id=expert_id,
496
- )
497
- break
498
- else:
499
- # Skip loading extra bias for GPTQ models.
500
- if name.endswith(".bias") and name not in params_dict:
501
- continue
502
-
503
- # Remapping the name of FP8 kv-scale.
504
- name = maybe_remap_kv_scale_name(name, params_dict)
505
- if name is None:
506
- continue
507
-
508
- if is_pp_missing_parameter(name, self):
509
- continue
510
-
511
- param = params_dict[name]
512
- weight_loader = getattr(
513
- param, "weight_loader", default_weight_loader
514
- )
515
- weight_loader(param, loaded_weight)
516
- loaded_params.add(name)
517
- return loaded_params
429
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
430
431
432
class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
433
+ hf_to_vllm_mapper = MiniMaxM2Model.hf_to_vllm_mapper
434
packed_modules_mapping = {
435
"qkv_proj": [
436
"q_proj",
@@ -577,17 +490,3 @@ class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
490
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
491
loader = AutoWeightsLoader(self)
492
return loader.load_weights(weights)
580
-
581
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
582
- return self.model.get_expert_mapping()
583
-
584
-
585
-def get_spec_layer_idx_from_weight_name(
586
- config: PretrainedConfig, weight_name: str
587
-) -> int | None:
588
- if hasattr(config, "num_mtp_modules") and (config.num_mtp_modules > 0):
589
- layer_idx = config.num_hidden_layers
590
- for i in range(config.num_mtp_modules):
591
- if weight_name.startswith(f"model.layers.{layer_idx + i}."):
592
- return layer_idx + i
593
- return None
vllm/model_executor/models/mixtral.py
+15
-113
@@ -24,8 +24,7 @@
24
# limitations under the License.
25
"""Inference-only Mixtral model."""
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
@@ -42,7 +41,6 @@ from vllm.distributed import (
41
from vllm.model_executor.layers.attention import Attention
42
from vllm.model_executor.layers.fused_moe import (
43
FusedMoE,
45
- fused_moe_make_expert_params_mapping,
44
)
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
@@ -57,17 +55,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
55
ParallelLMHead,
56
VocabParallelEmbedding,
57
)
60
-from vllm.model_executor.model_loader.weight_utils import (
61
- default_weight_loader,
62
- maybe_remap_kv_scale_name,
63
-)
58
from vllm.sequence import IntermediateTensors
59
60
from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
61
from .utils import (
62
AutoWeightsLoader,
63
PPMissingLayer,
70
- is_pp_missing_parameter,
64
+ WeightsMapper,
65
make_empty_intermediate_tensors_factory,
66
make_layers,
67
maybe_prefix,
@@ -142,6 +136,7 @@ class MixtralMoE(nn.Module):
136
prefix=f"{prefix}.experts",
137
enable_eplb=self.enable_eplb,
138
num_redundant_experts=self.n_redundant_experts,
139
+ ckpt_names=("w1", "w2", "w3"),
140
)
141
142
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -295,6 +290,15 @@ class MixtralDecoderLayer(nn.Module):
290
291
@support_torch_compile
292
class MixtralModel(nn.Module):
293
+ hf_to_vllm_mapper = WeightsMapper(
294
+ orig_to_new_stacked={
295
+ # weight_name: (param_name, shard_id)
296
+ ".q_proj": (".qkv_proj", "q"),
297
+ ".k_proj": (".qkv_proj", "k"),
298
+ ".v_proj": (".qkv_proj", "v"),
299
+ }
300
+ )
301
+
302
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
303
super().__init__()
304
@@ -364,114 +368,15 @@ class MixtralModel(nn.Module):
368
hidden_states, _ = self.norm(hidden_states, residual)
369
return hidden_states
370
367
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
368
- # Params for weights, fp8 weight scales, fp8 activation scales
369
- # (param_name, weight_name, expert_id, shard_id)
370
- return fused_moe_make_expert_params_mapping(
371
- self,
372
- ckpt_gate_proj_name="w1",
373
- ckpt_down_proj_name="w2",
374
- ckpt_up_proj_name="w3",
375
- num_experts=self.config.num_local_experts,
376
- num_redundant_experts=self.num_redundant_experts,
377
- )
378
-
371
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
380
- stacked_params_mapping = [
381
- # (param_name, shard_name, shard_id)
382
- ("qkv_proj", "q_proj", "q"),
383
- ("qkv_proj", "k_proj", "k"),
384
- ("qkv_proj", "v_proj", "v"),
385
- ]
386
-
387
- params_dict = dict(self.named_parameters())
388
- loaded_params: set[str] = set()
389
- expert_params_mapping = self.get_expert_mapping()
390
- for name, loaded_weight in weights:
391
- for param_name, weight_name, shard_id in stacked_params_mapping:
392
- if weight_name not in name:
393
- continue
394
- name = name.replace(weight_name, param_name)
395
- # Skip loading extra bias for GPTQ models.
396
- if (
397
- name.endswith(".bias") or name.endswith("_bias")
398
- ) and name not in params_dict:
399
- continue
400
- # Skip layers on other devices.
401
- if is_pp_missing_parameter(name, self):
402
- continue
403
- if name.endswith("scale"):
404
- # Remapping the name of FP8 kv-scale.
405
- name = maybe_remap_kv_scale_name(name, params_dict)
406
- if name is None:
407
- continue
408
- param = params_dict[name]
409
- weight_loader = param.weight_loader
410
- weight_loader(param, loaded_weight, shard_id)
411
- break
412
- else:
413
- is_expert_weight = False
414
- for mapping in expert_params_mapping:
415
- param_name, weight_name, expert_id, shard_id = mapping
416
-
417
- if weight_name not in name:
418
- continue
419
-
420
- is_expert_weight = True
421
- name_mapped = name.replace(weight_name, param_name)
422
-
423
- # Skip layers on other devices.
424
- if is_pp_missing_parameter(name_mapped, self):
425
- continue
426
-
427
- if (
428
- name_mapped.endswith(".bias") or name_mapped.endswith("_bias")
429
- ) and name_mapped not in params_dict:
430
- continue
431
-
432
- param = params_dict[name_mapped]
433
- weight_loader = typing.cast(
434
- Callable[..., bool], param.weight_loader
435
- )
436
- success = weight_loader(
437
- param,
438
- loaded_weight,
439
- name_mapped,
440
- shard_id=shard_id,
441
- expert_id=expert_id,
442
- return_success=True,
443
- )
444
- if success:
445
- name = name_mapped
446
- break
447
- else:
448
- if is_expert_weight:
449
- continue
450
- # Skip loading extra bias for GPTQ models.
451
- if (
452
- name.endswith(".bias") or name.endswith("_bias")
453
- ) and name not in params_dict:
454
- continue
455
- # Skip layers on other devices.
456
- if is_pp_missing_parameter(name, self):
457
- continue
458
- # Remapping the name of FP8 kv-scale.
459
- name = maybe_remap_kv_scale_name(name, params_dict)
460
- if name is None:
461
- continue
462
-
463
- param = params_dict[name]
464
- weight_loader = getattr(
465
- param, "weight_loader", default_weight_loader
466
- )
467
- weight_loader(param, loaded_weight)
468
- loaded_params.add(name)
469
- return loaded_params
372
+ loader = AutoWeightsLoader(self)
373
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
374
375
376
class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
377
fall_back_to_pt_during_load = False
378
379
+ hf_to_vllm_mapper = MixtralModel.hf_to_vllm_mapper
380
packed_modules_mapping = {
381
"qkv_proj": [
382
"q_proj",
@@ -582,6 +487,3 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
487
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
488
loader = AutoWeightsLoader(self)
489
return loader.load_weights(weights)
585
-
586
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
587
- return self.model.get_expert_mapping()
vllm/model_executor/models/molmo.py
+10
-32
@@ -672,6 +672,14 @@ class MolmoDecoderNormAfterLayer(MolmoDecoderLayer):
672
class MolmoVisionBackbone(nn.Module, SupportsQuant):
673
packed_modules_mapping = {"merged_linear": ["gate_proj", "up_proj"]}
674
675
+ hf_to_vllm_mapper = WeightsMapper(
676
+ orig_to_new_stacked={
677
+ # image_projector gate_up merge
678
+ "gate_proj": ("merged_linear", 0),
679
+ "up_proj": ("merged_linear", 1),
680
+ }
681
+ )
682
+
683
def __init__(
684
self,
685
config: PretrainedConfig,
@@ -800,38 +808,8 @@ class MolmoVisionBackbone(nn.Module, SupportsQuant):
808
return image_features
809
810
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
803
- stacked_params_mapping = [
804
- # (param_name, shard_name, shard_id)
805
- ("merged_linear", "gate_proj", 0),
806
- ("merged_linear", "up_proj", 1),
807
- ]
808
- params_dict = dict(self.named_parameters())
809
- loaded_params: set[str] = set()
810
-
811
- for name, loaded_weight in weights:
812
- for param_name, weight_name, shard_id in stacked_params_mapping:
813
- if weight_name not in name:
814
- continue
815
- name = name.replace(weight_name, param_name)
816
- # Skip loading extra bias for GPTQ models.
817
- if name.endswith(".bias") and name not in params_dict:
818
- continue
819
- if is_pp_missing_parameter(name, self):
820
- continue
821
- param = params_dict[name]
822
- weight_loader = param.weight_loader
823
- weight_loader(param, loaded_weight, shard_id)
824
- break
825
- else:
826
- if name.endswith(".bias") and name not in params_dict:
827
- continue
828
- if is_pp_missing_parameter(name, self):
829
- continue
830
- param = params_dict[name]
831
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
832
- weight_loader(param, loaded_weight)
833
- loaded_params.add(name)
834
- return loaded_params
811
+ loader = AutoWeightsLoader(self)
812
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
813
814
815
@support_torch_compile
vllm/model_executor/models/molmo2.py
+16
-37
@@ -709,6 +709,20 @@ class Molmo2VisionBackbone(nn.Module, SupportsQuant):
709
"merged_linear": ["gate_proj", "up_proj"],
710
}
711
712
+ # Runs after the top-level mapper, so image_pooling_2d/image_projector
713
+ # source names are already renamed to q/k/v_proj and gate/up_proj.
714
+ hf_to_vllm_mapper = WeightsMapper(
715
+ orig_to_new_stacked={
716
+ "wq": ("merged_qkv", "q"),
717
+ "wk": ("merged_qkv", "k"),
718
+ "wv": ("merged_qkv", "v"),
719
+ "k_proj": ("merged_kv", 0),
720
+ "v_proj": ("merged_kv", 1),
721
+ "gate_proj": ("merged_linear", 0),
722
+ "up_proj": ("merged_linear", 1),
723
+ },
724
+ )
725
+
726
def __init__(
727
self,
728
vit_config: VitConfig,
@@ -839,43 +853,8 @@ class Molmo2VisionBackbone(nn.Module, SupportsQuant):
853
]
854
855
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
842
- stacked_params_mapping = [
843
- # (param_name, shard_name, shard_id)
844
- ("merged_qkv", "wq", "q"),
845
- ("merged_qkv", "wk", "k"),
846
- ("merged_qkv", "wv", "v"),
847
- ("merged_kv", "k_proj", 0),
848
- ("merged_kv", "v_proj", 1),
849
- ("merged_linear", "gate_proj", 0),
850
- ("merged_linear", "up_proj", 1),
851
- ]
852
- params_dict = dict(self.named_parameters())
853
- loaded_params: set[str] = set()
854
-
855
- for name, loaded_weight in weights:
856
- for param_name, weight_name, shard_id in stacked_params_mapping:
857
- if weight_name not in name:
858
- continue
859
- name = name.replace(weight_name, param_name)
860
- # Skip loading extra bias for GPTQ models.
861
- if name.endswith(".bias") and name not in params_dict:
862
- continue
863
- if is_pp_missing_parameter(name, self):
864
- continue
865
- param = params_dict[name]
866
- weight_loader = param.weight_loader
867
- weight_loader(param, loaded_weight, shard_id)
868
- break
869
- else:
870
- if name.endswith(".bias") and name not in params_dict:
871
- continue
872
- if is_pp_missing_parameter(name, self):
873
- continue
874
- param = params_dict[name]
875
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
876
- weight_loader(param, loaded_weight)
877
- loaded_params.add(name)
878
- return loaded_params
856
+ loader = AutoWeightsLoader(self)
857
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
858
859
860
class Molmo2Attention(nn.Module):
vllm/model_executor/models/moss_audio.py
+5
-33
@@ -33,7 +33,6 @@ from vllm.model_executor.layers.linear import (
33
from vllm.model_executor.layers.logits_processor import LogitsProcessor
34
from vllm.model_executor.layers.quantization import QuantizationConfig
35
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
36
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
from vllm.multimodal import MULTIMODAL_REGISTRY
37
from vllm.multimodal.inputs import (
38
AudioItem,
@@ -742,37 +741,6 @@ class GatedMLP(nn.Module):
741
x, _ = self.down_proj(x)
742
return x
743
745
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
746
- stacked_params_mapping = [
747
- ("gate_up_proj", "gate_proj", 0),
748
- ("gate_up_proj", "up_proj", 1),
749
- ]
750
- params_dict = dict(self.named_parameters())
751
- loaded_params: set[str] = set()
752
-
753
- for name, loaded_weight in weights:
754
- target_name = name
755
- for param_name, weight_name, shard_id in stacked_params_mapping:
756
- components = target_name.split(".")
757
- if weight_name not in components:
758
- continue
759
-
760
- target_name = ".".join(
761
- param_name if component == weight_name else component
762
- for component in components
763
- )
764
- param = params_dict[target_name]
765
- weight_loader = param.weight_loader
766
- weight_loader(param, loaded_weight, shard_id)
767
- break
768
- else:
769
- param = params_dict[target_name]
770
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
771
- weight_loader(param, loaded_weight)
772
-
773
- loaded_params.add(target_name)
774
- return loaded_params
775
-
744
745
@support_torch_compile(
746
dynamic_arg_dims={
@@ -1478,7 +1446,11 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
1446
"language_model.embed_tokens.": "language_model.model.embed_tokens.",
1447
"language_model.layers.": "language_model.model.layers.",
1448
"language_model.norm.": "language_model.model.norm.",
1481
- }
1449
+ },
1450
+ orig_to_new_stacked={
1451
+ ".gate_proj": (".gate_up_proj", 0),
1452
+ ".up_proj": (".gate_up_proj", 1),
1453
+ },
1454
)
1455
1456
def get_mm_mapping(self) -> MultiModelKeys:
vllm/model_executor/models/olmoe.py
+10
-105
@@ -34,7 +34,6 @@ from vllm.logger import init_logger
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,
37
)
38
from vllm.model_executor.layers.layernorm import RMSNorm
39
from vllm.model_executor.layers.linear import (
@@ -49,13 +48,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
48
ParallelLMHead,
49
VocabParallelEmbedding,
50
)
52
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51
from vllm.sequence import IntermediateTensors
52
53
from .interfaces import SupportsLoRA, SupportsPP
54
from .utils import (
55
AutoWeightsLoader,
58
- is_pp_missing_parameter,
56
+ WeightsMapper,
57
make_empty_intermediate_tensors_factory,
58
make_layers,
59
maybe_prefix,
@@ -338,106 +336,16 @@ class OlmoeModel(nn.Module):
336
hidden_states = self.norm(hidden_states)
337
return hidden_states
338
341
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
342
- # Params for weights, fp8 weight scales, fp8 activation scales
343
- # (param_name, weight_name, expert_id, shard_id)
344
- return fused_moe_make_expert_params_mapping(
345
- self,
346
- ckpt_gate_proj_name="gate_proj",
347
- ckpt_down_proj_name="down_proj",
348
- ckpt_up_proj_name="up_proj",
349
- num_experts=self.config.num_experts,
350
- )
351
-
352
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
353
- stacked_params_mapping = [
354
- # (param_name, shard_name, shard_id)
355
- ("qkv_proj", "q_proj", "q"),
356
- ("qkv_proj", "k_proj", "k"),
357
- ("qkv_proj", "v_proj", "v"),
358
- ]
359
-
360
- params_dict = dict(self.named_parameters())
361
- loaded_params: set[str] = set()
362
- expert_params_mapping = self.get_expert_mapping()
363
- for name, loaded_weight in weights:
364
- for param_name, weight_name, shard_id in stacked_params_mapping:
365
- # Skip non-stacked layers and experts (experts handled below).
366
- if weight_name not in name:
367
- continue
368
- # We have mlp.experts[0].gate_proj in the checkpoint.
369
- # Since we handle the experts below in expert_params_mapping,
370
- # we need to skip here BEFORE we update the name, otherwise
371
- # name will be updated to mlp.experts[0].gate_up_proj, which
372
- # will then be updated below in expert_params_mapping
373
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
374
- if "mlp.experts" in name:
375
- continue
376
- name = name.replace(weight_name, param_name)
377
- # Skip loading extra bias for GPTQ models.
378
- if name.endswith(".bias") and name not in params_dict:
379
- continue
380
- # Skip layers on other devices.
381
- if is_pp_missing_parameter(name, self):
382
- continue
383
- if name not in params_dict:
384
- continue
385
-
386
- param = params_dict[name]
387
- weight_loader = param.weight_loader
388
- weight_loader(param, loaded_weight, shard_id)
389
- break
390
- else:
391
- for mapping in expert_params_mapping:
392
- param_name, weight_name, expert_id, shard_id = mapping
393
- if weight_name not in name:
394
- continue
395
- name = name.replace(weight_name, param_name)
396
- # Skip layers on other devices.
397
- if is_pp_missing_parameter(name, self):
398
- continue
399
- param = params_dict[name]
400
- weight_loader = param.weight_loader
401
- weight_loader(
402
- param,
403
- loaded_weight,
404
- name,
405
- shard_id=shard_id,
406
- expert_id=expert_id,
407
- )
408
- break
409
- else:
410
- # Skip loading extra bias for GPTQ models.
411
- if name.endswith(".bias") and name not in params_dict:
412
- continue
413
- # Skip layers on other devices.
414
- if is_pp_missing_parameter(name, self):
415
- continue
416
- # Remapping the name of FP8 kv-scale.
417
- if name.endswith("kv_scale"):
418
- remapped_kv_scale_name = name.replace(
419
- ".kv_scale", ".attn.kv_scale"
420
- )
421
- if remapped_kv_scale_name not in params_dict:
422
- logger.warning_once(
423
- "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501
424
- name,
425
- remapped_kv_scale_name,
426
- )
427
- continue
428
- else:
429
- name = remapped_kv_scale_name
430
-
431
- param = params_dict[name]
432
- weight_loader = getattr(
433
- param, "weight_loader", default_weight_loader
434
- )
435
- weight_loader(param, loaded_weight)
436
- loaded_params.add(name)
437
- return loaded_params
438
-
339
340
class OlmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
341
+ hf_to_vllm_mapper = WeightsMapper(
342
+ orig_to_new_stacked={
343
+ # weight_name: (param_name, shard_id)
344
+ ".q_proj": (".qkv_proj", "q"),
345
+ ".k_proj": (".qkv_proj", "k"),
346
+ ".v_proj": (".qkv_proj", "v"),
347
+ }
348
+ )
349
packed_modules_mapping = {
350
"qkv_proj": [
351
"q_proj",
@@ -496,7 +404,4 @@ class OlmoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
404
405
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
406
loader = AutoWeightsLoader(self)
499
- return loader.load_weights(weights)
500
-
501
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
502
- return self.model.get_expert_mapping()
407
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/paddleocr_vl.py
+21
-54
@@ -48,10 +48,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
48
from vllm.model_executor.layers.rotary_embedding.common import (
49
ApplyRotaryEmb,
50
)
51
-from vllm.model_executor.model_loader.weight_utils import (
52
- default_weight_loader,
53
- maybe_remap_kv_scale_name,
54
-)
51
from vllm.multimodal import MULTIMODAL_REGISTRY
52
from vllm.multimodal.inputs import (
53
MultiModalFeatureSpec,
@@ -81,7 +77,6 @@ from .utils import (
77
AutoWeightsLoader,
78
PPMissingLayer,
79
WeightsMapper,
84
- is_pp_missing_parameter,
80
maybe_prefix,
81
)
82
from .vision import get_vit_attn_backend
@@ -884,6 +879,14 @@ class SiglipVisionTransformer(nn.Module):
879
880
881
class SiglipVisionModel(nn.Module):
882
+ hf_to_vllm_mapper = WeightsMapper(
883
+ orig_to_new_stacked={
884
+ ".q_proj": (".qkv_proj", "q"),
885
+ ".k_proj": (".qkv_proj", "k"),
886
+ ".v_proj": (".qkv_proj", "v"),
887
+ }
888
+ )
889
+
890
def __init__(
891
self,
892
config,
@@ -928,55 +931,19 @@ class SiglipVisionModel(nn.Module):
931
)
932
933
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
931
- stacked_params_mapping = [
932
- ("qkv_proj", "q_proj", "q"),
933
- ("qkv_proj", "k_proj", "k"),
934
- ("qkv_proj", "v_proj", "v"),
935
- ]
936
- params_dict = dict(self.named_parameters(remove_duplicate=False))
937
- loaded_params: set[str] = set()
938
- for name, loaded_weight in weights:
939
- if "rotary_emb.inv_freq" in name:
940
- continue
941
- if "head.attention" in name or "head.layernorm" in name:
942
- continue
943
- if "head.mlp" in name or "head.probe" in name:
944
- continue
945
- if "packing_position_embedding" in name:
946
- continue
947
- for (
948
- param_name,
949
- weight_name,
950
- shard_id,
951
- ) in stacked_params_mapping:
952
- if weight_name not in name:
953
- continue
954
- name = name.replace(weight_name, param_name)
955
- if name.endswith(".bias") and name not in params_dict:
956
- continue
957
- if is_pp_missing_parameter(name, self):
958
- continue
959
- param = params_dict[name]
960
- weight_loader = param.weight_loader
961
- weight_loader(param, loaded_weight, shard_id)
962
- break
963
- else:
964
- if name.endswith(".bias") and name not in params_dict:
965
- continue
966
- name = maybe_remap_kv_scale_name(name, params_dict)
967
- if name is None:
968
- continue
969
- if is_pp_missing_parameter(name, self):
970
- continue
971
- param = params_dict[name]
972
- weight_loader = getattr(
973
- param,
974
- "weight_loader",
975
- default_weight_loader,
976
- )
977
- weight_loader(param, loaded_weight)
978
- loaded_params.add(name)
979
- return loaded_params
934
+ # Skip the SigLIP attention pooling head and packing pos embedding
935
+ # present in the checkpoint but absent from this vision tower.
936
+ loader = AutoWeightsLoader(
937
+ self,
938
+ skip_substrs=[
939
+ "head.attention",
940
+ "head.layernorm",
941
+ "head.mlp",
942
+ "head.probe",
943
+ "packing_position_embedding",
944
+ ],
945
+ )
946
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
947
948
949
@MULTIMODAL_REGISTRY.register_processor(
vllm/model_executor/models/phimoe.py
+11
-84
@@ -37,7 +37,6 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.fused_moe import (
39
FusedMoE,
40
- fused_moe_make_expert_params_mapping,
40
)
41
from vllm.model_executor.layers.linear import (
42
QKVParallelLinear,
@@ -51,16 +50,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
50
ParallelLMHead,
51
VocabParallelEmbedding,
52
)
54
-from vllm.model_executor.model_loader.weight_utils import (
55
- default_weight_loader,
56
- maybe_remap_kv_scale_name,
57
-)
53
from vllm.sequence import IntermediateTensors
54
55
from .interfaces import SupportsLoRA, SupportsPP
56
from .utils import (
57
AutoWeightsLoader,
63
- is_pp_missing_parameter,
58
+ WeightsMapper,
59
make_empty_intermediate_tensors_factory,
60
make_layers,
61
maybe_prefix,
@@ -289,6 +284,7 @@ class PhiMoE(nn.Module):
284
tp_size=tp_size,
285
custom_routing_function=phimoe_routing_function,
286
prefix=f"{prefix}.experts",
287
+ ckpt_names=("w1", "w2", "w3"),
288
)
289
290
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -516,84 +512,18 @@ class PhiMoEModel(nn.Module):
512
hidden_states = self.norm(hidden_states)
513
return hidden_states
514
519
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
520
- return fused_moe_make_expert_params_mapping(
521
- self,
522
- ckpt_gate_proj_name="w1",
523
- ckpt_down_proj_name="w2",
524
- ckpt_up_proj_name="w3",
525
- num_experts=self.config.num_local_experts,
526
- )
527
-
528
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
529
- stacked_params_mapping = [
530
- # (param_name, shard_name, shard_id)
531
- ("qkv_proj", "q_proj", "q"),
532
- ("qkv_proj", "k_proj", "k"),
533
- ("qkv_proj", "v_proj", "v"),
534
- ]
535
-
536
- params_dict = dict(self.named_parameters())
537
- loaded_params: set[str] = set()
538
- expert_params_mapping = self.get_expert_mapping()
539
- for name, loaded_weight in weights:
540
- for param_name, weight_name, shard_id in stacked_params_mapping:
541
- if weight_name not in name:
542
- continue
543
- name = name.replace(weight_name, param_name)
544
- # Skip loading extra bias for GPTQ models.
545
- if name.endswith(".bias") and name not in params_dict:
546
- continue
547
- # Skip layers on other devices.
548
- if is_pp_missing_parameter(name, self):
549
- continue
550
- param = params_dict[name]
551
- weight_loader = param.weight_loader
552
- weight_loader(param, loaded_weight, shard_id)
553
- break
554
- else:
555
- for mapping in expert_params_mapping:
556
- param_name, weight_name, expert_id, shard_id = mapping
557
- if weight_name not in name:
558
- continue
559
- name = name.replace(weight_name, param_name)
560
- # Skip layers on other devices.
561
- if is_pp_missing_parameter(name, self):
562
- continue
563
- param = params_dict[name]
564
- weight_loader = param.weight_loader
565
- weight_loader(
566
- param,
567
- loaded_weight,
568
- name,
569
- shard_id=shard_id,
570
- expert_id=expert_id,
571
- )
572
- break
573
- else:
574
- # Skip loading extra bias for GPTQ models.
575
- if name.endswith(".bias") and name not in params_dict:
576
- continue
577
- # Skip layers on other devices.
578
- if is_pp_missing_parameter(name, self):
579
- continue
580
- # Remapping the name of FP8 kv-scale.
581
- name = maybe_remap_kv_scale_name(name, params_dict)
582
- if name is None:
583
- continue
584
-
585
- param = params_dict[name]
586
- weight_loader = getattr(
587
- param, "weight_loader", default_weight_loader
588
- )
589
- weight_loader(param, loaded_weight)
590
- loaded_params.add(name)
591
- return loaded_params
592
-
515
516
class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
517
fall_back_to_pt_during_load = False
518
519
+ hf_to_vllm_mapper = WeightsMapper(
520
+ orig_to_new_stacked={
521
+ # weight_name: (param_name, shard_id)
522
+ ".q_proj": (".qkv_proj", "q"),
523
+ ".k_proj": (".qkv_proj", "k"),
524
+ ".v_proj": (".qkv_proj", "v"),
525
+ }
526
+ )
527
packed_modules_mapping = {
528
"qkv_proj": [
529
"q_proj",
@@ -654,7 +584,4 @@ class PhiMoEForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
584
585
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
586
loader = AutoWeightsLoader(self)
657
- return loader.load_weights(weights)
658
-
659
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
660
- return self.model.get_expert_mapping()
587
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/qwen2_5_vl.py
+12
-27
@@ -65,7 +65,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
65
from vllm.model_executor.layers.rotary_embedding.common import (
66
ApplyRotaryEmb,
67
)
68
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
68
from vllm.model_executor.models.module_mapping import MultiModelKeys
69
from vllm.multimodal import MULTIMODAL_REGISTRY
70
from vllm.multimodal.evs import (
@@ -613,6 +612,16 @@ class Qwen2_5_VisionPatchMerger(nn.Module):
612
613
614
class Qwen2_5_VisionTransformer(nn.Module):
615
+ hf_to_vllm_mapper = WeightsMapper(
616
+ orig_to_new_stacked={
617
+ ".attn.q.": (".attn.qkv.", "q"),
618
+ ".attn.k.": (".attn.qkv.", "k"),
619
+ ".attn.v.": (".attn.qkv.", "v"),
620
+ ".mlp.gate_proj.": (".mlp.gate_up_proj.", 0),
621
+ ".mlp.up_proj.": (".mlp.gate_up_proj.", 1),
622
+ }
623
+ )
624
+
625
def __init__(
626
self,
627
vision_config: Qwen2_5_VLVisionConfig,
@@ -1116,32 +1125,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
1125
return hidden_states
1126
1127
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1119
- stacked_params_mapping = [
1120
- # (param_name, shard_name, shard_id)
1121
- ("attn.qkv.", "attn.q.", "q"),
1122
- ("attn.qkv.", "attn.k.", "k"),
1123
- ("attn.qkv.", "attn.v.", "v"),
1124
- ("mlp.gate_up_proj.", "mlp.gate_proj.", 0),
1125
- ("mlp.gate_up_proj.", "mlp.up_proj.", 1),
1126
- ]
1127
- params_dict = dict(self.named_parameters(remove_duplicate=False))
1128
- loaded_params: set[str] = set()
1129
-
1130
- for name, loaded_weight in weights:
1131
- for param_name, weight_name, shard_id in stacked_params_mapping:
1132
- if weight_name not in name:
1133
- continue
1134
- name = name.replace(weight_name, param_name)
1135
- param = params_dict[name]
1136
- weight_loader = param.weight_loader
1137
- weight_loader(param, loaded_weight, shard_id)
1138
- break
1139
- else:
1140
- param = params_dict[name]
1141
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1142
- weight_loader(param, loaded_weight)
1143
- loaded_params.add(name)
1144
- return loaded_params
1128
+ loader = AutoWeightsLoader(self)
1129
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1130
1131
1132
class Qwen2_5_VLProcessingInfo(Qwen2VLProcessingInfo):
vllm/model_executor/models/qwen2_moe.py
+26
-122
@@ -42,7 +42,6 @@ from vllm.model_executor.layers.activation import SiluAndMul
42
from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.fused_moe import (
44
FusedMoE,
45
- fused_moe_make_expert_params_mapping,
45
)
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
@@ -58,14 +57,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
57
ParallelLMHead,
58
VocabParallelEmbedding,
59
)
61
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
60
from vllm.sequence import IntermediateTensors
61
62
from .interfaces import SupportsLoRA, SupportsPP
63
from .utils import (
64
AutoWeightsLoader,
65
+ WeightsMapper,
66
extract_layer_index,
68
- is_pp_missing_parameter,
67
make_empty_intermediate_tensors_factory,
68
make_layers,
69
maybe_prefix,
@@ -418,124 +416,22 @@ class Qwen2MoeModel(nn.Module):
416
hidden_states, _ = self.norm(hidden_states, residual)
417
return hidden_states
418
421
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
422
- # Params for weights, fp8 weight scales, fp8 activation scales
423
- # (param_name, weight_name, expert_id, shard_id)
424
- return fused_moe_make_expert_params_mapping(
425
- self,
426
- ckpt_gate_proj_name="gate_proj",
427
- ckpt_down_proj_name="down_proj",
428
- ckpt_up_proj_name="up_proj",
429
- num_experts=self.config.num_experts,
430
- )
431
-
432
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
433
- stacked_params_mapping = [
434
- # (param_name, shard_name, shard_id)
435
- ("qkv_proj", "q_proj", "q"),
436
- ("qkv_proj", "k_proj", "k"),
437
- ("qkv_proj", "v_proj", "v"),
438
- ("gate_up_proj", "gate_proj", 0),
439
- ("gate_up_proj", "up_proj", 1),
440
- ]
441
-
442
- params_dict = dict(self.named_parameters())
443
- loaded_params: set[str] = set()
444
- expert_params_mapping = self.get_expert_mapping()
445
- for name, loaded_weight in weights:
446
- for param_name, weight_name, shard_id in stacked_params_mapping:
447
- # Skip non-stacked layers and experts (experts handled below).
448
- if weight_name not in name:
449
- continue
450
- # We have mlp.experts[0].gate_proj in the checkpoint.
451
- # Since we handle the experts below in expert_params_mapping,
452
- # we need to skip here BEFORE we update the name, otherwise
453
- # name will be updated to mlp.experts[0].gate_up_proj, which
454
- # will then be updated below in expert_params_mapping
455
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
456
- if "mlp.experts" in name:
457
- continue
458
- name = name.replace(weight_name, param_name)
459
- # Skip loading extra bias for GPTQ models.
460
- if (
461
- name.endswith(".bias") or name.endswith("_bias")
462
- ) and name not in params_dict:
463
- continue
464
- # Skip layers on other devices.
465
- if is_pp_missing_parameter(name, self):
466
- continue
467
- if name not in params_dict:
468
- continue
469
-
470
- param = params_dict[name]
471
- weight_loader = param.weight_loader
472
- weight_loader(param, loaded_weight, shard_id)
473
- break
474
- else:
475
- for mapping in expert_params_mapping:
476
- param_name, weight_name, expert_id, shard_id = mapping
477
- if weight_name not in name:
478
- continue
479
- name = name.replace(weight_name, param_name)
480
-
481
- # Skip layers on other devices.
482
- if is_pp_missing_parameter(name, self):
483
- continue
484
- # Skip loading extra bias for GPTQ models.
485
- if (
486
- name.endswith(".bias") or name.endswith("_bias")
487
- ) and name not in params_dict:
488
- continue
489
- param = params_dict[name]
490
- weight_loader = param.weight_loader
491
- weight_loader(
492
- param,
493
- loaded_weight,
494
- name,
495
- shard_id=shard_id,
496
- expert_id=expert_id,
497
- )
498
- break
499
- else:
500
- # Skip loading extra bias for GPTQ models.
501
- if (
502
- name.endswith(".bias") or name.endswith("_bias")
503
- ) and name not in params_dict:
504
- continue
505
- # Skip layers on other devices.
506
- if is_pp_missing_parameter(name, self):
507
- continue
508
- # Remapping the name of FP8 kv-scale.
509
- if name.endswith("kv_scale"):
510
- remapped_kv_scale_name = name.replace(
511
- ".kv_scale", ".attn.kv_scale"
512
- )
513
- if remapped_kv_scale_name not in params_dict:
514
- logger.warning_once(
515
- "Found kv_scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv_scale is not loaded.", # noqa: E501
516
- name,
517
- remapped_kv_scale_name,
518
- )
519
- continue
520
- else:
521
- name = remapped_kv_scale_name
522
- # GGUF: make sure that shared_expert_gate is a 2D tensor.
523
- if (
524
- "mlp.shared_expert_gate" in name
525
- and len(loaded_weight.shape) == 1
526
- ):
527
- loaded_weight = loaded_weight[None, :]
528
- param = params_dict[name]
529
- weight_loader = getattr(
530
- param, "weight_loader", default_weight_loader
531
- )
532
- weight_loader(param, loaded_weight)
533
- loaded_params.add(name)
534
- return loaded_params
535
-
419
420
class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
421
fall_back_to_pt_during_load = False
422
+ hf_to_vllm_mapper = WeightsMapper(
423
+ orig_to_new_stacked={
424
+ # weight_name: (param_name, shard_id)
425
+ ".q_proj": (".qkv_proj", "q"),
426
+ ".k_proj": (".qkv_proj", "k"),
427
+ ".v_proj": (".qkv_proj", "v"),
428
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
429
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
430
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
431
+ ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0),
432
+ ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1),
433
+ }
434
+ )
435
packed_modules_mapping = {
436
"qkv_proj": [
437
"q_proj",
@@ -596,8 +492,16 @@ class Qwen2MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
492
return logits
493
494
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
599
- loader = AutoWeightsLoader(self)
600
- return loader.load_weights(weights)
495
+ def _maybe_reshape(
496
+ weights: Iterable[tuple[str, torch.Tensor]],
497
+ ) -> Iterable[tuple[str, torch.Tensor]]:
498
+ for name, loaded_weight in weights:
499
+ # GGUF: make sure that shared_expert_gate is a 2D tensor.
500
+ if "mlp.shared_expert_gate" in name and loaded_weight.dim() == 1:
501
+ loaded_weight = loaded_weight[None, :]
502
+ yield name, loaded_weight
503
602
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
603
- return self.model.get_expert_mapping()
504
+ loader = AutoWeightsLoader(self)
505
+ return loader.load_weights(
506
+ _maybe_reshape(weights), mapper=self.hf_to_vllm_mapper
507
+ )
vllm/model_executor/models/qwen2_vl.py
+11
-26
@@ -60,7 +60,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
60
from vllm.model_executor.layers.rotary_embedding.common import (
61
ApplyRotaryEmb,
62
)
63
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
63
from vllm.model_executor.models.module_mapping import MultiModelKeys
64
from vllm.multimodal import MULTIMODAL_REGISTRY
65
from vllm.multimodal.inputs import (
@@ -522,6 +521,15 @@ class Qwen2VisionPatchMerger(nn.Module):
521
522
523
class Qwen2VisionTransformer(nn.Module):
524
+ # Vision checkpoints store qkv pre-fused; merge separate q/k/v into qkv.
525
+ hf_to_vllm_mapper = WeightsMapper(
526
+ orig_to_new_stacked={
527
+ ".q_proj": (".qkv", "q"),
528
+ ".k_proj": (".qkv", "k"),
529
+ ".v_proj": (".qkv", "v"),
530
+ }
531
+ )
532
+
533
def __init__(
534
self,
535
vision_config: Qwen2VLVisionConfig,
@@ -742,31 +750,8 @@ class Qwen2VisionTransformer(nn.Module):
750
return x
751
752
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
745
- stacked_params_mapping = [
746
- # (param_name, shard_name, shard_id)
747
- ("qkv_proj", "q_proj", "q"),
748
- ("qkv_proj", "k_proj", "k"),
749
- ("qkv_proj", "v_proj", "v"),
750
- ]
751
- params_dict = dict(self.named_parameters(remove_duplicate=False))
752
- loaded_params: set[str] = set()
753
-
754
- for name, loaded_weight in weights:
755
- for param_name, weight_name, shard_id in stacked_params_mapping:
756
- if weight_name not in name:
757
- continue
758
- name = name.replace(weight_name, param_name)
759
-
760
- param = params_dict[name]
761
- weight_loader = param.weight_loader
762
- weight_loader(param, loaded_weight, shard_id)
763
- break
764
- else:
765
- param = params_dict[name]
766
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
767
- weight_loader(param, loaded_weight)
768
- loaded_params.add(name)
769
- return loaded_params
753
+ loader = AutoWeightsLoader(self)
754
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
755
756
757
def _create_qwen2vl_field_factory(
vllm/model_executor/models/qwen3_5.py
+30
-224
@@ -24,8 +24,7 @@
24
# limitations under the License.
25
"""Inference-only Qwen3.5 Series compatible with HuggingFace weights."""
26
27
-import typing
28
-from collections.abc import Callable, Iterable
27
+from collections.abc import Iterable
28
29
import torch
30
from torch import nn
@@ -37,12 +36,7 @@ from vllm.distributed import (
36
get_pp_group,
37
)
38
from vllm.logger import init_logger
40
-from vllm.model_executor.layers.fused_moe import (
41
- fused_moe_make_expert_params_mapping,
42
-)
43
-from vllm.model_executor.layers.layernorm import (
44
- GemmaRMSNorm as Qwen3_5RMSNorm,
45
-)
39
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm as Qwen3_5RMSNorm
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
42
QwenGatedDeltaNetAttention,
@@ -57,16 +51,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
51
ParallelLMHead,
52
VocabParallelEmbedding,
53
)
60
-from vllm.model_executor.model_loader.weight_utils import (
61
- default_weight_loader,
62
- maybe_remap_kv_scale_name,
63
-)
54
from vllm.multimodal import MULTIMODAL_REGISTRY
55
from vllm.sequence import IntermediateTensors
66
-from vllm.transformers_utils.configs.qwen3_5 import (
67
- Qwen3_5Config,
68
- Qwen3_5TextConfig,
69
-)
56
+from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig
57
from vllm.transformers_utils.configs.qwen3_5_moe import (
58
Qwen3_5MoeConfig,
59
Qwen3_5MoeTextConfig,
@@ -89,6 +76,7 @@ from .qwen3_next import (
76
Qwen3NextModel,
77
Qwen3NextSparseMoeBlock,
78
QwenNextMixtureOfExperts,
79
+ _is_shared_expert_fse_compatible,
80
)
81
from .qwen3_vl import (
82
Qwen3_VisionTransformer,
@@ -100,9 +88,9 @@ from .qwen3_vl import (
88
from .utils import (
89
AutoWeightsLoader,
90
PPMissingLayer,
91
+ WeightsMapper,
92
_merge_multimodal_embeddings,
93
extract_layer_index,
105
- is_pp_missing_parameter,
94
make_empty_intermediate_tensors_factory,
95
make_layers,
96
maybe_prefix,
@@ -210,6 +198,17 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
198
}
199
)
200
class Qwen3_5Model(Qwen3NextModel):
201
+ # Qwen3.5 ships the GDN in_proj checkpoints separately (qwen3-next
202
+ # pre-fuses them); fuse them on top of the qwen3-next QKV/gate_up mapping.
203
+ hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper | WeightsMapper(
204
+ orig_to_new_stacked={
205
+ ".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)),
206
+ ".in_proj_z": (".in_proj_qkvz", 3),
207
+ ".in_proj_b": (".in_proj_ba", 0),
208
+ ".in_proj_a": (".in_proj_ba", 1),
209
+ }
210
+ )
211
+
212
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
213
super(Qwen3NextModel, self).__init__()
214
@@ -222,6 +221,7 @@ class Qwen3_5Model(Qwen3NextModel):
221
self.num_redundant_experts = eplb_config.num_redundant_experts
222
223
self.config = config
224
+ self.quant_config = vllm_config.quant_config
225
226
self.vocab_size = config.vocab_size
227
@@ -251,212 +251,21 @@ class Qwen3_5Model(Qwen3NextModel):
251
252
self.aux_hidden_state_layers: tuple[int, ...] = ()
253
254
- def load_fused_expert_weights(
255
- self,
256
- name: str,
257
- params_dict: dict,
258
- loaded_weight: torch.Tensor,
259
- shard_id: str,
260
- num_experts: int,
261
- ) -> bool:
262
- param = params_dict[name]
263
- weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
264
- loaded_local_expert = False
265
- for expert_id in range(num_experts):
266
- curr_expert_weight = loaded_weight[expert_id]
267
- success = weight_loader(
268
- param,
269
- curr_expert_weight,
270
- name,
271
- shard_id=shard_id,
272
- expert_id=expert_id,
273
- return_success=True,
274
- )
275
- if success:
276
- loaded_local_expert = True
277
-
278
- return loaded_local_expert
279
-
254
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
281
- stacked_params_mapping = [
282
- # (param_name, shard_name, shard_id)
283
- # GDN
284
- ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
285
- ("in_proj_qkvz", "in_proj_z", 3),
286
- # self attention
287
- ("qkv_proj", "q_proj", "q"),
288
- ("qkv_proj", "k_proj", "k"),
289
- ("qkv_proj", "v_proj", "v"),
290
- # mlp
291
- ("gate_up_proj", "gate_proj", 0),
292
- ("gate_up_proj", "up_proj", 1),
293
- ("in_proj_ba", "in_proj_b", 0),
294
- ("in_proj_ba", "in_proj_a", 1),
295
- ]
296
-
297
- params_dict = dict(self.named_parameters())
298
- loaded_params: set[str] = set()
299
- expert_params_mapping = self.get_expert_mapping()
300
- is_fused_expert = False
301
- fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
302
- for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
303
- self,
304
- ckpt_gate_proj_name="gate_up_proj",
305
- ckpt_down_proj_name="down_proj",
306
- ckpt_up_proj_name="gate_up_proj",
307
- num_experts=1,
308
- ):
309
- if shard_id == "w3":
310
- continue
311
- parts = ckpt_name.split(".")
312
- fused_expert_params_mapping.append(
313
- (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
314
- )
315
- num_experts = (
316
- self.config.num_experts if hasattr(self.config, "num_experts") else 0
255
+ mapper = self.hf_to_vllm_mapper
256
+ # FSE must match construction (Qwen3NextSparseMoeBlock): reroute the
257
+ # shared expert into the extra fused slot only when AITER FSE is both
258
+ # requested and compatible with the quant spec.
259
+ is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and (
260
+ _is_shared_expert_fse_compatible(self.quant_config)
261
)
318
- from vllm.config import get_current_vllm_config
319
-
320
- from .qwen3_next import _is_shared_expert_fse_compatible
321
-
322
- is_fse = (
323
- rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
324
- and _is_shared_expert_fse_compatible(get_current_vllm_config().quant_config)
325
- )
326
-
327
- for name, loaded_weight in weights:
328
- if "rotary_emb.inv_freq" in name:
329
- continue
330
-
331
- if name.startswith("mtp."):
332
- continue
333
-
334
- # Remapping the name of FP8 kv-scale.
335
- if name.endswith("scale"):
336
- name = maybe_remap_kv_scale_name(name, params_dict)
337
- if name is None:
338
- continue
339
-
340
- # FSE: remap shared_expert weights to fused expert slot
341
- if is_fse and "mlp.shared_expert." in name:
342
- name = name.replace(
343
- "mlp.shared_expert.",
344
- f"mlp.experts.{num_experts}.",
345
- )
346
- is_fused_expert = False
347
- expert_params_mapping = self.get_expert_mapping()
348
-
349
- for param_name, weight_name, shard_id in stacked_params_mapping:
350
- if "experts.gate_up_proj" in name or "experts.down_proj" in name:
351
- is_fused_expert = True
352
- expert_params_mapping = fused_expert_params_mapping
353
-
354
- if weight_name not in name:
355
- continue
356
-
357
- if "mlp.experts" in name:
358
- continue
359
-
360
- name = name.replace(weight_name, param_name)
361
- # Skip loading extra bias for GPTQ models.
362
- if name.endswith(".bias") and name not in params_dict:
363
- continue
364
- # Skip layers on other devices.
365
- if is_pp_missing_parameter(name, self):
366
- continue
367
- # name = apply_attn_prefix(name, params_dict)
368
- if name not in params_dict:
369
- continue
370
- param = params_dict[name]
371
- weight_loader = param.weight_loader
372
- weight_loader(param, loaded_weight, shard_id)
373
- break
374
- else:
375
- is_expert_weight = False
376
- for mapping in expert_params_mapping:
377
- param_name, weight_name, expert_id, shard_id = mapping
378
- if weight_name not in name:
379
- continue
380
- is_expert_weight = True
381
- name_mapped = name.replace(weight_name, param_name)
382
- # Skip layers on other devices.
383
- if is_pp_missing_parameter(name_mapped, self):
384
- continue
385
- if is_fused_expert:
386
- # qwen3.5 no need to transpose
387
- # loaded_weight = loaded_weight.transpose(-1, -2)
388
- if "experts.gate_up_proj" in name:
389
- loaded_weight = loaded_weight.chunk(2, dim=-2)
390
- success_w1 = self.load_fused_expert_weights(
391
- name_mapped,
392
- params_dict,
393
- loaded_weight[0],
394
- "w1",
395
- num_experts,
396
- )
397
- success_w3 = self.load_fused_expert_weights(
398
- name_mapped,
399
- params_dict,
400
- loaded_weight[1],
401
- "w3",
402
- num_experts,
403
- )
404
- success = success_w1 and success_w3
405
- else:
406
- # down_proj
407
- success = self.load_fused_expert_weights(
408
- name_mapped,
409
- params_dict,
410
- loaded_weight,
411
- shard_id,
412
- num_experts,
413
- )
414
- if success:
415
- name = name_mapped
416
- break
417
- else:
418
- # Skip loading extra bias for GPTQ models.
419
- if (
420
- name_mapped.endswith(".bias")
421
- or name_mapped.endswith("_bias")
422
- ) and name_mapped not in params_dict:
423
- continue
424
- param = params_dict[name_mapped]
425
- weight_loader = param.weight_loader
426
- success = weight_loader(
427
- param,
428
- loaded_weight,
429
- name_mapped,
430
- shard_id=shard_id,
431
- expert_id=expert_id,
432
- return_success=True,
433
- )
434
- if success:
435
- name = name_mapped
436
- break
437
- else:
438
- if is_expert_weight:
439
- # We've checked that this is an expert weight
440
- # However it's not mapped locally to this rank
441
- # So we simply skip it
442
- continue
443
- # Skip loading extra bias for GPTQ models.
444
- if name.endswith(".bias") and name not in params_dict:
445
- continue
446
- if is_pp_missing_parameter(name, self):
447
- continue
448
- if name not in params_dict:
449
- logger.warning_once(
450
- f"Parameter {name} not found in params_dict, skip loading"
451
- )
452
- continue
453
- param = params_dict[name]
454
- weight_loader = getattr(
455
- param, "weight_loader", default_weight_loader
456
- )
457
- weight_loader(param, loaded_weight)
458
- loaded_params.add(name)
459
- return loaded_params
262
+ if is_fse:
263
+ num_routed = self.config.num_experts
264
+ mapper = mapper | WeightsMapper(
265
+ orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
266
+ )
267
+ loader = AutoWeightsLoader(self)
268
+ return loader.load_weights(weights, mapper=mapper)
269
270
271
class Qwen3_5ForCausalLMBase(
@@ -566,9 +375,6 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
375
# set MoE hyperparameters
376
self.set_moe_parameters()
377
569
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
570
- return self.model.get_expert_mapping()
571
-
378
379
########################################################
380
# Qwen3_5-Dense
vllm/model_executor/models/qwen3_5_mtp.py
+24
-189
@@ -2,29 +2,32 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only Qwen3_5 MTP model."""
4
5
-import typing
6
-from collections.abc import Callable, Iterable
5
+from collections.abc import Iterable
6
7
import torch
8
from torch import nn
9
10
+from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.compilation.decorators import support_torch_compile
12
-from vllm.config import VllmConfig
12
+from vllm.config import VllmConfig, get_current_vllm_config
13
from vllm.distributed.parallel_state import get_pp_group
14
from vllm.logger import init_logger
15
-from vllm.model_executor.layers.fused_moe import (
16
- fused_moe_make_expert_params_mapping,
17
-)
15
from vllm.model_executor.layers.linear import ColumnParallelLinear
16
from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.vocab_parallel_embedding import (
18
ParallelLMHead,
19
VocabParallelEmbedding,
20
)
24
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.model_executor.models.interfaces import LocalArgmaxMixin
26
-from vllm.model_executor.models.qwen3_5 import Qwen3_5DecoderLayer, Qwen3_5RMSNorm
27
-from vllm.model_executor.models.qwen3_next import QwenNextMixtureOfExperts
22
+from vllm.model_executor.models.qwen3_5 import (
23
+ Qwen3_5DecoderLayer,
24
+ Qwen3_5Model,
25
+ Qwen3_5RMSNorm,
26
+)
27
+from vllm.model_executor.models.qwen3_next import (
28
+ QwenNextMixtureOfExperts,
29
+ _is_shared_expert_fse_compatible,
30
+)
31
from vllm.sequence import IntermediateTensors
32
from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig
33
from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig
@@ -37,8 +40,8 @@ from .interfaces import (
40
from .utils import (
41
AutoWeightsLoader,
42
PPMissingLayer,
43
+ WeightsMapper,
44
_merge_multimodal_embeddings,
41
- is_pp_missing_parameter,
45
make_empty_intermediate_tensors_factory,
46
maybe_prefix,
47
)
@@ -58,6 +61,8 @@ logger = init_logger(__name__)
61
}
62
)
63
class Qwen3_5MultiTokenPredictor(nn.Module):
64
+ hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper
65
+
66
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
67
super().__init__()
68
@@ -159,188 +164,18 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
164
hidden_states, _ = self.norm(hidden_states, residual)
165
return hidden_states
166
162
- def load_fused_expert_weights(
163
- self,
164
- name: str,
165
- params_dict: dict,
166
- loaded_weight: torch.Tensor,
167
- shard_id: str,
168
- num_experts: int,
169
- ) -> bool:
170
- param = params_dict[name]
171
- weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
172
- loaded_local_expert = False
173
- for expert_id in range(num_experts):
174
- curr_expert_weight = loaded_weight[expert_id]
175
- success = weight_loader(
176
- param,
177
- curr_expert_weight,
178
- name,
179
- shard_id=shard_id,
180
- expert_id=expert_id,
181
- return_success=True,
182
- )
183
- if success:
184
- loaded_local_expert = True
185
-
186
- return loaded_local_expert
187
-
167
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
189
- stacked_params_mapping = [
190
- # (param_name, shard_name, shard_id)
191
- ("qkv_proj", "q_proj", "q"),
192
- ("qkv_proj", "k_proj", "k"),
193
- ("qkv_proj", "v_proj", "v"),
194
- ("gate_up_proj", "gate_proj", 0),
195
- ("gate_up_proj", "up_proj", 1),
196
- ]
197
-
198
- # Params for weights, fp8 weight scales, fp8 activation scales
199
- # (param_name, weight_name, expert_id, shard_id)
200
- expert_params_mapping = fused_moe_make_expert_params_mapping(
201
- self,
202
- ckpt_gate_proj_name="gate_proj",
203
- ckpt_down_proj_name="down_proj",
204
- ckpt_up_proj_name="up_proj",
205
- num_experts=self.config.num_experts
206
- if hasattr(self.config, "num_experts")
207
- else 0,
168
+ mapper = self.hf_to_vllm_mapper
169
+ is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and (
170
+ _is_shared_expert_fse_compatible(get_current_vllm_config().quant_config)
171
)
209
-
210
- params_dict = dict(self.named_parameters())
211
- loaded_params: set[str] = set()
212
- is_fused_expert = False
213
- fused_expert_params_mapping: list[tuple[str, str, int, str]] = []
214
- for param_name, ckpt_name, _, shard_id in fused_moe_make_expert_params_mapping(
215
- self,
216
- ckpt_gate_proj_name="gate_up_proj",
217
- ckpt_down_proj_name="down_proj",
218
- ckpt_up_proj_name="gate_up_proj",
219
- num_experts=1,
220
- ):
221
- if shard_id == "w3":
222
- continue
223
- parts = ckpt_name.split(".")
224
- fused_expert_params_mapping.append(
225
- (f"{param_name}weight", f"{parts[0]}.{parts[2]}", 0, shard_id)
172
+ if is_fse:
173
+ num_routed = getattr(self.config, "num_experts", 0)
174
+ mapper = mapper | WeightsMapper(
175
+ orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
176
)
227
- num_experts = (
228
- self.config.num_experts if hasattr(self.config, "num_experts") else 0
229
- )
230
- for name, loaded_weight in weights:
231
- if "rotary_emb.inv_freq" in name:
232
- continue
233
-
234
- for param_name, weight_name, shard_id in stacked_params_mapping:
235
- if "experts.gate_up_proj" in name or "experts.down_proj" in name:
236
- is_fused_expert = True
237
- expert_params_mapping = fused_expert_params_mapping
238
-
239
- if weight_name not in name:
240
- continue
241
-
242
- if "mlp.experts" in name:
243
- continue
244
-
245
- name = name.replace(weight_name, param_name)
246
- # Skip loading extra bias for GPTQ models.
247
- if name.endswith(".bias") and name not in params_dict:
248
- continue
249
- # Skip layers on other devices.
250
- if is_pp_missing_parameter(name, self):
251
- continue
252
- if name not in params_dict:
253
- continue
254
- param = params_dict[name]
255
- weight_loader = param.weight_loader
256
- weight_loader(param, loaded_weight, shard_id)
257
- break
258
- else:
259
- is_expert_weight = False
260
- for mapping in expert_params_mapping:
261
- param_name, weight_name, expert_id, shard_id = mapping
262
- if weight_name not in name:
263
- continue
264
- is_expert_weight = True
265
- name_mapped = name.replace(weight_name, param_name)
266
- # Skip layers on other devices.
267
- if is_pp_missing_parameter(name_mapped, self):
268
- continue
269
- if is_fused_expert:
270
- # qwen3.5 no need to transpose
271
- # loaded_weight = loaded_weight.transpose(-1, -2)
272
- if "experts.gate_up_proj" in name:
273
- loaded_weight = loaded_weight.chunk(2, dim=-2)
274
- success_w1 = self.load_fused_expert_weights(
275
- name_mapped,
276
- params_dict,
277
- loaded_weight[0],
278
- "w1",
279
- num_experts,
280
- )
281
- success_w3 = self.load_fused_expert_weights(
282
- name_mapped,
283
- params_dict,
284
- loaded_weight[1],
285
- "w3",
286
- num_experts,
287
- )
288
- success = success_w1 and success_w3
289
- else:
290
- # down_proj
291
- success = self.load_fused_expert_weights(
292
- name_mapped,
293
- params_dict,
294
- loaded_weight,
295
- shard_id,
296
- num_experts,
297
- )
298
- if success:
299
- name = name_mapped
300
- break
301
- else:
302
- # Skip loading extra bias for GPTQ models.
303
- if (
304
- name_mapped.endswith(".bias")
305
- or name_mapped.endswith("_bias")
306
- ) and name_mapped not in params_dict:
307
- continue
308
- param = params_dict[name_mapped]
309
- weight_loader = param.weight_loader
310
- success = weight_loader(
311
- param,
312
- loaded_weight,
313
- name_mapped,
314
- shard_id=shard_id,
315
- expert_id=expert_id,
316
- return_success=True,
317
- )
318
- if success:
319
- name = name_mapped
320
- break
321
- else:
322
- if is_expert_weight:
323
- # We've checked that this is an expert weight
324
- # However it's not mapped locally to this rank
325
- # So we simply skip it
326
- continue
327
- # Skip loading extra bias for GPTQ models.
328
- if name.endswith(".bias") and name not in params_dict:
329
- continue
330
- if is_pp_missing_parameter(name, self):
331
- continue
332
- if name not in params_dict:
333
- logger.warning_once(
334
- f"Parameter {name} not found in params_dict, skip loading"
335
- )
336
- continue
337
- param = params_dict[name]
338
- weight_loader = getattr(
339
- param, "weight_loader", default_weight_loader
340
- )
341
- weight_loader(param, loaded_weight)
342
- loaded_params.add(name)
343
- return loaded_params
177
+ loader = AutoWeightsLoader(self)
178
+ return loader.load_weights(weights, mapper=mapper)
179
180
181
@support_torch_compile(
vllm/model_executor/models/qwen3_eagle3.py
+14
-36
@@ -17,15 +17,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
17
ParallelLMHead,
18
VocabParallelEmbedding,
19
)
20
-from vllm.model_executor.model_loader.weight_utils import (
21
- default_weight_loader,
22
- maybe_remap_kv_scale_name,
23
-)
20
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3ForCausalLM
21
from vllm.multimodal.inputs import NestedTensors
22
23
from .utils import (
24
AutoWeightsLoader,
25
+ WeightsMapper,
26
get_draft_quant_config,
27
maybe_prefix,
28
process_eagle_weight,
@@ -132,6 +129,17 @@ class Qwen3Eagle3DecoderLayer(Qwen3DecoderLayer):
129
}
130
)
131
class Qwen3Eagle3Model(nn.Module):
132
+ hf_to_vllm_mapper = WeightsMapper(
133
+ orig_to_new_substr={"midlayer.": "layers.0."},
134
+ orig_to_new_stacked={
135
+ ".q_proj": (".qkv_proj", "q"),
136
+ ".k_proj": (".qkv_proj", "k"),
137
+ ".v_proj": (".qkv_proj", "v"),
138
+ ".gate_proj": (".gate_up_proj", 0),
139
+ ".up_proj": (".gate_up_proj", 1),
140
+ },
141
+ )
142
+
143
def __init__(
144
self,
145
*,
@@ -256,38 +264,8 @@ class Qwen3Eagle3Model(nn.Module):
264
return hidden_states, aux_output
265
266
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
259
- stacked_params_mapping = [
260
- # (param_name, shard_name, shard_id)
261
- (".qkv_proj", ".q_proj", "q"),
262
- (".qkv_proj", ".k_proj", "k"),
263
- (".qkv_proj", ".v_proj", "v"),
264
- (".gate_up_proj", ".gate_proj", 0),
265
- (".gate_up_proj", ".up_proj", 1),
266
- ]
267
- params_dict = dict(self.named_parameters())
268
- loaded_params: set[str] = set()
269
- for name, loaded_weight in weights:
270
- if "midlayer." in name:
271
- name = name.replace("midlayer.", "layers.0.")
272
- # Remapping the name FP8 kv-scale or zero point.
273
- if "scale" in name or "zero_point" in name:
274
- name = maybe_remap_kv_scale_name(name, params_dict)
275
- if name is None:
276
- continue
277
- for param_name, weight_name, shard_id in stacked_params_mapping:
278
- if weight_name not in name:
279
- continue
280
- name = name.replace(weight_name, param_name)
281
- param = params_dict[name]
282
- weight_loader = param.weight_loader
283
- weight_loader(param, loaded_weight, shard_id)
284
- break
285
- else:
286
- param = params_dict[name]
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
267
+ loader = AutoWeightsLoader(self)
268
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
269
270
271
class Eagle3Qwen3ForCausalLM(Qwen3ForCausalLM):
vllm/model_executor/models/qwen3_moe.py
+37
-154
@@ -23,8 +23,7 @@
23
# limitations under the License.
24
"""Inference-only Qwen3MoE 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
@@ -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,
@@ -61,10 +57,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
57
ParallelLMHead,
58
VocabParallelEmbedding,
59
)
64
-from vllm.model_executor.model_loader.weight_utils import (
65
- default_weight_loader,
66
- maybe_remap_kv_scale_name,
67
-)
60
from vllm.model_executor.models.utils import sequence_parallel_chunk
61
from vllm.sequence import IntermediateTensors
62
@@ -75,12 +67,13 @@ from .interfaces import (
67
SupportsEagle3,
68
SupportsLoRA,
69
SupportsPP,
70
+ SupportsQuant,
71
)
72
from .utils import (
73
AutoWeightsLoader,
74
PPMissingLayer,
75
+ WeightsMapper,
76
extract_layer_index,
83
- is_pp_missing_parameter,
77
make_empty_intermediate_tensors_factory,
78
make_layers,
79
maybe_prefix,
@@ -438,6 +431,20 @@ class Qwen3MoeDecoderLayer(nn.Module):
431
432
@support_torch_compile
433
class Qwen3MoeModel(nn.Module, EagleModelMixin):
434
+ hf_to_vllm_mapper = WeightsMapper(
435
+ orig_to_new_stacked={
436
+ # weight_name: (param_name, shard_id)
437
+ ".q_proj": (".qkv_proj", "q"),
438
+ ".k_proj": (".qkv_proj", "k"),
439
+ ".v_proj": (".qkv_proj", "v"),
440
+ # .experts.gate_up_proj must be handled by MoERunner.load_weights for EP
441
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
442
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
443
+ ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0),
444
+ ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1),
445
+ }
446
+ )
447
+
448
def __init__(
449
self,
450
*,
@@ -516,152 +523,31 @@ class Qwen3MoeModel(nn.Module, EagleModelMixin):
523
return hidden_states, aux_hidden_states
524
return hidden_states
525
519
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
520
- # Params for weights, fp8 weight scales, fp8 activation scales
521
- # (param_name, weight_name, expert_id, shard_id)
522
- return fused_moe_make_expert_params_mapping(
526
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
527
+ loader = AutoWeightsLoader(
528
self,
524
- ckpt_gate_proj_name="gate_proj",
525
- ckpt_down_proj_name="down_proj",
526
- ckpt_up_proj_name="up_proj",
527
- num_experts=self.config.num_experts,
528
- num_redundant_experts=self.num_redundant_experts,
529
+ ignore_unexpected_suffixes=[
530
+ ".bias",
531
+ "_bias",
532
+ ".weight_scale",
533
+ "_weight_scale",
534
+ ".input_scale",
535
+ "_input_scale",
536
+ ],
537
)
530
-
531
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
532
- stacked_params_mapping = [
533
- # (param_name, shard_name, shard_id)
534
- ("qkv_proj", "q_proj", "q"),
535
- ("qkv_proj", "k_proj", "k"),
536
- ("qkv_proj", "v_proj", "v"),
537
- ("gate_up_proj", "gate_proj", 0),
538
- ("gate_up_proj", "up_proj", 1),
539
- ]
540
-
541
- # Skip loading extra parameters for GPTQ/modelopt models.
542
- ignore_suffixes = (
543
- ".bias",
544
- "_bias",
545
- ".weight_scale",
546
- "_weight_scale",
547
- ".input_scale",
548
- "_input_scale",
549
- )
550
-
551
- params_dict = dict(self.named_parameters())
552
- loaded_params: set[str] = set()
553
- expert_params_mapping = self.get_expert_mapping()
554
- for name, loaded_weight in weights:
555
- if "scale" in name or "zero_point" in name:
556
- name = maybe_remap_kv_scale_name(name, params_dict)
557
- if name is None:
558
- continue
559
- for param_name, weight_name, shard_id in stacked_params_mapping:
560
- # Skip non-stacked layers and experts (experts handled below).
561
- if weight_name not in name:
562
- continue
563
- # We have mlp.experts[0].gate_proj in the checkpoint.
564
- # Since we handle the experts below in expert_params_mapping,
565
- # we need to skip here BEFORE we update the name, otherwise
566
- # name will be updated to mlp.experts[0].gate_up_proj, which
567
- # will then be updated below in expert_params_mapping
568
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
569
- if "mlp.experts" in name:
570
- continue
571
- name = name.replace(weight_name, param_name)
572
-
573
- # Skip loading extra parameters for GPTQ/modelopt models.
574
- if name.endswith(ignore_suffixes) and name not in params_dict:
575
- continue
576
-
577
- # Skip layers on other devices.
578
- if is_pp_missing_parameter(name, self):
579
- continue
580
- if name.endswith("scale"):
581
- # Remapping the name of FP8 kv-scale.
582
- name = maybe_remap_kv_scale_name(name, params_dict)
583
- if name is None:
584
- continue
585
- if name not in params_dict:
586
- continue
587
-
588
- param = params_dict[name]
589
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
590
- if weight_loader == default_weight_loader:
591
- weight_loader(param, loaded_weight)
592
- else:
593
- weight_loader(param, loaded_weight, shard_id)
594
- break
595
- else:
596
- is_expert_weight = False
597
- for mapping in expert_params_mapping:
598
- param_name, weight_name, expert_id, shard_id = mapping
599
- if weight_name not in name:
600
- continue
601
-
602
- # Anyway, this is an expert weight and should not be
603
- # attempted to load as other weights later
604
- is_expert_weight = True
605
-
606
- # Do not modify `name` since the loop may continue here
607
- # Instead, create a new variable
608
- name_mapped = name.replace(weight_name, param_name)
609
-
610
- if is_pp_missing_parameter(name_mapped, self):
611
- continue
612
-
613
- # Skip loading extra parameters for GPTQ/modelopt models.
614
- if (
615
- name_mapped.endswith(ignore_suffixes)
616
- and name_mapped not in params_dict
617
- ):
618
- continue
619
-
620
- param = params_dict[name_mapped]
621
- # We should ask the weight loader to return success or not
622
- # here since otherwise we may skip experts with other
623
- # available replicas.
624
- weight_loader = typing.cast(
625
- Callable[..., bool], param.weight_loader
626
- )
627
- success = weight_loader(
628
- param,
629
- loaded_weight,
630
- name_mapped,
631
- shard_id=shard_id,
632
- expert_id=expert_id,
633
- return_success=True,
634
- )
635
- if success:
636
- name = 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 parameters for GPTQ/modelopt models.
646
- if name.endswith(ignore_suffixes) and name not in params_dict:
647
- continue
648
- # Skip layers on other devices.
649
- if is_pp_missing_parameter(name, self):
650
- continue
651
- if name not in params_dict:
652
- continue
653
- param = params_dict[name]
654
- weight_loader = getattr(
655
- param, "weight_loader", default_weight_loader
656
- )
657
- weight_loader(param, loaded_weight)
658
- loaded_params.add(name)
659
- return loaded_params
538
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
539
540
541
class Qwen3MoeForCausalLM(
663
- nn.Module, SupportsPP, SupportsLoRA, SupportsEagle, SupportsEagle3, MixtureOfExperts
542
+ nn.Module,
543
+ SupportsPP,
544
+ SupportsLoRA,
545
+ SupportsEagle,
546
+ SupportsEagle3,
547
+ MixtureOfExperts,
548
+ SupportsQuant,
549
):
550
+ hf_to_vllm_mapper = Qwen3MoeModel.hf_to_vllm_mapper
551
packed_modules_mapping = {
552
"qkv_proj": [
553
"q_proj",
@@ -769,6 +655,3 @@ class Qwen3MoeForCausalLM(
655
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
656
loader = AutoWeightsLoader(self)
657
return loader.load_weights(weights)
772
-
773
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
774
- return self.model.get_expert_mapping()
vllm/model_executor/models/qwen3_next.py
+27
-136
@@ -10,12 +10,7 @@ from torch import nn
10
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.compilation.decorators import support_torch_compile
13
-from vllm.config import (
14
- CacheConfig,
15
- ModelConfig,
16
- VllmConfig,
17
- get_current_vllm_config,
18
-)
13
+from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import (
15
get_ep_group,
16
get_pp_group,
@@ -26,7 +21,6 @@ from vllm.logger import init_logger
21
from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.fused_moe import (
23
FusedMoE,
29
- fused_moe_make_expert_params_mapping,
24
)
25
from vllm.model_executor.layers.fused_qk_norm_rope import fused_qk_rmsnorm_rope_gate
26
from vllm.model_executor.layers.layernorm import (
@@ -53,10 +47,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
47
ParallelLMHead,
48
VocabParallelEmbedding,
49
)
56
-from vllm.model_executor.model_loader.weight_utils import (
57
- default_weight_loader,
58
- maybe_remap_kv_scale_name,
59
-)
50
from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
51
from vllm.model_executor.models.utils import sequence_parallel_chunk
52
from vllm.platforms import current_platform
@@ -76,8 +66,8 @@ from .interfaces import (
66
from .utils import (
67
AutoWeightsLoader,
68
PPMissingLayer,
69
+ WeightsMapper,
70
extract_layer_index,
80
- is_pp_missing_parameter,
71
make_empty_intermediate_tensors_factory,
72
make_layers,
73
maybe_prefix,
@@ -132,7 +122,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):
122
)
123
124
# Load balancing settings.
135
- vllm_config = get_current_vllm_config()
125
eplb_config = vllm_config.parallel_config.eplb_config
126
self.enable_eplb = parallel_config.enable_eplb
127
@@ -544,6 +533,19 @@ class Qwen3NextDecoderLayer(nn.Module):
533
534
@support_torch_compile
535
class Qwen3NextModel(nn.Module, EagleModelMixin):
536
+ hf_to_vllm_mapper = WeightsMapper(
537
+ orig_to_new_stacked={
538
+ # weight_name: (param_name, shard_id)
539
+ ".q_proj": (".qkv_proj", "q"),
540
+ ".k_proj": (".qkv_proj", "k"),
541
+ ".v_proj": (".qkv_proj", "v"),
542
+ ".mlp.gate_proj": (".mlp.gate_up_proj", 0),
543
+ ".mlp.up_proj": (".mlp.gate_up_proj", 1),
544
+ ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0),
545
+ ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1),
546
+ }
547
+ )
548
+
549
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
550
super().__init__()
551
@@ -625,123 +627,18 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
627
return hidden_states, aux_hidden_states
628
return hidden_states
629
628
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
629
- # Params for weights, fp8 weight scales, fp8 activation scales
630
- # (param_name, weight_name, expert_id, shard_id)
631
- num_experts = getattr(self.config, "num_experts", 0)
632
- if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
633
- num_experts += 1
634
- return fused_moe_make_expert_params_mapping(
635
- self,
636
- ckpt_gate_proj_name="gate_proj",
637
- ckpt_down_proj_name="down_proj",
638
- ckpt_up_proj_name="up_proj",
639
- num_experts=num_experts,
640
- num_redundant_experts=self.num_redundant_experts,
641
- )
642
-
630
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
644
- stacked_params_mapping = [
645
- # (param_name, shard_name, shard_id)
646
- ("qkv_proj", "q_proj", "q"),
647
- ("qkv_proj", "k_proj", "k"),
648
- ("qkv_proj", "v_proj", "v"),
649
- ("gate_up_proj", "gate_proj", 0),
650
- ("gate_up_proj", "up_proj", 1),
651
- ]
652
-
653
- params_dict = dict(self.named_parameters())
654
- loaded_params: set[str] = set()
655
- expert_params_mapping = self.get_expert_mapping()
656
-
657
- is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
658
- num_routed = getattr(self.config, "num_experts", 0)
659
-
660
- for name, loaded_weight in weights:
661
- if "rotary_emb.inv_freq" in name:
662
- continue
663
-
664
- if name.startswith("mtp."):
665
- continue
666
-
667
- # Remapping the name of FP8 kv-scale.
668
- if name.endswith("scale"):
669
- name = maybe_remap_kv_scale_name(name, params_dict)
670
- if name is None:
671
- continue
672
-
673
- # FSE: remap shared_expert weights to the fused expert slot
674
- if is_fse and "mlp.shared_expert." in name:
675
- name = name.replace(
676
- "mlp.shared_expert.",
677
- f"mlp.experts.{num_routed}.",
678
- )
679
-
680
- for param_name, weight_name, shard_id in stacked_params_mapping:
681
- if weight_name not in name:
682
- continue
683
-
684
- if "mlp.experts" in name:
685
- continue
686
-
687
- name = name.replace(weight_name, param_name)
688
- # Skip loading extra bias for GPTQ models.
689
- if name.endswith(".bias") and name not in params_dict:
690
- continue
691
- # Skip layers on other devices.
692
- if is_pp_missing_parameter(name, self):
693
- continue
694
- # name = apply_attn_prefix(name, params_dict)
695
- if name not in params_dict:
696
- continue
697
- param = params_dict[name]
698
- weight_loader = param.weight_loader
699
- weight_loader(param, loaded_weight, shard_id)
700
- break
701
- else:
702
- for mapping in expert_params_mapping:
703
- param_name, weight_name, expert_id, shard_id = mapping
704
- if weight_name not in name:
705
- continue
706
- name = name.replace(weight_name, param_name)
707
- # Skip layers on other devices.
708
- if is_pp_missing_parameter(name, self):
709
- continue
710
- # Skip loading extra bias for GPTQ models.
711
- if (
712
- name.endswith(".bias") or name.endswith("_bias")
713
- ) and name not in params_dict:
714
- continue
715
- if name not in params_dict:
716
- continue
717
- param = params_dict[name]
718
- weight_loader = param.weight_loader
719
- weight_loader(
720
- param,
721
- loaded_weight,
722
- name,
723
- shard_id=shard_id,
724
- expert_id=expert_id,
725
- )
726
- break
727
- else:
728
- # Skip loading extra bias for GPTQ models.
729
- if name.endswith(".bias") and name not in params_dict:
730
- continue
731
- if is_pp_missing_parameter(name, self):
732
- continue
733
- if name not in params_dict:
734
- logger.warning_once(
735
- f"Parameter {name} not found in params_dict, skip loading"
736
- )
737
- continue
738
- param = params_dict[name]
739
- weight_loader = getattr(
740
- param, "weight_loader", default_weight_loader
741
- )
742
- weight_loader(param, loaded_weight)
743
- loaded_params.add(name)
744
- return loaded_params
631
+ mapper = self.hf_to_vllm_mapper
632
+ if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
633
+ # AITER fused-shared-experts: route the shared_expert checkpoint
634
+ # weights into the extra fused expert slot. Merge (not mutate) so the
635
+ # shared class mapper isn't permanently altered.
636
+ num_routed = getattr(self.config, "num_experts", 0)
637
+ mapper = mapper | WeightsMapper(
638
+ orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
639
+ )
640
+ loader = AutoWeightsLoader(self)
641
+ return loader.load_weights(weights, mapper=mapper)
642
643
644
class QwenNextMixtureOfExperts(MixtureOfExperts):
@@ -901,11 +798,5 @@ class Qwen3NextForCausalLM(
798
return self.logits_processor(self.lm_head, hidden_states)
799
800
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
904
- loader = AutoWeightsLoader(
905
- self,
906
- skip_prefixes=["mtp."],
907
- )
801
+ loader = AutoWeightsLoader(self, skip_prefixes=["mtp."])
802
return loader.load_weights(weights)
909
-
910
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
911
- return self.model.get_expert_mapping()
vllm/model_executor/models/qwen3_next_mtp.py
+14
-100
@@ -12,18 +12,15 @@ from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
14
from vllm.logger import init_logger
15
-from vllm.model_executor.layers.fused_moe import (
16
- fused_moe_make_expert_params_mapping,
17
-)
15
from vllm.model_executor.layers.linear import ColumnParallelLinear
16
from vllm.model_executor.layers.logits_processor import LogitsProcessor
17
from vllm.model_executor.layers.vocab_parallel_embedding import (
18
ParallelLMHead,
19
VocabParallelEmbedding,
20
)
24
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.model_executor.models.qwen3_next import (
22
Qwen3NextDecoderLayer,
23
+ Qwen3NextModel,
24
Qwen3NextRMSNorm,
25
QwenNextMixtureOfExperts,
26
)
@@ -32,7 +29,7 @@ from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
29
30
from .utils import (
31
AutoWeightsLoader,
35
- is_pp_missing_parameter,
32
+ WeightsMapper,
33
make_empty_intermediate_tensors_factory,
34
maybe_prefix,
35
)
@@ -44,6 +41,8 @@ KVCache = tuple[torch.Tensor, torch.Tensor]
41
42
@support_torch_compile
43
class Qwen3NextMultiTokenPredictor(nn.Module):
44
+ hf_to_vllm_mapper = Qwen3NextModel.hf_to_vllm_mapper
45
+
46
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
47
super().__init__()
48
@@ -146,101 +145,16 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
145
return hidden_states
146
147
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
149
- stacked_params_mapping = [
150
- # (param_name, shard_name, shard_id)
151
- ("qkv_proj", "q_proj", "q"),
152
- ("qkv_proj", "k_proj", "k"),
153
- ("qkv_proj", "v_proj", "v"),
154
- ("gate_up_proj", "gate_proj", 0),
155
- ("gate_up_proj", "up_proj", 1),
156
- ]
157
-
158
- # Params for weights, fp8 weight scales, fp8 activation scales
159
- # (param_name, weight_name, expert_id, shard_id)
160
- is_fse = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
161
- num_experts = self.config.num_experts
162
- if is_fse:
163
- num_experts += 1
164
- expert_params_mapping = fused_moe_make_expert_params_mapping(
165
- self,
166
- ckpt_gate_proj_name="gate_proj",
167
- ckpt_down_proj_name="down_proj",
168
- ckpt_up_proj_name="up_proj",
169
- num_experts=num_experts,
170
- )
171
- num_routed = self.config.num_experts
172
-
173
- params_dict = dict(self.named_parameters())
174
- loaded_params: set[str] = set()
175
- for name, loaded_weight in weights:
176
- if "rotary_emb.inv_freq" in name:
177
- continue
178
-
179
- # FSE: remap shared_expert weights to the fused expert slot
180
- if is_fse and "mlp.shared_expert." in name:
181
- name = name.replace(
182
- "mlp.shared_expert.",
183
- f"mlp.experts.{num_routed}.",
184
- )
185
-
186
- for param_name, weight_name, shard_id in stacked_params_mapping:
187
- if weight_name not in name:
188
- continue
189
-
190
- if "mlp.experts" in name:
191
- continue
192
-
193
- name = name.replace(weight_name, param_name)
194
- # Skip loading extra bias for GPTQ models.
195
- if name.endswith(".bias") and name not in params_dict:
196
- continue
197
- # Skip layers on other devices.
198
- if is_pp_missing_parameter(name, self):
199
- continue
200
- if name not in params_dict:
201
- continue
202
- param = params_dict[name]
203
- weight_loader = param.weight_loader
204
- weight_loader(param, loaded_weight, shard_id)
205
- break
206
- else:
207
- for mapping in expert_params_mapping:
208
- param_name, weight_name, expert_id, shard_id = mapping
209
- if weight_name not in name:
210
- continue
211
- name = name.replace(weight_name, param_name)
212
- # Skip layers on other devices.
213
- if is_pp_missing_parameter(name, self):
214
- continue
215
- # Skip loading extra bias for GPTQ models.
216
- if (
217
- name.endswith(".bias") or name.endswith("_bias")
218
- ) and name not in params_dict:
219
- continue
220
- param = params_dict[name]
221
- weight_loader = param.weight_loader
222
- weight_loader(
223
- param,
224
- loaded_weight,
225
- name,
226
- shard_id=shard_id,
227
- expert_id=expert_id,
228
- )
229
- break
230
- else:
231
- # Skip loading extra bias for GPTQ models.
232
- if name.endswith(".bias") and name not in params_dict:
233
- continue
234
- if is_pp_missing_parameter(name, self):
235
- continue
236
-
237
- param = params_dict[name]
238
- weight_loader = getattr(
239
- param, "weight_loader", default_weight_loader
240
- )
241
- weight_loader(param, loaded_weight)
242
- loaded_params.add(name)
243
- return loaded_params
148
+ mapper = self.hf_to_vllm_mapper
149
+ if rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
150
+ # AITER fused-shared-experts: route the shared_expert checkpoint
151
+ # weights into the extra fused expert slot.
152
+ num_routed = self.config.num_experts
153
+ mapper = mapper | WeightsMapper(
154
+ orig_to_new_substr={"mlp.shared_expert.": f"mlp.experts.{num_routed}."}
155
+ )
156
+ loader = AutoWeightsLoader(self)
157
+ return loader.load_weights(weights, mapper=mapper)
158
159
160
@support_torch_compile
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+20
-55
@@ -63,7 +63,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
63
from vllm.model_executor.layers.quantization import QuantizationConfig
64
from vllm.model_executor.layers.rotary_embedding import get_rope
65
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
66
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
66
from vllm.model_executor.models.module_mapping import MultiModelKeys
67
from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
68
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -319,6 +318,14 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module):
318
class Qwen3OmniMoeAudioEncoder(nn.Module):
319
"""vLLM-native Qwen3-Omni Audio Encoder."""
320
321
+ hf_to_vllm_mapper = WeightsMapper(
322
+ orig_to_new_stacked={
323
+ ".self_attn.q_proj.": (".self_attn.qkv.", "q"),
324
+ ".self_attn.k_proj.": (".self_attn.qkv.", "k"),
325
+ ".self_attn.v_proj.": (".self_attn.qkv.", "v"),
326
+ }
327
+ )
328
+
329
def __init__(
330
self,
331
config: Qwen3OmniMoeAudioEncoderConfig,
@@ -531,35 +538,8 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):
538
return lengths
539
540
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
534
- """Load weights with mapping from HuggingFace format."""
535
- stacked_params_mapping = [
536
- # (param_name, shard_name, shard_id)
537
- ("self_attn.qkv.", "self_attn.q_proj.", "q"),
538
- ("self_attn.qkv.", "self_attn.k_proj.", "k"),
539
- ("self_attn.qkv.", "self_attn.v_proj.", "v"),
540
- ]
541
- params_dict = dict(self.named_parameters(remove_duplicate=False))
542
- loaded_params: set[str] = set()
543
-
544
- for name, loaded_weight in weights:
545
- for param_name, weight_name, shard_id in stacked_params_mapping:
546
- if weight_name not in name:
547
- continue
548
- name = name.replace(weight_name, param_name)
549
-
550
- param = params_dict[name]
551
- weight_loader = param.weight_loader
552
- weight_loader(param, loaded_weight, shard_id)
553
- break
554
- else:
555
- param = params_dict.get(name)
556
- if param is not None:
557
- weight_loader = getattr(
558
- param, "weight_loader", default_weight_loader
559
- )
560
- weight_loader(param, loaded_weight)
561
- loaded_params.add(name)
562
- return loaded_params
541
+ loader = AutoWeightsLoader(self)
542
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
543
544
545
class Qwen3_VisionPatchEmbed(nn.Module):
@@ -737,6 +717,14 @@ class Qwen3_VisionPatchMerger(nn.Module):
717
718
719
class Qwen3Omni_VisionTransformer(nn.Module):
720
+ hf_to_vllm_mapper = WeightsMapper(
721
+ orig_to_new_stacked={
722
+ ".attn.q.": (".attn.qkv.", "q"),
723
+ ".attn.k.": (".attn.qkv.", "k"),
724
+ ".attn.v.": (".attn.qkv.", "v"),
725
+ }
726
+ )
727
+
728
def __init__(
729
self,
730
vision_config,
@@ -1044,31 +1032,8 @@ class Qwen3Omni_VisionTransformer(nn.Module):
1032
return hidden_states
1033
1034
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1047
- stacked_params_mapping = [
1048
- # (param_name, shard_name, shard_id)
1049
- ("attn.qkv.", "attn.q.", "q"),
1050
- ("attn.qkv.", "attn.k.", "k"),
1051
- ("attn.qkv.", "attn.v.", "v"),
1052
- ]
1053
- params_dict = dict(self.named_parameters(remove_duplicate=False))
1054
- loaded_params: set[str] = set()
1055
-
1056
- for name, loaded_weight in weights:
1057
- for param_name, weight_name, shard_id in stacked_params_mapping:
1058
- if weight_name not in name:
1059
- continue
1060
- name = name.replace(weight_name, param_name)
1061
-
1062
- param = params_dict[name]
1063
- weight_loader = param.weight_loader
1064
- weight_loader(param, loaded_weight, shard_id)
1065
- break
1066
- else:
1067
- param = params_dict[name]
1068
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
1069
- weight_loader(param, loaded_weight)
1070
- loaded_params.add(name)
1071
- return loaded_params
1035
+ loader = AutoWeightsLoader(self)
1036
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1037
1038
1039
@support_torch_compile(
vllm/model_executor/models/qwen3_vl.py
+10
-26
@@ -67,7 +67,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
67
from vllm.model_executor.layers.quantization import QuantizationConfig
68
from vllm.model_executor.layers.rotary_embedding import get_rope
69
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
70
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
70
from vllm.model_executor.models.module_mapping import MultiModelKeys
71
from vllm.multimodal import MULTIMODAL_REGISTRY
72
from vllm.multimodal.evs import (
@@ -518,6 +517,14 @@ class Qwen3_VisionPatchMerger(nn.Module):
517
518
519
class Qwen3_VisionTransformer(nn.Module):
520
+ hf_to_vllm_mapper = WeightsMapper(
521
+ orig_to_new_stacked={
522
+ "attn.q.": ("attn.qkv.", "q"),
523
+ "attn.k.": ("attn.qkv.", "k"),
524
+ "attn.v.": ("attn.qkv.", "v"),
525
+ }
526
+ )
527
+
528
def __init__(
529
self,
530
vision_config: Qwen3VLVisionConfig,
@@ -834,31 +841,8 @@ class Qwen3_VisionTransformer(nn.Module):
841
return hidden_states
842
843
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
837
- stacked_params_mapping = [
838
- # (param_name, shard_name, shard_id)
839
- ("attn.qkv.", "attn.q.", "q"),
840
- ("attn.qkv.", "attn.k.", "k"),
841
- ("attn.qkv.", "attn.v.", "v"),
842
- ]
843
- params_dict = dict(self.named_parameters(remove_duplicate=False))
844
- loaded_params: set[str] = set()
845
-
846
- for name, loaded_weight in weights:
847
- for param_name, weight_name, shard_id in stacked_params_mapping:
848
- if weight_name not in name:
849
- continue
850
- name = name.replace(weight_name, param_name)
851
-
852
- param = params_dict[name]
853
- weight_loader = param.weight_loader
854
- weight_loader(param, loaded_weight, shard_id)
855
- break
856
- else:
857
- param = params_dict[name]
858
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
859
- weight_loader(param, loaded_weight)
860
- loaded_params.add(name)
861
- return loaded_params
844
+ loader = AutoWeightsLoader(self)
845
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
846
847
848
class Qwen3VLProcessingInfo(Qwen2VLProcessingInfo):
vllm/model_executor/models/qwen3_vl_moe.py
+1
-203
@@ -24,8 +24,6 @@
24
# limitations under the License.
25
"""Inference-only Qwen3-VL-MoE model compatible with HuggingFace weights."""
26
27
-import typing
28
-from collections.abc import Callable, Iterable
27
from itertools import islice
28
29
import torch
@@ -39,10 +37,6 @@ from vllm.distributed import get_pp_group
37
from vllm.logger import init_logger
38
from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
42
-from vllm.model_executor.model_loader.weight_utils import (
43
- default_weight_loader,
44
- maybe_remap_kv_scale_name,
45
-)
40
from vllm.multimodal import MULTIMODAL_REGISTRY
41
from vllm.sequence import IntermediateTensors
42
from vllm.tokenizers.registry import cached_tokenizer_from_config
@@ -60,7 +54,7 @@ from .qwen3_vl import (
54
Qwen3VLMultiModalProcessor,
55
Qwen3VLProcessingInfo,
56
)
63
-from .utils import is_pp_missing_parameter, maybe_prefix
57
+from .utils import maybe_prefix
58
59
logger = init_logger(__name__)
60
@@ -135,202 +129,6 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
129
return hidden_states, aux_hidden_states
130
return hidden_states
131
138
- def load_fused_expert_weights(
139
- self,
140
- name: str,
141
- params_dict: dict,
142
- loaded_weight: torch.Tensor,
143
- shard_id: str,
144
- num_experts: int,
145
- ) -> bool:
146
- param = params_dict[name]
147
- weight_loader = typing.cast(Callable[..., bool], param.weight_loader)
148
- loaded_local_expert = False
149
- for expert_id in range(num_experts):
150
- curr_expert_weight = loaded_weight[expert_id]
151
- success = weight_loader(
152
- param,
153
- curr_expert_weight,
154
- name,
155
- shard_id=shard_id,
156
- expert_id=expert_id,
157
- return_success=True,
158
- )
159
- if success:
160
- loaded_local_expert = True
161
-
162
- return loaded_local_expert
163
-
164
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
165
- stacked_params_mapping = [
166
- # (param_name, shard_name, shard_id)
167
- ("qkv_proj", "q_proj", "q"),
168
- ("qkv_proj", "k_proj", "k"),
169
- ("qkv_proj", "v_proj", "v"),
170
- ("gate_up_proj", "gate_proj", 0),
171
- ("gate_up_proj", "up_proj", 1),
172
- ]
173
- # Skip loading extra parameters for GPTQ/modelopt models.
174
- ignore_suffixes = (
175
- ".bias",
176
- "_bias",
177
- ".weight_scale",
178
- "_weight_scale",
179
- ".input_scale",
180
- "_input_scale",
181
- )
182
- params_dict = dict(self.named_parameters())
183
- loaded_params: set[str] = set()
184
- expert_params_mapping = self.get_expert_mapping()
185
- is_fused_expert = False
186
- base_layer = (
187
- "base_layer." if any(".base_layer." in name for name in params_dict) else ""
188
- )
189
- fused_expert_params_mapping = [
190
- (
191
- f"experts.routed_experts.{base_layer}w13_weight",
192
- "experts.gate_up_proj",
193
- 0,
194
- "w1",
195
- ),
196
- (
197
- f"experts.routed_experts.{base_layer}w2_weight",
198
- "experts.down_proj",
199
- 0,
200
- "w2",
201
- ),
202
- ]
203
- num_experts = self.config.num_experts
204
- for name, loaded_weight in weights:
205
- if "scale" in name or "zero_point" in name:
206
- name = maybe_remap_kv_scale_name(name, params_dict)
207
- if name is None:
208
- continue
209
-
210
- for param_name, weight_name, shard_id in stacked_params_mapping:
211
- if "experts.gate_up_proj" in name or "experts.down_proj" in name:
212
- is_fused_expert = True
213
- expert_params_mapping = fused_expert_params_mapping
214
-
215
- # Skip non-stacked layers and experts (experts handled below).
216
- if weight_name not in name:
217
- continue
218
- # We have mlp.experts[0].gate_proj in the checkpoint.
219
- # Since we handle the experts below in expert_params_mapping,
220
- # we need to skip here BEFORE we update the name, otherwise
221
- # name will be updated to mlp.experts[0].gate_up_proj, which
222
- # will then be updated below in expert_params_mapping
223
- # for mlp.experts[0].gate_gate_up_proj, which breaks load.
224
- if "mlp.experts" in name:
225
- continue
226
- name = name.replace(weight_name, param_name)
227
- # Skip loading extra parameters for GPTQ/modelopt models.
228
- if name.endswith(ignore_suffixes) and name not in params_dict:
229
- continue
230
- # Skip layers on other devices.
231
- if is_pp_missing_parameter(name, self):
232
- continue
233
- if name.endswith("scale"):
234
- # Remapping the name of FP8 kv-scale.
235
- name = maybe_remap_kv_scale_name(name, params_dict)
236
- if name is None:
237
- continue
238
- if name not in params_dict:
239
- continue
240
- param = params_dict[name]
241
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
242
- if weight_loader == default_weight_loader:
243
- weight_loader(param, loaded_weight)
244
- else:
245
- weight_loader(param, loaded_weight, shard_id)
246
- break
247
- else:
248
- is_expert_weight = False
249
- for mapping in expert_params_mapping:
250
- param_name, weight_name, expert_id, shard_id = mapping
251
- if weight_name not in name:
252
- continue
253
- # Anyway, this is an expert weight and should not be
254
- # attempted to load as other weights later
255
- is_expert_weight = True
256
- name_mapped = name.replace(weight_name, param_name)
257
- if is_pp_missing_parameter(name_mapped, self):
258
- continue
259
- if is_fused_expert:
260
- loaded_weight = loaded_weight.transpose(-1, -2) # no bias
261
- if "experts.gate_up_proj" in name:
262
- loaded_weight = loaded_weight.chunk(2, dim=-2)
263
- success_w1 = self.load_fused_expert_weights(
264
- name_mapped,
265
- params_dict,
266
- loaded_weight[0],
267
- "w1",
268
- num_experts,
269
- )
270
- success_w3 = self.load_fused_expert_weights(
271
- name_mapped,
272
- params_dict,
273
- loaded_weight[1],
274
- "w3",
275
- num_experts,
276
- )
277
- success = success_w1 and success_w3
278
- else:
279
- # down_proj
280
- success = self.load_fused_expert_weights(
281
- name_mapped,
282
- params_dict,
283
- loaded_weight,
284
- shard_id,
285
- num_experts,
286
- )
287
- else:
288
- # Skip loading extra parameters for GPTQ/modelopt models
289
- if (
290
- name_mapped.endswith(ignore_suffixes)
291
- and name_mapped not in params_dict
292
- ):
293
- continue
294
- param = params_dict[name_mapped]
295
- # We should ask the weight loader to return success or
296
- # not here since otherwise we may skip experts with
297
- # other available replicas.
298
- weight_loader = typing.cast(
299
- Callable[..., bool], param.weight_loader
300
- )
301
- success = weight_loader(
302
- param,
303
- loaded_weight,
304
- name_mapped,
305
- shard_id=shard_id,
306
- expert_id=expert_id,
307
- return_success=True,
308
- )
309
- if success:
310
- name = name_mapped
311
- break
312
- else:
313
- if is_expert_weight:
314
- # We've checked that this is an expert weight
315
- # However it's not mapped locally to this rank
316
- # So we simply skip it
317
- continue
318
- # Skip loading extra parameters for GPTQ/modelopt models.
319
- if name.endswith(ignore_suffixes) and name not in params_dict:
320
- continue
321
- # Skip layers on other devices.
322
- if is_pp_missing_parameter(name, self):
323
- continue
324
- if name not in params_dict:
325
- continue
326
- param = params_dict[name]
327
- weight_loader = getattr(
328
- param, "weight_loader", default_weight_loader
329
- )
330
- weight_loader(param, loaded_weight)
331
- loaded_params.add(name)
332
- return loaded_params
333
-
132
133
class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
134
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/siglip.py
+34
-126
@@ -33,10 +33,6 @@ from vllm.model_executor.layers.linear import (
33
from vllm.model_executor.layers.pooler import DispatchPooler
34
from vllm.model_executor.layers.quantization import QuantizationConfig
35
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
36
-from vllm.model_executor.model_loader.weight_utils import (
37
- default_weight_loader,
38
- maybe_remap_kv_scale_name,
39
-)
36
from vllm.multimodal import MULTIMODAL_REGISTRY
37
from vllm.multimodal.inputs import (
38
MultiModalFieldConfig,
@@ -62,7 +58,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
58
59
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
60
from .interfaces_base import default_pooling_type
65
-from .utils import AutoWeightsLoader, maybe_prefix
61
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
62
from .vision import (
63
VisionEncoderInfo,
64
VisionFeatureSelectStrategy,
@@ -571,6 +567,14 @@ class SiglipEncoder(nn.Module):
567
568
569
class SiglipTextTransformer(nn.Module):
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__(
579
self,
580
config: SiglipTextConfig,
@@ -615,30 +619,8 @@ class SiglipTextTransformer(nn.Module):
619
return last_hidden_state
620
621
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
618
- stacked_params_mapping = [
619
- # (param_name, shard_name, shard_id)
620
- ("qkv_proj", "q_proj", "q"),
621
- ("qkv_proj", "k_proj", "k"),
622
- ("qkv_proj", "v_proj", "v"),
623
- ]
624
- params_dict = dict(self.named_parameters())
625
- loaded_params: set[str] = set()
626
-
627
- for name, loaded_weight in weights:
628
- for param_name, weight_name, shard_id in stacked_params_mapping:
629
- if weight_name not in name:
630
- continue
631
- name = name.replace(weight_name, param_name)
632
- param = params_dict[name]
633
- weight_loader = param.weight_loader
634
- weight_loader(param, loaded_weight, shard_id)
635
- break
636
- else:
637
- param = params_dict[name]
638
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
639
- weight_loader(param, loaded_weight)
640
- loaded_params.add(name)
641
- return loaded_params
622
+ loader = AutoWeightsLoader(self)
623
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
624
625
626
class SiglipMultiheadAttentionPoolingHead(nn.Module):
@@ -682,6 +664,14 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):
664
665
666
class SiglipVisionTransformer(nn.Module):
667
+ hf_to_vllm_mapper = WeightsMapper(
668
+ orig_to_new_stacked={
669
+ ".q_proj": (".qkv_proj", "q"),
670
+ ".k_proj": (".qkv_proj", "k"),
671
+ ".v_proj": (".qkv_proj", "v"),
672
+ }
673
+ )
674
+
675
def __init__(
676
self,
677
config: SiglipVisionConfig,
@@ -806,46 +796,25 @@ class SiglipVisionTransformer(nn.Module):
796
return encoder_outputs
797
798
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
809
- stacked_params_mapping = [
810
- # (param_name, shard_name, shard_id)
811
- ("qkv_proj", "q_proj", "q"),
812
- ("qkv_proj", "k_proj", "k"),
813
- ("qkv_proj", "v_proj", "v"),
814
- ]
815
- params_dict = dict(self.named_parameters())
816
- loaded_params: set[str] = set()
817
- layer_count = len(self.encoder.layers)
818
-
819
- for name, loaded_weight in weights:
820
- # post_layernorm is not needed in SiglipVisionTransformer
821
- if name.startswith("post_layernorm") and self.post_layernorm is None:
822
- continue
799
+ skip_prefixes = []
800
+ if self.post_layernorm is None:
801
+ skip_prefixes.append("post_layernorm.")
802
+ if self.head is None:
803
+ skip_prefixes.append("head.")
804
+ loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
805
824
- # if the model configuration is not going to use
825
- # the pooling head for inference, don't load its weights
826
- if self.head is None and name.startswith("head"):
827
- continue
806
+ layer_count = len(self.encoder.layers)
807
829
- # omit layers when num_hidden_layers_override is set
830
- if name.startswith("encoder.layers"):
831
- layer_idx = int(name.split(".")[2])
832
- if layer_idx >= layer_count:
808
+ def _filter(ws):
809
+ for name, w in ws:
810
+ # omit layers when num_hidden_layers_override is set
811
+ if name.startswith("encoder.layers.") and (
812
+ int(name.split(".")[2]) >= layer_count
813
+ ):
814
continue
815
+ yield name, w
816
835
- for param_name, weight_name, shard_id in stacked_params_mapping:
836
- if weight_name not in name:
837
- continue
838
- name = name.replace(weight_name, param_name)
839
- param = params_dict[name]
840
- weight_loader = param.weight_loader
841
- weight_loader(param, loaded_weight, shard_id)
842
- break
843
- else:
844
- param = params_dict[name]
845
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
846
- weight_loader(param, loaded_weight)
847
- loaded_params.add(name)
848
- return loaded_params
817
+ return loader.load_weights(_filter(weights), mapper=self.hf_to_vllm_mapper)
818
819
820
class SiglipVisionModel(nn.Module):
@@ -896,67 +865,6 @@ class SiglipVisionModel(nn.Module):
865
feature_select_strategy=feature_select_strategy,
866
)
867
899
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
900
- stacked_params_mapping = [
901
- # (param_name, shard_name, shard_id)
902
- ("qkv_proj", "q_proj", "q"),
903
- ("qkv_proj", "k_proj", "k"),
904
- ("qkv_proj", "v_proj", "v"),
905
- ]
906
- params_dict = dict(self.named_parameters())
907
- loaded_params: set[str] = set()
908
- layer_count = len(self.vision_model.encoder.layers)
909
-
910
- for name, loaded_weight in weights:
911
- # post_layernorm is optional in SiglipVisionModel
912
- if (
913
- name.startswith("vision_model.post_layernorm")
914
- and self.vision_model.post_layernorm is None
915
- ):
916
- continue
917
-
918
- # if the model configuration is not going to use
919
- # the pooling head for inference, don't load its weights
920
- if self.vision_model.head is None and name.startswith("vision_model.head"):
921
- continue
922
-
923
- # omit layers when num_hidden_layers_override is set
924
- if name.startswith("vision_model.encoder.layers"):
925
- layer_idx = int(name.split(".")[3])
926
- if layer_idx >= layer_count:
927
- continue
928
-
929
- # Check if this is a scale parameter that needs remapping first
930
- if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
931
- # Try to remap the scale name first
932
- remapped_name = maybe_remap_kv_scale_name(name, params_dict)
933
- if remapped_name is not None and remapped_name in params_dict:
934
- # Successfully remapped, use the remapped name
935
- param = params_dict[remapped_name]
936
- weight_loader = getattr(
937
- param, "weight_loader", default_weight_loader
938
- )
939
- weight_loader(param, loaded_weight)
940
- loaded_params.add(remapped_name)
941
- continue
942
- # If remapping failed, continue with normal processing
943
-
944
- for param_name, weight_name, shard_id in stacked_params_mapping:
945
- if weight_name not in name:
946
- continue
947
- name = name.replace(weight_name, param_name)
948
-
949
- param = params_dict[name]
950
- weight_loader = param.weight_loader
951
- weight_loader(param, loaded_weight, shard_id)
952
- break
953
- else:
954
- param = params_dict[name]
955
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
956
- weight_loader(param, loaded_weight)
957
- loaded_params.add(name)
958
- return loaded_params
959
-
868
869
# Adapted from: https://github.com/huggingface/transformers/blob/v4.54.1/src/transformers/models/siglip/modeling_siglip.py#L200
870
class SiglipTextEmbeddings(nn.Module):
vllm/model_executor/models/siglip2navit.py
+11
-27
@@ -26,10 +26,9 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
26
from vllm.model_executor.layers.rotary_embedding.common import (
27
ApplyRotaryEmb,
28
)
29
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
29
from vllm.platforms import current_platform
30
32
-from .utils import maybe_prefix
31
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
32
from .vision import is_vit_use_data_parallel
33
34
@@ -588,6 +587,14 @@ class Siglip2VisionTransformer(nn.Module):
587
588
589
class Siglip2NavitModel(torch.nn.Module):
590
+ hf_to_vllm_mapper = WeightsMapper(
591
+ orig_to_new_stacked={
592
+ ".q_proj": (".qkv_proj", "q"),
593
+ ".k_proj": (".qkv_proj", "k"),
594
+ ".v_proj": (".qkv_proj", "v"),
595
+ }
596
+ )
597
+
598
def __init__(
599
self,
600
config: Siglip2VisionConfig,
@@ -613,28 +620,5 @@ class Siglip2NavitModel(torch.nn.Module):
620
)
621
622
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
616
- stacked_params_mapping = [
617
- # (param_name, shard_name, shard_id)
618
- ("qkv_proj", "q_proj", "q"),
619
- ("qkv_proj", "k_proj", "k"),
620
- ("qkv_proj", "v_proj", "v"),
621
- ]
622
- params_dict = dict(self.named_parameters())
623
- loaded_params: set[str] = set()
624
-
625
- for name, loaded_weight in weights:
626
- for param_name, weight_name, shard_id in stacked_params_mapping:
627
- if weight_name not in name:
628
- continue
629
- name = name.replace(weight_name, param_name)
630
-
631
- param = params_dict[name]
632
- weight_loader = param.weight_loader
633
- weight_loader(param, loaded_weight, shard_id)
634
- break
635
- else:
636
- param = params_dict[name]
637
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
638
- weight_loader(param, loaded_weight)
639
- loaded_params.add(name)
640
- return loaded_params
623
+ loader = AutoWeightsLoader(self)
624
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/transformers/moe.py
+15
-47
@@ -16,7 +16,6 @@
16
# limitations under the License.
17
"""Transformers modeling backend mixin for Mixture of Experts (MoE) models."""
18
19
-from collections.abc import Iterable
19
from dataclasses import dataclass
20
from functools import partial
21
from typing import TYPE_CHECKING, Any
@@ -82,11 +81,6 @@ class TransformersMoERunner(MoERunner):
81
) -> torch.Tensor:
82
return super().forward(hidden_states, topk_weights)
83
85
- def load_weights(
86
- self, weights: Iterable[tuple[str, torch.Tensor]]
87
- ) -> Iterable[str]:
88
- return self.routed_experts.load_weights(weights)
89
-
84
85
def _transformers_moe_forward(
86
hidden_states: torch.Tensor,
@@ -119,6 +113,20 @@ direct_register_custom_op(
113
)
114
115
116
+class TransformersRoutedExperts(RoutedExperts):
117
+ def get_expert_mapping(
118
+ self, include_fused: bool = False
119
+ ) -> list[tuple[str, str, int, str]]:
120
+ common_names = ("gate_proj", "down_proj", "up_proj")
121
+ common_map = super().get_expert_mapping(*common_names, include_fused)
122
+ mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused)
123
+ if not include_fused:
124
+ return common_map + mixtral_map
125
+ common_fused, common_unfused = common_map[:3], common_map[3:]
126
+ mixtral_fused, mixtral_unfused = mixtral_map[:3], mixtral_map[3:]
127
+ return common_fused + mixtral_fused + common_unfused + mixtral_unfused
128
+
129
+
130
class MoEMixin(MixtureOfExperts):
131
def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
132
self.check_version("5.0.0", "MoE models support")
@@ -140,43 +148,6 @@ class MoEMixin(MixtureOfExperts):
148
mlp.n_redundant_experts = self.num_redundant_experts
149
mlp.experts.update_expert_map()
150
143
- def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
144
- """
145
- Params for weights, fp8 weight scales, fp8 activation scales
146
- (param_name, weight_name, expert_id, shard_id)
147
- """
148
- # Models saved with fused experts. These are checkpoints released:
149
- # - After Transformers v5
150
- # - Before Transformers v5, but re-saved with save_original_format=False
151
- # In the fused experts case, we repurpose the expert_id as shard_idx for
152
- # deconcatenating w1 and w3 in FusedMoE.load_weights.
153
- expert_mapping = [
154
- ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"),
155
- ("experts.w13_weight", "experts.gate_up_proj", 1, "w3"),
156
- ("experts.w2_weight", "experts.down_proj", 0, "w2"),
157
- ]
158
- # Models saved with ModuleList experts
159
- ckpt_names = [
160
- # (ckpt_gate_proj_name, ckpt_down_proj_name, ckpt_up_proj_name)
161
- ("gate_proj", "down_proj", "up_proj"), # Most common MoE style
162
- ("w1", "w2", "w3"), # Granite, Mixtral, Phi MoE style
163
- ]
164
- num_experts = self.model_config.get_num_experts()
165
- num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
166
- for gate_proj, down_proj, up_proj in ckpt_names:
167
- expert_mapping.extend(
168
- RoutedExperts.make_expert_params_mapping(
169
- self,
170
- ckpt_gate_proj_name=gate_proj,
171
- ckpt_down_proj_name=down_proj,
172
- ckpt_up_proj_name=up_proj,
173
- num_experts=num_experts,
174
- num_redundant_experts=num_redundant_experts,
175
- routed_experts_prefix="",
176
- )
177
- )
178
- return expert_mapping
179
-
151
def recursive_replace(self):
152
"""Initialize the MoE layers."""
153
text_config = self.text_config
@@ -218,9 +189,6 @@ class MoEMixin(MixtureOfExperts):
189
if "gptoss" in wrapped_arch:
190
activation = "swigluoai"
191
221
- # Expert mapping for `AutoWeightsLoader`
222
- expert_mapping = self.get_expert_mapping()
223
-
192
# Expert parallel load balancing kwargs
193
enable_eplb = self.parallel_config.enable_eplb
194
num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
@@ -308,12 +276,12 @@ class MoEMixin(MixtureOfExperts):
276
enable_eplb=enable_eplb,
277
num_redundant_experts=num_redundant_experts,
278
has_bias=has_bias,
311
- expert_mapping=expert_mapping,
279
custom_routing_function=partial(
280
custom_routing_function,
281
moe_state=moe_state,
282
),
283
runner_cls=TransformersMoERunner,
284
+ routed_experts_cls=TransformersRoutedExperts,
285
runner_args={"moe_state": moe_state},
286
)
287
mlp.experts = fused_experts
vllm/model_executor/models/zamba2.py
+6
-30
@@ -43,7 +43,6 @@ 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 HasInnerState, IsHybrid, SupportsMambaPrefixCaching
@@ -804,34 +803,6 @@ class Zamba2Model(nn.Module):
803
hidden_states = self.final_layernorm(hidden_states)
804
return hidden_states
805
807
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
808
- stacked_params_mapping = [
809
- # (param_name, shard_name, shard_id)
810
- ("qkv_proj", "q_proj", "q"),
811
- ("qkv_proj", "k_proj", "k"),
812
- ("qkv_proj", "v_proj", "v"),
813
- ]
814
-
815
- params_dict = dict(self.named_parameters())
816
- loaded_params: set[str] = set()
817
- for chkpt_weight_name, loaded_weight in weights:
818
- for param_name, weight_name, shard_id in stacked_params_mapping:
819
- if weight_name not in chkpt_weight_name:
820
- continue
821
- chkpt_weight_name = chkpt_weight_name.replace(weight_name, param_name)
822
- param = params_dict[chkpt_weight_name]
823
- weight_loader = param.weight_loader
824
- weight_loader(param, loaded_weight, shard_id)
825
- break
826
- else:
827
- if chkpt_weight_name not in params_dict:
828
- continue
829
- param = params_dict[chkpt_weight_name]
830
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
831
- weight_loader(param, loaded_weight)
832
- loaded_params.add(chkpt_weight_name)
833
- return loaded_params
834
-
806
807
class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching):
808
"""Zamba2 model with causal language modeling head.
@@ -849,7 +820,12 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixC
820
"A_log": "A",
821
"0.weight": "A.weight",
822
"1.weight": "B.weight",
852
- }
823
+ },
824
+ orig_to_new_stacked={
825
+ ".self_attn.q_proj": (".self_attn.qkv_proj", "q"),
826
+ ".self_attn.k_proj": (".self_attn.qkv_proj", "k"),
827
+ ".self_attn.v_proj": (".self_attn.qkv_proj", "v"),
828
+ },
829
)
830
831
@classmethod
vllm/model_executor/utils.py
+10
-9
@@ -124,15 +124,16 @@ def get_packed_modules_mapping(model: torch.nn.Module) -> dict[str, list[str]]:
124
def get_moe_expert_mapping(
125
model: torch.nn.Module,
126
) -> list[tuple[str, str, int, str]]:
127
- if parent_map := getattr(model, "get_expert_mapping", None):
128
- return parent_map()
129
- else:
130
- # We only check main components instead of whole model submodules
131
- for child in model.children():
132
- child_map = getattr(child, "get_expert_mapping", None)
133
- if child_map is not None:
134
- return child_map()
135
- return []
127
+ """Get the expert mapping from a model.
128
+
129
+ It will be retrieved from the first module that has a `get_expert_mapping` method.
130
+ If the model manually implements `get_expert_mapping`, it will be used.
131
+ Otherwise, it will use the first RoutedExperts layer."""
132
+ for _, module in model.named_modules():
133
+ get_mapping = getattr(module, "get_expert_mapping", None)
134
+ if get_mapping is not None:
135
+ return get_mapping()
136
+ raise ValueError("No module in the model has a `get_expert_mapping` method.")
137
138
139
def maybe_disable_graph_partition(current_backend: str) -> dict[str, bool]:
vllm/models/deepseek_v4/nvidia/model.py
+1
-2
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import typing
4
-from collections.abc import Callable, Iterable, MutableSequence, Sequence
4
+from collections.abc import Callable, Iterable
5
from itertools import islice
6
7
import regex as re
@@ -1365,7 +1365,6 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts):
1365
self.set_moe_parameters()
1366
1367
def set_moe_parameters(self) -> None:
1368
- self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = []
1368
self.num_expert_groups = getattr(self.config, "n_group", 1)
1369
self.num_moe_layers = self.config.num_hidden_layers
1370
self.moe_layers: list[nn.Module] = []