Transformers DataCollatorForLanguageModeling 的 mlm 与 labels 构造方法

用 Hugging Face Transformers 训练语言模型时，DataCollatorForLanguageModeling 负责把不同长度的样本整理成批次，并按任务构造 labels。最容易出错的地方不是调用语法，而是把 mlm 选反：掩码语言模型需要只监督被选中的 token，因果语言模型则通常让标签复制输入，由模型内部完成错位预测。

官方文档：https://huggingface.co/docs/transformers/v4.57.1/en/main_classes/data_collator

mlm 开关决定 labels 的监督位置

mlm=True 的典型场景是 BERT 一类的掩码语言模型。collator 按 mlm_probability 选出候选 token，再按掩码替换比例生成输入；labels 只在这些被选位置保留原 token，其余位置使用 -100。PyTorch 的交叉熵会忽略 -100，所以 loss 只关注被遮住的词。

如果是 GPT 类的因果语言模型，应设置 mlm=False。此时 collator 不负责手工把标签左移，而是让标签与输入对齐；模型的 forward 通常在计算 loss 时完成 shift。手动再左移一次，反而会造成目标错位。

用 tokenizer 输出创建训练批次

下面示例使用掩码语言建模配置。输入样本应是 tokenizer 产生的字典列表，而不是已经堆叠好的单个张量；这样 collator 才能按当前批次最长长度动态补齐。

from transformers import AutoTokenizer, DataCollatorForLanguageModeling

# 选择带 pad_token 的 tokenizer，避免批处理时无法补齐序列
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
texts = ["数据整理器负责组成训练批次", "mlm 参数决定是否随机遮盖词元"]

# 返回特殊 token 标记，方便 collator 避开 [CLS]、[SEP] 等位置
features = [
    tokenizer(text, truncation=True, return_special_tokens_mask=True)
    for text in texts
]

# BERT 类任务使用掩码语言建模；0.15 是官方接口的默认掩码概率
collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15,
    return_tensors="pt",
)

# collator 会动态 padding，并返回 input_ids、attention_mask 与 labels
batch = collator(features)
print(batch["input_ids"].shape, batch["labels"].shape)

改做因果语言模型时，只需把 mlm=False，并确认 tokenizer 已配置可用的 padding token。不要同时保留 mlm_probability 作为“强制遮盖”开关；它只有在 mlm=True 时才参与掩码采样。

padding 与 special_tokens_mask 要一起检查

collator 会按批次最长样本动态补齐。对掩码语言建模来说，special_tokens_mask 能让特殊 token 不被随机替换；tokenizer 生成它时要显式传入 return_special_tokens_mask=True。如果输入是手工拼出的字典，缺少这项并不一定立即报错，却可能让掩码范围与预期不同。

pad_to_multiple_of 可以把长度补到指定倍数，适合需要 Tensor Core 对齐的硬件，但它改变的是批次形状，不改变监督目标。排查 loss 异常时，先打印 labels.eq(-100).sum()，再确认 tokenizer 的 pad_token_id 是否真实存在。

训练前的四项检查清单

1. 先根据模型目标确认 mlm，不要只照搬别人的 Trainer 配置。
2. 确认每个样本包含 input_ids，并让 tokenizer 负责生成 attention mask。
3. 检查 padding 位是否进入 -100，避免无意义 token 参与 loss。
4. 抽一个批次同时查看形状、有效标签数量和特殊 token 位置，再开始长时间训练。

常见问题

mlm=False 时为什么 labels 看起来和 input_ids 一样？

这是预期行为。因果模型通常由模型内部完成标签错位，collator 只负责对齐输入并屏蔽 padding。

mlm_probability 调大是不是一定更好？

不是。它改变每批被监督的 token 数量，过大可能让输入上下文被破坏，应结合语料和模型任务观察 loss。

special_tokens_mask 缺失会导致什么？

collator 可能无法准确避开特殊 token。使用 tokenizer 编码样本时显式返回该字段，最容易保持行为稳定。