高效大型語言模型知識蒸餾：離線Top-K Logits與融合分塊KL損失

摘要

小型語言模型通常是滿足嚴格延遲、成本與本地部署限制時的唯一選擇，但它們很少從零開始訓練：壓縮模型通常是透過知識蒸餾（KD）從較大模型中回復而來。這個回復步驟在很大程度上決定了最終品質，然而其成本高昂。我們提出一項實務導向研究，探討如何讓蒸餾訓練更有效率，並圍繞兩項系統貢獻進行組織。首先，我們展示離線知識蒸餾（將教師模型的頂層 K 個邏輯值快取一次，並讓學生模型針對該快取進行訓練）能在訓練損失幾乎相同的情況下媲美線上蒸餾，同時將教師模型移出記憶體，每次迭代速度快約 29%，且在單張 H200 GPU 上吞吐量最高可提升 41%。其次，我們引入一種融合的分塊 KL 損失，此損失從不實體化完整的詞彙量級邏輯張量，使峰值記憶體隨序列長度呈線性成長。這消除了原本會限制上下文長度的記憶體尖峰，讓我們能在單張 GPU 上以四倍上下文長度（32,768 個 token）進行訓練。一個獨立的僅輸出層小型基準測試將損失核心單獨隔離，並確認其在 4K 至 256K token 範圍內的記憶體與迭代速率擴展特性。兩者結合使大規模修復訓練與數百次消融實驗變得可負擔。我們亦報告關於損失設計與序列打包的輔助消融實驗。我們釋出分塊損失的實作：https://github.com/CompactifAI/Full-Chunked-KL-Loss。

English

Small language models are often the only option for deployment under tight latency, cost, and on-premises constraints, but they are rarely trained from scratch: a compressed model is usually recovered through knowledge distillation (KD). This recovery step largely decides the final quality, yet it is expensive. We present a practitioner's study of how to make distillation training efficient, organised around two systems contributions. First, we show that offline KD (caching the teacher's top-K logits once and training the student against the cache) matches online distillation at near-identical training loss while removing the teacher from memory, running about 29\% faster per iteration, and reaching up to 41\% higher throughput on a single H200 GPU. Second, we introduce a fused, chunked KL loss that never materialises the full vocabulary-sized logit tensor, making peak memory linear in the sequence length. This removes the memory spike that otherwise caps context length and lets us train at four times the context (32{,}768 tokens) on a single GPU. A separate output-head-only toy benchmark isolates the loss kernel and confirms its memory and iteration-rate scaling from 4K to 256K tokens. Together these make large-scale healing and hundreds of ablations affordable. We also report supporting ablations on loss design and sequence packing. We release our chunked-loss implementation: https://github.com/CompactifAI/Full-Chunked-KL-Loss.