TRL - Transformer 强化学习

用于微调和对齐大型语言模型的全栈库。

这是什么？

trl 库是一个全栈工具，用于使用监督微调步骤（SFT）、奖励建模（RM）和近端策略优化（PPO）以及直接偏好优化（DPO）等方法来微调和对齐 transformer 语言和扩散模型。

该库构建在 transformers 库之上，因此允许使用那里可用的任何模型架构。

亮点

高效且可扩展：
- accelerate 是 trl 的核心，它允许使用 DDP 和 DeepSpeed 等方法将模型训练从单个 GPU 扩展到大规模多节点集群。
- PEFT 已完全集成，允许通过量化和 LoRA 或 QLoRA 等方法在普通硬件上训练甚至最大的模型。
- unsloth 也已集成，允许使用专用内核显著加快训练速度。
命令行界面：使用命令行界面，您可以通过单个命令和灵活的配置系统微调和与 LLM 聊天，无需编写任何代码。
训练器：Trainer 类是一个抽象，可以轻松应用许多微调方法，如 SFTTrainer、DPOTrainer、RewardTrainer、PPOTrainer、CPOTrainer 和 ORPOTrainer。
自动模型：AutoModelForCausalLMWithValueHead 和 AutoModelForSeq2SeqLMWithValueHead 类为模型添加了一个额外的价值头，允许使用 PPO 等强化学习算法进行训练。
示例：使用 BERT 情感分类器训练 GPT2 生成正面电影评论，仅使用适配器进行完整的 RLHF，训练 GPT-j 减少毒性，StackLlama 示例等，参照示例。

安装

Python 包

使用 pip 安装库：

pip install trl

从源代码安装

如果您想在官方发布之前使用最新功能，可以从源代码安装：

pip install git+https://github.com/huggingface/trl.git

仓库

如果您想使用示例，可以使用以下命令克隆仓库：

git clone https://github.com/huggingface/trl.git

命令行界面（CLI）

您可以使用 TRL 命令行界面（CLI）快速开始监督微调（SFT）、直接偏好优化（DPO），并使用聊天 CLI 测试您对齐的模型：

SFT：

trl sft --model_name_or_path facebook/opt-125m --dataset_name imdb --output_dir opt-sft-imdb

DPO：

trl dpo --model_name_or_path facebook/opt-125m --dataset_name trl-internal-testing/hh-rlhf-helpful-base-trl-style --output_dir opt-sft-hh-rlhf

聊天：

trl chat --model_name_or_path Qwen/Qwen1.5-0.5B-Chat

在相关文档部分阅读更多关于 CLI 的信息，或使用 --help 获取更多详细信息。

如何使用

为了获得更大的灵活性和对训练的控制，您可以在 Python 中使用专用的训练器类来微调模型。

`SFTTrainer`

这是一个关于如何使用库中的SFTTrainer的基本示例。SFTTrainer是对transformers训练器的轻量级封装，用于在自定义数据集上轻松微调语言模型或适配器。

# 导入
from datasets import load_dataset
from trl import SFTTrainer

# 获取数据集
dataset = load_dataset("imdb", split="train")

# 获取训练器
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=512,
)

# 训练
trainer.train()

`RewardTrainer`

这是一个关于如何使用库中的RewardTrainer的基本示例。RewardTrainer是对transformers训练器的封装，用于在自定义偏好数据集上轻松微调奖励模型或适配器。

# 导入
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer

# 加载模型和数据集 - 数据集需要采用特定格式
model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=1)
tokenizer = AutoTokenizer.from_pretrained("gpt2")

...

# 加载训练器
trainer = RewardTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
)

# 训练
trainer.train()

`PPOTrainer`

这是一个关于如何使用库中的PPOTrainer的基本示例。基于查询，语言模型创建一个响应，然后对其进行评估。评估可以是人在循环中进行，也可以是另一个模型的输出。

# 导入
import torch
from transformers import AutoTokenizer
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead, create_reference_model
from trl.core import respond_to_batch

# 获取模型
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = create_reference_model(model)

tokenizer = AutoTokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token

# 初始化训练器
ppo_config = PPOConfig(batch_size=1, mini_batch_size=1)

# 编码查询
query_txt = "今天早上我去了"
query_tensor = tokenizer.encode(query_txt, return_tensors="pt")

# 获取模型响应
response_tensor  = respond_to_batch(model, query_tensor)

# 创建PPO训练器
ppo_trainer = PPOTrainer(ppo_config, model, ref_model, tokenizer)

# 为响应定义奖励
# （这可以是任何奖励，如人类反馈或另一个模型的输出）
reward = [torch.tensor(1.0)]

# 使用PPO训练模型一步
train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)

`DPOTrainer`

DPOTrainer是一个使用直接偏好优化算法的训练器。这是一个关于如何使用库中的DPOTrainer的基本示例。DPOTrainer是对transformers训练器的封装，用于在自定义偏好数据集上轻松微调奖励模型或适配器。

# 导入
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer

# 加载模型和数据集 - 数据集需要采用特定格式
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

...

# 加载训练器
trainer = DPOTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
)

# 训练
trainer.train()

开发

如果你想为trl做贡献或根据自己的需求进行定制，请确保阅读贡献指南并进行开发安装：

git clone https://github.com/huggingface/trl.git
cd trl/
make dev

参考文献

近端策略优化

PPO实现主要遵循D. Ziegler等人在论文**"Fine-Tuning Language Models from Human Preferences"**中介绍的结构。[论文, 代码]。

直接偏好优化

DPO基于E. Mitchell等人的**"Direct Preference Optimization: Your Language Model is Secretly a Reward Model"**的原始实现。[论文, 代码]

引用

@misc{vonwerra2022trl,
  author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang},
  title = {TRL: Transformer Reinforcement Learning},
  year = {2020},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/huggingface/trl}}
}