一个现实困境

你花了几周训练出一个效果不错的 MoE 模型。正准备上线，老板说：

本章地图

"Many insects成功优化了它们的幼虫形态，

以便从环境中汲取最大量的营养和能量。

类似地，我们可以训练一个庞大的模型，

然后将它的知识蒸馏到一个更适合部署的小模型中。"

—— Hinton et al., "Distilling the Knowledge in a Neural Network", 2015

Part 1: 核心概念

理解知识蒸馏为什么存在、解决什么问题、核心角色是谁。本节不涉及数学公式，只建立直觉。

什么是蒸馏

用大模型的知识来训练小模型，实现模型压缩

定义：知识蒸馏是一种模型压缩技术。用一个已训练好的大模型（教师）的输出概率分布来指导小模型（学生）训练，使学生获得接近教师的能力。

白话：好比你是个小学生，旁边坐着一位大学教授。你做题时，教授不只告诉你"答案是 C"，还告诉你"A 和 B 也有点道理但不太对，D 完全不靠谱"。这些额外信息就是教授的"思维方式"，你学到的远不止标准答案。

这个比喻只是为了帮你建立直观感受。在计算机内部，它实际上是教师模型输出 logits 作为额外监督信号，结合 KL 散度损失函数来训练小模型——没有真正的"教授"和"学生"在对话。

教师与学生

蒸馏中的两个角色分工明确，不能搞混。

蒸馏流程

同一份数据同时喂给教师和学生，只训练学生

蒸馏的核心流程：一个输入，两条路径，两个损失，一个优化目标。

硬标签与软标签

这是理解蒸馏的关键一步。两种标签决定了学生能从训练中获取多少信息。

输入: "猫坐在___"

硬标签: [0, 0, 0, 1, 0, 0]   ← 只有"垫子上"是1，其余全是0
         墙上 地上 桌上 垫子上 椅子 天上

软标签: [0.01, 0.05, 0.10, 0.70, 0.12, 0.02]
         墙上  地上  桌上  垫子上  椅子  天上

暗知识

Hinton 说的"Dark Knowledge"到底是什么

Part 2: 数学工具

蒸馏用到了两个关键的数学工具：KL 散度和温度缩放。理解它们，才能看懂蒸馏损失函数的设计。

KL 散度

衡量两个概率分布"距离"的工具

定义：KL 散度（Kullback-Leibler Divergence）衡量两个概率分布 P 和 Q 之间的"距离"。公式：

D_KL(P || Q) = Σ P(i) × log(P(i) / Q(i))

白话：教师输出概率分布 P，学生输出概率分布 Q。你想知道 Q 离 P 有多"远"——这就是 KL 散度。Q 越接近 P，KL 散度越小，说明学生学得越好。

代码：PyTorch 里用 F.kl_div 一行搞定。

# P = 教师概率分布, Q = 学生概率分布
# PyTorch 要求第一个参数传入 Q 的 log 形式
kl = F.kl_div(Q_log, P, reduction='batchmean')

温度缩放

让 softmax 变平滑，暴露暗知识

定义：在 softmax 中引入温度 T，公式变为 q_i = exp(z_i/T) / Σ exp(z_j/T)。

白话：标准 softmax 太"尖"了——最高概率接近 1，其余接近 0，软标签退化成硬标签。升温（T > 1）让分布变平滑，暗知识才能暴露出来。就好比教授讲课：T=1 时只说"答案是 C"；T=2 时会解释"每个选项的优劣"。

这个比喻只是为了帮你建立直观感受。在计算机内部，它实际上是温度参数通过 logits / T 运算来改变 softmax 输出的平滑程度——没有真正的"教授"在讲课。

代码：

# 标准 softmax（T=1）
probs = F.softmax(logits, dim=-1)

# 温度缩放 softmax（T=1.5，MiniMind 默认）
probs = F.softmax(logits / temperature, dim=-1)

温度数值对比

