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}}
}