奖励模型

TRL 支持自定义奖励模型，任何人都可以根据自己的数据集和模型进行奖励建模。

请在 examples/scripts/reward_modeling.py 查看一个完整且灵活的示例。

期望的数据集类型

RewardTrainer 需要一个隐式提示偏好数据集。这意味着数据集应该只包含 "chosen" 和 "rejected" 列（而不包含 "prompt"）。RewardTrainer 支持对话式和标准两种数据集格式。当提供对话式数据集时，训练器会自动将聊天模板应用于数据集。

您也可以使用预分词的数据集，在这种情况下，数据集应包含以下列：input_ids_chosen, attention_mask_chosen, input_ids_rejected 和 attention_mask_rejected。

使用 RewardTrainer

准备好数据集后，您可以像使用 🤗 Transformers 的 Trainer 类一样使用 RewardTrainer。您应该将一个 AutoModelForSequenceClassification 模型传递给 RewardTrainer，同时传递一个 RewardConfig 来配置训练的超参数。

利用 🤗 PEFT 训练奖励模型

只需在 RewardTrainer 的关键字参数中传递一个 peft_config，训练器就会自动将模型转换为 PEFT 模型！

from peft import LoraConfig, TaskType
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer, RewardConfig

model = AutoModelForSequenceClassification.from_pretrained("gpt2")
peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

...

trainer = RewardTrainer(
    model=model,
    args=training_args,
    processing_class=tokenizer,
    train_dataset=dataset,
    peft_config=peft_config,
)

trainer.train()

为损失函数添加边际

就像在 Llama 2 论文中一样，您可以通过向数据集中添加 margin 列来为损失函数添加边际。奖励整理器会自动传递它，并相应地计算损失。

def add_margin(row):
    # Assume you have a score_chosen and score_rejected columns that you want to use to compute the margin
    return {'margin': row['score_chosen'] - row['score_rejected']}

dataset = dataset.map(add_margin)

中心化奖励

在许多场景中，确保奖励模型的输出均值为零是更可取的。这通常通过首先计算模型的平均得分，然后减去它来实现。

[Eisenstein et al., 2023] 提出了一种辅助损失函数，旨在直接学习一个中心化的奖励模型。该辅助损失函数最小化奖励的平方和，从而鼓励模型自然地产生均值为零的输出。 $\Big( R(p, r_1) + R(p, r_2) \Big)^2$

这个辅助损失函数与主损失函数结合，权重由 [RewardConfig] 中的参数 center_rewards_coefficient 控制。默认情况下，此功能是禁用的（center_rewards_coefficient = None）。

training_args = RewardConfig(
    center_rewards_coefficient=0.01,
    ...
)

有关参考结果，请参阅 PR #1932。

RewardTrainer

class trl.RewardTrainer

< 源代码 >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module, NoneType] = None args: typing.Optional[trl.trainer.reward_config.RewardConfig] = None data_collator: typing.Optional[transformers.data.data_collator.DataCollator] = None train_dataset: typing.Optional[datasets.arrow_dataset.Dataset] = None eval_dataset: typing.Union[datasets.arrow_dataset.Dataset, dict[str, datasets.arrow_dataset.Dataset], NoneType] = None processing_class: typing.Union[transformers.tokenization_utils_base.PreTrainedTokenizerBase, transformers.image_processing_utils.BaseImageProcessor, transformers.feature_extraction_utils.FeatureExtractionMixin, transformers.processing_utils.ProcessorMixin, NoneType] = None model_init: typing.Optional[typing.Callable[[], transformers.modeling_utils.PreTrainedModel]] = None compute_metrics: typing.Optional[typing.Callable[[transformers.trainer_utils.EvalPrediction], dict]] = None callbacks: typing.Optional[list[transformers.trainer_callback.TrainerCallback]] = None optimizers: tuple = (None, None) preprocess_logits_for_metrics: typing.Optional[typing.Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None peft_config: typing.Optional[dict] = None )

train

< 源代码 >

( resume_from_checkpoint: typing.Union[str, bool, NoneType] = None trial: typing.Union[ForwardRef('optuna.Trial'), dict[str, typing.Any], NoneType] = None ignore_keys_for_eval: typing.Optional[list[str]] = None **kwargs )

参数

