Blackwellネイティブ8/4ビットRLへの道:MilesにおけるエンドツーエンドMXFP8とNVFP4の実践

Blackwellネイティブ8/4ビットRLへの道:MilesにおけるエンドツーエンドMXFP8とNVFP4の実践

はじめに

低精度RLでは、rollout・学習・チェックポイント変換・リアルタイム重み更新が統一された精度契約に従わなければならず、そうでなければサンプラーと学習器のポリシーが乖離する。MilesとSGLang RLエコシステムはすでに低精度方式を統合しており、本記事ではさらにMXFP8とNVFP4をBlackwellネイティブ形式へと拡張する。

なぜBlackwellネイティブ方式が必要なのか?

従来の低精度手法はMXFP8やNVFP4を想定して設計されたものではない。既存のMilesのパスはDeepSeek V3スタイルのブロックスケーリングFP8方式を採用しており、Blackwell上ではそのFP32スケールがネイティブのマイクロスケーリングハードウェアではなくソフトウェアで処理されている。一方、NVFP4はB200/B300のFP4 Tensor Coreのスループット優位性を最大限に活用できる。

フォーマットの背景

MXFP8

MXF P8はBlackwellネイティブのマイクロスケーリングFP8フォーマットであり、32個のE4M3値ごとに1つのE8M0ローカルスケールを共有する。

NVFP4

NVFP4は2段階スケーリングを採用する:FP32テンソルレベルスケールとE4M3ブロックレベルスケールである。

NVFP4のFP32テンソルスケールとE4M3ブロックスケールによる2段階スケーリング

方式1:エンドツーエンドMXFP8 RL

本方式では、rollout・順伝播・重み勾配・データ勾配のGEMMを通じてMXFP8を全面的に使用し、特定のテンソルはBF16のまま保持する。

エンドツーエンドMXFP8 RLレシピ

方式2:MoEエキスパートのPer-Token NVFP4 RL

NVFP4はより積極的な手法であり、MoEエキスパートのみを量化対象とし、その他の部分はデフォルトでBF16を使用する。

NVFP4事前学習の元レシピ高精度バックワードを用いたNVFP4逆量化バックワードを用いたNVFP4高精度バックワードを用いたMXFP8逆量化バックワードを用いたMXFP8

高精度レイヤー保護

最終レイヤーをBF16で保持することにより、学習の安定性が大幅に改善される。

最終レイヤーをBF16で保持した場合の効果第1レイヤーBF16と最終レイヤーBF16の比較共有エキスパートの高精度アブレーション