データセンター横断型強化学習へ

データセンターを越える強化学習：通信コストを劇的に削減する新アーキテクチャ

既存の強化学習（RL）トレーニングでは、モデルを学習させる「トレーナー」と行動データを生成する「ロールアウトワーカー」が同一データセンター内に固定されており、これが大規模化における最大のボトルネックとなっていました。しかし、Nan Jiang 氏らが提案する新アーキテクチャは、この物理的制約を打破し、地理的に分散した異種リソースを柔軟に統合する道を開きます。

従来の強化学習が抱える「物理的制約」

現在の標準的な RL トレーニングループでは、トレーナー（バックプロパゲーションを行う側）とロールアウトワーカー（環境とのインタラクションを行い、データを生成する側）は同じクラスタ内に存在することが前提となっています。この構成では、両者の通信が RDMA などの高速ファブリックで行われるため同期速度は速いものの、「十分な GPU 容量」「同一リージョン」「高速なネットワーク接続」「即時利用可能なリソース」のすべてを同時に満たすことが極めて困難です。

「IO（強化学習）はこれら 4 つの条件を同時に要求しますが、どれか一つでも欠けるとシステムが成り立ちません。特に GPU の弾力性がないため、トレーニング中に必要に応じてクラスタを拡大することはできません。」

この結果、ロールアウト側が必要以上に多くの計算リソースを求めても、トレーナー側の固定された容量に縛られてしまうというミスマッチが生じます。利用可能な計算資源は世界中に分散しているのに、強化学習のデフォルト形状が「一つの完璧なクラスタ」を要求しているため、大規模化が阻害されているのです。

島型アーキテクチャ：トレーナーとロールアウトの分離

解決策は、バックプロパゲーション（学習）を集中管理し、ロールアウト機能を地理的に分散した「島（Island）」として分離するというものです。この新しいアーキテクチャでは、トレーナーは特定のクラスタに留まり、複雑な集約処理を行います。一方、ロールアウト機能は世界各地の異なるプロバイダやオンプレミス環境にあるリソースを「島」として活用します。

「移動可能な単位は『バックプロパゲーション』ではなく、『ロールアウトサービングアイランド』です。島間にはグローバルな集約（All-reduce）のような重い依存関係はなく、単に『ポリシーバージョンの受け渡し』と『軌跡データの送信』のみが行われます。」

この分離により、トレーナー側は高速なファブリックを維持しつつ、ロールアウト側は価格や可用性が異なる多様なリソースプールを自由に組み合わせることが可能になります。グローバルな計算リソースプールを活用することで、地理的制約を取り払った弾力的なトレーニングが可能となるのです。

通信コスト劇減の鍵：「可視性閾値」と差分転送

アーキテクチャの変更だけでは不十分です。分散環境で最も懸念されるのは、モデルパラメータを頻繁に転送する際の通信帯域とストレージコストです。従来の考え方では、更新ごとに数ギガバイトから数百ギガバイトに及ぶフルチェックポイントの転送が必要でしたが、この新手法では「可視部分のみの変化（差分）」を送信することで、通信量を数百ギガバイトから数メガバイトレベルまで劇的に削減します。

これは以下の 2 つの技術的要素に基づいています。

1. BF16 の丸め誤差と閾値: オプティマイザは高精度（FP32）で重みを更新しますが、ロールアウトエンジンが実際に使用する際は BF16 にキャストされます。BF16 の数値表現には一定の「間隔（オップ）」があり、更新幅がこの間隔の半分未満であれば、BF16 形式での数値は変化しません。
2. Adam オプティマイザの小さな更新幅: Adam オプティマイザによるパラメータごとの更新幅（プッシュ）は、学習率に比例して非常に小さく制御されています。

「多くの重みにおいて、オプティマイザが更新する値は、BF16 の丸め誤差の閾値よりも千倍以上小さいのです。つまり、マスターウェイトは変化しても、ロールアウトエンジンが見る世界では『変化していない』と見なされる部分が多々あります。」

この現象を利用し、「ロールアウト側にとって実質的に変化しない重み部分は除外」し、変化が確認されたごく一部のパラメータのみを差分データとして送信します。これにより、転送するオブジェクトは数百ギガバイトから数メガバイトに縮小され、ネットワーク上の非同期環境でも秒単位での更新が可能になります。

Stitch プロトコルによる整合性と弾力性

この差分転送を実現するために採用されるのが「Stitch プロトコル」です。このプロトコルはバージョン管理と差分の適用を行い、異種バックエンドや非同期環境下でもロールアウトエンジン間の整合性を保ちます。

転送されるのは単なる浮動小数点の加算ではなく、ビットレベルでの等価なパッチデータです。ロールアウト側はこのパッチを正しく適用することで、フルチェックポイントを同期した場合と同じ状態をビット単位で再構築します。これにより、データの欠落やドリフトを防ぎつつ、世界中に散在する多様なリソースを安全にスケールさせることが可能になります。

まとめ

このアプローチは、大規模強化学習における計算リソースの制約を物理的な壁から解放し、地理的に分散した多様なクラウドプロバイダやオンプレミス環境のリソースを統合して活用する道を開きます。通信帯域とストレージコストの大幅削減により、より大規模かつ複雑な AI エージェントの開発が現実的なコストで可能になり、AI インフラの設計パラダイム自体を変革する可能性があります。