简体中文

English

简体中文

English

主题

🤗HuggingFace Trl ​

TRL (Transformers Reinforcement Learning，用强化学习训练Transformers模型) 是一个领先的Python库，旨在通过监督微调（SFT）、近端策略优化（PPO）和直接偏好优化（DPO）等先进技术，对基础模型进行训练后优化。TRL 建立在 🤗 Transformers 生态系统之上，支持多种模型架构和模态，并且能够在各种硬件配置上进行扩展。

logo

你可以使用Trl快速进行模型训练，同时使用SwanLab进行实验跟踪与可视化。

Demo

transformers>=4.50.0 的版本，已官方集成了SwanLab
如果你的版本低于4.50.0，请使用SwanLabCallback集成。

1. 一行代码集成 ​

只需要在你的训练代码中，找到HF的Config部分（比如SFTConfig、GRPOConfig等），添加report_to="swanlab"参数，即可完成集成。

python
from trl import SFTConfig, SFTTrainer

args = SFTConfig(
    ...,
    report_to="swanlab"
)

trainer = Trainer(..., args=args)

2. 自定义项目名 ​

默认下，项目名会使用你运行代码的目录名。

如果你想自定义项目名，可以设置SWANLAB_PROJECT环境变量：

python
import os
os.environ["SWANLAB_PROJECT"]="qwen2-sft"
bash
export SWANLAB_PROJECT="qwen2-sft"
bash
set SWANLAB_PROJECT="qwen2-sft"

3. 案例代码 ​

使用Qwen2.5-0.5B-Instruct模型，使用Capybara数据集进行SFT训练：

python
from trl import SFTConfig, SFTTrainer
from datasets import load_dataset

dataset = load_dataset("trl-lib/Capybara", split="train")

training_args = SFTConfig(
    output_dir="Qwen/Qwen2.5-0.5B-SFT",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    num_train_epochs=1,
    logging_steps=20,
    learning_rate=2e-5,
    report_to="swanlab", 
    )

trainer = SFTTrainer(
    args=training_args,
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)

trainer.train()

DPO、GRPO、PPO等同理，只需要将report_to="swanlab"传入对应的Config即可。

4. GUI效果展示 ​

超参数自动记录：

ig-hf-trl-gui-1

指标记录：

ig-hf-trl-gui-2

5.使用SwanLabCallback ​

如果你使用的是Transformers<4.50.0的版本，或者你希望更灵活地控制SwanLab的行为，则可以使用SwanLabCallback集成。

5.1 引入SwanLabCallback ​

python
from swanlab.integration.transformers import SwanLabCallback

SwanLabCallback是适配于Transformers的日志记录类。

SwanLabCallback可以定义的参数有：

  • project、experiment_name、description 等与 swanlab.init 效果一致的参数, 用于SwanLab项目的初始化。
  • 你也可以在外部通过swanlab.init创建项目，集成会将实验记录到你在外部创建的项目中。

5.2 传入Trainer ​

python
from swanlab.integration.transformers import SwanLabCallback
from trl import SFTConfig, SFTTrainer

...

# 实例化SwanLabCallback
swanlab_callback = SwanLabCallback(project="trl-visualization")

trainer = SFTTrainer(
    ...
    # 传入callbacks参数
    callbacks=[swanlab_callback],
)

trainer.train()

5.3 完整案例代码 ​

使用Qwen2.5-0.5B-Instruct模型，使用Capybara数据集进行SFT训练：

python
from trl import SFTConfig, SFTTrainer
from datasets import load_dataset
from swanlab.integration.transformers import SwanLabCallback

dataset = load_dataset("trl-lib/Capybara", split="train")

swanlab_callback = SwanLabCallback(
    project="trl-visualization",
    experiment_name="Qwen2.5-0.5B-SFT",
    description="测试使用trl框架sft训练"
)

training_args = SFTConfig(
    output_dir="Qwen/Qwen2.5-0.5B-SFT",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    num_train_epochs=1,
    logging_steps=20,
    learning_rate=2e-5,
    report_to="none",
    )

trainer = SFTTrainer(
    args=training_args,
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)

trainer.train()

DPO、GRPO、PPO等同理，只需要将SwanLabCallback传入对应的Trainer即可。

5. 环境变量 ​

参考：HuggingFace Docs: transformers.integrations.SwanLabCallback