784 lines
22 KiB
Markdown
784 lines
22 KiB
Markdown
# 第 20 课：Distillation 训练链路
|
||
|
||
这一课只解决一个问题：MiniMind 的 `train_distillation.py` 如何让 student 同时学习真实 assistant label 和 teacher 的 soft logits。
|
||
|
||
## 目录
|
||
|
||
- [0. 本节主线](#l20-mainline)
|
||
- [1. 本节要懂的 6 个原理](#l20-principles)
|
||
- [2. 蒸馏训练完整原理](#l20-complete-principle)
|
||
- [3. 源码阅读顺序图](#l20-reading-order)
|
||
- [4. MiniMind 源码走读](#l20-source-walkthrough)
|
||
- [5. 本节必须会写 / 暂时不要求](#l20-must-write)
|
||
- [6. 实验验证](#l20-experiment)
|
||
- [7. 阶段组装](#l20-stage-assembly)
|
||
- [8. 本节检查](#l20-check)
|
||
- [9. 下一课](#l20-next)
|
||
|
||
<a id="l20-mainline"></a>
|
||
## 0. 本节主线
|
||
|
||
蒸馏训练流程是：
|
||
|
||
```text
|
||
加载 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 直接给答案”，而是：
|
||
|
||
```text
|
||
teacher 给每个位置、每个 token 的概率分布；
|
||
student 学真实 label，也学 teacher 对其它 token 的相对偏好。
|
||
```
|
||
|
||
一句话：
|
||
|
||
```text
|
||
SFT 只告诉 student 正确 token 是什么，蒸馏还告诉 student teacher 认为哪些 token 也接近合理。
|
||
```
|
||
|
||
<a id="l20-principles"></a>
|
||
## 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` |
|
||
|
||
学完本节，你应该能说明：
|
||
|
||
```text
|
||
distillation_loss 的输入 shape 是什么；
|
||
teacher 为什么不 backward；
|
||
loss_mask 为什么来自 labels[..., 1:]；
|
||
alpha 和 temperature 分别控制什么；
|
||
最终保存的是 student，不是 teacher。
|
||
```
|
||
|
||
<a id="l20-complete-principle"></a>
|
||
## 2. 蒸馏训练完整原理
|
||
|
||
### 2.1 SFT 的监督信号太硬
|
||
|
||
普通 SFT 对每个位置只有一个目标 token。
|
||
|
||
假设当前位置的正确 token 是：
|
||
|
||
```text
|
||
北京
|
||
```
|
||
|
||
CE loss 看的是：
|
||
|
||
```text
|
||
模型给“北京”的概率够不够高
|
||
```
|
||
|
||
它不会告诉 student：
|
||
|
||
```text
|
||
“上海”是不是比“苹果”更接近合理；
|
||
“中国”是不是也有一定语义相关性；
|
||
标点、停用词、同义表达之间有什么相对偏好。
|
||
```
|
||
|
||
真实 label 是 hard target：
|
||
|
||
```text
|
||
正确 token = 1
|
||
其它 token = 0
|
||
```
|
||
|
||
蒸馏引入 teacher 的 soft target：
|
||
|
||
```text
|
||
teacher 对整个 vocab 的概率分布
|
||
```
|
||
|
||
这让 student 不只学习“唯一答案”，还学习 teacher 的分布形状。
|
||
|
||
### 2.2 teacher logits 是什么
|
||
|
||
teacher 和 student 都是 causal LM。
|
||
|
||
对同一个 `input_ids`，它们都会输出：
|
||
|
||
```text
|
||
student_logits: [batch, seq, vocab]
|
||
teacher_logits: [batch, seq, vocab]
|
||
```
|
||
|
||
第 10 课讲过 next-token shift，所以训练时真正使用：
|
||
|
||
```text
|
||
student_logits[..., :-1, :]
|
||
teacher_logits[..., :-1, :]
|
||
labels[..., 1:]
|
||
```
|
||
|
||
含义是：
|
||
|
||
```text
|
||
位置 i 的 logits 用来预测位置 i+1 的 label。
|
||
```
|
||
|
||
### 2.3 CE loss 和 distill loss 分工
|
||
|
||
MiniMind 的蒸馏总损失是：
|
||
|
||
```text
|
||
loss = alpha * CE + (1 - alpha) * Distill
|
||
```
|
||
|
||
CE loss 来自真实数据：
|
||
|
||
```text
|
||
student_logits -> labels
|
||
```
|
||
|
||
Distill loss 来自 teacher：
|
||
|
||
```text
|
||
student_logits -> teacher_logits
|
||
```
|
||
|
||
两者都只在 assistant label 区域上计算。这个区域由：
|
||
|
||
```python
|
||
labels[..., 1:] != -100
|
||
```
|
||
|
||
决定。
|
||
|
||
为什么不在 user/system 区域算 teacher KL？
|
||
|
||
因为 SFTDataset 把 user/system/prompt 区域设为 `-100`，表示这些 token 只是条件，不是训练目标。蒸馏训练沿用这个边界，避免 student 被训练去复述 prompt。
|
||
|
||
### 2.4 temperature 的作用
|
||
|
||
蒸馏时常把 logits 除以 temperature：
|
||
|
||
```text
|
||
softmax(logits / T)
|
||
```
|
||
|
||
当 `T > 1` 时，概率分布会更平滑：
|
||
|
||
```text
|
||
最大 token 的概率会降低；
|
||
次优 token 的概率会抬高；
|
||
student 更容易看到 teacher 的相对偏好。
|
||
```
|
||
|
||
MiniMind 默认：
|
||
|
||
```text
|
||
temperature = 1.5
|
||
```
|
||
|
||
源码最后乘回：
|
||
|
||
```text
|
||
temperature ** 2
|
||
```
|
||
|
||
这是常见蒸馏写法，用来补偿 temperature 对梯度尺度的影响。第 20 课只需要理解它保持 loss 尺度更稳定，不必推完整梯度公式。
|
||
|
||
### 2.5 teacher 为什么不更新
|
||
|
||
teacher 是训练目标的来源，不是被训练对象。
|
||
|
||
所以 teacher 必须：
|
||
|
||
```text
|
||
eval()
|
||
requires_grad_(False)
|
||
with torch.no_grad()
|
||
```
|
||
|
||
这三件事分别解决：
|
||
|
||
```text
|
||
eval(): 关闭 dropout 等训练态行为
|
||
requires_grad_(False): teacher 参数不需要梯度
|
||
no_grad(): teacher forward 不构建反向图
|
||
```
|
||
|
||
如果 teacher 也更新，目标本身会变化，student 学的分布不稳定；同时显存和计算都会浪费。
|
||
|
||
<a id="l20-reading-order"></a>
|
||
## 3. 源码阅读顺序图
|
||
|
||
建议按这个顺序读：
|
||
|
||
```text
|
||
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。
|
||
```
|
||
|
||
和前几课的关系：
|
||
|
||
```text
|
||
第 5 课：SFTDataset labels mask
|
||
第 10 课：logits/labels shift
|
||
第 14 课：lr、grad accumulation、AMP
|
||
第 16 课：Full SFT 训练入口
|
||
第 20 课：在 SFT 训练上额外加 teacher KL
|
||
```
|
||
|
||
<a id="l20-source-walkthrough"></a>
|
||
## 4. MiniMind 源码走读
|
||
|
||
### 4.1 脚本参数决定 student、teacher、alpha 和 temperature
|
||
|
||
#### 源码证据 A：student/teacher 参数
|
||
|
||
文件：`trainer/train_distillation.py:146-176`
|
||
|
||
看它是为了理解：蒸馏脚本和普通 SFT 脚本相比多了哪些训练对象和超参数。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 的平滑程度。
|
||
|
||
理解到这一步就够：
|
||
|
||
```text
|
||
蒸馏训练至少要同时关心两套模型配置，而不是只关心一个 lm_config。
|
||
```
|
||
|
||
暂时不要看：
|
||
|
||
```text
|
||
teacher 比 student 大多少效果最好；
|
||
alpha/temperature 搜参策略。
|
||
```
|
||
|
||
### 4.2 student 被训练，teacher 被冻结
|
||
|
||
#### 源码证据 A：构造两套 config
|
||
|
||
文件：`trainer/train_distillation.py:184-188`
|
||
|
||
看它是为了理解：student 和 teacher 在代码里是两套独立配置。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 目标。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 但不参与反向传播。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 维度对齐：
|
||
|
||
```text
|
||
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`。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 如何结合。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 对最终训练目标的影响。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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`
|
||
|
||
看它是为了理解：训练时应该观察哪些指标。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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`
|
||
|
||
看它是为了理解：蒸馏训练产物是什么。
|
||
|
||
代码摘录：
|
||
|
||
```python
|
||
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 权重名保存。
|
||
|
||
<a id="l20-must-write"></a>
|
||
## 5. 本节必须会写 / 暂时不要求
|
||
|
||
必须会写：
|
||
|
||
```text
|
||
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。
|
||
```
|
||
|
||
暂时不要求：
|
||
|
||
```text
|
||
1. 蒸馏效果评测。
|
||
2. teacher/student 架构搜索。
|
||
3. 多 teacher ensemble。
|
||
4. hidden states / attention map 蒸馏。
|
||
5. tokenizer 不一致时的 vocab 映射。
|
||
6. temperature 梯度尺度公式推导。
|
||
```
|
||
|
||
<a id="l20-experiment"></a>
|
||
## 6. 实验验证
|
||
|
||
本节新增实验：
|
||
|
||
```text
|
||
course/labs/trace_distillation_loss.py
|
||
```
|
||
|
||
这个实验不需要真实 `full_sft` 权重。它用随机初始化的 tiny student 和 tiny teacher 跑一次前向，只验证：
|
||
|
||
```text
|
||
input_ids / labels shape
|
||
student_logits / teacher_logits shape
|
||
loss_mask active token 数量
|
||
CE loss
|
||
distill loss
|
||
alpha 混合后的 total loss
|
||
```
|
||
|
||
运行命令：
|
||
|
||
```bash
|
||
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
|
||
```
|
||
|
||
重点记录：
|
||
|
||
```text
|
||
student_logits.shape =
|
||
teacher_logits.shape =
|
||
active_tokens =
|
||
ce_loss_raw =
|
||
distill_loss =
|
||
total_loss =
|
||
```
|
||
|
||
你应该看到：
|
||
|
||
```text
|
||
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 权重，例如：
|
||
|
||
```text
|
||
out/full_sft_768.pth
|
||
out/full_sft_768_moe.pth
|
||
dataset/sft_t2t_mini.jsonl
|
||
```
|
||
|
||
可以跑原项目蒸馏训练：
|
||
|
||
```bash
|
||
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 的张量和数值关系。
|
||
|
||
<a id="l20-stage-assembly"></a>
|
||
## 7. 阶段组装
|
||
|
||
第 20 课不新增 `course/impl/` 的手写训练阶段。原因是课程主线里的手写阶段已经覆盖：
|
||
|
||
```text
|
||
Pretrain
|
||
SFT
|
||
LoRA
|
||
DPO
|
||
PPO/GRPO
|
||
```
|
||
|
||
蒸馏在这里作为源码训练链路学习：
|
||
|
||
```text
|
||
理解 train_distillation.py
|
||
-> 用 lab 复现 CE + KL 的核心计算
|
||
-> 知道真实训练需要两套权重
|
||
```
|
||
|
||
如果后续要扩展教学版蒸馏，可以新增：
|
||
|
||
```text
|
||
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` 记录：
|
||
|
||
```text
|
||
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` 记录容易错的点：
|
||
|
||
```text
|
||
1. 忘记 teacher no_grad，导致 teacher 也进入反向图。
|
||
2. loss_mask 用 labels 而不是 shifted labels。
|
||
3. CE 和 KL 使用的 token 区域不一致。
|
||
4. student/teacher vocab 维度不一致时直接做 KL。
|
||
5. 把 alpha 理解反了：源码里 alpha 是 CE 权重。
|
||
```
|
||
|
||
<a id="l20-check"></a>
|
||
## 8. 本节检查
|
||
|
||
1. 蒸馏训练为什么要同时有 CE loss 和 distill loss？
|
||
2. `loss_mask = labels[..., 1:] != -100` 为什么要用 shifted labels？
|
||
3. teacher 为什么要 `eval()`、`requires_grad_(False)` 和 `torch.no_grad()`？
|
||
4. `temperature > 1` 对 teacher soft target 有什么影响？
|
||
5. 源码里的 `alpha` 是 CE 权重还是 distill 权重？
|
||
6. 为什么 KL 只在 assistant active token 上计算？
|
||
7. 蒸馏训练最终保存的是 student 还是 teacher？
|
||
|
||
<a id="l20-next"></a>
|
||
## 9. 下一课
|
||
|
||
第 21 课进入 DPO 数据和 reference model：
|
||
|
||
- `21_dpo_dataset_and_reference.md`
|
||
|
||
下一课要解决：
|
||
|
||
```text
|
||
DPO 数据为什么需要 chosen/rejected；
|
||
policy model 和 reference model 分别是什么；
|
||
DPODataset 如何构造 chosen/rejected 的 input、label 和 mask；
|
||
reference model 为什么冻结；
|
||
DPO 和 SFT/Distillation 的训练信号有什么不同。
|
||
```
|