Abstract
知識蒸留（KD）は、大規模言語モデル（LLMs）をより小さなモデルへ圧縮するための重要な手法です。しかし、KDにおいて生徒モデルと教師モデルがそれぞれ異なる役割を担うにもかかわらず、既存のほとんどのフレームワークでは両モデルに対して同種の学習バックエンド（例：FSDP や DeepSpeed）が用いられており、その結果、学習効率が最適でない状態になります。本論文では、
\textbf{KDFlow} と名付けたLLM蒸留のための新しいフレームワークを提案します。このフレームワークは分離されたアーキテクチャを特徴とし、教師推論には SGLang を使用します。FSDP2 の学習効率と SGLang の推論効率を橋渡しすることで、KDFlow は両方の利点を単一のシステムで完全に活用することを実現します。さらに、プロセス間で完全なロジットを転送するのではなく、教師の隠れ状態のみをゼロコピーのデータ転送で送信し、生徒側でロジットを再計算することで、通信コストとKD性能のバランスを効果的に取ります。加えて、本フレームワークはオフポリシー蒸留とオンポリシー蒸留の両方をサポートし、さらに、高度に拡張可能で使いやすい API を通じて、クロストークナイザーKDのためのKDアルゴリズムも組み込んでいます。実験の結果、KDFlow は既存のKDフレームワークと比べて
\textbf{1.44\times から 6.36\times} の速度向上を達成でき、最小限のエンジニアリング負荷で、研究者がLLM蒸留を迅速に試作しスケールできることを示しています。コードは以下で利用可能です： https://github.com/songmzhang/KDFlow
