fix

Author: gzy19990617

fastdeploy/rl/rollout_model.py

@@ -50,7 +50,7 @@ def _init_model(self) -> nn.Layer:
         model.eval()
         return model
 
-    def get_name_mappings_to_training(self, trainer_degree=1) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         """Get parameter name mappings between rollout and training models."""
         return getattr(self.rollout_model, "get_name_mappings_to_training", lambda: {})(trainer_degree)
 
@@ -92,9 +92,6 @@ def _complete_missing_mappings(self) -> None:
                 # Skip weight scale parameters in mapping. Train and infer have same key.
                 self.infer_to_train_mapping[key] = key
 
-        if getattr(self.fd_config.model_config, "tie_word_embeddings", False):
-            self.infer_to_train_mapping.pop("lm_head.linear.weight")
-    
     def get_quantization_infer_keys(self) -> list[str]:
         """Get quantization infer keys"""
         quant_weight_key = []
@@ -125,7 +122,7 @@ def name(self) -> str:
         """name"""
         return "Ernie4_5_MoeForCausalLMRL"
 
-    def get_name_mappings_to_training(self, trainer_degree=1) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         """Generate mapping between inference and training parameter for RL(donot delete!)."""
         # Prepare placeholders
         place_holders = ["weight"]
@@ -192,7 +189,7 @@ def name(self) -> str:
         """name"""
         return "Ernie4_5_VLMoeForConditionalGenerationRL"
 
-    def get_name_mappings_to_training(self, trainer_degree=2) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         """Generate mapping between inference and training parameter for RL(donot delete!)."""
         # Prepare placeholders
         place_holders = ["weight"]
@@ -255,6 +252,8 @@ def _generate_ranges(start, end, step=16, take=8):
 
         assert isinstance(self.fd_config.model_config.moe_num_experts, list)
         total_moe_num = sum(self.fd_config.model_config.moe_num_experts)
+        if not trainer_degree:
+            trainer_degree = self.fd_config.parallel_config.tensor_parallel_size
         expert_num_per_rank = self.fd_config.model_config.moe_num_experts[0] // trainer_degree
         # Process MoE layers
         for layer_idx in range(text_moe_layer_start_index, text_moe_layer_end_index):
@@ -285,7 +284,7 @@ def name(self) -> str:
         """name"""
         return "Qwen2ForCausalLMRL"
 
-    def get_name_mappings_to_training(self, trainer_degree=1) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         """Generate mapping between inference and training parameter for RL(donot delete!)."""
         # Prepare placeholders
         place_holders = ["weight"]
@@ -327,7 +326,7 @@ def name(self) -> str:
         """name"""
         return "Qwen3MoeForCausalLMRL"
 
-    def get_name_mappings_to_training(self, trainer_degree=1) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         """Generate mapping between inference and training parameter for RL(donot delete!)."""
         # Prepare placeholders
         place_holders = ["weight"]
@@ -394,5 +393,5 @@ def name(self) -> str:
         """name"""
         return "Qwen3ForCausalLMRL"
 
-    def get_name_mappings_to_training(self, trainer_degree=1) -> Dict[str, str]:
+    def get_name_mappings_to_training(self, trainer_degree=None) -> Dict[str, str]:
         pass