第 6 章 — 剪枝与知识蒸馏

《LLM Primer VI: Scaling AI Systems》章节走读的第六篇。这一章直接攻击权重数量 —— 先是把无关紧要的权重清零,然后把一个大模型的行为迁到一个更小的模型里。

这一章为什么存在

量化缩小每权重字节数。剪枝缩小真正参与相乘的那部分权重的数量。知识蒸馏则把一个大教师的行为迁到一个端到端更便宜的学生身上。剪枝是同架构里的编辑,削掉冗余;蒸馏是架构变更,拿一份训练期成本换一份永久性的推理期节省。第 6 章走这两条 —— 各自依赖的机制、实践者会踩到的失败模式,以及那个让结构剪枝在 Hopper 和 Blackwell 上重新变得经济有意思的硬件加速模式(2:4 稀疏)。

6.1 2:4 稀疏是硬件愿意奖励的剪枝模式

非结构化幅值剪枝 —— 把最小的 50% 权重清零,再短暂微调一下 —— 数值上可行,但根本不加速。GPU 执行的是稠密 GEMM;散落的零仍然要被乘一遍。结构化剪枝去掉的是硬件能跳过的整块单位:attention head、MLP 神经元,或者(重要那一档)权重矩阵里 N:M 模式。Hopper 和 Blackwell 原生支持 2:4 稀疏 —— 输入维度上连续四个权重里恰好两个是零。芯片把每一组四个权重存成两个非零值加一份两 bit 的 mask,稀疏 tensor core 跳过零乘,带宽和 FLOPs 都掉大约一半。实测 decoding 吞吐是 1.7–1.9× 的加速。窍门是 2:4 必须在剪枝时施加:算法一组四个地走过每一个矩阵,强制其中两个最小的清零,然后要么短暂微调,要么用 SparseGPT 做补偿。Mask 是元数据;运行时的收益是自动的。

6.2 蒸馏迁的是教师的分布,不是它的 argmax

训好的教师 LLM 在每一个位置产出一份关于词表的分布 —— 一份 128k 维的 softmax。argmax 是它贪婪解码时会输出的那个 token。但整份分布还编码了教师的不确定性、它在近义词之间的偏好结构、以及它的校准。训学生去匹配这份分布,能把上面这一切都迁过去:损失是软化后的教师和学生输出之间的 KL 散度,温度把梯度信号拉宽到低概率长尾。这样出来的学生,行为像是教师的一个更小、更快版本,而不是从零在同一份数据上训的一个小模型。Patient Knowledge Distillation 把这套配方扩展到中间层,在配对的师生层之间匹配隐状态。MiniLLM 处理一个更微妙的失败模式 —— mode-averaging 问题,即前向 KL 学生会在教师的多个合理续写之间打平均,产出乏味的输出 —— 靠切换到反向 KL 来解决,它偏好锐利地匹配教师的其中一个众数,而不是软性地覆盖所有。

6.3 三种压缩按特定顺序叠加

先蒸馏,再剪枝,最后量化。蒸馏最贵、对架构侵入最深;有道理先在一个稠密高精度模型上做一次,再让其他选择锁死。用 SparseGPT 把学生剪成 2:4,再做一次短的恢复微调。用 llm-compressor 把稀疏化后的学生量化到 FP8。一份走过的例子:70B BF16 教师 → 13B BF16 蒸馏学生 → 2:4 稀疏化 → FP8 量化。每次前向传播的权重字节数从 140 GB 掉到大约 6.5 GB —— 22 倍缩减 —— 而 Hopper 上的 FP8 和 2:4 通路把这份缩减变成每 GPU 8–10 倍的吞吐提升。这套叠加没缩减的是 KV cache,后者取决于隐维度、head 数、上下文长度,而不是权重数量。那是第 8 章的题目。而且模型级压缩是每次推理层面的;它给 GPU 腾出余量,但把余量填成并发用户吞吐的运行时杠杆,是 batching。

第 6 章为下面铺了什么

第 5、6 章把模型侧压缩讲完 —— 量化每权重字节、剪枝权重数量、蒸馏整个模型。GPU 每次前向传播都有了余量。第 7 章转向那个把这份余量填满有用工作的系统级杠杆:batching。静态 batching 差点就够用了,但被 fastest-finisher 问题拉垮。Continuous batching —— 迭代级调度 —— 才是生产 engine 真正在跑的,也是把带宽受限的 decode 阶段从一块空转的芯片变成一块同时服务多用户的芯片的东西。

下一章 — 第 7 章:进阶 Batching 策略。静态、动态、continuous batching —— 以及 continuous batching 给 KV cache 留下的债务。