22 KiB

Raw Blame History

第 20 课：Distillation 训练链路

这一课只解决一个问题：MiniMind 的 train_distillation.py 如何让 student 同时学习真实 assistant label 和 teacher 的 soft logits。

0. 本节主线

蒸馏训练流程是：

加载 student 模型
-> 加载 teacher 模型
-> 冻结 teacher
-> 用 SFTDataset 得到 input_ids / labels
-> student 前向得到 student_logits
-> teacher no_grad 前向得到 teacher_logits
-> labels 生成 CE loss
-> teacher soft distribution 生成 KL distillation loss
-> 总 loss = alpha * CE + (1 - alpha) * distill
-> 只反传和更新 student
-> 保存 student 权重

这节课的核心不是“teacher 直接给答案”，而是：

teacher 给每个位置、每个 token 的概率分布；
student 学真实 label，也学 teacher 对其它 token 的相对偏好。

一句话：

SFT 只告诉 student 正确 token 是什么，蒸馏还告诉 student teacher 认为哪些 token 也接近合理。

1. 本节要懂的 6 个原理

原理	要理解什么	源码位置
student/teacher 是两套模型	student 被训练，teacher 只提供 soft target	`trainer/train_distillation.py:204-214`
teacher 必须冻结	teacher 用 eval/no_grad/requires_grad_(False) 避免更新和节省显存	`trainer/train_distillation.py:43-45`, `trainer/train_distillation.py:207-209`
蒸馏仍用 SFTDataset	数据、chat template、assistant-only labels 仍沿用 SFT	`trainer/train_distillation.py:211`, `dataset/lm_dataset.py:88-119`
CE loss 学真实 label	student logits 与 labels shift 后计算 token-level CE	`trainer/train_distillation.py:68-80`
KL loss 学 teacher 分布	teacher softmax 后作为 target，student log_softmax 后做 KL	`trainer/train_distillation.py:25-36`, `trainer/train_distillation.py:82-90`
alpha/temperature 控制训练信号	alpha 混合 CE/KL，temperature 平滑 teacher 分布	`trainer/train_distillation.py:92-93`, `trainer/train_distillation.py:171-173`

学完本节，你应该能说明：

distillation_loss 的输入 shape 是什么；
teacher 为什么不 backward；
loss_mask 为什么来自 labels[..., 1:]；
alpha 和 temperature 分别控制什么；
最终保存的是 student，不是 teacher。

2. 蒸馏训练完整原理

2.1 SFT 的监督信号太硬

普通 SFT 对每个位置只有一个目标 token。

假设当前位置的正确 token 是：

北京

CE loss 看的是：

模型给“北京”的概率够不够高

它不会告诉 student：

“上海”是不是比“苹果”更接近合理；
“中国”是不是也有一定语义相关性；
标点、停用词、同义表达之间有什么相对偏好。

真实 label 是 hard target：

正确 token = 1
其它 token = 0

蒸馏引入 teacher 的 soft target：

teacher 对整个 vocab 的概率分布

这让 student 不只学习“唯一答案”，还学习 teacher 的分布形状。

2.2 teacher logits 是什么

teacher 和 student 都是 causal LM。

对同一个 input_ids，它们都会输出：

student_logits: [batch, seq, vocab]
teacher_logits: [batch, seq, vocab]

第 10 课讲过 next-token shift，所以训练时真正使用：

student_logits[..., :-1, :]
teacher_logits[..., :-1, :]
labels[..., 1:]

含义是：

位置 i 的 logits 用来预测位置 i+1 的 label。

2.3 CE loss 和 distill loss 分工

MiniMind 的蒸馏总损失是：

loss = alpha * CE + (1 - alpha) * Distill

CE loss 来自真实数据：

student_logits -> labels

Distill loss 来自 teacher：

student_logits -> teacher_logits

两者都只在 assistant label 区域上计算。这个区域由：

labels[..., 1:] != -100

决定。

为什么不在 user/system 区域算 teacher KL？

因为 SFTDataset 把 user/system/prompt 区域设为 -100，表示这些 token 只是条件，不是训练目标。蒸馏训练沿用这个边界，避免 student 被训练去复述 prompt。

2.4 temperature 的作用

蒸馏时常把 logits 除以 temperature：

softmax(logits / T)

当 T > 1 时，概率分布会更平滑：

最大 token 的概率会降低；
次优 token 的概率会抬高；
student 更容易看到 teacher 的相对偏好。

MiniMind 默认：

temperature = 1.5

源码最后乘回：

temperature ** 2

