From fa8c868a3c64994ac0f8b7a0e0bc6d575e378567 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Wed, 10 Jun 2026 13:40:45 -0400 Subject: [PATCH] [Bugfix] Fix Llama4 weight loading (#45047) Signed-off-by: Tyler Michael Smith Co-authored-by: Claude Opus 4.6 (1M context) --- vllm/model_executor/model_loader/weight_utils.py | 10 ++++++---- vllm/model_executor/models/lfm2_moe.py | 6 +++++- vllm/model_executor/models/llama4.py | 2 ++ vllm/model_executor/models/mllama4.py | 6 +++++- 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index 28423edeea..dd96e15261 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/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: diff --git a/vllm/model_executor/models/lfm2_moe.py b/vllm/model_executor/models/lfm2_moe.py index 55b00d2b9e..9ca7fb7aaa 100644 --- a/vllm/model_executor/models/lfm2_moe.py +++ b/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 diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py index 07ca2714ed..277848fb86 100644 --- a/vllm/model_executor/models/llama4.py +++ b/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 diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index 8fe1be721c..742dccc36f 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/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 (