refactor rl get_name_mappings_to_training (#2847)

* refactor rl get_name_mappings_to_training

* fix tp>1

* change variable name(ffn1->up_gate_proj/ffn2->down_proj)

* change variable name(linear_weight->weight/linear_bias->bias)

* add rl names mapping for vl

* fix ernie 0.3B error

* fix develop code

* fix

fastdeploy/rl/rollout_model.py

@@ -24,11 +24,14 @@ from fastdeploy.config import FDConfig
 from fastdeploy.model_executor.model_loader import ModelRegistry
 from fastdeploy.model_executor.models.ernie4_5_moe import \
     Ernie4_5_MoeForCausalLM
+from fastdeploy.model_executor.models.ernie4_5_vl.ernie4_5_vl_moe import \
+    Ernie4_5_VLMoeForConditionalGeneration
 from fastdeploy.model_executor.models.qwen2 import Qwen2ForCausalLM
 from fastdeploy.model_executor.models.qwen3 import Qwen3ForCausalLM
 from fastdeploy.model_executor.models.qwen3moe import Qwen3MoeForCausalLM
 from fastdeploy.rl.rollout_config import RolloutModelConfig
 
+
 class RolloutModel(nn.Layer):
     """Main model class for rollout operations, supports multimodal components for train."""
 
@@ -36,55 +39,26 @@ class RolloutModel(nn.Layer):
         """Initialize with FastDeploy configuration."""
         super(RolloutModel, self).__init__()
         self.fd_config = rollout_model_config.initialize()
-        self._init_models()
-
-    def _init_models(self):
-        """Initialize all model components including multimodal if needed."""
-        self.is_vl = "VL" in self.fd_config.model_config.architectures[0]
-        self.rollout_model = self._load_primary_model()
-        self.rollout_models = [self.rollout_model]
-
-        if self.is_vl:
-            self._init_multimodal_models()
-            self.rollout_models.extend(
-                [self.vision_model, self.resampler_model])
-
-    def _init_multimodal_models(self):
-        """Initialize vision and resampler components for multimodal models."""
-        # TODO:(gaoziyuan) Implement actual initialization
-        self.vision_model = nn.Layer()
-        self.resampler_model = nn.Layer()
-
-    def _load_primary_model(self):
-        """Load main model from loader based on config."""
-        if "VL" in self.fd_config.model_config.architectures[0]:
-            logger.error("Loaded Vision Language model, not support now")
+        self._init_model()
 
+    def _init_model(self):
+        """Load model from loader based on config."""
         context = paddle.LazyGuard()
         architectures = f"{self.fd_config.model_config.architectures[0]}RL"
         with context:
             model_cls = ModelRegistry.get_class(architectures)
             model = model_cls(self.fd_config)
 
-        model.eval()
-        return model
+        self.rollout_model = model.eval()
 
     def get_name_mappings_to_training(self) -> Dict[str, str]:
         """Get parameter name mappings between rollout and training models."""
-        mappings = {}
-        for model in self.rollout_models:
-            mappings.update(
-                getattr(model, "get_name_mappings_to_training", lambda: {})())
-        return mappings
+        return getattr(self.rollout_model, "get_name_mappings_to_training", lambda: {})()
 
     @paddle.no_grad()
     def state_dict(self):
         """state_dict"""
-        all_params = {}
-        for model in self.rollout_models:
-            for name, param in model.state_dict().items():
-                all_params[name] = param
-        return all_params
+        return self.rollout_model.state_dict()
 
 
 class Ernie4_5_MoeForCausalLMRL(Ernie4_5_MoeForCausalLM):
@@ -113,98 +87,159 @@ class Ernie4_5_MoeForCausalLMRL(Ernie4_5_MoeForCausalLM):
         # Initialize mapping dictionary
         infer_to_train = {}
 
-        infer_base_name = "model"
-        train_base_name = "ernie"
+        base_name = "ernie"
         # Static mappings (non-layer specific)
         static_mappings = {
-            f"{infer_base_name}.embeddings.word_embeddings.weight":
-            f"{train_base_name}.embed_tokens.weight",
-            f"{infer_base_name}.norm.ln_weight": f"{train_base_name}.norm.weight",
-            "lm_head.out_linear.weight": "lm_head.weight"
+            f"{base_name}.embed_tokens.embeddings.weight":
+            f"{base_name}.embed_tokens.weight",
+            "lm_head.linear.weight": "lm_head.weight"
         }
-        if self.fd_config.model_config.get("weight_sharing", False):
+        if self.fd_config.model_config.get("tie_word_embeddings", False):
             # Support tie_word_embeddings
             logger.debug("enable tie_word_embeddings")
-            static_mappings.pop("lm_head.out_linear.weight")
+            static_mappings.pop("lm_head.linear.weight")
         infer_to_train.update(static_mappings)
 
-        infer_base_name = infer_base_name + ".hidden_layers"
-        train_base_name = train_base_name + ".layers"
+        base_name = base_name + ".layers"
 
         # Helper function to add layer mappings
-        def _add_layer_mappings(layer_idx, is_moe_layer=False):
-            # Handle special case for layer 0's input layernorm
-            for ph in place_holders:
-                infer_key = f"{infer_base_name}.{layer_idx}.input_layernorm.ln_{ph}"
-                train_key = f"{train_base_name}.{layer_idx}.input_layernorm.{ph}"
-                infer_to_train[infer_key] = train_key
+        def _add_layer_mappings(layer_idx: int):
+            # MoE specific mappings
+            infer_to_train[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_weight"] = \
+                f"{base_name}.{layer_idx}.mlp.gate.weight"
 
-            # Common attention mappings
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.qkv_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.qkv_proj.{ph}"
+            if self.fd_config.model_config.moe_use_aux_free:
+                infer_to_train[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
+                    f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
 
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.o_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.o_proj.{ph}"
-
-            # Post-attention layernorm
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.post_attention_layernorm.ln_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.post_attention_layernorm.{ph}"
-
-            if not is_moe_layer:
-                # Dense FFN mappings
+            # MoE experts mappings
+            for expert_idx in range(self.fd_config.model_config.moe_num_experts):
                 for ph in place_holders:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.gate_up_proj.linear_{ph}"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.up_gate_proj.{ph}"
+                    # up_gate_proj (up_gate_proj)
+                    up_gate_proj_key = f"{base_name}.{layer_idx}.mlp.fused_moe.up_gate_proj_weight"
+                    if up_gate_proj_key not in infer_to_train:
+                        infer_to_train[up_gate_proj_key] = []
+                    infer_to_train[up_gate_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
+                    )
 
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.down_proj.linear_{ph}"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.down_proj.{ph}"
-            else:
-                # MoE specific mappings
-                infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.fused_moe.gate_weight"] = \
-                    f"{train_base_name}.{layer_idx}.mlp.gate.weight"
-
-                if self.fd_config.model_config.moe_use_aux_free:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
-
-                # Support shared experts
-                if self.fd_config.model_config.get(
-                        "moe_num_shared_experts") > 0:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.shared_experts.gate_up_proj.linear_weight"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.shared_experts.up_gate_proj.weight"
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.shared_experts.down_proj.linear_weight"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.shared_experts.down_proj.weight"
-
-                # MoE experts mappings
-                for expert_idx in range(self.fd_config.model_config.moe_num_experts):
-                    for ph in place_holders:
-                        # FFN1 (up_gate_proj)
-                        ffn1_key = f"{infer_base_name}.{layer_idx}.mlp.fused_moe.moe_ffn1_weight"
-                        if ffn1_key not in infer_to_train:
-                            infer_to_train[ffn1_key] = []
-                        infer_to_train[ffn1_key].append(
-                            f"{train_base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
-                        )
-
-                        # FFN2 (down_proj)
-                        ffn2_key = f"{infer_base_name}.{layer_idx}.mlp.fused_moe.moe_ffn2_weight"
-                        if ffn2_key not in infer_to_train:
-                            infer_to_train[ffn2_key] = []
-                        infer_to_train[ffn2_key].append(
-                            f"{train_base_name}.{layer_idx}.mlp.experts.{expert_idx}.down_proj.{ph}"
-                        )
-
-        # Process non-MoE layers
-        for layer_idx in range(
-                self.fd_config.model_config.moe_layer_start_index):
-            _add_layer_mappings(layer_idx, is_moe_layer=False)
+                    # down_proj (down_proj)
+                    down_proj_key = f"{base_name}.{layer_idx}.mlp.fused_moe.down_proj_weight"
+                    if down_proj_key not in infer_to_train:
+                        infer_to_train[down_proj_key] = []
+                    infer_to_train[down_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.down_proj.{ph}"
+                    )
 
+        assert isinstance(self.fd_config.model_config.moe_layer_start_index, int)
         # Process MoE layers
         for layer_idx in range(self.fd_config.model_config.moe_layer_start_index,
                                self.fd_config.model_config.num_hidden_layers):
-            _add_layer_mappings(layer_idx, is_moe_layer=True)
+            _add_layer_mappings(layer_idx)
+
+        return infer_to_train
+
+
+class Ernie4_5_VLMoeForConditionalGenerationRL(Ernie4_5_VLMoeForConditionalGeneration):
+    """
+    Ernie4_5_VLMoeForConditionalGenerationRL
+    """
+
+    def __init__(self, fd_config: FDConfig):
+        """
+        Args:
+            fd_config (FDConfig): Configurations for the LLM model.
+        """
+        super(Ernie4_5_VLMoeForConditionalGenerationRL, self).__init__(fd_config)
+
+    @classmethod
+    def name(self):
+        """name"""
+        return "Ernie4_5_VLMoeForConditionalGenerationRL"
+
+    def get_name_mappings_to_training(self):
+        """Generate mapping between inference and training parameter for RL(donot delete!)."""
+        have_bias = self.fd_config.model_config.get("have_norm_bias", False)
+        # Prepare placeholders
+        place_holders = ["weight"] + (["bias"] if have_bias else [])
+
+        # Initialize mapping dictionary
+        infer_to_train = {}
+
+        base_name = "ernie"
+        # Static mappings (non-layer specific)
+        static_mappings = {
+            f"{base_name}.embed_tokens.embeddings.weight":
+            f"{base_name}.embed_tokens.weight",
+            "lm_head.linear.weight": "lm_head.weight"
+        }
+        if self.fd_config.model_config.get("tie_word_embeddings", False):
+            # Support tie_word_embeddings
+            logger.debug("enable tie_word_embeddings")
+            static_mappings.pop("lm_head.linear.weight")
+        infer_to_train.update(static_mappings)
+
+        base_name = base_name + ".layers"
+
+        # Helper function to add layer mappings
+        def _add_layer_mappings(layer_idx: int, moe_tag: str):
+            # MoE specific mappings
+            infer_to_train[f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_weight"] = f"{base_name}.{layer_idx}.mlp.gate.weight" if moe_tag == "text" else f"{base_name}.{layer_idx}.mlp.gate.weight_1"
+
+            if self.fd_config.model_config.moe_use_aux_free:
+                infer_to_train[f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.gate_correction_bias"] = \
+                    f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
+
+            # MoE experts mappings
+            assert isinstance(self.fd_config.model_config.moe_num_experts, list)
+            if moe_tag == "text":
+                expert_idx_start = 0
+                expert_idx_end = self.fd_config.model_config.moe_num_experts[0]
+            else:
+                expert_idx_start = self.fd_config.model_config.moe_num_experts[0]
+                expert_idx_end = self.fd_config.model_config.moe_num_experts[1]
+
+            for expert_idx in range(expert_idx_start, expert_idx_end):
+                for ph in place_holders:
+                    # up_gate_proj (up_gate_proj)
+                    up_gate_proj_key = f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.up_gate_proj_weight"
+                    if up_gate_proj_key not in infer_to_train:
+                        infer_to_train[up_gate_proj_key] = []
+                    infer_to_train[up_gate_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
+                    )
+
+                    # down_proj (down_proj)
+                    down_proj_key = f"{base_name}.{layer_idx}.mlp.{moe_tag}_fused_moe.down_proj_weight"
+                    if down_proj_key not in infer_to_train:
+                        infer_to_train[down_proj_key] = []
+                    infer_to_train[down_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.down_proj.{ph}"
+                    )
+
+        moe_layer_start_index = self.fd_config.model_config.moe_layer_start_index
+        if isinstance(moe_layer_start_index, int):
+            text_moe_layer_start_index = moe_layer_start_index
+            image_moe_layer_start_index = moe_layer_start_index
+        else:
+            text_moe_layer_start_index = moe_layer_start_index[0]
+            image_moe_layer_start_index = moe_layer_start_index[1]
+
+        moe_layer_end_index = self.fd_config.model_config.moe_layer_end_index
+        if moe_layer_end_index is None:
+            text_moe_layer_end_index = self.fd_config.model_config.num_hidden_layers
+            image_moe_layer_end_index = self.fd_config.model_config.num_hidden_layers
+        elif isinstance(moe_layer_end_index, int):
+            text_moe_layer_end_index = moe_layer_end_index
+            image_moe_layer_end_index = moe_layer_end_index
+        else:
+            text_moe_layer_end_index = moe_layer_end_index[0]
+            image_moe_layer_end_index = moe_layer_end_index[1]
+        # Process MoE layers
+        for layer_idx in range(text_moe_layer_start_index, text_moe_layer_end_index):
+            _add_layer_mappings(layer_idx, "text")
+        for layer_idx in range(image_moe_layer_start_index, image_moe_layer_end_index):
+            _add_layer_mappings(layer_idx, "image")
 
         return infer_to_train
 
@@ -234,48 +269,23 @@ class Qwen2ForCausalLMRL(Qwen2ForCausalLM):
         # Initialize mapping dictionary
         infer_to_train = {}
 
-        infer_base_name = "model"
-        train_base_name = "qwen2"
+        base_name = "qwen2"
         # Static mappings (non-layer specific)
         static_mappings = {
-            f"{infer_base_name}.embeddings.word_embeddings.weight":
-            f"{train_base_name}.embed_tokens.weight",
-            f"{infer_base_name}.norm.ln_weight": f"{train_base_name}.norm.weight",
-            "lm_head.out_linear.weight": "lm_head.weight"
+            f"{base_name}.embed_tokens.embeddings.weight":
+            f"{base_name}.embed_tokens.weight",
+            "lm_head.linear.weight": "lm_head.weight"
         }
         infer_to_train.update(static_mappings)
 
-        infer_base_name = infer_base_name + ".layers"
-        train_base_name = train_base_name + ".layers"
+        base_name = base_name + ".layers"
 
         # Helper function to add layer mappings
         def _add_layer_mappings(layer_idx):
-            # Handle special case for layer 0's input layernorm and attn o_proj
-            for ph in place_holders:
-                infer_key = f"{infer_base_name}.{layer_idx}.input_layernorm.ln_{ph}"
-                train_key = f"{train_base_name}.{layer_idx}.input_layernorm.{ph}"
-                infer_to_train[infer_key] = train_key
-
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.o_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.o_proj.{ph}"
-
-            # qwen qkv proj need bias
-            for ph in ["weight", "bias"]:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.qkv_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.qkv_proj.{ph}"
-
-            # Post-attention layernorm
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.post_attention_layernorm.ln_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.post_attention_layernorm.{ph}"
-
             # FFN mappings
             for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.gate_up_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.mlp.gate_up_fused_proj.{ph}"
-
-                infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.down_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.mlp.down_proj.{ph}"
+                infer_to_train[f"{base_name}.{layer_idx}.mlp.up_gate_proj.{ph}"] = \
+                    f"{base_name}.{layer_idx}.mlp.gate_up_fused_proj.{ph}"
 
         for layer_idx in range(
                 self.fd_config.model_config.num_hidden_layers):
@@ -309,95 +319,49 @@ class Qwen3MoeForCausalLMRL(Qwen3MoeForCausalLM):
         # Initialize mapping dictionary
         infer_to_train = {}
 
-        infer_base_name = "model"
-        train_base_name = "model"
+        base_name = "model"
         # Static mappings (non-layer specific)
         static_mappings = {
-            f"{infer_base_name}.embeddings.word_embeddings.weight":
-            f"{train_base_name}.embed_tokens.weight",
-            f"{infer_base_name}.norm.ln_weight": f"{train_base_name}.norm.weight",
-            "lm_head.out_linear.weight": "lm_head.weight"
+            f"{base_name}.embed_tokens.embeddings.weight":
+            f"{base_name}.embed_tokens.weight",
+            "lm_head.linear.weight": "lm_head.weight"
         }
         infer_to_train.update(static_mappings)
 
-        infer_base_name = infer_base_name + ".layers"
-        train_base_name = train_base_name + ".layers"
+        base_name = base_name + ".layers"
 
         # Helper function to add layer mappings
-        def _add_layer_mappings(layer_idx, is_moe_layer=False):
-            # Handle special case for layer 0's input layernorm and attn o_proj
-            for ph in place_holders:
-                infer_key = f"{infer_base_name}.{layer_idx}.input_layernorm.ln_{ph}"
-                train_key = f"{train_base_name}.{layer_idx}.input_layernorm.{ph}"
-                infer_to_train[infer_key] = train_key
+        def _add_layer_mappings(layer_idx: int):
+            # MoE specific mappings
+            infer_to_train[f"{base_name}.{layer_idx}.mlp.gate_weight"] = \
+                f"{base_name}.{layer_idx}.mlp.gate.weight"
 
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.o_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.o_proj.{ph}"
+            if self.fd_config.moe_config.moe_use_aux_free:
+                infer_to_train[f"{base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
+                    f"{base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
 
-            # qwen q_norm/k_norm
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.q_norm.ln_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.q_norm.{ph}"
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.k_norm.ln_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.k_norm.{ph}"
-
-            # qwen qkv proj
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.self_attn.qkv_proj.linear_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.self_attn.qkv_proj.{ph}"
-
-            # Post-attention layernorm
-            for ph in place_holders:
-                infer_to_train[f"{infer_base_name}.{layer_idx}.post_attention_layernorm.ln_{ph}"] = \
-                    f"{train_base_name}.{layer_idx}.post_attention_layernorm.{ph}"
-
-            if not is_moe_layer:
-                # FFN mappings
+            # MoE experts mappings
+            for expert_idx in range(self.fd_config.moe_config.num_experts):
                 for ph in place_holders:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.gate_up_proj.linear_{ph}"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.gate_up_fused_proj.{ph}"
+                    # up_gate_proj (up_gate_proj)
+                    up_gate_proj_key = f"{base_name}.{layer_idx}.mlp.up_gate_proj_weight"
+                    if up_gate_proj_key not in infer_to_train:
+                        infer_to_train[up_gate_proj_key] = []
+                    infer_to_train[up_gate_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
+                    )
 
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.down_proj.linear_{ph}"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.down_proj.{ph}"
-            else:
-                # MoE specific mappings
-                infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.gate_weight"] = \
-                    f"{train_base_name}.{layer_idx}.mlp.gate.weight"
-
-                if self.fd_config.moe_config.moe_use_aux_free:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.fused_moe.gate_correction_bias"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.moe_statics.e_score_correction_bias"
-
-                # Support shared experts
-                if self.fd_config.model_config.get(
-                        "moe_num_shared_experts", 0) > 0:
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.shared_experts.gate_up_proj.linear_weight"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.shared_experts.up_gate_proj.weight"
-                    infer_to_train[f"{infer_base_name}.{layer_idx}.mlp.shared_experts.down_proj.linear_weight"] = \
-                        f"{train_base_name}.{layer_idx}.mlp.shared_experts.down_proj.weight"
-
-                # MoE experts mappings
-                for expert_idx in range(self.fd_config.moe_config.num_experts):
-                    for ph in place_holders:
-                        # FFN1 (up_gate_proj)
-                        ffn1_key = f"{infer_base_name}.{layer_idx}.mlp.moe_ffn1_weight"
-                        if ffn1_key not in infer_to_train:
-                            infer_to_train[ffn1_key] = []
-                        infer_to_train[ffn1_key].append(
-                            f"{train_base_name}.{layer_idx}.mlp.experts.{expert_idx}.up_gate_proj.{ph}"
-                        )
-
-                        # FFN2 (down_proj)
-                        ffn2_key = f"{infer_base_name}.{layer_idx}.mlp.moe_ffn2_weight"
-                        if ffn2_key not in infer_to_train:
-                            infer_to_train[ffn2_key] = []
-                        infer_to_train[ffn2_key].append(
-                            f"{train_base_name}.{layer_idx}.mlp.experts.{expert_idx}.down_proj.{ph}"
-                        )
+                    # down_proj (down_proj)
+                    down_proj_key = f"{base_name}.{layer_idx}.mlp.down_proj_weight"
+                    if down_proj_key not in infer_to_train:
+                        infer_to_train[down_proj_key] = []
+                    infer_to_train[down_proj_key].append(
+                        f"{base_name}.{layer_idx}.mlp.experts.{expert_idx}.down_proj.{ph}"
+                    )
 
         # Process MoE layers
         for layer_idx in range(self.fd_config.model_config.num_hidden_layers):
-            _add_layer_mappings(layer_idx, is_moe_layer=True)
+            _add_layer_mappings(layer_idx)
 
         return infer_to_train
 
@@ -417,4 +381,4 @@ class Qwen3ForCausalLMRL(Qwen3ForCausalLM):
     @classmethod
     def name(self):
         """name"""
-        return "Qwen3ForCausalLMRL"
+        return "Qwen3ForCausalLMRL"