megatron

Author: hjh0119

swift/megatron/trainers/grpo_trainer.py

@@ -43,7 +43,7 @@
 
 class MegatronGRPOTrainer(MegatronRolloutMixin, MegatronRLHFTrainer):
 
-    def __init__(self, args: MegatronRLHFArguments, template: Template, **kwargs):
+    def __init__(self, args: MegatronArguments, template: Template, **kwargs):
         self.vllm_client = kwargs.pop('vllm_client')
         super().__init__(args, template)
         self.args = args
@@ -145,14 +145,7 @@ def _prepare_rewards(self):
             for i, reward_func in enumerate(reward_funcs):
                 if reward_func in orms:
                     reward_func_class = orms[reward_func]
-                    reward_func_args = list(inspect.signature(reward_func_class.__init__).parameters)
-                    reward_func_kwargs = {
-                        key: getattr(args, key)
-                        for key in reward_func_args if key not in ['self', 'args', 'kwargs'] and hasattr(args, key)
-                    }
-                    if 'tokenizer' in reward_func_args:
-                        reward_func_kwargs['tokenizer'] = self.processing_class
-                    reward_funcs[i] = reward_func_class(**reward_func_kwargs)
+                    reward_funcs[i] = reward_func_class(args=self.args)
                 elif not callable(reward_func):
                     raise ValueError(f'reward_function {reward_func} is not implemented in swift.rewards')
 

swift/rewards/orm.py

@@ -9,6 +9,7 @@
 from swift.infer_engine import InferRequest
 
 if TYPE_CHECKING:
+    from swift.megatron.arguments import MegatronArguments
     from swift.rlhf_trainers import GRPOConfig
 
 
@@ -23,7 +24,7 @@ def __call__(self, completions, **kwargs) -> List[float]:
                 return [1.0 if len(c) > 100 else 0.0 for c in completions]
     """
 
-    def __init__(self, args: Optional['GRPOConfig'] = None, **kwargs):
+    def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
         self.args = args
 
     def __call__(self, **kwargs) -> List[float]:
@@ -58,7 +59,7 @@ async def score_single(session, text):
                     return list(rewards)
     """
 
-    def __init__(self, args: Optional['GRPOConfig'] = None, **kwargs):
+    def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
         self.args = args
 
     async def __call__(self, **kwargs) -> List[float]:
@@ -139,7 +140,7 @@ def __call__(self, completions, **kwargs) -> List[float]:
 
 class CosineReward(ORM):
     # https://arxiv.org/abs/2502.03373
-    def __init__(self, args: Optional['GRPOConfig'] = None, accuracy_orm=None):
+    def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, accuracy_orm=None):
         super().__init__(args)
         self.min_len_value_wrong = args.cosine_min_len_value_wrong
         self.max_len_value_wrong = args.cosine_max_len_value_wrong
@@ -174,7 +175,7 @@ def __call__(self, completions, solution, **kwargs) -> List[float]:
 
 class RepetitionPenalty(ORM):
     # https://arxiv.org/abs/2502.03373
-    def __init__(self, args: Optional['GRPOConfig'] = None, **kwargs):
+    def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
         super().__init__(args)
         self.ngram_size = args.repetition_n_grams
         self.max_penalty = args.repetition_max_penalty
@@ -214,7 +215,7 @@ def __call__(self, completions, **kwargs) -> List[float]:
 
 class SoftOverlong(ORM):
 
-    def __init__(self, args: Optional['GRPOConfig'] = None, **kwargs):
+    def __init__(self, args: Optional[Union['GRPOConfig', 'MegatronArguments']] = None, **kwargs):
         super().__init__(args)
         assert args.soft_cache_length < args.soft_max_length
         self.soft_max_length = args.soft_max_length