vLLM 使用 Ray Direct Transport 進行大規模分片權重傳輸

概述

vLLM 已實作了一種原生分片權重傳輸引擎，利用 Ray Direct Transport (RDT) 來優化在線上強化學習 (RL) 設定中，訓練器 (trainer) 與推理工作節點 (inference workers) 之間的模型權重同步。此系統取代了傳統的 NCCL broadcast 方法，後者在兆參數規模下經常面臨記憶體瓶頸與同步停頓問題，改採基於拉取 (pull-based) 的分片方法，從而降低峰值記憶體使用量並提升傳輸速度。

Broadcast-Based Sync 的限制

標準的權重同步通常依賴 NCCL broadcasts，其中訓練器會將參數以 HuggingFace 格式進行 all-gather 並 broadcast 到每個推理工作節點。這種方法對於大規模模型面臨兩個主要挑戰：

記憶體效率低下：在如 Tensor Parallelism 8 (TP8) 的配置中，工作節點會接收完整的模型，但僅保留 1/8 的權重，並捨棄其餘部分。對於大型混合專家模型 (MoE)，這會造成巨大的峰值記憶體開銷。
集體同步 (Collective Synchronization)：NCCL 要求所有 rank 必須同步參與。落後的 rank 或副本失效會導致整個集體操作停頓，這對於動態、大規模環境而言非常棘手。

技術實作：分片權重傳輸

Recording Tensor Dry Run

為了確保與各種模型架構（例如 Llama-4 的 fused experts 或各種模型中的 GQA）的相容性，vLLM 在初始化期間使用「記錄張量 (recording tensor)」進行 dry run。vLLM 的載入器 (loaders) 會獲得一個張量子類別，該子類別會回報其形狀 (shape) 與資料型別 (dtype)，但不包含實際數據。每一次轉換——包括 views、narrows、transposes 與 reshapes——都會被記錄為一連串的操作鏈（即「分片計畫 (sharding plan)」）。

此計畫允許訓練器執行初始佈局操作（融合、重新佈局、拆分與分片），並僅傳輸每個 vLLM rank 所需的 BF16 格式特定分片權重，從而確保過程在建構上即是正確的。

Ray Direct Transport (RDT) 與 NIXL

該引擎利用 Ray Direct Transport (RDT) 搭配 NIXL 後端，以實現 Ray actors 之間的直接 GPU-to-GPU 通訊。這種架構實現了基於拉取的系統，推理 rank 僅從對應的訓練器 rank 中拉取所需的特定分片張量。

初始化流程包含五個步驟：

訓練器 rank 進行 all-gather 以獲取所有權元數據 (ownership metadata，包括參數名稱、dtypes 與 shapes)。
Rank 0 將此元數據與訓練器 Ray actor 名稱傳送給推理工作節點。
Each vLLM worker 透過 recording-tensor dry run 建立其分片計畫。
工作節點以負載平衡的方式將自己映射到來源訓練器 rank。
兩端（生產者與消費者）預先配置並註冊 RDT buffers。

性能優化

vLLM 針對 Qwen3-235B-A22B 模型（TP4/PP2/EP8 訓練器對應 DP16/EP16 vLLM 伺服器）進行了三次引擎版本的迭代，以優化端到端延遲：

V1 (Simple Iterator)：逐一收集所有維度（TP, PP, EP）的參數。這導致了數千個微小的集體操作與冗餘的記憶體使用，同步時間為 25.02s。
V2 (PP/EP-Local)：實作了 PP-local gathers（僅在同一個 pipeline stage 內）與 EP-local transfers（專家模型不進行 gather；推理 rank 直接從持有該專家的 rank 中拉取）。這將同步時間降低至 5.61s。
V3 (Pipelined Execution)：引入了 all-gather、replay 操作與 RDMA 傳輸的重疊 (overlapping)。透過在 decoder block groups 中收集權重並在背景進行處理，同步延遲降低至 3.49s。

大規模驗證：Kimi K2

在 48 個 8xH100 節點（32 個訓練器節點，16 個推理節點）上的 Kimi K2 模型驗證結果如下：

Metric	Value
Bytes moved per sync	7.9 TB
Weight sync time	7.53s
Aggregate bandwidth	1,049 GB/s

考慮到 vLLM 的逐層重新載入邏輯限制，此效能約為該特定設定下預期「光速 (SoL)」傳輸時間的 1.5 倍。

容錯與整合

透過使用 NIXL 而非 broadcast collectives，系統本身對失效具有更高的韌性。如果推理引擎失效，路由器會繼續將流量導向其餘引擎，而訓練器在下次同步時僅與運作中的引擎進行通訊。一旦失效的副本被修復，它會在下次同步邊界重新加入，並接收更新後的權重，而不會影響整體的收斂性。

框架整合

該引擎已整合至 SkyRL。其他 RL 框架可以透過實作一個提供參數元數據並產出實體化張量的 WeightSource 迭代器來採用此技術，並可選擇實作 held_names 方法以啟用 EP/PP-local 優化。

目前的限制

Loader Constraints：載入器必須使用可記錄的操作；那些在載入期間檢查實際數值的載入器將會失敗。
Memory Budget：RDT 目標緩衝區 (destination buffers) 存在於 vLLM 的 gpu_memory_utilization 預算之外。
Compatibility：目前的實作與 vLLM 中的 EPLB 不相容。
Serialization：傳輸目前在訓練器 PP groups 之間是序列化的，以防止在逐層重新載入期間發生 OOM。