假设教师模型对某位置的 logits：[2.0, 1.0, 0.1, 4.0, 0.5]，对应词汇：猫、跑了、睡了、坐下、跳了。

为什么乘 T²

蒸馏中最容易被忽略的细节

看 distillation_loss 函数的返回语句：

return (temperature ** 2) * kl

组合损失

硬标签和软标签各取所长

定义：蒸馏总损失由两部分加权组合：

白话：学生不能只跟教师学（会偏离真实数据），也不能只看标准答案（浪费了教师的知识）。α 就是"听老师讲"和"看标准答案"的平衡。

代码（来自 train_distillation.py）：

# 总损失 = alpha * CE + (1-alpha) * Distill
loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps

为什么双损失

只用一个损失行不行？

Part 3：从数学到实现

前面的 KL 散度和温度缩放都是数学公式。现在把这些公式变成能跑的代码——看看蒸馏训练究竟是怎么运转的。

训练脚本三件事

蒸馏训练只需做三件事：定义损失 → 同时跑两个模型 → 更新学生

# train_distillation.py 全局结构

# 1. 核心损失函数
def distillation_loss(student_logits, teacher_logits, temperature, reduction):
    ...

# 2. 训练主循环
def train_epoch(epoch, loader, iters, teacher_model, ...):
    ...

# 3. 主程序入口
if __name__ == "__main__":
    # 3.1 解析参数（alpha, temperature, 学生/教师配置）
    # 3.2 初始化模型（学生可训练，教师冻结）
    # 3.3 训练循环
    # 3.4 清理

蒸馏损失函数

整个文件最核心的 12 行代码

def distillation_loss(student_logits, teacher_logits,
                      temperature=1.0, reduction='batchmean'):
    # 教师端：关闭梯度 + 温度缩放 + softmax
    with torch.no_grad():
        teacher_probs = F.softmax(
            teacher_logits / temperature, dim=-1
        ).detach()

    # 学生端：温度缩放 + log_softmax
    student_log_probs = F.log_softmax(
        student_logits / temperature, dim=-1
    )

    # KL 散度 + T² 补偿
    kl = F.kl_div(
        student_log_probs, teacher_probs,
        reduction=reduction
    )
    return (temperature ** 2) * kl

数据怎么流

冻结教师模型

train_epoch 的第一步操作

train_epoch 函数开头，先冻结教师：

if teacher_model is not None:
    teacher_model.eval()              # 切到评估模式
    teacher_model.requires_grad_(False)  # 禁用所有参数梯度

双模型前向

同一个输入，两条路径

# 学生前向传播（在混合精度上下文中）
with autocast_ctx:
    res = model(input_ids)
    student_logits = res.logits[..., :-1, :].contiguous()

# 教师前向传播（纯推理，不计算梯度）
if teacher_model is not None:
    with torch.no_grad():
        teacher_logits = teacher_model(input_ids)
        teacher_logits = teacher_logits.logits[..., :-1, :].contiguous()

词表对齐

教师和学生词表大小不同时的处理

如果教师模型的词表比学生大，logits 的最后一维不匹配。MiniMind 用一行代码解决：

# 对齐词表：如果教师词表更大，截断到学生词表大小
vocab_size_student = student_logits.size(-1)
teacher_logits = teacher_logits[..., :vocab_size_student]

三重损失计算

CE + KL + 辅助损失，组合出击

# 1) 硬标签 CE 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,       # 忽略 padding 位置
    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

# 2) 软标签蒸馏 Loss（仅在有效 token 位置计算）
if teacher_model is not None:
    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
    )

# 3) 总损失 = alpha * CE + (1-alpha) * Distill
loss = (alpha * ce_loss + (1 - alpha) * distill_loss) / args.accumulation_steps

损失计算流程

loss_mask 解析

为什么需要遮罩？遮住了什么？

# loss_mask 来自 labels 中非 -100 的位置
loss_mask = (labels[..., 1:] != -100).float()
loss_mask_flat = loss_mask.view(-1)