这是常见蒸馏写法，用来补偿 temperature 对梯度尺度的影响。第 20 课只需要理解它保持 loss 尺度更稳定，不必推完整梯度公式。

2.5 teacher 为什么不更新

teacher 是训练目标的来源，不是被训练对象。

所以 teacher 必须：

eval()
requires_grad_(False)
with torch.no_grad()

这三件事分别解决：

eval(): 关闭 dropout 等训练态行为
requires_grad_(False): teacher 参数不需要梯度
no_grad(): teacher forward 不构建反向图

如果 teacher 也更新，目标本身会变化，student 学的分布不稳定；同时显存和计算都会浪费。

3. 源码阅读顺序图

建议按这个顺序读：

1. train_distillation.py:146-176
   看脚本参数：student/teacher 配置、alpha、temperature。

2. train_distillation.py:184-214
   看 student 和 teacher 如何初始化、冻结、建 optimizer。

3. train_distillation.py:39-67
   看一个 batch 里 student/teacher logits 如何取得。

4. train_distillation.py:68-93
   看 CE loss、distill loss、总 loss 如何组合。

5. train_distillation.py:25-36
   单独看 KL 蒸馏损失函数。

6. train_distillation.py:124-143
   看保存 student 和尾批 gradient accumulation。

和前几课的关系：

第 5 课：SFTDataset labels mask
第 10 课：logits/labels shift
第 14 课：lr、grad accumulation、AMP
第 16 课：Full SFT 训练入口
第 20 课：在 SFT 训练上额外加 teacher KL

4. MiniMind 源码走读

4.1 脚本参数决定 student、teacher、alpha 和 temperature

源码证据 A：student/teacher 参数

文件：trainer/train_distillation.py:146-176

看它是为了理解：蒸馏脚本和普通 SFT 脚本相比多了哪些训练对象和超参数。

代码摘录：

parser.add_argument('--student_hidden_size', default=768, type=int)
parser.add_argument('--student_num_layers', default=8, type=int)
parser.add_argument('--teacher_hidden_size', default=768, type=int)
parser.add_argument('--teacher_num_layers', default=8, type=int)
parser.add_argument('--student_use_moe', default=0, type=int, choices=[0, 1])
parser.add_argument('--teacher_use_moe', default=1, type=int, choices=[0, 1])
parser.add_argument('--from_student_weight', default='full_sft', type=str)
parser.add_argument('--from_teacher_weight', default='full_sft', type=str)
parser.add_argument('--alpha', default=0.5, type=float)
parser.add_argument('--temperature', default=1.5, type=float)

这段代码说明：

student 和 teacher 可以有不同结构。
默认场景是 teacher 用 MoE，student 用 dense。
两者默认都从 full_sft 权重开始。
alpha 控制 CE 和 KL 的权重。
temperature 控制 teacher soft target 的平滑程度。

理解到这一步就够：

蒸馏训练至少要同时关心两套模型配置，而不是只关心一个 lm_config。

暂时不要看：

teacher 比 student 大多少效果最好；
alpha/temperature 搜参策略。

4.2 student 被训练，teacher 被冻结

源码证据 A：构造两套 config

文件：trainer/train_distillation.py:184-188

看它是为了理解：student 和 teacher 在代码里是两套独立配置。

代码摘录：

lm_config_student = MiniMindConfig(hidden_size=args.student_hidden_size, num_hidden_layers=args.student_num_layers, use_moe=bool(args.student_use_moe))
lm_config_teacher = MiniMindConfig(hidden_size=args.teacher_hidden_size, num_hidden_layers=args.teacher_num_layers, use_moe=bool(args.teacher_use_moe))
ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None

这段代码说明：

student 和 teacher 都用 MiniMindConfig 构建。
resume checkpoint 按 student 配置保存。
最终被训练和恢复的是 student。

源码证据 B：加载并冻结 teacher

文件：trainer/train_distillation.py:204-214

看它是为了理解：optimizer 为什么只更新 student。

代码摘录：

model, tokenizer = init_model(lm_config_student, args.from_student_weight, device=args.device)
teacher_model, _ = init_model(lm_config_teacher, args.from_teacher_weight, device=args.device)
teacher_model.eval()
teacher_model.requires_grad_(False)
train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

这段代码说明：

model 是 student。
teacher_model 单独加载。
teacher 被设置为 eval 并冻结参数。
optimizer 只接收 model.parameters()，也就是 student 参数。
数据仍然用 student tokenizer 构造。

4.3 一个 batch 同时跑 student 和 teacher

