AI・機械学習
4ビットの教訓:NVFP4 RLにおける安定性とパフォーマンスのバランス
The 4-Bitter Lesson: Balancing Stability and Performance in NVFP4 RL (humansand.ai)
要約
この記事では、非同期強化学習(RL)システムにおける安定性とパフォーマンスのトレードオフについて論じています。特に、NVFP4のような低精度フォーマットを使用する際の課題と、それを克服するための技術について解説しています。目標は、ポリシーのドリフトが加速する閾値以下に保ちつつ、スループットを最大化する設定を見つけることです。
全文翻訳
はじめに RLトレーニングシミュレーター。
シミュレーターは、サンプラーが継続的にロールアウトを生成し、トレーナーがポリシーを更新する非同期RLシステムをモデル化します。
ポリシーの不一致は、オフポリシー(古いロールアウト)と量子化誤差の両方から生じ、ポリシーのドリフトとして蓄積され、最終的にはオプティマイザーの補正能力を超えると報酬が低下します。
RLと効率化のノブ。
オフポリシー、ウェイト同期、バッチサイズ、ホライゾンは、非同期の度合いとポリシーの古さを制御します。
MXFP8とNVFP4は、トレーニングとロールアウトの効率を向上させます。一方、デ量子化バックワード、BF16ラスト15%、共有エキスパートは、数値的安定性を向上させます。
使用方法。
アルゴリズムおよびシステムノブを変化させることで、スループットと安定性のトレードオフを探求します。
非同期性を高めたり、精度を下げたりして利用率を向上させ、次に安定化技術が安定性マージンをどのように回復するかを観察します。
目標は、ドリフトが加速し報酬が崩壊するクリティカルな閾値を下回るポリシーの不一致を維持しながら、スループットを最大化する構成を特定することです。
RLの本質は、行動と結果を通じてモデルを教えることです。humans&では、人との相互作用の長期的な影響を理解するモデルをトレーニングするためにRLを使用しています。
RLループ1(特にポリシー勾配アルゴリズム)では、モデルは行動し、報酬を得て、次にアクションの確率を更新します。
しかし、LMのリアルワールドRLトレーニングでは、他のロールアウトがまだサンプリングされている間でも、利用可能になり次第、観察結果からポリシーを更新したいと考えます。
私たちのミッションの中心である長ホライゾンマルチプレイヤーロールアウトでは、モデルが1つのロールアウトを完了するのに数十回のトレーニングステップを費やすことさえあります。
これにより、RLにおけるスループットと安定性の間の綱引きが生じます。
一方では、できるだけ多くの例をできるだけ早くトレーニングしたいと考えています2(各トレーニングステップを速くする、各ロールアウトを速くする、またはトレーナーとサンプラーをより多くオーバーラップさせる)。
他方では、スループットを向上させるほとんどの技術は、サンプリングされたポリシーとトレーニングされたポリシーを乖離させ、学習を遅らせたりトレーニングを不安定にしたりする可能性があります。
量子化は、このトレードオフの例です。低精度フォーマットはハードウェア上での高速な通信と計算を可能にしますが、量子化は安定性を損ないます。
NVFP4のような高速で正確な低精度フォーマットは、ハードウェアサポート3(NVIDIA Rubin GPUで16ビットトレーニングよりも最大9倍の演算速度を可能にする)と組み合わされることで、モデルトレーニングと推論の両方で大幅なスループット向上を推進してきました。
しかし、オープンソースコミュニティには、主にサンプリングとトレーニングの不安定性がRLで複合するため、安定したハードウェアネイティブの4ビット4(INT4 QATレシピは4ビット量子化をシミュレートするために使用されていますが、これらは16ビットアクティベーション(Kimi K2 Thinking)を使用し、DSv4のようなMXFP4レシピはMXFP8アクティベーション(DeepSeek-V4 on Day 0)を使用しています。私たちの方法は、ウェイトとアクティベーションの両方にNVFP4量子化を使用します)のRLレシピはほとんどありません。
私たちは、オープンソースコミュニティとの長年の協力により、高精度トレーニングダイナミクスを維持する低精度RLレシピを開発し、共有しました。
このレシピでは、ポリシー量子化エラーによるフォワードパスの不安定性、勾配の不一致によるバックワードパスの不安定性、そして特に敏感な少数のウェイトによる両者の交差点での不安定性に対処する必要がありました。
以下に、それぞれをどのように対処し、最終的なレシピを検証したかを説明します。
この取り組みは、RadixArkとNVIDIAの素晴らしい協力者、およびトレーニング、推論、RLスタック全体での彼らの仕事なしには実現できなかったことに注意してください。
ベースライン:安定したトレーニングダイナミクスを持つ開始レシピ。
一貫した比較のため、特に断りのない限り、このレポートのすべての実験では、DAPO-math-17kデータセットで8kシーケンス長でトレーニングされたQwen3-30B-A3Bモデルを使用します。
私たちは、NVFP4フォーマットを使用するベースラインレシピから始めます。
NVFP4は、新しいNVIDIA GPUアーキテクチャでスループットとメモリ効率を向上させる4ビット浮動小数点フォーマットです。
これは階層的なブロックスケーリングを使用します。16個の値の各ブロックはFP8 E4M3スケールを持ち、テンソル全体はグローバルFP32スケールを持ちます。
これらのスケールを組み合わせることで、より高精度の値を回復できます。
NVFP4事前トレーニングレシピがRLに十分でないのはなぜですか?
NVIDIAのNVFP4事前トレーニングレシピは、この作業の自然な出発点です。
これは混合精度戦略を使用しており、ほとんどの操作はFP4精度で実行されますが、数値的に敏感なコンポーネントはより高い精度で維持されます。
しかし、事前トレーニングとRLでは、失敗のモードが異なります。
事前トレーニングでは、勾配信号は密であり、多数のトークンにわたって繰り返し平均化されます。
主な目標は、最適化の方向を摂動させる量子化バイアスを回避することです。
確率的丸めは、勾配量子化をほぼバイアスなくすることで役立ちます。
RL設定では、バイアス分散トレードオフが異なります。
ポリシー勾配は、サンプリングされたロールアウト、アドバンテージ推定、報酬推定、KL正則化、およびポリシーの古さに依存するため、すでにノイズの多い推定値です。
したがって、量子化方法が単にバイアスがないだけでは十分ではありません。
量子化ノイズは、更新ごとの真のポリシー勾配信号を低下させない程度に小さくなければなりません。
したがって、ポリシーの量子化誤差を減らし、勾配計算の精度を向上させる介入を優先します。
ベースラインレシピ。
ベースラインレシピでは、MoEレイヤーのみを量子化し、他のすべてのレイヤーはより高精度のBF16で維持します。
DeepSeek-V3スタイルのアーキテクチャでは、MoEエキスパートは総パラメータの97%5(各エキスパートは7168 x 2048の3つの射影行列を含み、エキスパートあたり約44Mパラメータです。256のエキスパートと各MOEレイヤーあたり1つの追加の共有エキスパートで、レイヤーあたり約11.32Bパラメータです。58レイヤーにわたって、656Bパラメータとなり、総パラメータの約97.8%に相当します)を占めるため、MoEレイヤーを積極的に量子化すると、メモリのメリットのほとんどが得られます。
フォワードパスはNVFP4で実行し、バックワードパスはBF16精度で維持します。
これは、FP4のロールアウトとメモリのメリットを享受しつつ、バックワードパスでのFP4の使用を避けるための保守的な出発点です。
ウェイトについては、FP8のブロックごとスケールと単一のFP32グローバルスケールを持つ標準的なNVFP4フォーマットを使用します。
アクティベーションについては、テンソル全体にわたるグローバルFP32スケールは使用しません。
Cursor Composer 2の技術レポートが指摘しているように、グローバルスケールは2つの問題を引き起こす可能性があります。
同じトークンが、バッチ内の他のトークンがどのようなものかによって、異なる方法で量子化される可能性があります。
テンソルに同じシーケンスの複数の位置が含まれている場合、後続のトークンが先行トークンによって共有されるスケールに影響を与え、将来のトークンから過去のトークンへのリークパスを作成する可能性があります。
この問題を解決するために、Cursor Composer 2と同様に、各トークンが隠れ次元全体で独自のFP32アクティベーションスケールを計算する、トークンごとのアクティベーションスケーリングを使用します。
これにより、量子化は各トークンにローカルに保たれ、スケールはフォワードパス中に観察されたアクティベーションから直接計算できるため、別のキャリブレーションステップが不要になります。
また、このきめ細かなFP32スケールは、トークンごとのFP32スケールと比較して、量子化誤差が低いのに役立ちます。
ロールアウト中、このトークンごとのスケール計算は、アクティベーション量子化カーネルに融合されます。
各トークン行について、カーネルはFP32スケール、16値ブロックのFP8 E4M3スケールを計算し、これらの量子化値をFP4にパックします。
この融合実装は、追加のメモリ移動とカーネル起動のオーバーヘッドを削減します。
実装。
トークンごとのNVFP4レシピは、スタック全体での変更を必要とし、実装をオープンソース化しました。
TransformerEngine行スケーリングレシピ。
行スケーリンググループ化GEMM量子化用のcuDNNカーネル。
MoEトークンごとの量子化用のFlashInfer推論カーネル。
SGLang。