总目录›TRL›第 3 章
Chapter 03

快速上手：用 SFTTrainer 跑 LoRA SFT

📌 commit 8221a9f30 行代码训完一个对话模型

TRL 最常用的就是 SFTTrainer（Supervised Fine-Tuning）。它在 transformers.Trainer 之上加了两层封装：自动 chat template、自动 packing。本章给一个能跑通的最小例子。

3.1最小完整脚本

from trl import SFTConfig, SFTTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from datasets import load_dataset

MODEL = "Qwen/Qwen2.5-0.5B"

# 1) 模型 + tokenizer
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="bfloat16")
tok   = AutoTokenizer.from_pretrained(MODEL)

# 2) 数据（标准 chat 格式）
ds = load_dataset("trl-lib/Capybara", split="train[:1000]")
# ds[0] = {"messages": [{"role": "user", ...}, {"role": "assistant", ...}]}

# 3) 配置
cfg = SFTConfig(
    output_dir="./sft_out",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    learning_rate=2e-4,
    bf16=True,
    logging_steps=10,
    max_length=2048,
    packing=True,           # ★ 把多条短样本打包成长序列
)

# 4) LoRA（可选）
peft_cfg = LoraConfig(r=16, lora_alpha=32, target_modules="all-linear")

# 5) 训
trainer = SFTTrainer(
    model=model, processing_class=tok,
    args=cfg, train_dataset=ds,
    peft_config=peft_cfg,
)
trainer.train()
trainer.save_model("./sft_out/final")

跑完 ./sft_out/final/ 是标准 HF + PEFT 目录（base 配置 + adapter）。

3.2SFTTrainer 在做什么

看 trl/trainer/sft_trainer.py 的核心逻辑：

  1. 数据预处理：自动判断 dataset 列结构（messages / text / prompt+completion），应用 chat template；
  2. Packing：把多条短样本拼成 max_length 的序列，加 EOS 分隔，节省 padding；
  3. Label mask：可以选只对 assistant 部分算 loss（completion-only loss）；
  4. 包装 PEFT：如果传了 peft_config，自动 get_peft_model；
  5. 调底层 Trainer：剩下的训练循环复用 transformers.Trainer。

3.3三种数据格式

格式样本示例用途
conversational{"messages": [{"role": "user", ...}, ...]}多轮对话（推荐）
prompt-completion{"prompt": "...", "completion": "..."}单轮
plain text{"text": "..."}continued pretraining

SFTTrainer 自动识别三种格式，不需要手动 tokenize。

3.4completion-only loss

SFT 的标准做法是不对 user 部分算 loss，只学 assistant。在 SFTConfig 里加：

cfg = SFTConfig(
    ...,
    assistant_only_loss=True,    # ★ 仅对 assistant 部分算 loss
)

实现：通过 chat template 的 generation_prompt 找到 assistant 起点，前面的 token 全部 label=-100。

3.5验证训练效果

from transformers import pipeline
gen = pipeline("text-generation", model="./sft_out/final", torch_dtype="bfloat16")
print(gen([{"role":"user","content":"用一句话介绍 SFTTrainer"}], max_new_tokens=64))

3.6SFTTrainer 性能基线

配置显存速度
0.5B + LoRA + bf16~5 GB~3.5k tok/s
7B + LoRA + bf16~22 GB~5k tok/s
7B + QLoRA + bf16~14 GB~3.8k tok/s
70B + QLoRA + bf16~48 GB~600 tok/s

3.7不写一行代码：trl sft CLI 路径

如果只是想跑标准 SFT，根本不需要 3.1 那 30 行代码。trl sft 子命令做了完全等价的事。来源 trl/scripts/sft.py：

# 等价于 3.1 的最简单形式
trl sft \
    --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name trl-lib/Capybara \
    --output_dir ./sft_out \
    --num_train_epochs 1 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
    --learning_rate 2e-4 \
    --bf16 \
    --max_length 2048 \
    --packing

背后的调用栈：

步位置做什么
1trl/cli/main.py argparse dispatcher，按子命令名分发
2trl/cli/commands/__init__.py把 sft 绑到 trl/scripts/sft.py:main()
3trl/scripts/sft.py:main() make_parser() 拼好 args，用 HfArgumentParser([SFTScriptArguments, SFTConfig, ModelConfig]) 解析
4同上 调 SFTTrainer(model=..., args=cfg, train_dataset=ds)
5SFTTrainer.train() 转给 transformers.Trainer.train()

配合 yaml 时更干净：

# sft.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: trl-lib/Capybara
output_dir: ./sft_out
num_train_epochs: 1
per_device_train_batch_size: 2
gradient_accumulation_steps: 4
learning_rate: 2.0e-4
bf16: true
max_length: 2048
packing: true
trl sft --config sft.yaml                       # 单卡
accelerate launch -m trl.scripts.sft \
    --config sft.yaml --num_processes 4        # 多卡，accelerate 接管

"何时该写代码？" → 需要自定义数据预处理、自定义 callback、自定义 collator 时；否则 CLI 路径足够覆盖 80% 业务场景。

3.8这章你需要带走的