wehub-resource-sync
124 lines
4.0 KiB
Python
124 lines
4.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import trl
|
|
from dataclasses import dataclass
|
|
from packaging import version
|
|
from transformers.utils.versions import require_version
|
|
|
|
if version.parse(trl.__version__) <= version.parse('0.28'):
|
|
from trl import CPOConfig as HfCPOConfig
|
|
from trl import GKDConfig as HfGKDConfig
|
|
from trl import ORPOConfig as HfORPOConfig
|
|
from trl import PPOConfig as HfPPOConfig
|
|
else:
|
|
from trl.experimental.cpo import CPOConfig as HfCPOConfig
|
|
from trl.experimental.gkd import GKDConfig as HfGKDConfig
|
|
from trl.experimental.orpo import ORPOConfig as HfORPOConfig
|
|
from trl.experimental.ppo import PPOConfig as HfPPOConfig
|
|
|
|
from trl import DPOConfig as HfDPOConfig
|
|
from trl import GRPOConfig as HfGRPOConfig
|
|
from trl import KTOConfig as HfKTOConfig
|
|
from trl import RewardConfig as HfRewardConfig
|
|
from typing import Optional
|
|
|
|
from swift.trainers import TrainArgumentsMixin
|
|
from .args_mixin import GRPOArgumentsMixin, RolloutTrainerArgumentsMixin
|
|
|
|
|
|
@dataclass
|
|
class DPOConfig(TrainArgumentsMixin, HfDPOConfig):
|
|
ld_alpha: Optional[float] = None # compat trl==0.15
|
|
# Fields removed in trl 0.29, kept here for backward compatibility
|
|
rpo_alpha: Optional[float] = None
|
|
ref_adapter_name: Optional[str] = None
|
|
reference_free: Optional[bool] = None
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfDPOConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class CPOConfig(TrainArgumentsMixin, HfCPOConfig):
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfCPOConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class ORPOConfig(TrainArgumentsMixin, HfORPOConfig):
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfORPOConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class KTOConfig(TrainArgumentsMixin, HfKTOConfig):
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfKTOConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class RewardConfig(TrainArgumentsMixin, HfRewardConfig):
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfRewardConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class PPOConfig(TrainArgumentsMixin, HfPPOConfig):
|
|
|
|
def __post_init__(self):
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfPPOConfig.__post_init__(self)
|
|
|
|
|
|
@dataclass
|
|
class GKDConfig(RolloutTrainerArgumentsMixin, TrainArgumentsMixin, HfGKDConfig):
|
|
sft_alpha: float = 0
|
|
|
|
offload_teacher_model: bool = False
|
|
max_completion_length: int = 512
|
|
log_completions: bool = False
|
|
|
|
def __post_init__(self):
|
|
RolloutTrainerArgumentsMixin.__post_init__(self)
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
HfGKDConfig.__post_init__(self)
|
|
self._init_generation_batch_params()
|
|
|
|
|
|
@dataclass
|
|
class GRPOConfig(GRPOArgumentsMixin, TrainArgumentsMixin, HfGRPOConfig):
|
|
|
|
offload_teacher_model: bool = False
|
|
|
|
def __post_init__(self):
|
|
require_version('trl>=0.26')
|
|
GRPOArgumentsMixin.__post_init__(self)
|
|
TrainArgumentsMixin.__post_init__(self)
|
|
# Skip trl GRPOConfig.__post_init__ (hard-requires num_generations>=2); keep TrainingArguments init.
|
|
super(HfGRPOConfig, self).__post_init__()
|
|
|
|
if self.vllm_reasoning_parser is not None:
|
|
raise ValueError('vllm_reasoning_parser is not supported for GRPO Training, please unset it.')
|
|
|
|
if self.cosine_max_len is None:
|
|
self.cosine_max_len = self.max_completion_length
|
|
|
|
if self.deepspeed and 'zero_optimization' in self.deepspeed and self.deepspeed['zero_optimization'][
|
|
'stage'] == 3:
|
|
# https://github.com/modelscope/ms-swift/issues/3237
|
|
self.deepspeed['zero_optimization']['stage3_prefetch_bucket_size'] = 0
|
|
self.deepspeed_plugin.hf_ds_config.config['zero_optimization']['stage3_prefetch_bucket_size'] = 0
|
|
|
|
# https://github.com/modelscope/ms-swift/issues/3863
|
|
self.dataloader_drop_last = True
|
|
|
|
self._init_generation_batch_params()
|