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)
 
@@ -125,7 +125,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 +192,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 +255,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 +287,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 +329,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 +396,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