技术博客

teacher forcing 的实现细节:标签位移的正确做法与常见错误

第 107 / 147 篇 · 2026-01-22 · ai · 预训练训练工程自回归

teacher forcing 的概念一句话就能说完:训练时输入真实前缀、预测下一个 token(暴露偏差见本系列概率篇)。但把它写成能跑的工程代码,坑密度高得不成比例——我在这上面烧过不止一个下午。这篇按实现顺序过一遍标准做法,把每一步对应的典型错误列出来。参考实现是 PyTorch 2.x + transformers 4.x 的组合,Loss 用交叉熵。

标准做法:错一位,加边界

输入序列 $x_1,\dots,x_T$ 本身不变,标签是它的左移:$\hat y_t = x_{t+1}$。等价于在序列头部插 BOS、尾部补 EOS:

tokens:  [BOS] 我 们 吃 饭 [EOS]
input:   [BOS] 我 们 吃 饭        ← 网络输入
labels:   我  们 吃 饭 [EOS]      ← 第 t 步的监督目标

第 $t$ 个位置的 logits 经过 softmax 后与 $\hat y_t$ 算交叉熵。实现上几乎从不真的做 shift 拷贝,而是让 labels 与 input_ids 等长、整体错位,第一个位置(或最后一个)填 ignore index:

labels = input_ids.clone()
labels[:, :-1] = input_ids[:, 1:]   # 错一位
labels[:, -1] = -100                # 最后一步没有下一个 token

transformers 的 ForCausalLM 系列在模型内部做同样的事:forward(input_ids, labels=...) 时会把 labels 左移再算 loss,第一个位置自动置 -100。所以给 HF 模型传 labels 时不要再手动移位——这是最高频的错误,移两次等于让模型用 $w_t$ 预测 $w_t$,loss 会漂亮地降到接近 ln V(均匀分布)但模型什么也没学到。

标签位移一行:网络输入与监督目标错开一位,首位置不参与 loss input: BOS 我 们 吃 饭 pad labels: 我 们 吃 饭 EOS -100 causal mask: 位置 t 只能 attend 到 ≤ t:mask 与标签错位同向,不能只对齐一边

六个高频错误

一、移位两次。 上面说过,HF 模型内部已移。判断方法很简单:把训练 loss 和 model(input_ids).logits 手动对齐 labels 算一次交叉熵对比,不一致就说明有人移了两遍。

二、pad 的 mask 漏了。 batch 里有变长序列,右对齐 pad 到最长。attention mask 遮住 pad 是对的,但 loss mask 也得同步把 pad 位置置 -100——CrossEntropyLoss 默认 mean 聚合,一个没遮住的 pad 位置就是一份垃圾监督,还会污染有效 token 的平均。我自己实践里会同时检查 attention_mask.sum() == (labels != -100).sum() 在两个维度是否分别匹配。

三、EOS 该喂不该喂的两难。 训练时 EOS 必须出现在 labels 里(模型得学会停),input 里最后一个有效 token 的预测目标就是 EOS。但推理时模型自己生成 EOS 后就要停——所以生成侧的 tokenizer 配置(eos_token_id)与训练侧的必须一致,改了特殊 token 表没重训,模型学到停止信号对应的 id 就不是你推理时查的那个 id 了。

四、logits 的 gather 写错维度。 logits 形状 [B, T, V],labels 是 [B, T]。交叉熵要求先把 logits 展平成 [B*T, V]、labels 展成 [B*T],F.crossEntropy 的 mean 模式默认会把 -100 去掉再平均。手写 gather 时常见错误是沿 T 维 gather(应该用 V 维),或者忘了 .contiguous() 展平时交错错位。我现在的习惯是直接交给 fused kernel(如 Liger 的 FLCE),展平、移位、mask 一步完成,还能省掉 logits 的显存峰值。

五、packing 时文档边界没隔离。 把多篇文档拼进一条定长序列能消灭 pad 浪费(见本系列 packing 篇),但如果直接拼,前一篇的结尾会通过 attention 泄漏给后一篇的预测——模型学到"某篇维基条目结束后下一个 token 是另一篇 Reddit 帖子的开头",这不是任何真实文档的分布。两种修法:在文档之间放 BOS 并训练模型把它当硬边界;或者用 FlashAttention 的 varlen 接口把每条文档作为独立序列传入,attention 不跨文档。新代码一律用第二种,第一种在边界位置 loss 有系统性偏差。

六、右 padding 用于训练、左 padding 用于生成,两边混用。 训练惯例右 pad(loss 在左侧对齐);批量生成用左 pad(最新 token 在右侧,position id 好算)。同一份代码两个分支如果共享 attention mask 的构造逻辑,很容易在训练时用了左 pad 的 mask——肉眼看不出来,loss 也不炸,但 causal 结构的右侧约束被 mask 反转了,训练出来的模型生成质量明显差。排查方法:用一条两个文档拼接的输入手动打印 attention 权重热力图,看可见区域是不是下三角。

为什么这些细节值得专门写一篇

teacher forcing 是"显而易见"的操作,但上述每个错误都不会让训练崩溃——loss 正常下降、曲线看着健康,只在下游 eval 或人工看输出时才暴露。预训练里这种"静默错误"是最贵的:几周的算力烧完才发现模型把 pad 当目标学了。把移位、mask、EOS、packing 四条线做对,本质是保证训练时的条件分布与推理时的自回归采样严格同构——这正是下一篇 UL2 要在目标函数层面做的事。

小结

标签位移的正确做法:labels 左移一位、首位置置 -100、HF 模型内部移位不要再手动移、pad 双遮、EOS 在训练目标里但不在推理输入里、packing 用 varlen 隔离文档。六个高频错误都不报错、不炸曲线,只腐蚀模型质量,只能靠代码审查和对照验证拦住。