wip

swift/llm/argument/rlhf_args.py

@@ -211,8 +211,15 @@ def _check_rlhf(self):
     def _check_grpo(self):
         if self.rlhf_type != 'grpo':
             return
-
         from packaging import version
+
+        if self.use_liger_loss:
+            import liger_kernel
+            liger_kernel_version = version.parse(liger_kernel.__version__)
+            assert liger_kernel_version >= version.parse('0.5.8'), (
+                'Your current version of `liger-kernel` is outdated. '
+                'Please update it by running: pip install -U liger-kernel')
+
         import trl
         trl_version = version.parse(trl.__version__)
         assert trl_version >= version.parse('0.17'), ('Your current version of `trl` is outdated. '
@@ -278,10 +285,3 @@ def _external_vllm_warning(self):
                 "Configuration conflict: 'vllm_max_model_len=%s' is ignored for external vLLM. "
                 'Please specify it when launching the inference service: '
                 '`swift deploy --max_model_len <value>`', self.vllm_max_model_len)
-
-        if self.use_liger_loss:
-            assert self.mini_batch_size is None, 'liger loss is not compatible with mini batch currently'
-            try:
-                from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss
-            except ImportError:
-                raise ImportError('liger_kernel is not available. Run `pip install -U liger-kernel`.')

swift/trainers/rlhf_trainer/grpo_trainer.py

@@ -166,7 +166,8 @@ def __init__(self,
         else:
             self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32)
 
-        self.num_generations = args.num_generations
+        self.num_generations = args.num_generations  # = G in the GRPO paper
+        self.max_completion_length = args.max_completion_length  # = |o_i| in the GRPO paper
         self.temperature = args.temperature
         self.loss_type = args.loss_type
         model.warnings_issued['estimate_tokens'] = True
@@ -175,6 +176,8 @@ def __init__(self,
 
         use_vllm = args.use_vllm
         use_lmdeploy = args.use_lmdeploy
+
+        # we initialize vllm_client in RLHFArguments._init_external_vllm (swift/llm/rlhf_args)
         vllm_client = kwargs.pop('vllm_client')  # for external vllm
         self.use_vllm = args.use_vllm
         self.use_lmdeploy = args.use_lmdeploy
@@ -209,6 +212,8 @@ def __init__(self,
                 epsilon_high=self.epsilon_high,
                 temperature=self.temperature,
                 use_ref_model=self.ref_model is not None,
+                loss_type=self.loss_type,
+                max_completion_length=self.max_completion_length,
             )
             self._forward_redirection = _ForwardRedirection()