源码证据 A：student 前向

文件：trainer/train_distillation.py:47-60

看它是为了理解：student logits 如何对齐 next-token 目标。

代码摘录：

input_ids = input_ids.to(args.device)
labels = labels.to(args.device)
loss_mask = (labels[..., 1:] != -100).float()

with autocast_ctx:
    res = model(input_ids)
    student_logits = res.logits[..., :-1, :].contiguous()

这段代码说明：

loss_mask 来自 shifted labels。
student forward 没传 labels，因为脚本要手写 CE 和 KL。
student_logits[..., :-1, :] 去掉最后一个位置。
这和第 10 课的 next-token shift 一致。

源码证据 B：teacher no_grad 前向

文件：trainer/train_distillation.py:61-67

看它是为了理解：teacher logits 如何参与 loss 但不参与反向传播。

代码摘录：

if teacher_model is not None:
    with torch.no_grad():
        teacher_logits = teacher_model(input_ids).logits[..., :-1, :].contiguous()
        vocab_size_student = student_logits.size(-1)
        teacher_logits = teacher_logits[..., :vocab_size_student]

这段代码说明：

teacher forward 包在 torch.no_grad() 里。
teacher logits 和 student logits 一样做 [..., :-1, :]。
teacher vocab 会被截断到 student vocab 大小。
KL loss 需要 student 和 teacher 最后一维一致。

注意这个截断只解决 vocab 维度对齐：

teacher_logits[..., :vocab_size_student]

它不解决 tokenizer 不一致问题。真正训练时 student 和 teacher 应该使用兼容 tokenizer，否则同一个 token id 的含义可能不同。

4.4 CE loss 只在 assistant 区域上算

源码证据：手写 CE loss

文件：trainer/train_distillation.py:68-80

看它是为了理解：为什么 distillation 脚本没有直接用 model(input_ids, labels=labels) 的 res.loss。

代码摘录：

shift_labels = labels[..., 1:].contiguous()
loss_mask_flat = loss_mask.view(-1)
ce_loss = F.cross_entropy(
    student_logits.view(-1, student_logits.size(-1)),
    shift_labels.view(-1),
    ignore_index=-100,
    reduction='none'
)
ce_loss_raw = torch.sum(ce_loss * loss_mask_flat) / (loss_mask_flat.sum() + 1e-8)
if lm_config_student.use_moe: ce_loss = ce_loss_raw + res.aux_loss
else: ce_loss = ce_loss_raw

这段代码说明：

CE loss 的 logits 和 labels 仍然 shift 对齐。
reduction='none' 保留每个 token 的 loss。
loss_mask_flat 只保留 assistant 训练区域。
ce_loss_raw 是 active token 上的平均 CE。
如果 student 是 MoE，还要加 res.aux_loss。

为什么不用 res.loss？

因为蒸馏脚本后面还要对同一批 active token 计算 KL。手写 CE 能让 CE 和 KL 使用同一个 loss_mask_flat，日志里也能分别打印 ce_loss 和 distill_loss。

4.5 KL distillation loss 学 teacher 分布

源码证据 A：distillation_loss 函数

文件：trainer/train_distillation.py:25-36

看它是为了理解：teacher logits 如何变成 student 的 soft target。

代码摘录：

def distillation_loss(student_logits, teacher_logits, temperature=1.0, reduction='batchmean'):
    with torch.no_grad():
        teacher_probs = F.softmax(teacher_logits / temperature, dim=-1).detach()

    student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

    kl = F.kl_div(
        student_log_probs,
        teacher_probs,
        reduction=reduction
    )
    return (temperature ** 2) * kl

这段代码说明：

teacher logits 先除以 temperature，再 softmax 成概率。
teacher 概率 detach，不参与梯度。
student logits 变成 log probability。
F.kl_div 的输入是 student log probs，target 是 teacher probs。
最后乘 temperature ** 2。

源码证据 B：只在 active token 上算 KL

文件：trainer/train_distillation.py:82-90

看它是为了理解：teacher KL 和 SFT labels mask 如何结合。

代码摘录：

distill_loss = distillation_loss(
    student_logits.view(-1, student_logits.size(-1))[loss_mask_flat == 1],
    teacher_logits.view(-1, teacher_logits.size(-1))[loss_mask_flat == 1],
    temperature=temperature
)

这段代码说明：

logits 被展平成 [batch * (seq - 1), vocab]。
只取 loss_mask_flat == 1 的 token。
user/system/padding 区域不会进入 KL。
CE 和 KL 使用同一批 active token。

