TRL项目中的监督微调训练器(SFTTrainer)使用指南

概述

监督微调(Supervised Fine-Tuning, SFT)是强化学习人类反馈(RLHF)流程中的关键步骤。TRL项目提供的SFTTrainer让开发者能够用简洁的API快速实现模型微调。本文将全面介绍SFTTrainer的核心功能和使用方法。

基础用法

快速开始

使用SFTTrainer进行基础微调非常简单：

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

# 加载数据集
dataset = load_dataset("imdb", split="train")

# 配置训练参数
sft_config = SFTConfig(
    dataset_text_field="text",  # 指定文本字段
    max_seq_length=512,        # 最大序列长度
    output_dir="/tmp",         # 输出目录
)

# 初始化训练器
trainer = SFTTrainer(
    "facebook/opt-350m",  # 模型名称
    train_dataset=dataset,
    args=sft_config,
)

# 开始训练
trainer.train()

关键参数说明：

• max_seq_length：必须正确设置，默认值为tokenizer.model_max_length和1024中的较小值
• dataset_text_field：指定数据集中包含文本的字段名

外部模型加载

你也可以先加载模型再传递给训练器：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
trainer = SFTTrainer(
    model,  # 传入预加载的模型
    train_dataset=dataset,
    args=sft_config,
)

高级功能

仅训练完成部分

使用DataCollatorForCompletionOnlyLM可以只训练模型生成的部分：

from trl import DataCollatorForCompletionOnlyLM

response_template = " ### Answer:"
collator = DataCollatorForCompletionOnlyLM(response_template, tokenizer=tokenizer)

trainer = SFTTrainer(
    model,
    data_collator=collator,
    # 其他参数...
)

对于对话数据，需要同时指定指令和响应模板：

instruction_template = "### Human:"
response_template = "### Assistant:"
collator = DataCollatorForCompletionOnlyLM(
    instruction_template=instruction_template,
    response_template=response_template,
    tokenizer=tokenizer
)

特殊令牌处理

对话模型需要添加特殊令牌来区分不同角色：

from trl import setup_chat_format

model, tokenizer = setup_chat_format(model, tokenizer)

此函数会：

1. 添加对话开始/结束标记
2. 调整模型嵌入层大小
3. 设置tokenizer的聊天模板

数据集格式支持

SFTTrainer原生支持两种常见格式：

1. 对话格式：

{"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

2. 指令格式：

{"prompt": "...", "completion": "..."}

输入提示格式化

自定义提示格式示例：

def formatting_prompts_func(example):
    return [f"问题：{q}\n回答：{a}" for q, a in zip(example['question'], example['answer'])]

trainer = SFTTrainer(
    formatting_func=formatting_prompts_func,
    # 其他参数...
)

数据集打包

启用打包功能可提高训练效率：

sft_config = SFTConfig(packing=True)
trainer = SFTTrainer(
    packing=True,  # 启用打包
    # 其他参数...
)

模型控制

精度控制

sft_config = SFTConfig(
    model_init_kwargs={"torch_dtype": "bfloat16"}
)

适配器训练

结合PEFT库进行适配器训练：

from peft import LoraConfig

peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    task_type="CAUSAL_LM",
)

trainer = SFTTrainer(
    peft_config=peft_config,  # 传入LoRA配置
    # 其他参数...
)

8位模型训练

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-125m",
    load_in_8bit=True,
    device_map="auto",
)

trainer = SFTTrainer(
    model,  # 传入8位模型
    peft_config=peft_config,
    # 其他参数...
)

性能优化

Flash Attention支持

安装最新依赖：

pip install -U flash-attn

使用Flash Attention 2：

model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-350m",
    attn_implementation="flash_attention_2"
)

性能对比（NVIDIA T4 16GB）：

启用Flash Attention	模型	序列长度	批次大小	每步时间
是	opt-350m	2048	8	~59.1s
否	opt-350m	2048	8	OOM
是	opt-350m	2048	4	~30.3s
否	opt-350m	2048	4	~148.9s

总结

TRL的SFTTrainer提供了从基础到高级的完整监督微调解决方案。通过合理配置，开发者可以高效地微调各种规模的模型，并利用现代加速技术提升训练效率。无论是简单的指令微调还是复杂的对话模型训练，SFTTrainer都能提供简洁而强大的支持。