Hugging Face 的 logo

TRL 文档

XPO 训练器

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

XPO 训练器

概述

探索性偏好优化 (XPO) 在论文 探索性偏好优化：利用隐式 Q*-近似进行样本高效 RLHF 中提出，作者为谢腾阳、Dylan J. Foster、Akshay Krishnamurthy、Corby Rosset、Ahmed Awadallah 和 Alexander Rakhlin。它是一种基于 DPO 损失和奖励模型 (RM) 的简单在线偏好调整方法。XPO 通过添加探索奖励来增强 DPO 目标，使该方法能够探索初始模型和人类反馈数据支持范围之外的内容。

论文摘要如下：

来自人类反馈的强化学习 (RLHF) 已成为语言模型对齐的核心工具。我们考虑 RLHF 中的在线探索，它通过故意鼓励模型生成多样化、信息量最大的响应来利用对人类或 AI 反馈的交互式访问。通过允许 RLHF 有信心地偏离预训练模型，在线探索提供了新颖的、可能超越人类能力的可能性，但由于直接调整现有强化学习技术的计算和统计瓶颈，其作为语言模型训练范式的全部潜力尚未实现。我们提出了一种用于 RLHF 中在线探索的新算法，即探索性偏好优化 (XPO)，它简单实用——对（在线）直接偏好优化 (DPO；Rafailov 等人，2023) 进行了一行代码的修改——但具有已知的最佳可证明保证和良好的经验性能。XPO 通过新颖且有原则的探索奖励增强了 DPO 目标，使算法能够探索初始模型和人类反馈数据支持范围之外的内容。在理论上，我们证明了 XPO 在自然探索条件下具有可证明的样本效率，并且会收敛到接近最优的语言模型策略，无论初始模型是否具有良好的覆盖率。我们的分析建立在 DPO 隐式执行某种形式的 Q*-近似（或贝尔曼误差最小化）的观察结果之上，以 KL 正则化马尔可夫决策过程的角度，以偶然的方式结合了语言建模和理论强化学习中以前不同的技术。在实践中，我们发现 XPO 在初步评估中比非探索性 DPO 变体更具样本效率。

此训练后方法由 Kashif Rasul、Quentin Gallouédec 和 Lewis Tunstall 贡献。

快速入门

此示例演示了如何使用 XPO 方法训练模型。我们使用 Qwen 0.5B 模型 作为基础模型，并使用 Qwen 0.5B 奖励模型 作为奖励模型。我们使用来自 UltraFeedback 数据集 的提示。您可以在此处查看数据集中的提示。

以下是训练模型的脚本

# train_xpo.py
from datasets import load_dataset
from trl import XPOConfig, XPOTrainer
from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
reward_model = AutoModelForSequenceClassification.from_pretrained("trl-lib/Qwen2-0.5B-Reward", num_labels=1)
train_dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

args = XPOConfig(output_dir="nash-md-qwen2", logging_steps=10)
trainer = XPOTrainer(
    model=model,
    reward_model=reward_model,
    args=args,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
)
trainer.train()

使用以下命令执行脚本

accelerate launch train_xpo.py

预期数据集格式

XPO 需要一个 仅提示数据集。 XPOTrainer 支持 对话 和 标准 数据集格式。当提供对话数据集时，训练器会自动将聊天模板应用于数据集。

使用技巧

⚠️ 使用相同的聊天模板

确保 SFT 模型和奖励模型使用相同的聊天模板。否则，您可能会发现模型完成在训练期间得分不正确。

鼓励 EOS 标记生成

我们希望模型在给定长度内生成完成。在学习过程中，模型将生成完成，直到 XPOConfig 的max_new_tokens参数中指定的最大完成长度。如果您希望惩罚在最大完成长度之前没有生成 EOS 标记的情况，您可以使用 XPOConfig 的missing_eos_penalty参数。

args = XPOConfig(..., max_new_tokens=128, missing_eos_penalty=1.0)

记录完成

为了更好地了解模型在训练期间的行为，您可以使用 LogCompletionsCallback 定期记录样本完成情况。

trainer = XPOTrainer(..., eval_dataset=eval_dataset)
completions_callback = LogCompletionsCallback(trainer, num_prompts=8)
trainer.add_callback(completions_callback)

此回调将模型生成的完成直接记录到 Weights & Biases。

Logged Completions

示例脚本

我们提供了一个使用 XPO 方法训练模型的示例脚本。该脚本位于 examples/scripts/nash_md.py

要使用 Pythia 14M 模型 在 TL;DR 摘要任务上测试 XPO 脚本，请运行以下命令

