vLLM Ray Direct Transport による大規模シャード化された重みの転送

概要

vLLM は、オンライン強化学習 (RL) 設定において、トレーナーと推論ワーカー間のモデル重みの同期を最適化するために、Ray Direct Transport (RDT) を活用したネイティブなシャード化された重み転送エンジンを実装しました。このシステムは、従来の NCCL ブロードキャスト方式（兆単位のパラメータ規模ではメモリのボトルネックや同期の停止が発生しがちです）に代わり、ピークメモリ使用量を削減し転送速度を向上させる、プルベースのシャード化されたアプローチを採用しています。

ブロードキャストベースの同期の限界

標準的な重みの同期は、通常 NCCL ブロードキャストに依存しています。ここでは、トレーナーがパラメータを HuggingFace 形式で all-gather し、それらをすべての推論ワーカーにブロードキャストします。このアプローチは、大規模モデルにおいて主に 2 つの課題を提示します。

メモリの非効率性: Tensor Parallelism 8 (TP8) のような構成では、ワーカーはモデル全体を受け取りますが、重みの 1/8 だけを保持し、残りは破棄します。大規模な Mixture-of-Experts (MoE) モデルの場合、これは膨大なピークメモリのオーバーヘッドを生み出します。
集団同期 (Collective Synchronization): NCCL はすべてのランクが同期的に参加することを要求します。遅延のあるランク (straggler ranks) やレプリカの失敗は、集団全体の同期を停止させることがあり、動的な大規模環境においては問題となります。

技術的実装: シャード化された重みの転送

Recording Tensor によるドライラン

多様なモデルアーキテクチャ（Llama-4 の fused experts や、さまざまなモデルにおける GQA など）との互換性を確保するため、vLLM は初期化中に「recording tensor」によるドライランを行います。vLLM のローダーには、形状 (shape) とデータ型 (dtype) を報告しますがデータは含まないテンソル・サブクラスが提供されます。view、narrow、transpose、reshape を含むすべての変換は、一連の操作（「sharding plan」）として記録されます。

このプランにより、トレーナーは初期のレイアウト操作（fusion、relayout、splitting、sharding）を実行し、各 vLLM ランクが必要とする特定のシャード化された重みを BF16 形式で転送することだけが可能になり、プロセスが構造的に正しいことが保証されます。

Ray Direct Transport (RDT) と NIXL

このエンジンは、Ray Direct Transport (RDT) と NIXL バックエンドを利用して、Ray actor 間の直接的な GPU-to-GPU 通信を可能にします。このアーキテクチャにより、推論ランクが、マッピングされたトレーナーランクから必要なシャード化されたテンソルのみをプルするプルベースのシステムが実現します。

初期化フローは 5 つのステップで構成されます:

トレーナーランクが所有権のメタデータ（パラメータ名、dtype、shape）を all-gather します。
Rank 0 がこのメタデータとトレーナーの Ray actor 名を推論ワーカーに送信します。
各 vLLM ワーカーは、recording-tensor ドライランを通じて自身の sharding plan を作成します。
ワーカーは、ロードバランスが取れた方法でソースとなるトレーナーランクに自身をマッピングします。
プロデューサーとコンシューマーの両方が、事前に RDT バッファを割り当てて登録します。

パフォーマンス最適化

vLLM は、Qwen3-235B-A22B モデル（TP4/PP2/EP8 トレーナーから DP16/EP16 vLLM サーバー）のエンドツーエンドのレイテンシを最適化するために、エンジンの 3 つのバージョンを反復開発しました:

V1 (Simple Iterator): すべての次元 (TP, PP, EP) にわたってパラメータを一つずつ gather します。これにより、数千の小さな集団操作が発生し、冗長なメモリ使用が生じ、同期時間は 25.02s となりました。
V2 (PP/EP-Local): PP-local gather（同じパイプラインステージ内のみ）と EP-local 転送（エキスパートは gather しません。推論ランクがエキスパートを保持するランクから直接プルします）を実装しました。これにより、同期時間は 5.61s に減少しました。
V3 (Pipelined Execution): all-gather、replay 操作、および RDMA 転送のオーバーラップを導入しました。デコーダーブロックのグループごとに重みを gather し、バックグラウンドで処理することで、同期レイテンシは 3.49s まで低下しました。

大規模検証: Kimi K2

48 ノードの 8xH100（32 トレーナーノード、16 推論ノード）にわたる Kimi K2 モデルでの検証では、以下の結果が得られました:

メトリック	値
同期あたりの転送バイト数	7.9 TB
重みの同期時間	7.53s
合計帯域幅	1,049 GB/s

このパフォーマンスは、vLLM のレイヤーごとのリロードロジックの制約を考慮すると、この特定のセットアップにおける期待される「光速 (SoL)」転送時間の約 1.5 倍です。

障害耐性と統合

ブロードキャスト集団操作の代わりに NIXL を利用することで、システムは本質的に障害に対してより強靭です。推論エンジンが失敗した場合、ルーターは残りのエンジンにトラフィックを継続的に誘導し、トレーナーは次の同期の際に、稼働中のエンジンのみと通信します。失敗したレプリカが復旧した際、次の同期境界で再参加し、全体の収束に影響を与えることなく更新された重みを送受信します。

フレームワーク統合

このエンジンは SkyRL に統合されています。他の RL フレームワークは、パラメータのメタデータを提供し、実体化されたテンソルを yield する WeightSource イテレータを実装することで、このエンジンを採用できます（オプションとして EP/PP-local 最適化を有効にするための held_names メソッドを備えることができます）。

現在の制限事項

ローダーの制約: ローダーは記録可能な操作を使用する必要があります。ロード中に実際の値を確認する操作は失敗します。
メモリ予算: RDT 転送先バッファは、vLLM の gpu_memory_utilization 予算の外部に存在します。
互換性: 現在の実装は vLLM の EPLB と互換性がありません。
シリアル化: レイヤーごとのリロード中の OOM を防ぐため、現在の転送はトレーナーの PP グループ間でシリアルに行われます。