A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix Llama4 weight loading (#45047)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Tyler Michael Smith committed
Jun 10, 2026 at 13:40 UTC
fa8c868a3c64994ac0f8b7a0e0bc6d575e378567
4 files changed
+18
-6
vllm/model_executor/model_loader/weight_utils.py
+6
-4
@@ -1648,8 +1648,10 @@ def maybe_remap_moe_expert_param_name(
1648
1649
Checkpoint weights have names like:
1650
layers.0.mlp.experts.w13_weight
1651
+ layers.0.feed_forward.experts.w2_input_scale
1652
But actual parameters are now:
1653
layers.0.mlp.experts.routed_experts.w13_weight
1654
+ layers.0.feed_forward.experts.routed_experts.w2_input_scale
1655
1656
This function inserts 'routed_experts.' into the path when needed.
1657
@@ -1662,11 +1664,11 @@ def maybe_remap_moe_expert_param_name(
1664
otherwise the original name
1665
"""
1666
# Only remap if this looks like an expert parameter
1665
- if ".mlp.experts." not in name:
1667
+ if ".experts." not in name:
1668
return name
1669
1670
# Skip if already has routed_experts
1669
- if ".mlp.experts.routed_experts." in name:
1671
+ if ".experts.routed_experts." in name:
1672
return name
1673
1674
# Expert parameter patterns to check
@@ -1700,8 +1702,8 @@ def maybe_remap_moe_expert_param_name(
1702
if not is_expert_param:
1703
return name
1704
1703
- # Try inserting routed_experts
1704
- new_name = name.replace(".mlp.experts.", ".mlp.experts.routed_experts.")
1705
+ # Try inserting routed_experts after .experts.
1706
+ new_name = name.replace(".experts.", ".experts.routed_experts.", 1)
1707
1708
# Only use the new name if it exists in the model
1709
if new_name in params_dict:
vllm/model_executor/models/lfm2_moe.py
+5
-1
@@ -40,7 +40,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
40
ParallelLMHead,
41
VocabParallelEmbedding,
42
)
43
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
43
+from vllm.model_executor.model_loader.weight_utils import (
44
+ default_weight_loader,
45
+ maybe_remap_moe_expert_param_name,
46
+)
47
from vllm.sequence import IntermediateTensors
48
from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig
49
@@ -572,6 +575,7 @@ class Lfm2MoeModel(nn.Module):
575
# Skip layers on other devices.
576
if is_pp_missing_parameter(name, self):
577
continue
578
+ name = maybe_remap_moe_expert_param_name(name, params_dict)
579
param = params_dict[name]
580
weight_loader = getattr(
581
param, "weight_loader", default_weight_loader
vllm/model_executor/models/llama4.py
+2
@@ -51,6 +51,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
51
from vllm.model_executor.model_loader.weight_utils import (
52
default_weight_loader,
53
maybe_remap_kv_scale_name,
54
+ maybe_remap_moe_expert_param_name,
55
)
56
from vllm.model_executor.models.interfaces import MixtureOfExperts
57
from vllm.model_executor.models.utils import sequence_parallel_chunk
@@ -662,6 +663,7 @@ class Llama4Model(LlamaModel):
663
if "experts." in name and any(
664
scale_name in name for scale_name in scale_names
665
):
666
+ name = maybe_remap_moe_expert_param_name(name, params_dict)
667
param = params_dict[name]
668
weight_loader = getattr(
669
param, "weight_loader", default_weight_loader
vllm/model_executor/models/mllama4.py
+5
-1
@@ -52,7 +52,10 @@ from vllm.model_executor.layers.linear import (
52
from vllm.model_executor.layers.quantization import QuantizationConfig
53
from vllm.model_executor.layers.rotary_embedding import get_rope
54
from vllm.model_executor.model_loader.utils import initialize_model
55
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
55
+from vllm.model_executor.model_loader.weight_utils import (
56
+ default_weight_loader,
57
+ maybe_remap_moe_expert_param_name,
58
+)
59
from vllm.model_executor.models.module_mapping import MultiModelKeys
60
from vllm.multimodal import MULTIMODAL_REGISTRY
61
from vllm.multimodal.inputs import (
@@ -1021,6 +1024,7 @@ class Llama4ForConditionalGeneration(
1024
and "scale" in name
1025
and ".shared_expert" not in name
1026
):
1027
+ name = maybe_remap_moe_expert_param_name(name, params_dict)
1028
if name in params_dict:
1029
param = params_dict[name]
1030
if (