技術的概要：通信効率の高い分散 RL のための重み更新スパース性の理解と活用

問題提起

大規模言語モデル（LLM）の分散強化学習（RL）におけるポストトレーニングは、特に帯域幅が制約された環境において、深刻な通信ボトルネックに直面しています。著者らは、これらのボトルネックを引き起こす 2 つの主要なチャネルを特定しました：

1. トレーナーから推論ワーカーへの同期：更新されたモデル重みを用いて、推論ワーカー上のロールアウトポリシーを刷新すること。
2. トレーナー間同期：分散トレーナー間で勾配または疑似勾配を交換すること。

地理的に分散された環境では、32B パラメータモデル（BF16 で 62 GB）の同期は、一般的な帯域幅では 1 ラウンドあたり約 14 分を要し、GPU の利用率が大幅に低下します。先行研究では、RL 微調整が粗い比較においてパラメータのごく一部のみを変更することが観察されてきましたが、既存の手法は RL パイプラインにおける効率的な同期に必要な「ステップごとの」スパース性を十分に活用できていません。

核心的な観察：計算可視的なスパース性

本論文の中心的な発見は、標準的な RL パイプラインにおいて、ステップごとの重み更新の約**99% が「計算不可視」**であることです。この現象は、以下の相互作用によって生じます：

• BF16 精度：標準的なトレーニングと推論はフォワードパスに BF16 を使用し、オプティマイザは FP32 のマスター重みを維持します。
• RL 学習率：一般的な RL ポストトレーニングの学習率（例：$\eta \approx 3 \times 10^{-6}$）は、教師あり微調整で使用されるものよりも著しく小さいです。
• Adam 更新の境界：これらの低い学習率において、Adam 更新の大きさは、しばしば BF16 の丸め閾値（約 $|w|/256$）を下回ります。

その結果、勾配はほぼ 100% 密ですが、重みの BF16 表現に対する結果的な更新はスパースになります。更新が「可視的」になるのは、次のフォワードパスで使用される BF16 値を変更する場合に限られます。丸めセルより小さい更新はキャストによって吸収され、基盤となる FP32 マスター重みが変更されても計算には影響しません。

手法：PULSE

著者らは、計算可視的なスパース化の原則に基づいたフレームワークであるPULSE（Precision-gated Updates for Low-precision Sparse Exchange：低精度スパース交換のための精度ゲート付き更新）を提案します。PULSE は 2 つのアルゴリズムで構成されます：

1. PULSESync（トレーナーから推論ワーカーへ）

• メカニズム：トレーナーから推論ワーカーへ重みを同期するための損失なし手法。
• プロセス：トレーナーは連続する BF16 チェックポイントを比較します。BF16 視点で変更されたパラメータのインデックスと新しい値のみを送信します。
• 再構築：推論ワーカーは、これらのスパースパッチをローカルの BF16 重みに直接適用します。パッチは算術的な差分ではなく絶対値を格納するため、再構築はフルチェックポイント転送とビット同一であり、エラーフィードバックは不要です。
• 圧縮：スパースパッチは、汎用コーデック（例：zstd）を用いてさらに圧縮されます。

2. PULSELoCo（トレーナーからトレーナーへ）

• メカニズム：トレーナー間で疑似勾配を同期する手法であり、DiLoCo（Distributed Low-Communication）アプローチを拡張したものです。
• プロセス：
  • ワーカーは、共有されたグローバルチェックポイントから開始して $H$ ステップのローカル Adam 更新を実行します。
  • 各ワーカーは、$H$ ステップにわたるパラメータ更新である疑似勾配を形成します。
  • これらの FP32 疑似勾配に計算可視性ゲートを適用します。BF16 表現を変更するエントリのみを同期します。
  • ゲートを通過しないエントリはFP32 エラーフィードバックバッファに格納され、次のラウンドで再検討されます。
