2 秒內完成 RL 後訓練的權重傳輸

我們最近實現了 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 都可以重建完整參數，因此都可以作為權重傳輸的來源。

多個不相交的 DeviceMeshes 形成一個 網格群組 (mesh group)。由於同一群組中的 DeviceMeshes 是不相交的，因此它們的傳輸不會互相干擾，並且可以完全並行執行。在網格群組之間，我們插入了一個全域屏障來強制執行順序。

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

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

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

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

RDMA 傳輸 — 零拷貝寫入遠端推理 GPU 記憶體

全域屏障 — 在所有 full_tensor() 呼叫完成後，透過乙太網路使用 GLOO 跨網格群組進行同步。

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

full_tensor() 和其他 GPU 操作會引入額外的 GPU 記憶體使用量。為了避免記憶體不足 (OOM) 錯誤，只有當當前進行中的任務所佔用的臨時 GPU 記憶體小於可配置的水位線時，我們才開始執行任務。

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

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

相比之下，我們的點對點設定允許每個訓練 GPU 直接向每個推理 GPU 發送，從而填滿整個網路架構，而不是單個鏈路。

基於 RPC 的傳輸會使用推理端緩衝區與 worker 呼叫；而單邊傳輸則會在不需要推理 worker 參與的情況下直接寫入權重。

某些系統依賴於為每個張量呼叫推理引擎的 update_weight() 方法。這意味著對推理程式碼進行侵入性更改，以及來自 RPC、序列化和控制平面協調的額外負荷。

借助 RDMA WRITE 原語，我們可以在推理 GPU 記憶體上悄悄更新權重，而無需額外拷貝。不涉及控制平面訊息和 CPU 控制邏輯。不需要修改推理引擎。

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

我們的設計將權重傳輸任務拆分為管線階段，從而便於在不同的硬體資源之間進行重疊。

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

我們的排程在初始化時計算一次。每個訓練迭代只需重播計劃：控制器發出「開始」信號，GPU 遵循其預先指派的路徑。執行是可預測且輕量級的。

人們很容易將整個權重更新過程糾纏在一個單體函數中：收集中繼資料、名稱匹配、節點內收集、投影融合、量化、子切片通訊世界、節點間網路傳輸。這很難正確編程，而且更難優化。

在我們的工程中，我們將這些步驟分離為單獨的元件。每個元件都可以獨立進行單元測試、推導和優化。

快速、可靠的權重傳輸是大規模 RL 微調的關鍵建構區塊。透過結合 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 Groups)

任務管線

GPU 記憶體使用量控制

為什麼它既快又簡單

點對點通訊

單向資料傳輸

管線化

靜態排程

清晰的分離

結論

引文