第3章：TRL 与强化学习实战 (SFT / DPO / PPO)#

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig
from datasets import load_dataset

# 1. 准备模型与数据
model_id = "Qwen/Qwen2-0.5B-Instruct"
# 使用 HuggingFaceH4 的精选数据集 (Ultrafachat)
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:1%]")

tokenizer = AutoTokenizer.from_pretrained(model_id)
# ⚠️ Qwen 的 pad_token 有时需要手动指定，避免 loss 为 NaN
tokenizer.pad_token = tokenizer.eos_token

# 2. 配置参数
args = SFTConfig(
    output_dir="./qwen-sft",
    max_seq_length=2048,
    packing=True,                # 核心加速
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-5,          # SFT 典型学习率
    lr_scheduler_type="cosine",
    logging_steps=10,
    fp16=True,                   # T4 用 fp16, A100 用 bf16
)

# 3. 开始训练
trainer = SFTTrainer(
    model=model_id,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="messages", # 数据集中的列名
    args=args,
)

trainer.train()
trainer.save_model("./qwen-sft-final")

from trl import DPOTrainer, DPOConfig

# 1. 加载 SFT 后的模型 (作为 Policy Model)
model_id = "./qwen-sft-final"

# 2. 加载数据
dataset = load_dataset("HuggingFaceH4/ultrafeedback_binarized", split="train_prefs[:1000]")

# 3. DPO特殊的配置
# 注意：DPO 的学习率通常比 SFT 低一个数量级 (5e-6 vs 2e-5)
dpo_args = DPOConfig(
    output_dir="./qwen-dpo",
    beta=0.1,
    learning_rate=5e-6,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    fp16=True,
)

# 4. 初始化 Trainer
# TRL 会自动加载 ref_model (也就是 model 的一份拷贝，冻结参数)
trainer = DPOTrainer(
    model=model_id,
    ref_model=None, # 自动处理
    args=dpo_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    max_length=1024,
    max_prompt_length=512,
)

trainer.train()

from trl import RewardTrainer, RewardConfig

# 定义模型：AutoModelForSequenceClassification (num_labels=1)
model = AutoModelForSequenceClassification.from_pretrained(
    "Qwen/Qwen2-0.5B", num_labels=1
)

trainer = RewardTrainer(
    model=model,
    args=RewardConfig(output_dir="./qwen-rm", learning_rate=1e-5),
    train_dataset=dataset, # 包含 chosen/rejected
    tokenizer=tokenizer
)
trainer.train()

# 伪代码流程
ppo_trainer = PPOTrainer(...)

for batch in dataloader:
    query = batch["input_ids"]

    # 1. Actor 生成回复
    response = ppo_trainer.generate(query)

    # 2. RM 打分
    reward = reward_model(query, response)

    # 3. PPO Update Step
    # 这一步会综合 reward 和 KL(actor, ref) 来更新 actor
    stats = ppo_trainer.step(query, response, reward)