1.1 最小可用示例#

1
from trl import SFTTrainer
2
from datasets import load_dataset
3

4
trainer = SFTTrainer(
5
    model="Qwen/Qwen3-0.6B",
6
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
7
)
8
trainer.train()

SFTTrainer 自动处理 tokenizer 加载、padding、truncation、损失 mask 等。

1.2 数据格式#

支持三种主流格式：

1
# 语言建模（标准）
2
{"text": "The sky is blue."}
3

4
# 对话格式
5
{"messages": [
6
    {"role": "user", "content": "What color is the sky?"},
7
    {"role": "assistant", "content": "It is blue."}
8
]}
9

10
# Prompt-Completion
11
{"prompt": "The sky is", "completion": " blue."}

1.3 损失函数#

SFT 是 token 级的交叉熵：

因果语言模型用当前位置之前的 token 预测下一个 token；训练实现通常在计算损失时对预测值和标签做错位对齐。Padding 位置的标签设为 -100，在交叉熵中被忽略；对话数据还可以进一步只计算助手回复部分的损失。

图：SFT 训练中，预测值与下一 token 的标签对齐；padding 位置不参与损失计算。

2.1 为什么需要#

batch 内序列长度不齐时，所有序列要 pad 到最长那条。如果 1% 的序列特别长，剩下 99% 都被拖着 pad 到那个长度——大量 padding token 浪费显存与算力。

图：序列长度分布不均时，不加截断会造成大量 padding。

2.2 配置#

1
from trl import SFTConfig
2

3
training_args = SFTConfig(max_length=2048)

2.3 选 max_length 的取舍#

• 太小：超长样本被截掉尾部，信息丢失
• 太大：padding 占比高，显存浪费

实操做法：先看数据集的长度分布（TRL 提供在线工具 https://trl-lib-dataset-length-profiler.hf.space），选 P90 ~ P95 分位数附近的值。

但 truncation 终究是”用信息损失换显存”的笨办法。下面的 Packing 是更优雅的解。

3.1 思路#

Truncation 砍尾巴，Packing 是 拼接：把多条短序列首尾相接成一条接近 max_length 的长序列，几乎消除 padding。

图：Packing 把多条序列拼到一个训练样本中，最大化利用上下文窗口。

打包后的注意力边界、position ID 和跨样本语义由所选 packing strategy、模型及 FlashAttention 实现共同决定；不要用仅按 EOS 扫描的手写 mask 取代 trainer 的实现。自定义 collator 时，应以当前 TRL/模型文档和单元测试确认每个样本的可见范围与位置编码约定。

3.2 TRL 配置#

1
training_args = SFTConfig(
2
    packing=True,
3
    packing_strategy="bfd",  # "bfd" | "bfd_split" | "wrapped"
4
    max_length=2048,
5
)

三种策略：

这里使用的 TRL 默认 bfd 策略会自动开启 padding_free，因此需要兼容的 FlashAttention 2/3。不能据此推断所有 Packing 实现都必须使用 FlashAttention。

所有序列都短于 max_length 时，bfd 和 bfd_split 行为完全一致。

4.1 QLoRA：量化基模型 + LoRA#

QLoRA 在 LoRA 的基础上，把基模型量化到 4-bit（NF4 格式），只在前向时反量化到 BF16 计算。显存收益叠加：

• 基模型权重：NF4 编码本体为 4 bit/参数；实际占用还包括量化常数、分组元数据、对齐与运行时工作区
• LoRA 适配器：通常仍用 BF16/FP16，参数量取决于 rank 与目标层
• 优化器状态：只针对 LoRA 参数

三个关键技巧（来自 QLoRA 论文 ¹）：

1. NF4（NormalFloat4）：一种非均匀量化格式，专门为「正态分布」的权重设计——量化级别在 0 附近更密、两端更疏，比均匀 INT4 精度高
2. 双重量化（Double Quantization）：连量化常数（scale）自己也做一次 FP8 量化，每个参数额外省 ~0.4 bit
3. 分页优化器（Paged Optimizers）：用 unified memory 把优化器状态分页到 CPU，OOM 时自动换出

1
import torch
2
from datasets import load_dataset
3
from peft import LoraConfig
4
from transformers import BitsAndBytesConfig
5
from trl import SFTConfig, SFTTrainer
6

7
bnb_config = BitsAndBytesConfig(
8
    load_in_4bit=True,
9
    bnb_4bit_quant_type="nf4",
10
    bnb_4bit_compute_dtype=torch.bfloat16,
11
    bnb_4bit_use_double_quant=True,
12
)
13

14
trainer = SFTTrainer(
15
    model="Qwen/Qwen3-8B",
16
    args=SFTConfig(max_length=2048),
17
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
18
    quantization_config=bnb_config,
19
    peft_config=LoraConfig(r=16, lora_alpha=32),
20
)

这里由 BitsAndBytesConfig 在加载时量化基模型，并由 peft_config 只训练 LoRA adapter；所需显存仍取决于序列长度、batch、target modules、dtype 和内核。量化原理详见 第 3 篇 — 量化深入。

6.1 大词表的麻烦#

LM head 输出的 logits 形状是 [batch, seq_len, vocab]。Qwen3 的 vocab ≈ 152K，Llama-3 是 128K。算一下：

