15 KiB

Raw Blame History

第 10 课：Causal LM forward 与 loss shift

这一课只解决一个问题：MiniMindForCausalLM.forward 如何把 hidden_states 变成 logits，再把 logits 和 labels shift 后算出 next-token loss。

0. 本节主线
1. 本节要懂的 5 个原理
2. 变量流转
3. 原理一：lm_head 把 hidden 映射成 logits
4. 原理二：next-token loss 必须 shift
5. 原理三：-100 决定哪些位置不参与 loss
6. 原理四：Pretrain 和 SFT 共享同一个 forward loss
7. 原理五：训练目标差异来自 labels
8. 实验验证
9. 本节检查
10. 下一课

0. 本节主线

Causal LM loss 的本质是：

input_ids
-> MiniMindModel
-> hidden_states [batch, seq, hidden_size]
-> lm_head
-> logits [batch, seq, vocab_size]
-> logits[..., :-1, :]
-> labels[..., 1:]
-> cross_entropy(ignore_index=-100)
-> loss

最关键的一句话：

logits 在位置 i 的输出，用来预测 labels 在位置 i+1 的 token。

所以模型不是预测“当前位置 token”，而是用当前位置之前的信息预测“下一个 token”。

1. 本节要懂的 5 个原理

原理	要理解什么	源码证据
`lm_head` 把 hidden 映射成 logits	主干输出 hidden_states，外层 CausalLM 头输出词表分数	`model/model_minimind.py:245-249`
next-token loss 必须 shift	`logits[..., :-1, :]` 对齐 `labels[..., 1:]`	`model/model_minimind.py:249-252`
`-100` 决定哪些位置不参与 loss	cross entropy 用 `ignore_index=-100` 忽略被 mask 的 label	`model/model_minimind.py:250-252`
Pretrain 和 SFT 共享同一个 forward loss	训练脚本都只是传入 `(input_ids, labels)`	`trainer/train_pretrain.py:35-40`, `trainer/train_full_sft.py:35-40`
训练目标差异来自 labels	Pretrain 基本全段监督，SFT 只监督 assistant 区域	`dataset/lm_dataset.py:47-55`, `dataset/lm_dataset.py:88-119`

学完本节，你应该能解释：为什么 SFT 仍然是 causal LM next-token loss，以及为什么 user token 的 label 要设成 -100。

2. 变量流转

先看 shape：

input_ids:
  [batch, seq]

hidden_states:
  [batch, seq, hidden_size]

logits:
  [batch, seq, vocab_size]

labels:
  [batch, seq]

再看 shift：

logits[..., :-1, :]
  [batch, seq - 1, vocab_size]
  位置 0 到 seq-2 的预测分数

labels[..., 1:]
  [batch, seq - 1]
  位置 1 到 seq-1 的目标 token

对齐关系是：

logits[:, 0, :] -> labels[:, 1]
logits[:, 1, :] -> labels[:, 2]
logits[:, 2, :] -> labels[:, 3]
...

所以：

labels[:, 0] 不参与 loss
logits[:, -1, :] 不参与 loss

3. 原理一：`lm_head` 把 hidden 映射成 logits

原理讲解

第 9 课讲到，MiniMindModel 的输出还是 hidden states：

[batch, seq, hidden_size]

但训练语言模型时，我们需要在每个位置预测下一个 token。预测 token 就要输出词表大小的分数：

[batch, seq, vocab_size]

这一步由 lm_head 完成：

hidden_size -> vocab_size

每个位置都会得到一组长度为 vocab_size 的 logits。logits 还不是概率，它只是未归一化分数；cross entropy 内部会处理 softmax。

源码证据 A：forward 先得到 hidden，再调用 lm_head

文件：model/model_minimind.py:245-249

看它是为了理解：logits 是从哪里来的。

源码：

def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
    hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    loss = None

这段代码说明：

self.model(...) 返回的是 hidden_states。
self.lm_head(...) 把 hidden 映射成词表 logits。
logits 的最后一维是 vocab_size。
loss 初始为 None，只有传入 labels 才计算训练损失。

理解到这一步就够

你应该能说清楚：

MiniMindModel 负责产生 hidden 表示；
MiniMindForCausalLM 负责把 hidden 变成 logits，并在有 labels 时计算 loss。

暂时不用看：

logits_to_keep 的推理节省计算用途。
softmax 的数值稳定实现。

4. 原理二：next-token loss 必须 shift

原理讲解

Causal LM 的训练目标是：

看到前面的 token，预测下一个 token。

如果输入是：

input_ids = [A, B, C, D]

那么监督关系应该是：

看到 A      -> 预测 B
看到 A B    -> 预测 C
看到 A B C  -> 预测 D

这就是为什么要 shift：

logits positions: 0, 1, 2
labels positions: 1, 2, 3

模型输出的 logits[:, i, :] 不是用来预测 input_ids[:, i]，而是用来预测下一位 labels[:, i+1]。

