NeMo-DCR、兆パラメータ級エージェントRLの重み同期を約35倍高速化

背景

エージェント強化学習では、方策を更新する学習クラスタと、現在の方策でデータを生成するrolloutクラスタを分離する構成が使われます。計算資源を柔軟に拡張できる一方、学習で更新された重みを次のデータ生成に間に合うよう転送しなければなりません。モデルが兆パラメータ規模になると、完全なチェックポイントを毎回送る処理が反復全体のボトルネックになります。

NeMo-DCR（Delta-Compressed Refit）は、この重み同期を差分転送として再設計する手法です。論文の観測では、BF16学習の1ステップで保存値が変化する重み要素は約1%にとどまります。ただし、変更箇所だけを送ればよいわけではありません。学習側とサービス側ではシャード構成やテンソル配置が異なるため、最終的なビット列が完全ロード時と一致する必要があります。

仕組みの要点

• 正規座標への投影: 固定されたアフィン写像を用いて、学習シャード上の変更をチェックポイントの正規座標へ移します。約96%は直接処理し、残りは残差変換で扱います。
• XORマスクと上書き: ロード経路を通っても保存ビットを保てる変更には、圧縮しやすいXORマスクを使います。それ以外は絶対値の上書きとして送信し、算術的な再構成による誤差を避けます。
• ネイティブローダーの活用: モデルごとの配置ルールを新たに実装したり、完全なテンソルを組み立てたりせず、rolloutランタイムの既存ローダーによるメモリコピーを介入して、受信側の領域へ更新を直接適用します。
• 失敗時の再試行: 部分的な書き込みが残っても上書きによって再試行できます。共同コミットにより、インストールした方策と次の差分が依存するベースラインを結び付けます。
• 柔軟な転送経路: オブジェクトストレージや中継ツリーを利用でき、クラスタ間collectiveを必須としません。差分の生成、転送、適用を重ね合わせる構成も可能です。

結果と意義

1兆パラメータモデルで変更率3%を設定したテストでは、完全チェックポイントの転送に87.5分かかったのに対し、NeMo-DCRのrefitは2.5分で完了し、約35倍の高速化となりました。提供された素材では、変更率3%および5%、30Bから1Tまでのモデルを対象とした評価も説明されています。

この成果のポイントは、単なる通信量削減ではありません。差分更新、異なる配置への変換、ビット単位の正確性、インプレース適用、途中失敗からの復旧を一つの同期経路にまとめています。実運用では、差分が小さくても配置を誤ったり、丸めの違いでビット列が変わったり、失敗後に半端な状態が残ったりすれば利用できません。

NeMo-DCRは、rolloutワーカーが古い方策を使い続ける時間を短縮し、学習パイプラインの待ち時間を抑える可能性があります。一方で、効果は変更率、ネットワークとストレージの帯域、ローダー実装に左右されます。差分転送を安全に運用するには、常にベースラインとバージョンの対応関係を維持することが重要です。

出典：Hugging Face Daily Papers