• 差別化：密な勾配を直接量子化またはスパース化する標準的な勾配圧縮器とは異なり、PULSELoCo はローカル最適化から導出されたスパースな更新を同期し、同じ BF16 吸収メカニズムを活用します。

主要な貢献

1. スパース性のメカニズム的説明：RL における重み更新のスパース性は、特定のモデルアーキテクチャの性質ではなく、低い学習率と BF16 精度の自然な帰結であることを実証しました。一般的な RL 学習率において、ほとんどの更新が BF16 の可視性閾値を下回ることを示す理論的限界を提供します。
2. PULSESync：損失なく、ビット同一の重み同期手法を導入しました。これは、あらゆる RL パイプラインにおける密な同期のドロップイン代替手段です。
3. PULSELoCo：ローカル更新とエラーフィードバックスパース化を組み合わせた疑似勾配同期手法を導入しました。これは、マルチトレーナー RL ポストトレーニング用に特別に設計された最初の圧縮器です。
4. 実証的検証：DiLoCo 型のローカル更新トレーニングが、RL ポストトレーニング設定における LLM において機能することを初めて実証的に示しました。

実験結果

PULSESync の性能

• 展開：grail フレームワークを用いて、公衆インターネット上のライブ地理分散環境で評価されました。
• 帯域幅削減：7B モデルの場合、約 14 GB から約 108 MB へ、100 倍以上のペイロードサイズ削減を達成し、ビット同一の再構築を維持しました。
• トレーニング品質：検証性能（MATH および MBPP における Pass@1）は着実に向上し、標準的なトレーニング動作と一致しました。これにより、同期の損失な性質が確認されました。

PULSELoCo の性能

• 設定：Qwen2.5（1.5B、3B、7B）および Llama-3.2-3B モデルを MATH データセット上で、DDP および DiLoCo ベースラインに対して評価しました。
• 収束：PULSELoCo は、すべてのモデルにおいて DiLoCo の学習動作と一致しました。エラーフィードバックの蓄積による初期の遅れが観察されましたが、最終チェックポイントによって解消されました。
• 通信効率：
  • DiLoCo と比較して、トレーナー間の通信を17 倍以上削減しました（例：7B モデルで 1.77 GB 対 30.5 GB）。
  • 同一のローカル更新ウィンドウにおいて、標準的な DDP（各ステップで同期）と比較して100 倍以上削減しました。
• スパース性：疑似勾配ペイロードで 94.8–96.4% のスパース性を達成し、1 ラウンドあたりに送信される FP32 値は約 3.6–5.2% であることを意味します。

意義と主張

本論文は、計算可視的なスパース化が、分散 RL のための根本的かつデータ型に依存しない原則であると主張します。

• 帯域幅効率：これらの手法は、地理分散 RL の帯域幅要件を劇的に削減し、一般的なネットワーク（例：PULSESync で約 0.2 Gbit/s、PULSELoCo で約 2.6 Gbit/s）であっても高い GPU 利用率（90%）を可能にします。
• 損失な性質：PULSESync は近似誤差なしに圧縮を達成し、損失のある勾配圧縮技術と区別されます。
• スケーラビリティ：このアプローチは、モデルファミリー（Qwen、Llama、Gemma）およびサイズ（0.5B から 7B）全体に拡張可能であり、この現象は特定のモデルアーキテクチャではなく、RL の最適化ダイナミクスに起因することを示唆しています。
• 将来の可能性：著者らは、より低い精度の形式（例：FP8、MXFP4）は、より粗い丸めセルによりスパース性をさらに増加させる可能性があると指摘していますが、これは現在の作業では推測されたものであり、実証的に測定されたものではありません。

この研究は、低い学習率と低精度キャストによって引き起こされる自然なスパース性を活用することで、分散 RL が、現在 LLM の RL ポストトレーニングの規模と展開を制限している通信ボトルネックを克服できるという結論に至ります。