resume_from_checkpoint (str 或 bool, 可选) — 如果是 str，则为之前 Trainer 实例保存的检查点的本地路径。如果是 bool 且等于 True，则加载之前 Trainer 实例在 *args.output_dir* 中保存的最后一个检查点。如果提供，训练将从此处加载的模型/优化器/调度器状态恢复。
trial (optuna.Trial 或 dict[str, Any], 可选) — 用于超参数搜索的试验运行或超参数字典。
ignore_keys_for_eval (list[str], 可选) — 模型输出（如果为字典）中在训练期间收集评估预测时应忽略的键列表。
kwargs (dict[str, Any], 可选) — 用于隐藏已弃用参数的附加关键字参数。

主训练入口点。

save_model

< 源代码 >

( output_dir: typing.Optional[str] = None _internal_call: bool = False )

将保存模型，以便您可以使用 `from_pretrained()` 重新加载它。

仅从主进程保存。

push_to_hub

< 源代码 >

( commit_message: typing.Optional[str] = 'End of training' blocking: bool = True token: typing.Optional[str] = None revision: typing.Optional[str] = None **kwargs )

参数

commit_message (str, 可选, 默认为 "End of training") — 推送时提交的消息。
blocking (bool, 可选, 默认为 True) — 函数是否应在 git push 完成后才返回。
token (str, 可选, 默认为 None) — 具有写入权限的令牌，用于覆盖 Trainer 的原始参数。
revision (str, 可选) — 提交的 git 版本。默认为“main”分支的头部。
kwargs (dict[str, Any], 可选) — 传递给 ~Trainer.create_model_card 的附加关键字参数。

将 `self.model` 和 `self.processing_class` 上传到 🤗 模型中心的 `self.args.hub_model_id` 存储库。

RewardConfig

class trl.RewardConfig

< 源代码 >

