Hugging Face 的 logo

TRL 文档

奖励建模

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

奖励建模

TRL 支持自定义奖励建模，任何人都可以在其数据集和模型上执行奖励建模。

请查看 examples/scripts/reward_modeling.py 中的完整灵活示例。

预期数据集格式

由于模型将针对示例对进行训练以预测哪个示例更受欢迎，因此 RewardTrainer 对数据集的格式有非常具体的要求。我们提供以下来自 Anthropic/hh-rlhf 数据集的示例

因此，如果您使用默认的 RewardDataCollatorWithPadding 数据整理器，则最终数据集对象至少应包含两个 4 个条目。条目应命名为

  • 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,
    tokenizer=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 等人，2023] 提出了一个旨在直接学习中心化奖励模型的辅助损失函数。该辅助损失函数最小化奖励的平方和，鼓励模型自然地产生均值为零的输出(R(p,r1)+R(p,r2))2\Big( R(p, r_1) + R(p, r_2) \Big)^2

该辅助损失函数与主损失函数结合，由 [RewardConfig] 中的参数 center_rewards_coefficient 加权。默认情况下，此功能处于禁用状态 (center_rewards_coefficient = None)。

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

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

RewardTrainer

类 trl.RewardTrainer

 < >

( model: Union = None args: Optional = None data_collator: Optional = None train_dataset: Optional = None eval_dataset: Union = None tokenizer: Optional = None model_init: Optional = None compute_metrics: Optional = None callbacks: Optional = None optimizers: Tuple = (None, None) preprocess_logits_for_metrics: Optional = None max_length: Optional = None peft_config: Optional = None )

RewardTrainer 可用于训练自定义奖励模型。它是 transformers.Trainer 类的子类，并继承其所有属性和方法。建议使用 AutoModelForSequenceClassification 作为奖励模型。奖励模型应在配对示例的数据集上进行训练，其中每个示例都是两个序列的元组。奖励模型应该被训练来预测配对中哪个示例与手头的任务更相关。

奖励训练器期望数据集具有非常特定的格式。如果未使用默认的 RewardDataCollatorWithPadding 数据收集器，则数据集应至少包含两个 4 个条目。条目应命名为

  • input_ids_chosen
  • attention_mask_chosen
  • input_ids_rejected
  • attention_mask_rejected

可选地，您还可以将 margin 条目传递给数据集。此条目应包含用于调节奖励模型损失的边际值，如 https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/ 中所述。如果您不传递边际值，则不会使用任何边际值。

可视化样本

 < >

( num_print_samples: int )

参数

  • num_print_samples (int，默认为 4) — 要打印的样本数量。设置为 -1 以打印所有样本。

可视化奖励模型 logits 预测

RewardConfig

类 trl.RewardConfig

< >

( output_dir: str overwrite_output_dir: bool = False do_train: bool = False do_eval: bool = False do_predict: bool = False eval_strategy: Union = '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: Optional = None per_gpu_eval_batch_size: Optional = None gradient_accumulation_steps: int = 1 eval_accumulation_steps: Optional = None eval_delay: Optional = 0 torch_empty_cache_steps: Optional = 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: Union = 'linear' lr_scheduler_kwargs: Union = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: Optional = 'passive' log_level_replica: Optional = 'warning' log_on_each_node: bool = True logging_dir: Optional = None logging_strategy: Union = 'steps' logging_first_step: bool = False logging_steps: float = 500 logging_nan_inf_filter: bool = True save_strategy: Union = 'steps' save_steps: float = 500 save_total_limit: Optional = None save_safetensors: Optional = 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: Optional = None jit_mode_eval: bool = False use_ipex: bool = False bf16: bool = False fp16: bool = False fp16_opt_level: str = 'O1' half_precision_backend: str = 'auto' bf16_full_eval: bool = False fp16_full_eval: bool = False tf32: Optional = None local_rank: int = -1 ddp_backend: Optional = None tpu_num_cores: Optional = None tpu_metrics_debug: bool = False debug: Union = '' dataloader_drop_last: bool = False eval_steps: Optional = None dataloader_num_workers: int = 0 dataloader_prefetch_factor: Optional = None past_index: int = -1 run_name: Optional = None disable_tqdm: Optional = None remove_unused_columns: Optional = True label_names: Optional = None load_best_model_at_end: Optional = False metric_for_best_model: Optional = None greater_is_better: Optional = None ignore_data_skip: bool = False fsdp: Union = '' fsdp_min_num_params: int = 0 fsdp_config: Union = None fsdp_transformer_layer_cls_to_wrap: Optional = None accelerator_config: Union = None deepspeed: Union = None label_smoothing_factor: float = 0.0 optim: Union = 'adamw_torch' optim_args: Optional = None adafactor: bool = False group_by_length: bool = False length_column_name: Optional = 'length' report_to: Union = None ddp_find_unused_parameters: Optional = None ddp_bucket_cap_mb: Optional = None ddp_broadcast_buffers: Optional = 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: Optional = None hub_model_id: Optional = None hub_strategy: Union = 'every_save' hub_token: Optional = None hub_private_repo: bool = False hub_always_push: bool = False gradient_checkpointing: bool = False gradient_checkpointing_kwargs: Union = None include_inputs_for_metrics: bool = False include_for_metrics: List = <factory> eval_do_concat_batches: bool = True fp16_backend: str = 'auto' evaluation_strategy: Union = None push_to_hub_model_id: Optional = None push_to_hub_organization: Optional = None push_to_hub_token: Optional = None mp_parameters: str = '' auto_find_batch_size: bool = False full_determinism: bool = False torchdynamo: Optional = None ray_scope: Optional = 'last' ddp_timeout: Optional = 1800 torch_compile: bool = False torch_compile_backend: Optional = None torch_compile_mode: Optional = None dispatch_batches: Optional = None split_batches: Optional = None include_tokens_per_second: Optional = False include_num_input_tokens_seen: Optional = False neftune_noise_alpha: Optional = None optim_target_modules: Union = None batch_eval_metrics: bool = False eval_on_start: bool = False use_liger_kernel: Optional = False eval_use_gather_object: Optional = False max_length: Optional = None dataset_num_proc: Optional = None center_rewards_coefficient: Optional = None )

参数

  • max_length (Optional[int], 可选, 默认为 None) — 批次中序列（提示 + 补全）的最大长度。如果您想使用默认的数据整理器，则需要此参数。
  • dataset_num_proc (int, 可选, 默认为 None) — 用于处理数据集的进程数。
  • center_rewards_coefficient (float, 可选, 默认为 None) — 鼓励奖励模型输出均值为零的奖励的系数（由 https://huggingface.co/papers/2312.09244，公式 2 提议）。推荐值：0.01。

RewardTrainer 的配置类。

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

< > 在 GitHub 上更新

←PPOv2 RLOO→