让知识蒸馏便宜到可以规模化运行
来自 Hugging Face · 期号 2026-08-10
摘要
知识蒸馏是压缩大模型的标准做法，但也是最贵的一步：同时加载教师和学生模型，为每个 token 在全词表上算概率分布，VRAM 轻松突破单卡上限。 Hugging Face 团队用两个系统改动解决了这个问题：把教师模型的 top-K logits 离线缓存一次，训练时不再加载教师；用分块 KL 散度损失避免实例化完整的词表×序列长度矩阵。 改动后 VRAM 峰值从约 250GB（超过 H200 的 141GB）降到约 128GB，长上下文蒸馏可以在单张 GPU 上完成。论文附了开源实现。
深读
导读 知识蒸馏是压缩大模型的标准做法，但也是最贵的一步：同时加载教师和学生模型、为每个 token 在全词表上算概率分布，VRAM 轻松突破单卡上限。Hugging Face 团队用两个系统改动解决了这个问题，离线缓存教师模型的 top-K logits 和分块 KL 散度损失，使长上下文蒸馏能在单张 GPU 上完成，并开源了实现。读完能带走一套可落地的蒸馏降本方案及其背后的权衡。 关键信息 - 在线蒸馏时，以 gpt-oss-120b（词表 201,088 tokens）为例，序列长度 32K、batch size 4 下，教师概率张量本身约 50GB，加上梯度、激活值、优化器状态后单次迭代峰值约 250GB，超过 H200 或 B200 的 141GB 容量。 - 离线蒸馏方案将教师模型的 top-100 logits 缓存一次，训练时不再加载教师，同一缓存可复用多次消融实验。 - 三种离线 KL 损失实现中，fused chunked KL 是主要贡献：将模型输出投影直接融合进损失计算，不产生完整学生 logits 网格，按 chunk 端到端处理序列，反向传播时按需重算。 - 在单张 H200 上以 Llama 3.1 8B Instruct 为教师、3.2B Llama 模型为学生、8K 上下文测试，四种方法（在线蒸馏、离线 dense KL、离线 forward-chunked KL、离线 fused chunked KL）训练损失曲线几乎完全重合，证明离线 top-100 蒸馏相对于在线蒸馏是无损的。 - 在 32K 上下文的 GPT-OSS 20B 蒸馏中，fused chunked loss 将配置从四个 GPU 节点缩减到一个，步时间从 57.0 秒降至 12.23 秒（约 5 倍加速），每 GPU 吞吐量从 74.2 TFLOP/s 升至 345.7 TFLOP/s。 逐节讲解 为什么蒸馏恢复很贵 标准在线蒸馏同时加载教师和学生模型，每步教师都要做完整前向传播。以 gpt-oss-120b 为例，词表 201,088，序列长度 32K、batch size 4 时，教师概率张量形状为 4 × 201,088 × 32,768，bfloat16 下约 50GB。加上其他状态后峰值约 250GB，超过单张 H200 的 141GB 容量。文章指出，将 KL 损失改为分块处理可大幅削减这一成本。 两个系统改动 离线蒸馏：教师输出只计算一次，缓存每个位置 top-100 最可能 token 的 logits，训练时教师不再驻留内存，同一缓存可跨多次消融实验复用。 分块 KL 损失：文章比较了三种数学等价的实现方式。Dense KL 从缓存的 top-100 logits 重建完整教师概率网格并与学生网格比较，作为正确性基线但内存开销最大。For…