Nash-MD 训练器

概述

Nash-MD 在 Rémi Munos、Michal Valko、Daniele Calandriello、Mohammad Gheshlaghi Azar、Mark Rowland、Daniel Guo、Yunhao Tang、Matthieu Geist、Thomas Mésnard 和 Andrea Michi 发表的论文从人类反馈中学习纳什均衡中提出。

论文摘要如下：

从人类反馈中强化学习 (RLHF) 已成为将大型语言模型 (LLM) 与人类偏好对齐的主要范式。通常，RLHF 包括学习从人类反馈中获得的奖励模型的初始步骤，这些反馈通常表示为对预训练 LLM 生成的文本生成对的偏好。随后，通过使用强化学习算法优化 LLM 的策略以最大化奖励模型来对其进行微调。然而，当前奖励模型的固有局限性在于它们无法完全代表人类偏好的丰富性及其对采样分布的依赖性。在本研究中，我们介绍了使用成对人类反馈对 LLM 进行微调的替代流程。我们的方法包括初始学习偏好模型，该模型以给定提示的两个输入为条件，然后追求一个策略，该策略始终生成比任何竞争策略生成的响应更受青睐的响应，从而定义此偏好模型的纳什均衡。我们将这种方法称为从人类反馈中学习纳什均衡 (NLHF)。在表格策略表示的上下文中，我们提出了一种新颖的算法解决方案 Nash-MD，它以镜像下降的原理为基础。该算法产生一系列策略，其中最后一次迭代收敛到正则化的纳什均衡。此外，我们探索了策略的参数表示，并引入了深度学习架构的梯度下降算法。为了证明我们方法的有效性，我们展示了涉及对 LLM 进行文本摘要任务微调的实验结果。我们相信 NLHF 为偏好学习和策略优化提供了一条引人注目的途径，有可能推动将 LLM 与人类偏好对齐领域的进步。

此训练后方法由 Kashif Rasul 和 Daniil Tiapkin、Pierre Ménard、Daniele Calandriello 和 Quentin Gallouédec 贡献。

快速入门

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

以下是训练模型的脚本。

# train_nash_md.py
from datasets import load_dataset
from trl import NashMDConfig, NashMDTrainer
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 = NashMDConfig(output_dir="nash-md-qwen2", logging_steps=10)
trainer = NashMDTrainer(
    model=model,
    reward_model=reward_model,
    args=args,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
)
trainer.train()

使用以下命令执行脚本。

accelerate launch train_nash_md.py

预期数据集格式

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

使用技巧

⚠️ 使用相同的聊天模板

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

鼓励 EOS 标记生成

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

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

记录完成情况

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

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

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

Logged Completions

示例脚本

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

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

python examples/scripts/nash_md.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-nash-md \
    --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/kl：模型与参考数据之间的平均 KL 散度。
objective/entropy：模型与参考数据的平均熵。
loss/score：平均强化得分损失。
rewards/chosen：模型完成情况的平均得分（根据奖励模型）。
rewards/rejected：混合完成情况的平均得分（根据奖励模型）。
rewards/accuracies：Nash-MD 隐式奖励模型的准确率。
rewards/margins：选定完成情况和混合完成情况之间平均奖励差值（根据奖励模型）。
logps/chosen：选定完成情况的平均对数概率。
logps/rejected：参考完成情况的平均对数概率。
val/model_contain_eos_token：模型输出包含 eos 标记的次数。
val/ref_contain_eos_token：混合输出包含 eos 标记的次数。
beta：控制表示与参考模型偏差的损失项权重的参数。通常是固定的，但可以通过传递列表到 NashMDConfig 来使其动态变化。
mixture_coef：模型和参考模型的对数混合系数。通常是固定的，但可以通过传递列表到 NashMDConfig 来使其动态变化。

NashMDTrainer

类 trl.NashMDTrainer

< 源代码 >

( 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 (NashMDConfig) — 用于训练的 NashMD 配置参数。
data_collator (transformers.DataCollator) — 用于训练的数据整理器。如果未指定，则将使用默认数据整理器 (DPODataCollatorWithPadding)，它将根据成对序列的数据集，将序列填充到批次中序列的最大长度。
train_dataset (datasets.Dataset) — 用于训练的数据集。
eval_dataset (datasets.Dataset) — 用于评估的数据集。
tokenizer (transformers.PreTrainedTokenizerBase) — 用于训练的分词器。如果您想使用默认数据整理器，则需要此参数。
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 的函数。

将 NashMDTrainer 初始化为 OnlineDPOConfig 的子类。

NashMDConfig

类 trl.NashMDConfig

< 源代码 >

( 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 mixture_coef: List = <factory> )

参数

mixture_coef (float 或 list[float]，可选，默认为 0.5) — 模型和参考模型的logit混合系数。如果提供了一个浮点数列表，则为每个新 epoch 选择混合系数，最后一个系数将用于剩余的 epoch。

NashMDTrainer 的配置类。

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

< > 在 GitHub 上更新

TRL

Nash-MD 训练器

概述

快速入门

预期数据集格式

使用技巧

⚠️ 使用相同的聊天模板

鼓励 EOS 标记生成

记录完成情况

示例脚本

记录的指标

NashMDTrainer

类 trl.NashMDTrainer

NashMDConfig

类 trl.NashMDConfig