이미지: METAL LAB 생성
摘要
- Multiverse Computing公司公开了两项系统改进,可大幅降低知识蒸馏训练中的显存占用
- 通过预先缓存教师模型的top-K logits,消除了需要与学生模型同时载入内存的问题
- 应用了不生成完整词表×序列长度矩阵的新型KL散度损失函数
- 논문 제목
- Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss
- 발표
- Multiverse Computing, Hugging Face Blog, 2026-08-10
- 예시 모델
- gpt-oss-120b, 어휘 201,088 토큰
- 메모리 사례
- 시퀀스 32K·배치 4 기준 교사 확률 텐서 1개만 bf16으로 약 50GB
- 기존 증류 전체 소요
- 학습 1스텝 당 최대 약 250GB VRAM
Multiverse Computing公司通过Hugging Face博客公开了一篇论文,内容涉及降低大型语言模型知识蒸馏(knowledge distillation)训练成本的两项系统改进。论文标题为《Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss》。
知识蒸馏是让较小的学生(student)模型学习模仿大型教师(teacher)模型性能的技术。传统方法需要在训练的每一步都将教师模型同时载入内存以计算整个词表的概率分布,因此显存负担很大。例如,词表规模为201,088的gpt-oss-120b模型,在序列长度32K、批大小为4的条件下,仅教师概率张量在bf16精度下就约达50GB,若再加上梯度、激活值和优化器状态等,训练单步的显存占用最高可能飙升至约250GB。
两项改进
研究团队通过预先一次性计算并缓存教师模型的top-K logits,消除了训练过程中需要将教师模型与学生模型同时保留在内存中的必要性。此外,团队还引入了一种新型的内存高效KL散度损失函数,该函数不会直接生成词表规模×序列长度大小的矩阵,相比PyTorch或NVIDIA Megatron-Bridge等库的默认实现,大幅降低了显存占用。
