強化學習後訓練的權重傳輸，不到 2 秒即可完成

我們近期實現了 Kimi-K2 (1T 參數) 的 1.3 秒跨機器參數更新，將權重從 256 個訓練 GPU (BF16) 傳輸至 128 個推理 GPU (FP8)。

在非同步強化學習微調中，訓練與推理任務在不同的 GPU 上執行。在每個訓練步驟後，必須將新權重推送到推理節點。對於兆級參數模型，許多現有的框架需要數秒甚至數分鐘的時間。

透過利用 RDMA 點對點通訊，我們能夠在不更動推理引擎的情況下，極速完成權重傳輸，並簡化程式碼的編寫與維護工作。

我們的解決方案基於 RDMA WRITE，這是一種單向原語，允許源端直接寫入目的端的 GPU 記憶體。

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 開始傳輸。

定義了高階工作流程後，關鍵挑戰在於如何以兆級參數規模高效執行權重傳輸。我們在此說明執行路徑的細節。

訓練中的參數根據 FSDP 配置進行分佈。使用 full_tensor()，DeviceMesh 中的所有 GPU 都能重建完整參數，因此皆可作為權重傳輸的源端。

多個不相交的 DeviceMesh 組成一個 mesh group。由於同一群組中的 DeviceMesh 互不相交，其傳輸不會相互干擾，並可完全平行執行。我們在各 mesh group 之間插入一個全域屏障 (global barrier) 以強制執行順序。

權重傳輸管線會重疊主機到裝置的複製、張量組裝、融合與量化，以及跨 MeshGroups 的 RDMA 傳輸。

我們將每個參數張量的傳輸視為一個任務。權重傳輸過程利用了多種類型的硬體資源，因此我們將權重傳輸任務拆分為在時間上重疊的不同管線階段：

主機至裝置記憶體複製 (Host to device memcpy) — 若 FSDP 將權重卸載至 CPU

參數準備 — 使用 full_tensor() 重建完整權重，套用投影融合，並根據需要進行量化。

RDMA 傳輸 — 零複製寫入遠端推理 GPU 記憶體

全域屏障 (Global barrier) — 在所有 full_tensor() 呼叫完成後，透過乙太網路使用 GLOO 在各 mesh group 間進行同步。

在實作中，我們為每個管線階段維護一個 FIFO 任務佇列。當佇列頭部的任務完成該階段後，即會移動至下一個階段佇列的尾端。

full_tensor() 及其他 GPU 操作會引入額外的 GPU 記憶體佔用。為避免記憶體不足錯誤，只有在目前執行中任務所佔用的臨時 GPU 記憶體低於可設定的閾值時，我們才會開始執行任務。

幾項設計選擇使我們的系統比常見的開源解決方案執行速度顯著更快，且更易於維護。

常見的模式是透過 rank-0 GPU 匯聚所有參數：在訓練 rank-0 收集，發送至推理 rank-0，然後再次進行分散。這會迅速成為瓶頸，受限於單一 GPU 的 PCIe 頻寬和網卡（例如 400 Gbps ≈ 50 GB/s）。

相比之下，我們的點對點設定允許每個訓練 GPU 直接發送至每個推理 GPU，從而飽和整個網路架構，而非單一鏈路。

基於 RPC 的傳輸使用推理端緩衝區和工作程序呼叫，而單邊傳輸則可直接寫入權重，無需推理工作程序參與。

有些系統依賴於為每個張量呼叫推理引擎的 update_weight() 方法。這意味著對推理程式碼進行侵入式修改，外加 RPC、序列化與控制平面協調的開銷。

利用 RDMA WRITE 原語，我們可以在不進行額外複製的情況下，悄悄地更新推理 GPU 記憶體中的權重。過程中不涉及控制平面訊息與 CPU 控制邏輯，也無需修改推理引擎。

權重傳輸過程可利用四種類型的硬體資源：(1) 主機與裝置間的資料移動 (2) 用於投影融合與量化的 GPU 計算 (3) 用於資料平面的 RDMA 網路 (4) 用於控制平面的乙太網路。

我們的設計將權重傳輸任務拆分為管線階段，實現不同硬體資源間的輕鬆重疊。

有些實作會在每個訓練步驟重新計算傳輸排程，重複收集中繼資料並分發指令，這增加了不必要的控制平面延遲。

我們的排程僅在初始化時計算一次。每次訓練迭代只需重複該計劃：控制器發出「開始」訊號，GPU 則遵循其預先分配的路徑。執行過程可預測且輕量。

人們很容易試圖將整個權重更新過程糾纏在一個單體函數中：收集中繼資料、名稱匹配、節點內收集、投影融合、量化、拆分通訊組、節點間網路傳輸。這不僅難以正確編寫，更難以優化。

在工程實作中，我們將這些步驟拆分為獨立組件。每個組件都可以單獨進行單元測試、邏輯驗證與優化。

快速且可靠的權重傳輸是大規模強化學習微調的關鍵建構模組。透過結合 RDMA WRITE 原語、靜態傳輸排程與 管線化執行，我們將 Kimi-K2 上的兆級參數更新時間縮短至僅 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 Group

任務管線

GPU 記憶體使用控制

為何快速且簡單

點對點通訊

單向資料傳輸

管線化

靜態排程

清晰的解耦

結論

引用