言語モデルNemotron 3.5 Lightningの重みを量子化し、元のモデルから蒸留することで、精度低下を回復しながらチェックポイント容量を66GBから22GBに縮めた。NVIDIAはNVFP4版で最大4倍のスループットを得たとしている。
量子化したモデルの容量は22GBに
NVIDIAが開発した言語モデルNemotron 3.5 LightningのNVFP4チェックポイントは22GBで、フル精度版の66GBから容量が縮んだ。重みの多くを4ビットへ量子化した結果である。
NVIDIAは、このチェックポイントで最大4倍のスループットを得たとしている。容量と処理性能を両立させるため、量子化後に元のモデルの出力を学習させる工程を加えた。
Mambaの線形層をW4A16へ量子化
量子化を意識した蒸留は、まずBF16の元モデルを学習後量子化し、重みを4ビット、演算時の値を16ビットで扱うW4A16の生徒モデルを作る。NVIDIAはMambaの線形層について、FP8ではなくW4A16を採用し、大幅な精度低下を伴わずにスループットを高めたとしている。
検討した量子化設定では、出力層のlm_headもW4A16にする一方、Attentionの射影層はBF16のまま残した。このように層ごとに精度を選んだチェックポイントが、次の蒸留工程の出発点となる。
教師モデルの出力で量子化後の精度を回復
蒸留では、固定したBF16の元モデルを教師に使う。生徒モデルの出力値であるlogitを教師のlogitと比較し、KLダイバージェンスによる損失を小さくするよう生徒モデルを学習させる。
生徒モデルは各順伝播で量子化をシミュレートし、推論時に生じる量子化の影響を学習中にも受ける。NVIDIAは、この工程によって積極的な量子化で低下した精度を回復させたとしている。
較正方法に合わせて蒸留時のスケールを選択
蒸留時のスケールの扱いは、最初の学習後量子化で使った較正方法に合わせる。最大値による較正から作ったチェックポイントには動的スケールを、平均二乗誤差（MSE）に基づく設定には固定スケールを使う。
NVIDIAは検討した量子化設定を、単一のDGX B300で1000サンプルを使って較正した。そのうえで実施したエージェント関連の評価では、量子化を意識した蒸留が学習後量子化のみの場合を上回ったとしている。