11月7日、NVIDIAが「Accelerating Reinforcement Learning Weight Refit with the Reshard Collective」と題した記事を公開した。RL後学習では学習GPUで更新した重みを推論GPUへ定期的に同期する「refit」処理が必要だが、現行の主流実装では論理的な通信ペイロードがモデルサイズの約2倍に膨らむという根本的な非効率がある。NVIDIAはこの問題をPyTorchのDTensor抽象で再定式化し、新たな集合通信プリミティブ「Reshard Collective」として解決した。ベンチマークでは従来比最大5.7倍の高速化を確認している。

問題の核心：refitのペイロードは「モデルサイズの2倍」になる

RL後学習では、推論と学習を別々のGPUで並列実行するdisaggregated構成が一般的だ。学習GPUで更新されたモデルの重みを、推論GPUへ定期的に同期する処理が「refit」である。

現在の主流実装はgather-and-broadcast方式をとる。各学習GPUがローカルのシャードをall-gatherして完全なパラメータテンソルを復元し、それを推論GPUへブロードキャストする仕組みだ。推論GPUは受け取った完全テンソルのうち自分に必要なシャードだけを保持する。

問題は、推論GPUが必要としているのはシャードの一部にもかかわらず、全体を受信しなければならないことによる冗長な通信だ。all-gather段階とbroadcast段階の合計で、論理的な通信ペイロードはモデルサイズのおよそ2倍に膨らむ。元記事によれば、現在のオープンウェイトモデルの最大規模クラスでは1回のrefitで約10 TBの通信が発生しうるとされており、これを学習ステップごとに繰り返すコストは無視できないと指摘されている。

Reshard Collective：DTensor-to-DTensor転送への再定式化

NVIDIAのチームはこの問題を、PyTorchの**DTensor**（分散テンソル）抽象を使って解決するアプローチをとった。

DTensorはテンソルのグローバルな形状、GPUの論理的な配置（device mesh）、シャーディング・レプリケーションの方法（placements）で分散レイアウトを記述できる。この抽象を使うと、refitは「学習レイアウトのDTensorを推論レイアウトのDTensorへreshardする問題」として統一的に記述できる。

具体例として、8×8の重み行列を学習側2GPUで行方向にシャーディング、推論側4GPUで列方向にシャーディングするケースを考える。各学習GPUは4×8シャードを持ち、各推論GPUには8×2シャードが必要だ。この場合、各推論GPUは各学習GPUから4×2の領域を受け取るだけで済む——完全なテンソルを再構成せずに直接転送できる（下図参照）。

この原理に基づいた新しい集合通信操作がReshard Collectiveだ。all-reduceやall-gatherと同様に、アプリケーションから1回の操作として呼び出せる。送信元・送信先のDTensor記述を渡すだけで、必要な転送パターンをライブラリ側が導出する。

NVLinkを最大活用するNVIDIAの実装

オープンソースの参照実装として**NCCL M2N Reshard**が公開されている（NVIDIA/nccl-extensions リポジトリ）。M（学習GPU数）からN（推論GPU数）への転送を担う実装だ。

通信効率化の要点はNVLinkドメインを意識したデータ配置にある。同一NVLinkドメイン内のレプリカ群に対して、ネットワーク経由の受信データをできるだけ均等に分散する。各レプリカは受け取ったサブセットをNVLink経由で交換し合うことで、ネットワーク上の重複転送を排除する。

gather-and-broadcast方式では、各ソースシャードをDP（データ並列）レプリカ数分だけネットワーク越しに送信していた。NCCL M2N Reshardはシャードをネットワーク経由で1回だけ転送し、残りはNVLink経由の内部交換で補完する。これによりネットワーク帯域の利用効率が大幅に改善される。

実測値：最大5.7倍の高速化

NVIDIAはNeMo RLフレームワーク（RL後学習向けオープンソースフレームワーク）に参照実装を組み込み、ベンチマークを実施した。なお実装上の工夫として、Transformerブロック内のfeed-forward射影の重みがrefit全体のペイロードの約95%を占めるという観測に基づき、これらはReshard経由、QKV射影など残りの小さい重みは従来のgather-and-broadcast経由とするハイブリッド方式を採用している。

各ベンチマークで2.7倍〜5.7倍の高速化を達成している。

まとめ

Reshard Collectiveのポイントをまとめると以下の通りだ。

• 通信量の削減：シャード間の差分だけを転送し、gather-and-broadcastの冗長性を排除
• 統一的な抽象：DTensor記述を渡すだけでよく、並列構成の違いをライブラリ側で吸収
• NVLink活用：ドメイン内交換によりネットワーク帯域を無駄なく使い切る
• 既存フレームワークへの統合：NeMo RLでの実装はDTensor記述を提供するだけで済む

今回の実装はNeMo RLを対象としているが、DTensor抽象はPyTorchの標準機能であるため、同様のdisaggregated RL構成をとる他フレームワークへの応用も技術的には自然な延長線上にある。RL後学習においてrefit待ちが学習スループットのボトルネックになっている環境では、実際的な改善手段として注目に値する。

詳細はAccelerating Reinforcement Learning Weight Refit with the Reshard Collectiveを参照していただきたい。