[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>

vllm/model_executor/model_loader/weight_utils.py

@@ -1648,8 +1648,10 @@ def maybe_remap_moe_expert_param_name(
 
     Checkpoint weights have names like:
         layers.0.mlp.experts.w13_weight
+        layers.0.feed_forward.experts.w2_input_scale
     But actual parameters are now:
         layers.0.mlp.experts.routed_experts.w13_weight
+        layers.0.feed_forward.experts.routed_experts.w2_input_scale
 
     This function inserts 'routed_experts.' into the path when needed.
 
@@ -1662,11 +1664,11 @@ def maybe_remap_moe_expert_param_name(
         otherwise the original name
     """
     # Only remap if this looks like an expert parameter
-    if ".mlp.experts." not in name:
+    if ".experts." not in name:
         return name
 
     # Skip if already has routed_experts
-    if ".mlp.experts.routed_experts." in name:
+    if ".experts.routed_experts." in name:
         return name
 
     # Expert parameter patterns to check
@@ -1700,8 +1702,8 @@ def maybe_remap_moe_expert_param_name(
     if not is_expert_param:
         return name
 
-    # Try inserting routed_experts
-    new_name = name.replace(".mlp.experts.", ".mlp.experts.routed_experts.")
+    # Try inserting routed_experts after .experts.
+    new_name = name.replace(".experts.", ".experts.routed_experts.", 1)
 
     # Only use the new name if it exists in the model
     if new_name in params_dict:

vllm/model_executor/models/lfm2_moe.py

@@ -40,7 +40,10 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_moe_expert_param_name,
+)
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.lfm2_moe import Lfm2MoeConfig
 
@@ -572,6 +575,7 @@ class Lfm2MoeModel(nn.Module):
                     # Skip layers on other devices.
                     if is_pp_missing_parameter(name, self):
                         continue
+                    name = maybe_remap_moe_expert_param_name(name, params_dict)
                     param = params_dict[name]
                     weight_loader = getattr(
                         param, "weight_loader", default_weight_loader

vllm/model_executor/models/llama4.py

@@ -51,6 +51,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    maybe_remap_moe_expert_param_name,
 )
 from vllm.model_executor.models.interfaces import MixtureOfExperts
 from vllm.model_executor.models.utils import sequence_parallel_chunk
@@ -662,6 +663,7 @@ class Llama4Model(LlamaModel):
                 if "experts." in name and any(
                     scale_name in name for scale_name in scale_names
                 ):
+                    name = maybe_remap_moe_expert_param_name(name, params_dict)
                     param = params_dict[name]
                     weight_loader = getattr(
                         param, "weight_loader", default_weight_loader

vllm/model_executor/models/mllama4.py

@@ -52,7 +52,10 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.model_loader.utils import initialize_model
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_moe_expert_param_name,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
@@ -1021,6 +1024,7 @@ class Llama4ForConditionalGeneration(
                 and "scale" in name
                 and ".shared_expert" not in name
             ):
+                name = maybe_remap_moe_expert_param_name(name, params_dict)
                 if name in params_dict:
                     param = params_dict[name]
                     if (