• batch=4, seq=2048, vocab=152K, BF16
• logits 张量 = $4 \times 2048 \times 152000 \times 2 \approx 2.5\,\text{GB}$

这只是 logits 自身的理论大小。反向传播还可能需要保存或生成与 logits 同量级的张量，但具体峰值取决于交叉熵实现、精度、张量是否复用及计算是否融合，不能固定写成“必占 5 GB”。

6.2 思路#

观察：交叉熵 loss 一次只需要每个位置的”自己那个 label 对应的 logit”和归一化项 $\log \sum_v \exp(\text{logit}_v)$，不需要同时持有全部位置的全 vocab logits。

Chunked NLL：

1. 在 LM head 矩阵乘之前 先过滤掉 labels == -100 的位置（padding/prompt 部分）
2. 把剩下的有效 token 切成 chunk
3. 每个 chunk 单独算 logits → 单独算 loss → 释放
4. 配合 gradient checkpointing 在反向重新计算

峰值显存：batch × seq × vocab → chunk_size × vocab。

6.3 配置#

1
training_args = SFTConfig(loss_type="chunked_nll")

TRL 官方 benchmark（Qwen3-1.7B，词表约 152K）报告：

• 单 GPU 峰值 -30%
• FSDP2 × 4 GPU 峰值 -50%
• 训练时间基本持平或略快

当前 TRL 中，chunked_nll 是 SFTTrainer 的默认 loss；use_liger_kernel=True 时会自动改用 nll，两者不兼容。chunked_nll 也不兼容 PEFT 和 VLM。

如何取舍：在非 PEFT 场景下，如果大词表产生的 logits 张量是显存瓶颈，可以使用 Chunked NLL。Liger Kernel 优化的是一组算子，是否更合适，要看模型支持情况和实际测试结果。

30%/50% 是上述官方 benchmark 的结果，不是普遍收益；FSDP 选项、模型结构和词表大小都会改变结果。

10.1 pad_to_multiple_of：硬件友好对齐#

1
training_args = SFTConfig(pad_to_multiple_of=2048)

把序列 pad 到 64/128/2048 的倍数。Tensor Core 在某些维度对齐下吞吐显著提升。SFT 和 Reward 两个 trainer 支持。

10.2 ZeRO-3 + 在线 RL#

使用 DeepSpeed ZeRO-3 的在线方法在生成时可能临时 gather 模型权重，大模型因此可能 OOM。对 GRPO 可关闭该 gather（生成速度可能下降）：

1
from trl import GRPOConfig
2
training_args = GRPOConfig(ds3_gather_for_generation=False)

10.3 vLLM Sleep Mode#

在线 RL 用 vLLM 做生成时，把 vLLM 的参数和缓存在优化步暂时卸载到 CPU：

1
training_args = GRPOConfig(vllm_enable_sleep_mode=True)

显存腾出来给训练，生成时再加载回 GPU。

单卡约 7/8B QLoRA 配置示例（显存以实测为准）#

1
import torch
2
from datasets import load_dataset
3
from peft import LoraConfig
4
from transformers import BitsAndBytesConfig
5
from trl import SFTConfig, SFTTrainer
6

7
dataset = load_dataset("trl-lib/Capybara", split="train")
8

9
quantization_config = BitsAndBytesConfig(
10
    load_in_4bit=True,
11
    bnb_4bit_quant_type="nf4",
12
    bnb_4bit_compute_dtype=torch.bfloat16,
13
    bnb_4bit_use_double_quant=True,
14
)
15

16
training_args = SFTConfig(
17
    max_length=2048,
18
    packing=True,
19
    model_init_kwargs={
20
        "attn_implementation": "kernels-community/flash-attn2",
21
    },
22
)
23

24
trainer = SFTTrainer(
25
    model="Qwen/Qwen3-8B",
26
    args=training_args,
27
    train_dataset=dataset,
28
    quantization_config=quantization_config,
29
    peft_config=LoraConfig(r=16, lora_alpha=32),
30
)

该示例同时传入 4-bit BitsAndBytesConfig 和 LoraConfig，因此是 QLoRA；packing 还要求安装并验证兼容的 FlashAttention。它没有声称固定显存可运行：开始训练前应以目标 GPU、实际 batch 和数据长度记录峰值显存。

多卡训 70B+#

1
training_args = SFTConfig(
2
    max_length=4096,
3
    packing=True,
4
    use_liger_kernel=True,
5
    gradient_checkpointing=True,
6
)
7
# 启动时配 FSDP / DeepSpeed ZeRO-3
8
# 如果显存瓶颈在 logits 且非 PEFT，可换 loss_type="chunked_nll"（需关掉 use_liger_kernel）

在线 RL（GRPO）#

1
training_args = GRPOConfig(
2
    use_vllm=True,
3
    vllm_enable_sleep_mode=True,
4
    ds3_gather_for_generation=False,
5
)

vllm_enable_sleep_mode 与 ds3_gather_for_generation 是在线生成的显存取舍；前者会引入主机与设备间传输，后者可能降低生成速度。不要把 padding_free 或 activation_offloading 当作当前 GRPOConfig 的开关叠加进去。