[Model] Migrate MistralLarge3ForCausalLM to AutoWeightsLoader (#48153)
Signed-off-by: Yuchen Fan <functionhx@gmail.com>
FAN YUCHEN committed
Jul 10, 2026 at 20:27 UTC
fabec87f63cdf0725d48e3eb6fdef9799837bcaf
2 files changed
+89
-57
vllm/model_executor/models/mistral_large_3.py
+80
-51
@@ -2,62 +2,91 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable
4
5
-import regex as re
5
+import regex
6
import torch
7
8
from vllm.model_executor.models.deepseek_v2 import DeepseekV3ForCausalLM
9
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
10
11
12
class MistralLarge3ForCausalLM(DeepseekV3ForCausalLM):
12
- # fmt: off
13
- remapping = {
14
- r"layers\.(\d+)\.attention_norm\.weight": r"model.layers.\1.input_layernorm.weight", # noqa: E501
15
- r"layers\.(\d+)\.attention\.wq_a\.(\w+)": r"model.layers.\1.self_attn.q_a_proj.\2", # noqa: E501
16
- r"layers\.(\d+)\.attention\.q_a_norm\.weight": r"model.layers.\1.self_attn.q_a_layernorm.weight", # noqa: E501
17
- r"layers\.(\d+)\.attention\.wq_b\.(\w+)": r"model.layers.\1.self_attn.q_b_proj.\2", # noqa: E501
18
- r"layers\.(\d+)\.attention\.wkv_a_with_mqa\.(\w+)": r"model.layers.\1.self_attn.kv_a_proj_with_mqa.\2", # noqa: E501
19
- r"layers\.(\d+)\.attention\.kv_a_norm\.weight": r"model.layers.\1.self_attn.kv_a_layernorm.weight", # noqa: E501
20
- r"layers\.(\d+)\.attention\.wkv_b\.(\w+)": r"model.layers.\1.self_attn.kv_b_proj.\2", # noqa: E501
21
- r"layers\.(\d+)\.attention\.wo\.(\w+)": r"model.layers.\1.self_attn.o_proj.\2", # noqa: E501
22
- r"layers\.(\d+)\.ffn_norm\.weight": r"model.layers.\1.post_attention_layernorm.weight", # noqa: E501
23
- r"layers\.(\d+)\.feed_forward\.w1\.(\w+)": r"model.layers.\1.mlp.gate_proj.\2", # noqa: E501
24
- r"layers\.(\d+)\.feed_forward\.w2\.(\w+)": r"model.layers.\1.mlp.down_proj.\2", # noqa: E501
25
- r"layers\.(\d+)\.feed_forward\.w3\.(\w+)": r"model.layers.\1.mlp.up_proj.\2", # noqa: E501
26
- r"layers\.(\d+)\.gate\.weight": r"model.layers.\1.mlp.gate.weight", # noqa: E501
27
- r"layers\.(\d+)\.shared_experts\.w1\.(\w+)": r"model.layers.\1.mlp.shared_experts.gate_proj.\2", # noqa: E501
28
- r"layers\.(\d+)\.shared_experts\.w2\.(\w+)": r"model.layers.\1.mlp.shared_experts.down_proj.\2", # noqa: E501
29
- r"layers\.(\d+)\.shared_experts\.w3\.(\w+)": r"model.layers.\1.mlp.shared_experts.up_proj.\2", # noqa: E501
30
- r"layers\.(\d+)\.experts\.(\d+)\.w1\.(\w+)": r"model.layers.\1.mlp.experts.\2.gate_proj.\3", # noqa: E501
31
- r"layers\.(\d+)\.experts\.(\d+)\.w2\.(\w+)": r"model.layers.\1.mlp.experts.\2.down_proj.\3", # noqa: E501
32
- r"layers\.(\d+)\.experts\.(\d+)\.w3\.(\w+)": r"model.layers.\1.mlp.experts.\2.up_proj.\3", # noqa: E501
33
- r"norm\.weight": "model.norm.weight", # noqa: E501
34
- r"tok_embeddings\.weight": "model.embed_tokens.weight", # noqa: E501
35
- r"output\.weight": "lm_head.weight", # noqa: E501
36
- }
37
- # fmt: on
13
+ # WeightsMapper applies all matching patterns sequentially (no break on first
14
+ # match). This is safe here because every pattern is anchored at both ends
15
+ # (\A...\Z) and after substitution the resulting key always starts with
16
+ # "model." or "lm_head.", so no later pattern can accidentally match again.
17
+ hf_to_vllm_mapper = WeightsMapper(
18
+ orig_to_new_regex={ # noqa: B950
19
+ regex.compile(
20
+ r"\Alayers\.(\d+)\.attention_norm\.weight\Z"
21
+ ): r"model.layers.\1.input_layernorm.weight",
22
+ regex.compile(
23
+ r"\Alayers\.(\d+)\.attention\.wq_a\.(\w+)\Z"
24
+ ): r"model.layers.\1.self_attn.q_a_proj.\2",
25
+ regex.compile(
26
+ r"\Alayers\.(\d+)\.attention\.q_a_norm\.weight\Z"
27
+ ): r"model.layers.\1.self_attn.q_a_layernorm.weight",
28
+ regex.compile(
29
+ r"\Alayers\.(\d+)\.attention\.wq_b\.(\w+)\Z"
30
+ ): r"model.layers.\1.self_attn.q_b_proj.\2",
31
+ regex.compile(
32
+ r"\Alayers\.(\d+)\.attention\.wkv_a_with_mqa\.(\w+)\Z"
33
+ ): r"model.layers.\1.self_attn.kv_a_proj_with_mqa.\2",
34
+ regex.compile(
35
+ r"\Alayers\.(\d+)\.attention\.kv_a_norm\.weight\Z"
36
+ ): r"model.layers.\1.self_attn.kv_a_layernorm.weight",
37
+ regex.compile(
38
+ r"\Alayers\.(\d+)\.attention\.wkv_b\.(\w+)\Z"
39
+ ): r"model.layers.\1.self_attn.kv_b_proj.\2",
40
+ regex.compile(
41
+ r"\Alayers\.(\d+)\.attention\.wo\.(\w+)\Z"
42
+ ): r"model.layers.\1.self_attn.o_proj.\2",
43
+ regex.compile(
44
+ r"\Alayers\.(\d+)\.ffn_norm\.weight\Z"
45
+ ): r"model.layers.\1.post_attention_layernorm.weight",
46
+ regex.compile(
47
+ r"\Alayers\.(\d+)\.feed_forward\.w1\.(\w+)\Z"
48
+ ): r"model.layers.\1.mlp.gate_proj.\2",
49
+ regex.compile(
50
+ r"\Alayers\.(\d+)\.feed_forward\.w2\.(\w+)\Z"
51
+ ): r"model.layers.\1.mlp.down_proj.\2",
52
+ regex.compile(
53
+ r"\Alayers\.(\d+)\.feed_forward\.w3\.(\w+)\Z"
54
+ ): r"model.layers.\1.mlp.up_proj.\2",
55
+ regex.compile(
56
+ r"\Alayers\.(\d+)\.gate\.weight\Z"
57
+ ): r"model.layers.\1.mlp.gate.weight",
58
+ regex.compile(
59
+ r"\Alayers\.(\d+)\.shared_experts\.w1\.(\w+)\Z"
60
+ ): r"model.layers.\1.mlp.shared_experts.gate_proj.\2",
61
+ regex.compile(
62
+ r"\Alayers\.(\d+)\.shared_experts\.w2\.(\w+)\Z"
63
+ ): r"model.layers.\1.mlp.shared_experts.down_proj.\2",
64
+ regex.compile(
65
+ r"\Alayers\.(\d+)\.shared_experts\.w3\.(\w+)\Z"
66
+ ): r"model.layers.\1.mlp.shared_experts.up_proj.\2",
67
+ regex.compile(
68
+ r"\Alayers\.(\d+)\.experts\.(\d+)\.w1\.(\w+)\Z"
69
+ ): r"model.layers.\1.mlp.experts.\2.gate_proj.\3",
70
+ regex.compile(
71
+ r"\Alayers\.(\d+)\.experts\.(\d+)\.w2\.(\w+)\Z"
72
+ ): r"model.layers.\1.mlp.experts.\2.down_proj.\3",
73
+ regex.compile(
74
+ r"\Alayers\.(\d+)\.experts\.(\d+)\.w3\.(\w+)\Z"
75
+ ): r"model.layers.\1.mlp.experts.\2.up_proj.\3",
76
+ regex.compile(r"\Anorm\.weight\Z"): "model.norm.weight",
77
+ regex.compile(r"\Atok_embeddings\.weight\Z"): "model.embed_tokens.weight",
78
+ regex.compile(r"\Aoutput\.weight\Z"): "lm_head.weight",
79
+ },
80
+ orig_to_new_suffix={
81
+ ".qscale_act": ".input_scale",
82
+ ".qscale_weight": ".weight_scale",
83
+ },
84
+ )
85
86
+ # Bypass super().load_weights() and construct AutoWeightsLoader(self)
87
+ # directly (same pattern as Qwen2ForCausalLM). Any logic in the parent
88
+ # class's load_weights is a thin wrapper around AutoWeightsLoader, and
89
+ # we must apply hf_to_vllm_mapper before the loader walks the tree.
90
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
40
- return super().load_weights(map(self._remap_mistral_to_ds, weights))
41
-
42
- def _remap_mistral_to_ds(
43
- self, weight: tuple[str, torch.Tensor]
44
- ) -> tuple[str, torch.Tensor]:
45
- """Remap Mistral parameters to DeepseekV2 parameters."""
46
- name, loaded_weight = weight
47
-
48
- for k, v in self.remapping.items():
49
- match = re.fullmatch(k, name)
50
- if match:
51
- name = re.sub(k, v, name)
52
- break
53
- else:
54
- raise ValueError(f"Cannot remap {name}")
55
-
56
- # Remapping scale names. We could do this in the regex above but it
57
- # would triple the number of lines for most layers.
58
- if name.endswith(".qscale_act"):
59
- name = re.sub(r"\.qscale_act$", ".input_scale", name)
60
- elif name.endswith(".qscale_weight"):
61
- name = re.sub(r"\.qscale_weight$", ".weight_scale", name)
62
-
63
- return name, loaded_weight
91
+ loader = AutoWeightsLoader(self)
92
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/mistral_large_3_eagle.py
+9
-6
@@ -5,6 +5,7 @@ import copy
5
from collections.abc import Iterable
6
from functools import partial
7
8
+import regex
9
import torch
10
import torch.nn as nn
11
@@ -22,7 +23,7 @@ from vllm.model_executor.models.deepseek_v2 import (
23
from vllm.model_executor.models.mistral_large_3 import MistralLarge3ForCausalLM
24
25
from .interfaces import SupportsMultiModal
25
-from .utils import make_empty_intermediate_tensors_factory, maybe_prefix
26
+from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix
27
28
logger = init_logger(__name__)
29
@@ -107,11 +108,13 @@ class EagleMistralLarge3Model(DeepseekV2Model):
108
109
110
class EagleMistralLarge3ForCausalLM(MistralLarge3ForCausalLM):
110
- remapping = MistralLarge3ForCausalLM.remapping | {
111
- r"eagle_linear\.weight": r"model.fc.weight",
112
- r"eagle_linear\.qscale_act": r"model.fc.input_scale",
113
- r"eagle_linear\.qscale_weight": r"model.fc.weight_scale",
114
- }
111
+ hf_to_vllm_mapper = MistralLarge3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
112
+ orig_to_new_regex={
113
+ regex.compile(r"\Aeagle_linear\.weight\Z"): r"model.fc.weight",
114
+ regex.compile(r"\Aeagle_linear\.qscale_act\Z"): r"model.fc.input_scale",
115
+ regex.compile(r"\Aeagle_linear\.qscale_weight\Z"): r"model.fc.weight_scale",
116
+ },
117
+ )
118
119
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
120
target_layer_num = vllm_config.model_config.get_num_layers(