4.6 总 loss 混合 CE 和 distill

源码证据 A：loss 组合

文件：trainer/train_distillation.py:92-102

看它是为了理解：alpha 对最终训练目标的影响。

代码摘录：

loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps

scaler.scale(loss).backward()

if step % args.accumulation_steps == 0:
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

这段代码说明：

alpha=1.0 时接近普通 SFT。
alpha=0.0 时只学 teacher 分布。
默认 alpha=0.5，CE 和 KL 各占一半。
梯度累积、AMP、clip、optimizer step 沿用普通训练脚本。
optimizer 只管理 student，所以只更新 student。

源码证据 B：日志同时打印 CE 和 distill

文件：trainer/train_distillation.py:104-122

看它是为了理解：训练时应该观察哪些指标。

代码摘录：

Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, ce: {current_ce_loss:.4f}, aux_loss: {current_aux_loss:.4f}, distill: {distill_loss.item():.4f}, learning_rate: {current_lr:.8f}, epoch_time: {eta_min:.3f}min')

这段代码说明：

loss 是混合后的总损失。
ce 是真实 label 的监督损失。
distill 是 teacher 分布的 KL 损失。
MoE student 时还要观察 aux_loss。

4.7 保存的是 student 权重

源码证据：保存 checkpoint

文件：trainer/train_distillation.py:124-143

看它是为了理解：蒸馏训练产物是什么。

代码摘录：

