集成SwanLab与HuggingFace TRL：跟踪与优化强化学习实验

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

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

Demo

1. 引入SwanLabCallback

from swanlab.integration.transformers import SwanLabCallback

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

SwanLabCallback可以定义的参数有：

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

2. 传入Trainer

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()

3. 完整案例代码

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

from trl import SFTConfig, SFTTrainer
from datasets import load_dataset
from swanlab.integration.transformers import SwanLabCallbackdataset = 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,)trainer = SFTTrainer(args=training_args,model="Qwen/Qwen2.5-0.5B-Instruct",train_dataset=dataset,callbacks=[swanlab_callback]
)trainer.train()

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

4. GUI效果展示

超参数自动记录：

指标记录：