trl - 微调、对齐大模型的全栈工具

最新推荐文章于 2025-04-16 16:14:37 发布

编程乐园

最新推荐文章于 2025-04-16 16:14:37 发布

阅读量2.5k

点赞数 29

文章标签： trl SFTTrainer RewardTrainer 微调对齐 transformers sft

本文链接：https://blog.csdn.net/lovechris00/article/details/140488897

版权

trl

文章目录

一、关于 TRL

TRL : Transformer Reinforcement Learning

Full stack library to fine-tune and align large language models.

Train transformer language models with reinforcement learning.

github : https://github.com/huggingface/trl
文档：https://huggingface.co/docs/trl/index

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

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

亮点

Efficient and scalable
- accelerate是trl的支柱，它允许使用DDP和DeepSpeed等方法将模型训练从单个GPU扩展到大规模多节点集群。
- PEFT是完全集成的，即使是最大的模型也可以通过量化和LoRA或QLoRA等方法在适度的硬件上训练。
- unsloth也是集成的，允许使用专用内核显着加快训练速度。
CLI：使用CLI，您可以使用单个命令和灵活的配置系统微调LLM并与之聊天，而无需编写任何代码。
Trainers：培训师类是一个抽象，可以轻松应用许多微调方法，如SFTTrainer、DPOTrainer、RewardTrainer、PPOTrainer、CPOTrainer和ORPOTrainer。
AutoModels：AutoModelForCausalLMWithValueHead & AutoModelForSeq2SeqLMWithValueHead 类为模型添加了一个额外的值头，允许使用RL算法（如PPO）训练它们。
Examples：使用BERT情感分类器训练GPT2以生成积极的电影评论，仅使用适配器的完整RLHF，训练GPT-j毒性更小，StackLlama示例等。以下是示例。

二、安装

1、Python包

使用pip安装库：

pip install trl

2、从源码安装

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

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

3、存储库

如果您想使用这些示例，您可以使用以下命令克隆存储库：

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

在 relevant documentation section 阅读有关CLI的更多信息，或使用--help获取更多详细信息。

四、如何使用

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

1、`SFTTrainer`

这是如何使用库中的SFTTrainer的基本示例。

SFTTrainer 是围绕transformersTrainer的轻型包装器，可轻松微调自定义数据集上的语言模型或适配器。

# imports
from datasets import load_dataset
from trl import SFTTrainer

# get dataset
dataset = load_dataset("imdb", split="train")

# get trainer
trainer = SFTTrainer(
    "facebook/opt-350m",
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=512,
)

# train
trainer.train()

2、`RewardTrainer`

这是如何使用库中的RewardTrainer的基本示例。

RewardTrainer 是 transformers Trainer 的包装器，可轻松微调自定义偏好数据集上的奖励模型或适配器。

# imports
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from trl import RewardTrainer

# load model and dataset - dataset needs to be in a specific format
model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=1)
tokenizer = AutoTokenizer.from_pretrained("gpt2")

...

# load trainer
trainer = RewardTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
)

# train
trainer.train()

3、`PPOTrainer`

这是如何使用库中的PPOTrainer的基本示例。

基于查询，语言模型创建一个响应，然后对其进行评估。评估可以是循环中的人或另一个模型的输出。

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

# get models
model = AutoModelForCausalLMWithValueHead.from_pretrained('gpt2')
ref_model = create_reference_model(model)

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

# initialize trainer
ppo_config = PPOConfig(batch_size=1, mini_batch_size=1)

# encode a query
query_txt = "This morning I went to the "
query_tensor = tokenizer.encode(query_txt, return_tensors="pt")

# get model response
response_tensor  = respond_to_batch(model, query_tensor)

# create a ppo trainer
ppo_trainer = PPOTrainer(ppo_config, model, ref_model, tokenizer)

# define a reward for response
# (this could be any reward such as human feedback or output from another model)
reward = [torch.tensor(1.0)]

# train model for one step with ppo
train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)

4、`DPOTrainer`

DPOTrainer是使用直接偏好优化算法的培训师，这是如何使用库中的DPOTrainer的基本示例DPOTrainer是transformersTrainer的包装器，可轻松微调自定义偏好数据集上的奖励模型或适配器。

# imports
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import DPOTrainer

# load model and dataset - dataset needs to be in a specific format
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

...

# load trainer
trainer = DPOTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
)

# train
trainer.train()

五、其它

开发 & 贡献

如果您想为trl做出贡献或根据您的需求对其进行定制，请务必阅读贡献指南并确保您进行了开发安装：

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

参考文献

最近策略优化 PPO

PPO实现在很大程度上遵循D. Ziegler等人的**“来自人类偏好的微调语言模型”**论文中介绍的结构。[论文，代码]。

直接偏好优化 DPO

DPO基于E. Mitchell等人的**《直接偏好优化：您的语言模型是秘密的奖励模型》**的原始实现。[论文，代码]

2024-07-17(三)