fix generation-batch-size&steps_per_generation check (#8048)

Author: hjh0119

swift/llm/argument/deploy_args.py

@@ -115,19 +115,11 @@ class RolloutArguments(DeployArguments):
     context_manager: Optional[str] = None
 
     def __post_init__(self):
-        self._check_trl_version()
         self._set_default_engine_type()
         super().__post_init__()
         self._check_args()
         self._check_device_count()
 
-    def _check_trl_version(self):
-        try:
-            from trl.scripts.vllm_serve import WeightSyncWorkerExtension
-        except ImportError as e:
-            raise ImportError("Could not import 'WeightSyncWorkerExtension' from 'trl.scripts.vllm_serve'. "
-                              "Please upgrade your 'trl' package by 'pip install -U trl'") from e
-
     def _set_default_engine_type(self):
         if self.vllm_use_async_engine is None:
             if self.multi_turn_scheduler:

swift/llm/argument/rlhf_args.py

@@ -451,14 +451,13 @@ def _check_grpo(self):
 
         import trl
         trl_version = version.parse(trl.__version__)
-        assert trl_version >= version.parse('0.17'), ('Your current version of `trl` is outdated. '
+        assert trl_version >= version.parse('0.20'), ('Your current version of `trl` is outdated. '
                                                       'Please update it by running: pip install -U trl')
         if is_mp() and self.use_vllm:
             raise ValueError('GRPO with vLLM is not compatible with `device_map`. '
                              'Please set NPROC_PER_NODE equal to num_processes.')
         if self.use_liger_kernel:
             liger_kernel_version = version.parse(importlib.metadata.version('liger-kernel'))
-            assert trl_version >= version.parse('0.18')
             if self.delta is not None:
                 raise ValueError('Liger loss does not support two-sided GRPO loss yet.')
             if self.sequence_parallel_size > 1:
@@ -485,12 +484,6 @@ def _check_grpo(self):
         if self.async_generate and self.multi_turn_scheduler is not None:
             raise NotImplementedError('Currently, async_generate is not supported with multi-turn functionality.')
 
-        if self.generation_batch_size or self.steps_per_generation:
-            from trl.trainer.grpo_config import GRPOConfig
-            assert 'generation_batch_size' in GRPOConfig.__dict__, (
-                'generation_batch_size or steps_per_generation needs trl >= 0.18, '
-                'please install trl `pip install trl>=0.18')
-
     def _external_vllm_warning(self):
         if self.rlhf_type not in rlhf_support_vllm_types or not self.vllm_server_host:
             return

swift/trainers/rlhf_arguments.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
+from transformers.utils.versions import require_version
 from trl import CPOConfig as HfCPOConfig
 from trl import DPOConfig as HfDPOConfig
 from trl import GKDConfig as HfGKDConfig
@@ -58,6 +59,7 @@ def __post_init__(self):
 class GRPOConfig(GRPOArgumentsMixin, SwiftArgumentsMixin, HfGRPOConfig):
 
     def __post_init__(self):
+        require_version('trl>=0.20')
         GRPOArgumentsMixin.__post_init__(self)
         SwiftArgumentsMixin.__post_init__(self)
         if self.vllm_reasoning_parser is not None:
@@ -75,25 +77,6 @@ def __post_init__(self):
         # https://github.com/modelscope/ms-swift/issues/3863
         self.dataloader_drop_last = True
 
-        # from trl https://github.com/huggingface/trl/blob/7a39ff3995f2f8b7cb4f8ca29a09390ac587a43d/trl/trainer/grpo_config.py#L843 # noqa: E501
-        num_processes = self.world_size
-        # The current default effective batch size
-        if self.generation_batch_size is None and self.steps_per_generation is None:
-            self.steps_per_generation = self.gradient_accumulation_steps
-            self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation
-        elif self.generation_batch_size is not None and self.steps_per_generation is None:
-            # Just ensure the value is divisible by the global batch size
-            if self.generation_batch_size % (self.per_device_train_batch_size * num_processes) != 0:
-                raise ValueError(
-                    f'generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size '
-                    f'({self.per_device_train_batch_size * num_processes}).')
-            self.steps_per_generation = self.generation_batch_size // (self.per_device_train_batch_size * num_processes)
-        elif self.generation_batch_size is None and self.steps_per_generation is not None:
-            self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation
-        else:
-            raise ValueError(
-                "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time")
-
         self.check_num_generations()
 
     def check_num_generations(self):