# CE Loss：用 ignore_index=-100 忽略 padding
ce_loss = F.cross_entropy(..., ignore_index=-100, reduction='none')
ce_loss_raw = torch.sum(ce_loss * loss_mask_flat) / (loss_mask_flat.sum() + 1e-8)

# Distill Loss：用布尔索引只选有效位置
distill_loss = distillation_loss(
    student_logits.view(-1, ...)[loss_mask_flat == 1],
    teacher_logits.view(-1, ...)[loss_mask_flat == 1],
    temperature=temperature
)

梯度与更新

梯度累积、裁剪、更新的完整流程

# 反向传播（混合精度）
scaler.scale(loss).backward()

# 每 accumulation_steps 步执行一次更新
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)  # 清空梯度

Part 4: 实战与调优

知道原理后，怎么实际运行蒸馏？参数怎么选？不同场景的策略有何不同？

参数配置

蒸馏相关的所有可配置参数

# 学生模型配置（小模型 / Dense）
lm_config_student = MiniMindConfig(
    hidden_size=args.student_hidden_size,      # 默认 768
    num_hidden_layers=args.student_num_layers, # 默认 8
    use_moe=bool(args.student_use_moe)         # 默认 False（Dense）
)

# 教师模型配置（大模型 / MoE）
lm_config_teacher = MiniMindConfig(
    hidden_size=args.teacher_hidden_size,      # 默认 768
    num_hidden_layers=args.teacher_num_layers, # 默认 8
    use_moe=bool(args.teacher_use_moe)         # 默认 True（MoE）
)

MoE 到 Dense

MiniMind 的默认蒸馏场景

MoE 效果好但推理慢，蒸馏成 Dense 模型后部署更轻便：

python trainer/train_distillation.py \
    --student_hidden_size 768 \
    --student_num_layers 8 \
    --student_use_moe 0 \
    --teacher_hidden_size 768 \
    --teacher_num_layers 8 \
    --teacher_use_moe 1 \
    --alpha 0.5 \
    --temperature 1.5 \
    --from_student_weight full_sft \
    --from_teacher_weight full_sft

大模型到小模型

部署到资源受限环境的场景

需要部署到手机、边缘设备时，用大模型蒸馏出小模型：

python trainer/train_distillation.py \
    --student_hidden_size 512 \
    --student_num_layers 8 \
    --student_use_moe 0 \
    --teacher_hidden_size 1024 \
    --teacher_num_layers 16 \
    --teacher_use_moe 0 \
    --alpha 0.5 \
    --temperature 2.0 \
    --from_student_weight full_sft \
    --from_teacher_weight full_sft

参数选择指南

日志解读

训练时你会看到这样的输出

Epoch:[1/6](100/500), loss: 2.3456, ce: 1.8923,
aux_loss: 0.0000, distill: 2.7989,
learning_rate: 0.00000450, epoch_time: 3.200min

初始化流程

从参数解析到模型加载的完整过程

Part 5: 练习与排错

动手实践才能把知识变成技能。以下练习从简单到困难，建议按顺序完成。

动手练习

进阶练习

# 提示：KLDivLoss 的用法
criterion = torch.nn.KLDivLoss(
    reduction='batchmean'
)
# 注意第一个参数需要 log 概率

def distillation_loss_v2(student_logits, teacher_logits,
                         temperature=1.0):
    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
    )
    criterion = torch.nn.KLDivLoss(reduction='batchmean')
    kl = criterion(student_log_probs, teacher_probs)
    return (temperature ** 2) * kl

常见错误诊断

RuntimeError: CUDA out of memory

Epoch:[1/6](100/500), distill: 3.2456
Epoch:[2/6](100/500), distill: 3.2431

更多错误诊断

distill: -0.5234

RuntimeError: The size of tensor a (32000)
must match the size of tensor b (32005)

调优速查表

核心回顾

延伸阅读