XPO 训练器

概述

探索性偏好优化 (XPO) 在论文 Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF 中提出，作者为 Tengyang Xie、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 模型作为基础模型，并使用 PairRMJudge 作为评判器。我们使用来自 UltraFeedback 数据集的提示。您可以在此处查看数据集中的提示

以下是训练模型的脚本

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

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
judge = PairRMJudge()
train_dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

training_args = XPOConfig(output_dir="Qwen2-0.5B-XPO", logging_steps=10)
trainer = XPOTrainer(
    model=model, judge=judge, args=training_args, processing_class=tokenizer, train_dataset=train_dataset
)
trainer.train()

使用以下命令执行脚本

accelerate launch train_xpo.py

在 8 个 GPU 上分布式运行，训练大约需要 1 小时。

要了解已训练模型的性能，您可以使用 Transformers Chat CLI。

$ transformers-cli chat --model_name_or_path trl-lib/Qwen2-0.5B-XPO
<quentin_gallouedec>:
What is the best programming language?

<trl-lib/Qwen2-0.5B-XPO>:
The best programming language depends on individual preferences and familiarity with coding concepts. Some popular languages include Python, Java, C++, and JavaScript.

预期数据集类型

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

使用技巧

使用奖励模型

您可以选择使用奖励模型而不是评判器——请参阅 Reward Bench，其中包含您可以使用的公共模型排行榜。以下代码示例展示了如何使用 trl-lib/Qwen2-0.5B-Reward 模型替换评判器

- from trl import PairRMJudge
+ from transformers import AutoModelForSequenceClassification

- judge = PairRMJudge()
+ reward_model = AutoModelForSequenceClassification.from_pretrained("trl-lib/Qwen2-0.5B-Reward", num_labels=1)

  trainer = XPOTrainer(
      ...
-     judge=judge,
+     reward_model=reward_model,
  )

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

鼓励生成 EOS 令牌

当使用奖励模型时，我们可能希望模型在给定长度内生成补全。在训练期间，模型将生成补全，直到 XPOConfig 的 max_new_tokens 参数中指定的最大长度。如果您想惩罚模型在达到最大长度之前未生成 EOS 令牌，您可以使用 XPOConfig 的 missing_eos_penalty 参数

training_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/xpo.py 中找到

要在 UltraFeedback 数据集上使用 Qwen2.5 0.5B 模型测试 XPO 脚本，请运行以下命令

python examples/scripts/xpo.py \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --judge pair_rm \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 5.0e-7 \
    --logging_steps 25 \
    --output_dir Qwen2.5-0.5B-XPO-PairRM \
    --warmup_ratio 0.1 \
    --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

< source >

( model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] = None ref_model: typing.Union[transformers.modeling_utils.PreTrainedModel, torch.nn.modules.module.Module] = None reward_model: typing.Optional[torch.nn.modules.module.Module] = None judge: typing.Optional[trl.trainer.judges.BasePairwiseJudge] = None args: typing.Optional[trl.trainer.xpo_config.XPOConfig] = None data_collator: typing.Optional[typing.Callable] = None train_dataset: typing.Union[datasets.arrow_dataset.Dataset, datasets.iterable_dataset.IterableDataset, NoneType] = 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 peft_config: typing.Optional[dict] = 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 )

参数

model (transformers.PreTrainedModel) — 要训练的模型，最好是 AutoModelForCausalLM。
ref_model (PreTrainedModelWrapper) — Hugging Face transformers 模型，带有因果语言建模头。用于隐式奖励计算和损失。如果没有提供参考模型，训练器将创建一个与要优化的模型具有相同架构的参考模型。
reward_model (transformers.PreTrainedModel) — 用于对补全结果进行评分的奖励模型，最好是 AutoModelForSequenceClassification。
judge (BasePairwiseJudge) — 用于模型补全结果成对比较的评判器。
args (XPOConfig) — 用于训练的 XPO 配置参数。
data_collator (transformers.DataCollator) — 用于训练的数据整理器。如果未指定，将使用默认数据整理器 (DPODataCollatorWithPadding)，它将成对序列数据集中的序列填充到批次中最长序列的长度。
train_dataset (datasets.Dataset) — 用于训练的数据集。
eval_dataset (datasets.Dataset) — 用于评估的数据集。
processing_class (PreTrainedTokenizerBase 或 BaseImageProcessor 或 FeatureExtractionMixin 或 ProcessorMixin, 可选) — 用于处理数据的处理类。如果提供，将用于自动处理模型的输入，并将与模型一起保存，以便更容易地重新运行中断的训练或重用微调后的模型。
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

class trl.XPOConfig

< source >

( 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-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: typing.Union[transformers.trainer_utils.SchedulerType, str] = 'linear' lr_scheduler_kwargs: typing.Union[dict, str, NoneType] = <factory> warmup_ratio: float = 0.0 warmup_steps: int = 0 log_level: typing.Optional[str] = 'passive' log_level_replica: typing.Optional[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 = 500 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: 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: 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: typing.Optional[bool] = True 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, NoneType] = None tp_size: typing.Optional[int] = 0 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 gradient_checkpointing: bool = False gradient_checkpointing_kwargs: typing.Union[dict, 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: typing.Optional[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 eval_use_gather_object: typing.Optional[bool] = False average_tokens_across_devices: typing.Optional[bool] = False reward_model_path: typing.Optional[str] = None judge: typing.Optional[str] = None max_new_tokens: int = 64 max_length: int = 512 temperature: float = 0.9 missing_eos_penalty: typing.Optional[float] = None beta: list = <factory> loss_type: str = 'sigmoid' dataset_num_proc: typing.Optional[int] = None disable_dropout: bool = True use_vllm: bool = False ds3_gather_for_generation: bool = True alpha: list = <factory> )

参数

alpha (float or list[float], optional, defaults to 1e-5) — XPO 损失项的权重。如果提供浮点数列表，则为每个新 epoch 选择 alpha，最后一个 alpha 用于其余的 epochs。

为 XPOTrainer 配置类。

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

< > 在 GitHub 上更新

TRL

XPO 训练器

概述

快速开始

预期数据集类型

使用技巧

使用奖励模型

鼓励生成 EOS 令牌

记录补全

示例脚本

记录的指标

XPOTrainer

class trl.XPOTrainer

XPOConfig

class trl.XPOConfig