2秒未満で完了するRLポストトレーニング用ウェイト転送

Kimi-K2（1Tパラメータ）において、256基の学習用GPU（BF16）から128基の推論用GPU（FP8）へウェイトを転送する際、機械間パラメータ更新を1.3秒で達成しました。

非同期強化学習のファインチューニングでは、学習と推論が個別のGPUで実行されます。各学習ステップの終了後、新しいウェイトを推論ノードにプッシュする必要があります。既存の多くのフレームワークでは、1兆パラメータ規模のモデルにおいて、これに数秒から数分を要しています。

RDMAポイントツーポイント通信を活用することで、推論エンジンを変更することなく極めて高速なウェイト転送を実現し、コードの記述と保守を簡素化しました。

当社のソリューションは、ソースから宛先のGPUメモリへ直接書き込む片方向プリミティブであるRDMA WRITEに基づいています。

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

宛先側には転送の通知すら行われません。これにより、推論ノード側に制御ロジックを必要とせず、学習ノード主導による低遅延、高スループット、ゼロコピーの転送が実現します。

メタデータ収集 – コントローラーがすべての学習用および推論用GPUからパラメータのメタデータを収集します。

スケジュール計算 – コントローラーが、どの学習用GPUがどのパラメータをどの推論用GPUに、どのような順序で送信するかをマッピングする静的なウェイト転送スケジュールを計算します。

スケジュールの配布 – コントローラーがスケジュールをすべての学習用GPUに送信します。

実行 – 各学習ステップの後、コントローラーが学習用GPUに対して転送開始のシグナルを送ります。

高レベルなワークフローを定義した上で、重要な課題は、1兆パラメータ規模でいかに効率的にウェイト転送を実行するかという点です。以下に実行パスの詳細を説明します。

学習時のパラメータはFSDPプレースメントに従って分散されます。full_tensor()を使用することで、DeviceMesh内の全GPUが完全なパラメータを再構成できるため、すべてのGPUがウェイト転送のソースとして機能します。

複数の分離されたDeviceMeshがmeshグループを形成します。同一グループ内のDeviceMeshは分離されているため、転送が干渉せず完全に並列実行可能です。Meshグループ間には、順序を強制するためにグローバルバリアを挿入します。

重み転送パイプラインは、ホストからデバイスへのコピー、テンソルの組み立て、融合と量子化、そしてMeshGroup間のRDMA転送を重ね合わせます。

各パラメータテンソルの転送を1つのタスクとして扱います。ウェイト転送プロセスは複数種類のハードウェアリソースを利用するため、ウェイト転送タスクを時間的にオーバーラップする以下のパイプラインステージに分割します。

ホストからデバイスへのメモリコピー — FSDPがウェイトをCPUにオフロードする場合

パラメータ準備 — full_tensor()による完全なウェイトの再構成、プロジェクションフュージョンの適用、必要に応じた量子化を行います。

RDMA転送 — リモートの推論用GPUメモリへのゼロコピー書き込みを行います。

グローバルバリア — すべてのfull_tensor()呼び出し終了後、イーサネット経由でGLOOを使用し、Meshグループ間で同期をとります。

実装においては、各パイプラインステージでタスクのFIFOキューを維持しています。キューの先頭のタスクがステージを完了すると、次のステージのキューの末尾に移動されます。

full_tensor()およびその他のGPU操作は追加のGPUメモリ消費を招きます。メモリ不足エラーを回避するため、実行中のタスクが使用する一時的なGPUメモリ量が設定可能なウォーターマークを下回っている場合にのみ、タスクの実行を開始します。

いくつかの設計上の選択により、一般的なオープンソースソリューションよりもシステムを大幅に高速化し、保守を容易にしています。

ポイントツーポイント通信は、両方の学習ノードから推論ノードのネットワークインターフェースへ重みを転送します。

一般的なパターンとして、すべてのパラメータをランク0のGPUに集約する方法（学習ランク0で収集し、推論ランク0に送信、その後再分配）があります。これは単一GPUのPCIe帯域幅やNIC（例：400 Gbps ≈ 50 GB/s）に制限されるため、すぐにボトルネックとなります。

これに対し、当社のポイントツーポイント構成では、すべての学習用GPUがすべての推論用GPUへ直接送信できるため、単一リンクではなくネットワークファブリック全体を飽和させることができます。

RPCベースの転送では推論側バッファーとワーカー呼び出しを使用し、一方、一方向転送では推論ワーカーの関与なしに重みを直接書き込みます。

一部のシステムでは、各テンソルに対して推論エンジンのupdate_weight()メソッドを呼び出す手法に依存しています。これは推論コードへの介入的な変更を意味し、さらにRPC、シリアライズ、コントロールプレーンの調整によるオーバーヘッドが発生します。

RDMA WRITEプリミティブを使用することで、余分なコピーなしで推論用GPUメモリ上のウェイトを静かに更新します。コントロールプレーンメッセージやCPU制御ロジックは関与せず、推論エンジンの修正も不要です。

ウェイト転送プロセスは、(1) ホスト-デバイス間のデータ移動、(2) プロジェクションフュージョンと量子化のためのGPU計算、(3) データプレーンのためのRDMAネットワーク、(4) コントロールプレーンのためのイーサネット、という4種類のハードウェアリソースを活用できます。

当社の設計ではウェイト転送タスクをパイプラインステージに分割しており、異なるハードウェアリソース間での容易なオーバーラップを可能にしています。

一部の実装では、学習ステップごとに転送スケジュールを再計算しており、メタデータの収集と命令の配布を繰り返しています。これは不要なコントロールプレーンの遅延を追加します。

当社のスケジュールは初期化時に一度計算されます。各学習反復では単にプランを再実行するだけであり、コントローラーが「開始」シグナルを送り、GPUが事前に割り当てられたルートに従います。実行は予測可能かつ軽量です。

メタデータ収集、名前照合、ノード内集約、プロジェクションフュージョン、量子化、通信世界のサブスライス、ノード間ネットワーク転送といったウェイト更新プロセス全体を、1つのモノリシックな関数に詰め込みたくなる誘惑があります。しかし、そのような実装は正確にプログラムするのが難しく、最適化はさらに困難です。

当社ではエンジニアリングにおいて、これらのステップを個別のコンポーネントとして分離しています。各コンポーネントは分離された状態で単体テスト、推論、最適化が可能です。

高速で信頼性の高いウェイト転送は、大規模なRLファインチューニングにおける重要な構成要素です。RDMA WRITEプリミティブ、静的転送スケジュール、パイプライン化された実行を組み合わせることで、Kimi-K2において1兆パラメータの更新をわずか1.3秒に短縮しました。このアプローチは論理的にシンプルで保守しやすく、従来型設計のボトルネックを回避します。

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE: 片方向転送

高レベルワークフロー

ウェイト転送の実行

DeviceMeshとMeshグループ

タスクパイプライン

GPUメモリ使用量の制御

高速かつシンプルな理由

ポイントツーポイント通信

片方向データ転送

パイプライン処理

静的スケジュール

明確な分離

結論

引用