源码证据 A：shift logits 和 labels

文件：model/model_minimind.py:249-252

看它是为了理解：next-token 对齐关系在源码里是哪一行实现的。

源码：

loss = None
if labels is not None:
    x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
    loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)

这段代码说明：

logits[..., :-1, :] 去掉最后一个位置的预测。
labels[..., 1:] 去掉第一个位置的目标。
x 和 y 的 seq 维都变成 seq-1。
这样 x[:, i, :] 就对齐 y[:, i]，也就是原始的 labels[:, i+1]。

理解到这一步就够

你应该能回答：

为什么 labels[0] 不参与 loss？
因为没有前一个位置的 logits 去预测它。

为什么 logits 最后一个位置不参与 loss？
因为序列里没有下一个 label 给它预测。

暂时不用看：

teacher forcing 的训练范式。
推理时如何一步步 autoregressive generate。

5. 原理三：`-100` 决定哪些位置不参与 loss

原理讲解

labels 不只是目标 token id，也是一张 loss mask。

规则是：

label 是正常 token id -> 这个位置参与 loss
label 是 -100       -> 这个位置被忽略

PyTorch 的 cross entropy 支持 ignore_index。MiniMind 设置：

ignore_index=-100

所以只要 Dataset 把某些位置的 label 写成 -100，这些位置就不会对 loss 和梯度产生贡献。

注意 shift 后才判断是否忽略：

真正参与 loss 的是 labels[..., 1:] 里的非 -100 位置。

源码证据 A：cross entropy 忽略 `-100`

文件：model/model_minimind.py:250-252

看它是为了理解：-100 为什么不是普通 token id，而是 loss mask。

源码：

if labels is not None:
    x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
    loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)

这段代码说明：

y.view(-1) 里有正常 token id，也可能有 -100。
ignore_index=-100 让 -100 位置不参与 loss。
Dataset 决定哪些位置是 -100，模型 forward 只负责执行这个规则。

理解到这一步就够

你应该能说清楚：

-100 不是 tokenizer 词表里的特殊 token；
它是 PyTorch loss 的忽略标记。

暂时不用看：

reduction='mean' 时忽略位置如何影响平均值。
label smoothing。

6. 原理四：Pretrain 和 SFT 共享同一个 forward loss

原理讲解

从训练循环看，Pretrain 和 SFT 都是同一种调用：

model(input_ids, labels=labels)

训练脚本没有写两套 loss。模型 forward 也没有判断“现在是 pretrain 还是 SFT”。

这说明一个关键事实：

阶段差异不在 loss 函数；
阶段差异在 labels 是怎么构造的。

Pretrain 传进来的 labels 和 SFT 传进来的 labels 不一样，所以同一个 causal LM loss 会产生不同训练目标。

源码证据 A：Pretrain 训练循环

文件：trainer/train_pretrain.py:35-40

看它是为了理解：pretrain 训练时如何调用模型 loss。

源码：

with autocast_ctx:
    res = model(input_ids, labels=labels)
    loss = res.loss + res.aux_loss
    loss = loss / args.accumulation_steps

scaler.scale(loss).backward()

这段代码说明：

pretrain 阶段直接把 input_ids 和 labels 交给模型。
res.loss 来自 MiniMindForCausalLM.forward。
训练循环本身没有手写 next-token loss。

源码证据 B：SFT 训练循环

文件：trainer/train_full_sft.py:35-40

看它是为了理解：SFT 阶段是否用了另一套 loss。

源码：

with autocast_ctx:
    res = model(input_ids, labels=labels)
    loss = res.loss + res.aux_loss
    loss = loss / args.accumulation_steps

scaler.scale(loss).backward()

这段代码说明：

SFT 调用方式和 pretrain 一样。
模型 forward 不知道当前阶段名称。
训练目标差异要去 Dataset 的 labels 里找。

理解到这一步就够

你应该能回答：

为什么 SFT 仍然可以使用 causal LM next-token loss？
因为 SFT 只是把对话 prompt 的 assistant 区域设成监督目标，其余位置设成 -100。

暂时不用看：

SFT 超参数选择。
SFT 是否会注入新知识。

7. 原理五：训练目标差异来自 labels

原理讲解

Pretrain 和 SFT 的 input_ids 都是 token 序列，但 labels 的含义不同。

Pretrain：

labels = input_ids.clone()
padding labels = -100

这表示普通文本几乎每个 token 都参与 next-token prediction。

SFT：

labels 默认全是 -100
只把 assistant 回答区改成 token id

这表示模型只在 assistant 回答区域被监督。user/system/tool prompt 只是条件上下文，不作为要模仿的答案。

关键点：

SFT 不是让模型学习“复读 user 问题”；
SFT 是让模型在 user/system 上下文之后，学习 assistant 应该怎么回答。

源码证据 A：Pretrain labels

文件：dataset/lm_dataset.py:47-55

看它是为了理解：pretrain 为什么几乎全段参与 loss。