( output_dir: typing.Optional[str] = None overwrite_output_dir: bool = False do_train: bool = False do_eval: bool = False do_predict: bool = False eval_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'no' prediction_loss_only: bool = False per_device_train_batch_size: int = 8 per_device_eval_batch_size: int = 8 per_gpu_train_batch_size: typing.Optional[int] = None per_gpu_eval_batch_size: typing.Optional[int] = None gradient_accumulation_steps: int = 1 eval_accumulation_steps: typing.Optional[int] = None eval_delay: typing.Optional[float] = 0 torch_empty_cache_steps: typing.Optional[int] = None learning_rate: float = 5e-05 weight_decay: float = 0.0 adam_beta1: float = 0.9 adam_beta2: float = 0.999 adam_epsilon: float = 1e-08 max_grad_norm: float = 1.0 num_train_epochs: float = 3.0 max_steps: int = -1 lr_scheduler_type: typing.Union[transformers.trainer_utils.SchedulerType, str] = 'linear' lr_scheduler_kwargs: typing.Union[dict[str, typing.Any], str, NoneType] = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: str = 'passive' log_level_replica: str = 'warning' log_on_each_node: bool = True logging_dir: typing.Optional[str] = None logging_strategy: typing.Union[transformers.trainer_utils.IntervalStrategy, str] = 'steps' logging_first_step: bool = False logging_steps: float = 10 logging_nan_inf_filter: bool = True save_strategy: typing.Union[transformers.trainer_utils.SaveStrategy, str] = 'steps' save_steps: float = 500 save_total_limit: typing.Optional[int] = None save_safetensors: typing.Optional[bool] = True save_on_each_node: bool = False save_only_model: bool = False restore_callback_states_from_checkpoint: bool = False no_cuda: bool = False use_cpu: bool = False use_mps_device: bool = False seed: int = 42 data_seed: typing.Optional[int] = None jit_mode_eval: bool = False use_ipex: bool = False bf16: typing.Optional[bool] = None fp16: bool = False fp16_opt_level: str = 'O1' half_precision_backend: str = 'auto' bf16_full_eval: bool = False fp16_full_eval: bool = False tf32: typing.Optional[bool] = None local_rank: int = -1 ddp_backend: typing.Optional[str] = None tpu_num_cores: typing.Optional[int] = None tpu_metrics_debug: bool = False debug: typing.Union[str, list[transformers.debug_utils.DebugOption]] = '' dataloader_drop_last: bool = False eval_steps: typing.Optional[float] = None dataloader_num_workers: int = 0 dataloader_prefetch_factor: typing.Optional[int] = None past_index: int = -1 run_name: typing.Optional[str] = None disable_tqdm: typing.Optional[bool] = None remove_unused_columns: bool = False label_names: typing.Optional[list[str]] = None load_best_model_at_end: typing.Optional[bool] = False metric_for_best_model: typing.Optional[str] = None greater_is_better: typing.Optional[bool] = None ignore_data_skip: bool = False fsdp: typing.Union[list[transformers.trainer_utils.FSDPOption], str, NoneType] = '' fsdp_min_num_params: int = 0 fsdp_config: typing.Union[dict[str, typing.Any], str, NoneType] = None fsdp_transformer_layer_cls_to_wrap: typing.Optional[str] = None accelerator_config: typing.Union[dict, str, NoneType] = None deepspeed: typing.Union[dict, str, NoneType] = None label_smoothing_factor: float = 0.0 optim: typing.Union[transformers.training_args.OptimizerNames, str] = 'adamw_torch' optim_args: typing.Optional[str] = None adafactor: bool = False group_by_length: bool = False length_column_name: typing.Optional[str] = 'length' report_to: typing.Union[NoneType, str, list[str]] = None ddp_find_unused_parameters: typing.Optional[bool] = None ddp_bucket_cap_mb: typing.Optional[int] = None ddp_broadcast_buffers: typing.Optional[bool] = None dataloader_pin_memory: bool = True dataloader_persistent_workers: bool = False skip_memory_metrics: bool = True use_legacy_prediction_loop: bool = False push_to_hub: bool = False resume_from_checkpoint: typing.Optional[str] = None hub_model_id: typing.Optional[str] = None hub_strategy: typing.Union[transformers.trainer_utils.HubStrategy, str] = 'every_save' hub_token: typing.Optional[str] = None hub_private_repo: typing.Optional[bool] = None hub_always_push: bool = False hub_revision: typing.Optional[str] = None gradient_checkpointing: bool = False gradient_checkpointing_kwargs: typing.Union[dict[str, typing.Any], str, NoneType] = None include_inputs_for_metrics: bool = False include_for_metrics: list = <factory> eval_do_concat_batches: bool = True fp16_backend: str = 'auto' push_to_hub_model_id: typing.Optional[str] = None push_to_hub_organization: typing.Optional[str] = None push_to_hub_token: typing.Optional[str] = None mp_parameters: str = '' auto_find_batch_size: bool = False full_determinism: bool = False torchdynamo: typing.Optional[str] = None ray_scope: typing.Optional[str] = 'last' ddp_timeout: int = 1800 torch_compile: bool = False torch_compile_backend: typing.Optional[str] = None torch_compile_mode: typing.Optional[str] = None include_tokens_per_second: typing.Optional[bool] = False include_num_input_tokens_seen: typing.Optional[bool] = False neftune_noise_alpha: typing.Optional[float] = None optim_target_modules: typing.Union[NoneType, str, list[str]] = None batch_eval_metrics: bool = False eval_on_start: bool = False use_liger_kernel: typing.Optional[bool] = False liger_kernel_config: typing.Optional[dict[str, bool]] = None eval_use_gather_object: typing.Optional[bool] = False average_tokens_across_devices: bool = True max_length: typing.Optional[int] = 1024 disable_dropout: bool = True dataset_num_proc: typing.Optional[int] = None center_rewards_coefficient: typing.Optional[float] = None )

参数

max_length (int 或 None，可选，默认为 1024) — 批次中序列（提示+补全）的最大长度，会过滤掉超过此限制的条目。如果要使用默认的数据整理器，则此参数是必需的。
disable_dropout (bool, 可选，默认为 True) — 是否在模型中禁用 dropout。
dataset_num_proc (int, 可选，默认为 None) — 用于处理数据集的进程数。
center_rewards_coefficient (float, 可选，默认为 None) — 激励奖励模型输出均值为零的奖励的系数（由 https://huggingface.ac.cn/papers/2312.09244 提出，公式 2）。推荐值：0.01。
remove_unused_columns (bool, 可选，默认为 False) — 是否移除模型前向传播中未使用的列。仅当数据集已预先分词时，此项才能为 True。

用于 RewardTrainer 的配置类。

此类仅包含特定于奖励（Reward）训练的参数。有关训练参数的完整列表，请参阅 TrainingArguments 文档。请注意，此类中的默认值可能与 TrainingArguments 中的不同。

使用 HfArgumentParser，我们可以将此类别转换为可在命令行上指定的 argparse 参数。

< > 在 GitHub 上更新