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

Multiverse Computing 推出了一種全新的知識蒸餾方法，能顯著降低 VRAM 使用量與訓練成本，使單張 GPU 即可進行長文本上下文蒸餾。透過結合離線 Logit 快取與融合分塊 KL 散度（KL-divergence）損失，研究人員能夠將大型教師模型蒸餾至較小的學生模型，而不會產生通常與全詞表機率分佈相關的大量記憶體開銷。

標準知識蒸餾的高昂成本

傳統使用 Kullback-Leibler (KL) 散度損失的「在線」知識蒸餾在計算上非常昂貴，因為它需要同時將教師模型與學生模型載入至 VRAM 中。在每個訓練步驟中，教師模型必須進行完整的正向傳播，以便為每個 Token 位置生成整個詞表的機率分佈。

對於大型模型而言，這會造成巨大的記憶體瓶頸。例如，使用具有 201,088 個 Token 詞表的 gpt-oss-120b 模型，在序列長度為 32K 且 Batch Size 為 4 的情況下，僅教師機率張量（teacher-probability tensor）在 bfloat16 精度下就需要約 50GB 的 VRAM。若計入梯度、激活值與優化器狀態，單次迭代的 VRAM 使用量峰值可能達到約 250GB，甚至超過了 H200 或 B200 等高端 GPU 的容量。

實現可擴展蒸餾的兩項系統變革

為了克服這些記憶體限制，Multiverse Computing 實施了兩項主要的技術變革：

1. 透過 Top-K Logit 快取進行離線蒸餾

與其在每個步驟重新計算教師模型的輸出，該系統會計算一次教師模型的輸出，並針對每個位置快取前 100 個最可能的 Token。接著，學生模型會針對此快取進行訓練。這消除了在訓練過程中將教師模型保留在記憶體中的需求，並允許在不同的實驗消融研究中重複使用相同的快取。

2. 融合分塊 KL 損失

標準的 KL 損失實作方式是建立一個完整的「詞表 x 序列」網格來計算學生與教師模型之間的不一致性。Multiverse Computing 的「融合分塊 KL 損失」透過將模型的輸出投影直接融合到損失計算中來優化此過程。

該過程並非實體化完整的 Logits 網格，而是如下運作：

它一次處理序列的一個分塊（chunk）。
它僅針對該特定分塊將隱藏狀態（hidden states）投影至 Logits。
它將結果摺疊（fold）進持續累加的損失中，並立即捨棄該分塊。
反向傳播時會即時重新計算每個分塊。

這種方法確保了峰值記憶體會隨著序列長度線性增長，而非根據全詞表大小而出現激增。

性能與記憶體基準測試

在單張 H200 GPU 上，使用 Llama 3.1 8B Instruct（教師）與 3.2B Llama 模型（學生）進行 8K Token 上下文的比較測試，研究人員發現所有方法的訓練損失都接近一致。這證實了僅使用前 100 個快取的 Logits 進行離線蒸餾，相對於在線蒸餾而言，幾乎是無損的。

方法 (8K 上下文, 單張 H200)	峰值記憶體	迭代時間	吞吐量
在線蒸餾	102.8 GB	25.9 s	237 TFLOP/s
離線, 密集型 KL	78.3 GB	18.5 s	331 TFLOP/s
離線, 前向分塊 KL	61.8 GB	18.4 s	335 TFLOP/s
離線, 融合分塊 KL	58.3 GB	20.2 s	304 TFLOP/s

擴展至長文本上下文

隨著上下文長度增加，融合分塊損失的優勢變得更加明顯。在針對小型輸出投影網路的隔離基準測試中：

在 32K Token 時： 峰值記憶體從 85.2 GiB（密集型損失）降至 5.45 GiB（完全分塊），降幅達 15.6 倍。
在 64K Token 時： 密集型損失實作完全失效。
在 256K Token 時： 完全分塊損失僅使用 11.6 GiB，而次優的分塊變體則需使用 134.2 GiB，且每次迭代的速度快了 3.3 倍。

在一個真實場景中，針對 GPT-OSS 20B 模型進行 32,768 Token 上下文的蒸餾，融合損失將所需的硬體從四個 GPU 節點減少到一個。這使得單步時間從 57.0 秒縮短至 12.23 秒（5 倍加速），並將每張 GPU 的吞吐量從 74.2 提升至 345.7 TFLOP/s。

學生模型能力

此流程的高效率使得大規模的蒸餾工作得以進行，並產出了一個從 Llama 3.1 8B Instruct 蒸餾而來的 3.2B 參數學生模型。儘管參數數量不到教師模型的一半，這個精簡的學生模型在 BoolQ 和 HellaSwag 任務上保留了教師模型大部分的準確度，且在 MMLU 基準測試中與教師模型的差距保持在約 9 分以內。