基于 Python 的自然语言处理系列（82）：Transformer Reinforcement Learning

本文将介绍一套完整的 RLHF（Reinforcement Learning with Human Feedback）训练流程，基于 Hugging Face 推出的 trl 库，从 SFT（Supervised Fine-tuning）、RM（Reward Modeling）、到 PPO（Proximal Policy Optimization）三大阶段，逐步实现对 Transformer 模型

冷基栋_攻城师

617人浏览 · 2025-04-19 13:00:00

冷基栋_攻城师 · 2025-04-19 13:00:00 发布

🔗 本文所用工具：trl、transformers、peft、bitsandbytes
📘 官方文档参考：https://huggingface.co/docs/trl

一、引言：从有监督微调到 RLHF 全流程

随着语言大模型的发展，如何在大规模预训练模型基础上更精细地对齐人类偏好，成为了研究与应用的热点。本文将介绍一套完整的 RLHF（Reinforcement Learning with Human Feedback）训练流程，基于 Hugging Face 推出的 trl 库，从 SFT（Supervised Fine-tuning）、RM（Reward Modeling）、到 PPO（Proximal Policy Optimization）三大阶段，逐步实现对 Transformer 模型的强化学习优化。

本篇聚焦于 SFT 阶段的实现，并以 Hugging Face 提供的 instruction-dataset 为例，介绍如何使用 trl 和 PEFT（参数高效微调）技术训练一个高效对齐指令的语言模型。

二、安装与环境准备

确保安装以下库（建议使用 PyTorch + CUDA 环境）：

pip install trl transformers datasets peft bitsandbytes accelerate

三、加载并准备数据集

本例使用 HuggingFaceH4 团队整理的 instruction-dataset：

from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/instruction-dataset")
dataset = dataset.remove_columns("meta")  # 移除无用字段
dataset

四、构建模型及量化配置（4-bit）

使用 BitsAndBytesConfig 对模型进行 4-bit 量化，可大幅降低显存占用：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import prepare_model_for_kbit_training

model_name = "lmsys/fastchat-t5-3b-v1.0"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)

base_model.config.use_cache = False
base_model = prepare_model_for_kbit_training(base_model)

五、注入 LoRA 参数高效微调机制

首先识别所有 4-bit 线性模块并定义 LoRA 参数配置：

import bitsandbytes as bnb
from peft import get_peft_model, LoraConfig

def find_all_linear_names(model):
    cls = bnb.nn.Linear4bit
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split(".")
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
    return list(lora_module_names)

peft_config = LoraConfig(
    r=128,
    lora_alpha=16,
    target_modules=find_all_linear_names(base_model),
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

base_model = get_peft_model(base_model, peft_config)

打印可训练参数占比：

def print_trainable_parameters(model):
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"Trainable params: {trainable} / {total} ({trainable / total:.2%})")

print_trainable_parameters(base_model)

六、定义 Prompt 格式化函数

将数据集中的 prompt 和 completion 格式化为统一格式：

def formatting_prompts_func(example):
    return [
        f"### Input: ```{prompt}```\n ### Output: {completion}"
        for prompt, completion in zip(example["prompt"], example["completion"])
    ]

七、训练参数设置与 SFTTrainer 训练器

使用 SFTTrainer 执行指令微调训练，支持 gradient checkpointing、cosine 学习率调度等高级策略：

from transformers import TrainingArguments
from trl import SFTTrainer

output_dir = "./results"

training_args = TrainingArguments(
    output_dir=output_dir,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    max_grad_norm=0.3,
    num_train_epochs=15,
    learning_rate=2e-4,
    bf16=True,
    save_total_limit=3,
    logging_steps=10,
    optim="paged_adamw_32bit",
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

trainer = SFTTrainer(
    model=base_model,
    train_dataset=dataset,
    tokenizer=tokenizer,
    max_seq_length=2048,
    formatting_func=formatting_prompts_func,
    args=training_args
)

执行训练：

trainer.train()
trainer.save_model(output_dir)

保存最终模型权重与 tokenizer：

import os
final_output_dir = os.path.join(output_dir, "final_checkpoint")
trainer.model.save_pretrained(final_output_dir)
tokenizer.save_pretrained(final_output_dir)