源码：

def __getitem__(self, index):
    sample = self.samples[index]
    tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids
    tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
    input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
    input_ids = torch.tensor(input_ids, dtype=torch.long)
    labels = input_ids.clone()
    labels[input_ids == self.tokenizer.pad_token_id] = -100
    return input_ids, labels

这段代码说明：

labels = input_ids.clone() 是 pretrain 的核心。
padding 位置改成 -100，不参与 loss。
非 padding 文本 token 基本都作为 next-token 监督目标。

源码证据 B：SFT labels

文件：dataset/lm_dataset.py:88-119

看它是为了理解：SFT 如何只监督 assistant 区域。

源码：

def generate_labels(self, input_ids):
    labels = [-100] * len(input_ids)
    i = 0
    while i < len(input_ids):
        if input_ids[i:i + len(self.bos_id)] == self.bos_id:
            start = i + len(self.bos_id)
            end = start
            while end < len(input_ids):
                if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                    break
                end += 1
            for j in range(start, min(end + len(self.eos_id), self.max_length)):
                labels[j] = input_ids[j]
            i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
        else:
            i += 1
    return labels
...
labels = self.generate_labels(input_ids)
return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long)

这段代码说明：

SFT labels 初始全部是 -100。
代码找到 assistant 回答区域后，才把对应 label 改成 token id。
user/system/padding 默认保持 -100，所以不参与 loss。

理解到这一步就够

你应该能说清楚：

同一个 forward loss；
不同 Dataset 构造出来的 labels；
最终得到不同训练目标。

暂时不用看：

tool call 样本的完整标签细节。
多轮对话中多个 assistant span 的边界情况。

8. 实验验证

实验 A：手动复现 SFT 的 shifted loss

这个实验验证：

model.forward 的 loss
等于
logits[..., :-1, :] 与 labels[..., 1:] 手动 cross_entropy 的 loss

运行：

cd /home/sun/minimind
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_loss_shift.py --mode sft --show_rows 70

记录：

logits.shape =
shift_logits.shape =
shift_labels.shape =
model_forward_loss =
manual_shift_loss =
abs_diff =
active_loss_positions =
ignored_loss_positions =

你应该看到：

abs_diff 接近 0
shift_logits.shape = [1, seq-1, vocab_size]
shift_labels.shape = [1, seq-1]
SFT 里大量位置 participates=False

这说明模型源码里的 loss 就是 shifted next-token cross entropy。

实验 B：对比 Pretrain 的参与位置

运行：

PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_loss_shift.py --mode pretrain --show_rows 50

记录：

active_loss_positions =
ignored_loss_positions =
哪些位置是 padding 导致的 ignored

你应该看到：

pretrain 的非 padding 位置基本都参与 loss。
SFT 的 user/system/prompt 大量不参与 loss。
两者 manual_shift_loss 都和 model_forward_loss 一致。

9. 本节检查

如果你真懂了本节，应该能不看答案说清楚：

为什么 logits[..., :-1, :] 要对齐 labels[..., 1:]。
logits[:, i, :] 预测的是哪个位置的 token。
为什么 labels[:, 0] 不参与 loss。
ignore_index=-100 在 MiniMind 里起什么作用。
为什么 SFT 不需要单独写一个 SFT loss。
Pretrain 和 SFT 的训练目标差异到底来自哪里。

10. 下一课

第 11 课进入 Attention 原理与源码：我们会看 q_proj / k_proj / v_proj 如何把 [batch, seq, hidden_size] 拆成多头 Q/K/V，并开始解释 causal mask 与注意力权重。

15 KiB Raw Blame History Unescape Escape

第 10 课：Causal LM forward 与 loss shift

目录

0. 本节主线

1. 本节要懂的 5 个原理

2. 变量流转

3. 原理一：lm_head 把 hidden 映射成 logits

原理讲解

源码证据 A：forward 先得到 hidden，再调用 lm_head

理解到这一步就够

4. 原理二：next-token loss 必须 shift

原理讲解

源码证据 A：shift logits 和 labels

理解到这一步就够

5. 原理三：-100 决定哪些位置不参与 loss

原理讲解

源码证据 A：cross entropy 忽略 -100

理解到这一步就够

6. 原理四：Pretrain 和 SFT 共享同一个 forward loss

原理讲解

源码证据 A：Pretrain 训练循环

源码证据 B：SFT 训练循环

理解到这一步就够

7. 原理五：训练目标差异来自 labels

原理讲解

源码证据 A：Pretrain labels

源码证据 B：SFT labels

理解到这一步就够

8. 实验验证

实验 A：手动复现 SFT 的 shifted loss

实验 B：对比 Pretrain 的参与位置

9. 本节检查

10. 下一课

15 KiB

Raw Blame History

3. 原理一：`lm_head` 把 hidden 映射成 logits

5. 原理三：`-100` 决定哪些位置不参与 loss

源码证据 A：cross entropy 忽略 `-100`