python examples/scripts/xpo.py \
    --model_name_or_path EleutherAI/pythia-14m  \
    --reward_model_path EleutherAI/pythia-14m \
    --dataset_name trl-lib/tldr \
    --learning_rate 5.0e-7 \
    --output_dir pythia-14m-tldr-xpo \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 32 \
    --num_train_epochs 3 \
    --max_new_tokens 64 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --push_to_hub

记录的指标

记录的指标如下

  • loss/xpo：完整损失的 XPO 部分的平均值。
  • loss/dpo：完整损失的 DPO 部分的平均值。
  • objective/kl：模型与参考数据之间的平均 KL 散度。
  • objective/entropy：模型与参考数据的平均熵。
  • objective/model_scores：模型完成的平均得分（根据奖励模型）。
  • objective/ref_scores：参考完成的平均得分（根据奖励模型）。
  • objective/scores_margin：选定完成和拒绝完成之间平均得分差（根据外部奖励模型）。
  • rewards/chosen：选定完成的平均奖励（根据 XPO 的 DPO 隐式奖励模型）。
  • rewards/rejected：拒绝完成的平均奖励（根据 XPO 的 DPO 隐式奖励模型）。
  • rewards/accuracies：XPO 的隐式奖励模型的准确率。
  • rewards/margins：选定完成和拒绝完成之间平均奖励差（根据在线 DPO 的隐式奖励模型）。
  • logps/chosen：选定完成的平均对数概率。
  • logps/rejected：拒绝完成的平均对数概率。
  • val/model_contain_eos_token：模型输出包含 eos 标记的次数。
  • val/ref_contain_eos_token：参考输出包含 eos 标记的次数。
  • alpha：XPO 损失项的权重。通常是固定的，但可以通过传递列表到 XPOConfig 来使其动态变化。
  • beta：控制表示与参考模型偏差的损失项权重的参数。通常是固定的，但可以通过传递列表到 XPOConfig 来使其动态变化。

XPOTrainer

class trl.XPOTrainer

< >

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

参数

  • model (transformers.PreTrainedModel) — 用于训练的模型，最好是 AutoModelForCausalLM。
  • ref_model (PreTrainedModelWrapper) — 具有因果语言建模头的 Hugging Face Transformer 模型。用于隐式奖励计算和损失。如果未提供参考模型，则训练器将创建与要优化的模型具有相同架构的参考模型。
  • reward_model (transformers.PreTrainedModel) — 用于对完成情况进行评分的奖励模型，最好是 AutoModelForSequenceClassification。
  • judge (BasePairwiseJudge) — 用于模型完成情况的成对比较的评判器。
  • args (XPOConfig) — 用于训练的 XPO 配置参数。
  • data_collator (transformers.DataCollator) — 用于训练的数据整理器。如果未指定，则将使用默认数据整理器 (DPODataCollatorWithPadding)，它将在给定配对序列数据集的情况下，将序列填充到批次中序列的最大长度。
  • train_dataset (datasets.Dataset) — 用于训练的数据集。
  • eval_dataset (datasets.Dataset) — 用于评估的数据集。
  • peft_config (Dict) — 用于训练的 peft 配置。
  • compute_metrics (Callable[[EvalPrediction], Dict], 可选) — 用于计算指标的函数。必须接收一个 EvalPrediction 并返回一个字典，其中键为指标名称，值为指标值。
  • callbacks (List[transformers.TrainerCallback]) — 用于训练的回调函数。
  • optimizers (Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]) — 用于训练的优化器和调度器。
  • preprocess_logits_for_metrics (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]) — 用于在计算指标之前预处理 logits 的函数。

将 XPOTrainer 初始化为 OnlineDPOConfig 的子类。

XPOConfig

类 trl.XPOConfig

< >

( 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-07 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 reward_model_path: Optional = None max_new_tokens: int = 64 temperature: float = 0.9 missing_eos_penalty: Optional = None beta: List = <factory> loss_type: Literal = 'sigmoid' dataset_num_proc: Optional = None disable_dropout: bool = True alpha: List = <factory> )

参数

  • alpha (float 或 List[float], 可选, 默认为 1e-5) — XPO 损失项的权重。如果提供浮点数列表，则为每个新纪元选择 alpha，并将最后一个 alpha 用于剩余纪元。

XPOTrainer 的配置类。

OnlineDPOConfig 的子类，我们可以使用它的所有参数并添加以下内容

< > 在 GitHub 上更新

←迭代 SFT 模型类→