wehub-resource-sync
30 lines
942 B
Python
30 lines
942 B
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from dataclasses import dataclass
|
|
from typing import Literal, Optional
|
|
|
|
from .sft_args import MegatronSftArguments
|
|
|
|
|
|
@dataclass
|
|
class MegatronRLHFArguments(MegatronSftArguments):
|
|
rlhf_type: Literal['dpo', 'kto', 'grpo', 'gkd', 'rm'] = 'dpo'
|
|
loss_scale: str = 'last_round'
|
|
truncation_strategy: Optional[Literal['delete', 'left', 'right', 'split', None]] = None
|
|
|
|
calculate_per_token_loss: bool = False
|
|
|
|
def __post_init__(self):
|
|
if self.rlhf_type == 'rm':
|
|
self.task_type = 'seq_cls'
|
|
self.num_labels = 1
|
|
self._init_truncation_strategy()
|
|
super().__post_init__()
|
|
|
|
def _init_truncation_strategy(self):
|
|
if self.truncation_strategy is not None:
|
|
return
|
|
if self.rlhf_type == 'grpo':
|
|
self.truncation_strategy = 'left'
|
|
else:
|
|
self.truncation_strategy = 'delete'
|