ckp = f'{args.save_dir}/{args.save_weight}_{lm_config_student.hidden_size}{moe_suffix}.pth'
raw_model = model.module if isinstance(model, DistributedDataParallel) else model
raw_model = getattr(raw_model, '_orig_mod', raw_model)
state_dict = raw_model.state_dict()
torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
lm_checkpoint(lm_config_student, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')

这段代码说明：

保存路径使用 student config。
保存的是当前训练中的 model，也就是 student。
teacher 不会被保存成训练产物。
resume checkpoint 也按 student 权重名保存。

5. 本节必须会写 / 暂时不要求

必须会写：

1. 说明 distillation_loss 的输入输出：
   student_logits: [active_tokens, vocab]
   teacher_logits: [active_tokens, vocab]
   output: scalar KL loss

2. 手动写出 teacher soft target：
   teacher_probs = softmax(teacher_logits / temperature)

3. 手动写出 student log probability：
   student_log_probs = log_softmax(student_logits / temperature)

4. 说明 loss_mask 的来源：
   loss_mask = labels[..., 1:] != -100

5. 说明总 loss：
   loss = alpha * ce_loss + (1 - alpha) * distill_loss

6. 说明 teacher 为什么 eval/no_grad/frozen。

暂时不要求：

1. 蒸馏效果评测。
2. teacher/student 架构搜索。
3. 多 teacher ensemble。
4. hidden states / attention map 蒸馏。
5. tokenizer 不一致时的 vocab 映射。
6. temperature 梯度尺度公式推导。

6. 实验验证

本节新增实验：

course/labs/trace_distillation_loss.py

这个实验不需要真实 full_sft 权重。它用随机初始化的 tiny student 和 tiny teacher 跑一次前向，只验证：

input_ids / labels shape
student_logits / teacher_logits shape
loss_mask active token 数量
CE loss
distill loss
alpha 混合后的 total loss

运行命令：

cd /home/sun/minimind
python course/labs/trace_distillation_loss.py \
  --data_path course/labs/tiny_sft.jsonl \
  --max_length 96 \
  --alpha 0.5 \
  --temperature 1.5

重点记录：

student_logits.shape =
teacher_logits.shape =
active_tokens =
ce_loss_raw =
distill_loss =
total_loss =

你应该看到：

student_logits.shape 和 teacher_logits.shape 的最后一维都是 vocab_size；
active_tokens 大于 0；
ce_loss_raw、distill_loss、total_loss 都是标量；
total_loss 约等于 alpha * ce_loss_raw + (1-alpha) * distill_loss。

这个实验验证的是计算路径，不验证模型效果。随机 teacher 不会教出好 student，但足够让你看懂源码里的张量流。

6.1 真实训练命令

如果本地已经有 student 和 teacher 权重，例如：

out/full_sft_768.pth
out/full_sft_768_moe.pth
dataset/sft_t2t_mini.jsonl

可以跑原项目蒸馏训练：

cd /home/sun/minimind/trainer
python train_distillation.py \
  --from_student_weight full_sft \
  --from_teacher_weight full_sft \
  --student_use_moe 0 \
  --teacher_use_moe 1 \
  --alpha 0.5 \
  --temperature 1.5 \
  --epochs 1

当前课程环境如果缺默认权重和数据，就不要把真实训练作为本节必跑验收。先跑 tiny trace 实验，确认 CE/KL 的张量和数值关系。

7. 阶段组装

第 20 课不新增 course/impl/ 的手写训练阶段。原因是课程主线里的手写阶段已经覆盖：

Pretrain
SFT
LoRA
DPO
PPO/GRPO

蒸馏在这里作为源码训练链路学习：

理解 train_distillation.py
-> 用 lab 复现 CE + KL 的核心计算
-> 知道真实训练需要两套权重

如果后续要扩展教学版蒸馏，可以新增：

course/impl/core/losses.py::distillation_loss
course/impl/train_distillation_impl.py
course/impl/tests/test_distillation_loss.py

但第 20 课先不强制做这件事，避免在 DPO 之前引入太多额外阶段。

7.1 和前面课程的连接

已学内容	在蒸馏里的作用
SFTDataset	提供 assistant-only labels 和 loss mask
loss shift	让 logits[..., :-1, :] 对齐 labels[..., 1:]
MoE aux loss	student 是 MoE 时加到 CE loss
AMP / grad accumulation	蒸馏训练循环继续复用
checkpoint	保存和恢复 student 训练状态

7.2 Portfolio 记录

完成本节后，可以在 course/portfolio/experiments.md 记录：

Distillation trace:
- 构造 tiny student/teacher。
- 对同一批 SFT input_ids 计算 student_logits 和 teacher_logits。
- 用 labels[..., 1:] != -100 只保留 assistant active tokens。
- 计算 CE loss、KL distillation loss 和 alpha 混合总 loss。

在 course/notes/mistakes.md 记录容易错的点：

1. 忘记 teacher no_grad，导致 teacher 也进入反向图。
2. loss_mask 用 labels 而不是 shifted labels。
3. CE 和 KL 使用的 token 区域不一致。
4. student/teacher vocab 维度不一致时直接做 KL。
5. 把 alpha 理解反了：源码里 alpha 是 CE 权重。

8. 本节检查

蒸馏训练为什么要同时有 CE loss 和 distill loss？
loss_mask = labels[..., 1:] != -100 为什么要用 shifted labels？
teacher 为什么要 eval()、requires_grad_(False) 和 torch.no_grad()？
temperature > 1 对 teacher soft target 有什么影响？
源码里的 alpha 是 CE 权重还是 distill 权重？
为什么 KL 只在 assistant active token 上计算？
蒸馏训练最终保存的是 student 还是 teacher？

9. 下一课

第 21 课进入 DPO 数据和 reference model：

21_dpo_dataset_and_reference.md

下一课要解决：

DPO 数据为什么需要 chosen/rejected；
policy model 和 reference model 分别是什么；
DPODataset 如何构造 chosen/rejected 的 input、label 和 mask；
reference model 为什么冻结；
DPO 和 SFT/Distillation 的训练信号有什么不同。

22 KiB Raw Blame History Unescape Escape

第 20 课：Distillation 训练链路

目录

0. 本节主线

1. 本节要懂的 6 个原理

2. 蒸馏训练完整原理

2.1 SFT 的监督信号太硬

2.2 teacher logits 是什么

2.3 CE loss 和 distill loss 分工

2.4 temperature 的作用

2.5 teacher 为什么不更新

3. 源码阅读顺序图

4. MiniMind 源码走读

4.1 脚本参数决定 student、teacher、alpha 和 temperature

源码证据 A：student/teacher 参数

4.2 student 被训练，teacher 被冻结

源码证据 A：构造两套 config

源码证据 B：加载并冻结 teacher

4.3 一个 batch 同时跑 student 和 teacher

源码证据 A：student 前向

源码证据 B：teacher no_grad 前向

4.4 CE loss 只在 assistant 区域上算

源码证据：手写 CE loss

4.5 KL distillation loss 学 teacher 分布

源码证据 A：distillation_loss 函数

源码证据 B：只在 active token 上算 KL

4.6 总 loss 混合 CE 和 distill

源码证据 A：loss 组合

源码证据 B：日志同时打印 CE 和 distill

4.7 保存的是 student 权重

源码证据：保存 checkpoint

5. 本节必须会写 / 暂时不要求

6. 实验验证

6.1 真实训练命令

7. 阶段组装

7.1 和前面课程的连接

7.2 Portfolio 记录

8. 本节检查

9. 下一课

22 KiB

Raw Blame History