Address review comments

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

tests/utils/test_vllm_weight_name_normalization_on_cpu.py

@@ -31,6 +31,12 @@ class _FakeModel:
     def __init__(self):
         self.hf_to_vllm_mapper = _FakeMapper(
             {
+                "model.language_model.layers.0.mlp.experts.base_layer.w13_weight": (
+                    "language_model.model.layers.0.mlp.experts.base_layer.w13_weight"
+                ),
+                "model.language_model.layers.0.mlp.experts.base_layer.w2_weight": (
+                    "language_model.model.layers.0.mlp.experts.base_layer.w2_weight"
+                ),
                 "model.language_model.layers.0.self_attn.qkv_proj.base_layer.weight": (
                     "language_model.model.layers.0.self_attn.qkv_proj.base_layer.weight"
                 ),
@@ -89,6 +95,23 @@ def test_normalize_base_sync_weight_names_handles_bridge_inserted_base_layer_on_
     ]
 
 
+def test_normalize_base_sync_weight_names_handles_fused_expert_leaf_params():
+    worker = _make_worker(_FakeModel())
+    tensor = torch.empty(0)
+
+    normalized_weights = worker._normalize_base_sync_weight_names(
+        [
+            ("model.language_model.layers.0.mlp.experts.w13_weight", tensor),
+            ("model.language_model.layers.0.mlp.experts.base_layer.w2_weight", tensor),
+        ]
+    )
+
+    assert [name for name, _ in normalized_weights] == [
+        "model.language_model.layers.0.mlp.experts.base_layer.w13_weight",
+        "model.language_model.layers.0.mlp.experts.base_layer.w2_weight",
+    ]
+
+
 def test_update_weights_from_ipc_accumulates_lora_tensors_across_buckets(monkeypatch):
     import verl.workers.rollout.vllm_rollout.bucketed_weight_transfer as bucketed_weight_transfer
 

verl/utils/megatron_utils.py

@@ -1345,10 +1345,16 @@ def check_mtp_config(model_config: HFModelConfig, engine_config: McoreEngineConf
     Check and configure MTP (Multi-Token Prediction) settings.
 
     Cases:
-        - mtp.enable == False and no MTP layers: return directly
-        - mtp.enable == False and has MTP layers: set num_nextn_predict_layers = 0
-        - mtp.enable == True and has MTP layers: configure override_transformer_config
-        - mtp.enable == True and no MTP layers: raise ValueError
+        - mtp.enable == False and neither ``num_nextn_predict_layers`` nor
+          ``mtp_num_hidden_layers`` is enabled on ``hf_config`` /
+          ``hf_config.text_config``: return directly.
+        - mtp.enable == False and MTP layers are configured: zero the first
+          supported MTP layer-count field (``num_nextn_predict_layers`` when
+          present, otherwise ``mtp_num_hidden_layers``).
+        - mtp.enable == True and MTP layers are configured: keep the existing
+          layer counts and populate ``override_transformer_config`` as needed.
+        - mtp.enable == True and no MTP layers are configured: raise
+          ``ValueError``.
     """
     text_hf_config = getattr(model_config.hf_config, "text_config", model_config.hf_config)
     has_mtp = (

verl/workers/megatron_workers.py

@@ -157,7 +157,7 @@ def _init_hf_config_and_tf_config(
             assert (
                 getattr(text_hf_config, "num_nextn_predict_layers", 0) > 0
                 or getattr(text_hf_config, "mtp_num_hidden_layers", 0) > 0
-            ), "MTP requires at least one nextn_predict_layer"
+            ), "MTP requires at least one MTP layer (num_nextn_predict_layers or mtp_num_hidden_layers)"
             assert megatron_config.use_mbridge, "MTP requires use_mbridge to be True"
             override_transformer_config["mtp_loss_scaling_factor"] = self.config.model.mtp.mtp_loss_scaling_factor
         else:

verl/workers/rollout/vllm_rollout/utils.py

@@ -195,7 +195,8 @@ def _iter_model_weight_name_candidates(weight_name: str):
 
     @staticmethod
     def _is_leaf_weight_or_bias_name(weight_name: str) -> bool:
-        return weight_name.rsplit(".", 1)[-1] in {"weight", "bias"}
+        leaf = weight_name.rsplit(".", 1)[-1]
+        return leaf in {"weight", "bias"} or leaf.endswith(("_weight", "_bias"))
 
     @classmethod
     def _strip_bridge_base_layer_from_expert_alias(cls, weight_name: str) -> str: