[BugFix] LoRA: Support loading base_layer of experts

Signed-off-by: Hollow Man <[email protected]>

Author: HollowMan6

vllm/model_executor/layers/fused_moe/layer.py

@@ -42,6 +42,9 @@
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     is_flashinfer_supporting_global_sf,
 )
+from vllm.model_executor.model_loader.weight_utils import (
+    remap_expert_weight_name,
+)
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 from vllm.utils.math_utils import cdiv, round_up
@@ -1376,7 +1379,9 @@ def load_weights(
             for param_name, weight_name, expert_id, shard_id in expert_mapping:
                 if weight_name not in qual_name:
                     continue
-                weight_name = qual_name.replace(weight_name, param_name)
+                weight_name = remap_expert_weight_name(
+                    qual_name, weight_name, param_name
+                )
                 param_name = weight_name.removeprefix(f"{self.layer_name}.")
                 param = getattr(self, param_name)
                 success = self.weight_loader(

vllm/model_executor/model_loader/weight_utils.py

@@ -1178,3 +1178,38 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
 
     # If there were no matches, return the untouched param name
     return name
+
+
+def remap_expert_weight_name(
+    name: str,
+    weight_name: str,
+    param_name: str,
+) -> str:
+    """Remap expert weight names, handling base_layer prefix for LoRA.
+
+    When loading expert weights, this function maps from checkpoint weight
+    names to model parameter names. It handles the special case where
+    LoRA wraps the original layer with a `base_layer` prefix.
+
+    For example:
+        - Input: name="model.layers.0.mlp.experts.0.up_proj.base_layer.weight"
+                 weight_name="experts.0.up_proj."
+                 param_name="experts.w13_"
+        - Output: "model.layers.0.mlp.experts.base_layer.w13_weight"
+
+    Args:
+        name: The full checkpoint weight name.
+        weight_name: The weight name pattern to match (e.g., "experts.0.up_proj.").
+        param_name: The parameter name to substitute (e.g., "experts.w13_").
+
+    Returns:
+        The remapped weight name with proper base_layer handling.
+    """
+    prefix, _, suffix = name.partition(weight_name)
+    middle = param_name
+    base = "base_layer"
+    if suffix.startswith(f"{base}."):
+        param_list = param_name.split(".", 1)
+        param_list.insert(1, base)
+        middle = ".".join(param_list)
+    return prefix + middle + suffix.removeprefix(f"{base}.")

vllm/model_executor/models/afmoe.py

@@ -36,6 +36,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
 from vllm.model_executor.models.llama import LlamaMLP as AfmoeMLP
@@ -533,7 +534,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
 
                     # Do not modify `name` since the loop may continue here
                     # Instead, create a new variable
-                    name_mapped = name.replace(weight_name, param_name)
+                    # Remap expert weight name (handles base_layer prefix)
+                    name_mapped = remap_expert_weight_name(
+                        name, weight_name, param_name
+                    )
 
                     if is_pp_missing_parameter(name_mapped, self):
                         continue

vllm/model_executor/models/arctic.py

@@ -38,7 +38,10 @@
     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,
+    remap_expert_weight_name,
+)
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -609,7 +612,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     for param_name, weight_name, shard_id in expert_params_mapping:
                         if weight_name not in name:
                             continue
-                        name = name.replace(weight_name, param_name)
+                        # Remap expert weight name (handles base_layer prefix)
+                        name = remap_expert_weight_name(name, weight_name, param_name)
                         if is_pp_missing_parameter(name, self):
                             continue
                         param = params_dict[name]

vllm/model_executor/models/bailing_moe.py

@@ -55,7 +55,10 @@
     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,
+    remap_expert_weight_name,
+)
 from vllm.sequence import IntermediateTensors
 
 from .interfaces import SupportsLoRA, SupportsPP
@@ -524,7 +527,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     param_name, weight_name, expert_id, shard_id = mapping
                     if weight_name not in name:
                         continue
-                    name = name.replace(weight_name, param_name)
+                    # Remap expert weight name (handles base_layer prefix)
+                    name = remap_expert_weight_name(name, weight_name, param_name)
 
                     if is_pp_missing_parameter(name, self):
                         continue

vllm/model_executor/models/dbrx.py

@@ -31,6 +31,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.sequence import IntermediateTensors
 
@@ -411,7 +412,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             for param_name, weight_name in expert_params_mapping:
                 if weight_name not in name:
                     continue
-                name = name.replace(weight_name, param_name)
+                # Remap expert weight name (handles base_layer prefix)
+                name = remap_expert_weight_name(name, weight_name, param_name)
                 if is_pp_missing_parameter(name, self):
                     continue
                 param = params_dict[name]

vllm/model_executor/models/deepseek_eagle.py

@@ -18,6 +18,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.model_executor.models.deepseek_v2 import (
     DeepseekV2DecoderLayer,
@@ -155,7 +156,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     param_name, weight_name, expert_id, shard_id = mapping
                     if weight_name not in name:
                         continue
-                    name = name.replace(weight_name, param_name)
+                    # Remap expert weight name (handles base_layer prefix)
+                    name = remap_expert_weight_name(name, weight_name, param_name)
 
                     param = params_dict[name]
                     weight_loader = param.weight_loader

vllm/model_executor/models/deepseek_mtp.py

@@ -22,6 +22,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -357,9 +358,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         # attempted to load as other weights later
                         is_expert_weight = True
 
-                        # Do not modify `name` since the loop may continue here
-                        # Instead, create a new variable
-                        name_mapped = chunk_name.replace(weight_name, param_name)
+                        # Remap expert weight name (handles base_layer prefix)
+                        name_mapped = remap_expert_weight_name(
+                            chunk_name, weight_name, param_name
+                        )
 
                         param = params_dict[name_mapped]
                         # We should ask the weight loader to return success or

vllm/model_executor/models/deepseek_v2.py

@@ -72,6 +72,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.platforms import current_platform
@@ -1635,9 +1636,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                         # attempted to load as other weights later
                         is_expert_weight = True
 
-                        # Do not modify `name` since the loop may continue here
-                        # Instead, create a new variable
-                        name_mapped = chunk_name.replace(weight_name, param_name)
+                        # Remap expert weight name (handles base_layer prefix)
+                        name_mapped = remap_expert_weight_name(
+                            chunk_name, weight_name, param_name
+                        )
 
                         if is_pp_missing_parameter(name_mapped, self):
                             continue

vllm/model_executor/models/dots1.py

@@ -59,6 +59,7 @@
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    remap_expert_weight_name,
 )
 from vllm.sequence import IntermediateTensors
 
@@ -464,7 +465,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
                     param_name, weight_name, expert_id, shard_id = mapping
                     if weight_name not in name:
                         continue
-                    name = name.replace(weight_name, param_name)
+                    # Remap expert weight name (handles base_layer prefix)
+                    name = remap_expert_weight_name(name, weight_name, param_name)
 
                     if is_pp_missing_parameter(name, self):
                         continue