Kimi開源又雙叒放大招了！

一個中間件，就能讓Kimi K2的萬億模型參數進入“秒更時代”。

圖片

不僅支持一次性把更新完的權重從一個節點同時發送給所有節點，還能實現點對點動態更新。

網友也算是大開眼界了，可謂頻頻驚嘆。

圖片

下面讓我們一起看看這個中間件到底是如何發揮大作用的。

20秒更新萬億參數

此中間件名為checkpoint-engine（檢查點引擎），主要用于強化學習中的一個關鍵步驟——在大語言模型推理過程中更新模型權重。

借助此組件，Kimi-K2只需約20秒就可實現在數千個GPU上更新1萬億參數。

與Kimi k1. 5類似，K2在同步強化學習訓練中采用了混合共置架構，即訓練引擎和推理引擎部署在同一組工作節點上。

當一個引擎處于活躍工作狀態時，另一個引擎會釋放或卸載其GPU資源以以配合資源調配。

在每一次強化學習訓練迭代中，集中式控制器會先調用推理引擎生成新的訓練數據，隨后通知訓練引擎基于這些數據進行訓練，并將更新后的參數發送至推理引擎，供下一輪迭代使用。

因此，每個引擎都針對高吞吐量進行了深度優化。

然而，隨著模型規模擴展至K2級別，引擎切換與故障恢復的延遲變得尤為顯著。

所以，研究團隊思考如何才能更高效地更新模型參數？

首先，在rollout階段，訓練引擎的參數會被卸載至DRAM（動態隨機存取存儲器）中，因此啟動訓練引擎僅需執行一次簡單的H2D數據傳輸。

但在這個過程中，啟動推理引擎會面臨更大挑戰，因為它必須從訓練引擎獲取更新后的參數，而兩者的分片范式并不相同。

研究團隊又考慮到K2的規模及龐大的設備數量，使用網絡文件系統來重新切分并廣播參數是不現實的。因為將系統開銷保持在較低水平，所需的總帶寬高達每秒數PB（千萬億字節）。

在上述背景下，檢查點引擎應運而生。

圖片

研究團隊選擇在訓練節點上部署分布式檢查點引擎來管理參數狀態。

執行參數更新時，每個檢查點工作節點先從訓練引擎獲取本地參數副本，然后將完整參數集廣播到所有檢查點節點。

隨后，推理引擎僅從檢查點引擎中獲取自己所需的參數分片即可。

為了支持1萬億參數的模型更新，他們還選擇采用參數逐條更新的流水線方式，將內存占用降至最低。

理論上的3階段流水線如下所示：

• H2D階段：將最新權重的一個分片異步復制到 H2D 緩沖區；
• 廣播階段：一旦復制完成，該分片會被復制到其中一個IPC緩沖區，并廣播到所有GPU；
• 重載階段：推理引擎同時從另一個IPC緩沖區加載參數。

圖片

但需注意的是，這種理想的3階段流水線目前尚未實現，K2應用的更簡單的兩階段方案。

• 所有設備先進行一次同步的H2D傳輸；
• 廣播和重載操作隨后并行進行。

圖片

他們選擇將完整參數集廣播到整個集群，而不考慮每個推理工作節點的具體切分方式。

雖然這種方式傳輸的數據量會比理論最優方案更多，但它可以簡化系統設計，對訓練和推理引擎的侵入性更低。

研究團隊認為，通過犧牲這一點微小的開銷，實現訓練引擎與推理引擎的完全解耦，大大簡化了維護和測試流程。

除了上述問題外，像Kimi K2這樣的大模型，優化啟動時間也至關重要。

啟動訓練引擎時，他們讓每個訓練工作節點選擇性地從磁盤讀取部分或不讀取任何參數，并將必要參數廣播至其他對等節點。

這么做的目的是確保所有工作節點只需集體讀取一次檢查點，從而最大限度地減少昂貴的磁盤IO開銷。

除此之外，由于推理引擎是獨立副本，研究團隊希望避免在它們之間引入額外的同步屏障。

因此，他們選擇在啟動階段復用檢查點引擎。

讓檢查點引擎先像訓練引擎啟動時一樣，集體從磁盤讀取檢查點，然后更新尚未初始化的推理引擎狀態。

值得一提的是，通過利用專門的檢查點引擎，系統還可以抵御單點故障，因為某個推理副本可以獨立重啟，而無需與其他副本通信。

這么一看，這一中間件真在Kimi K2中起了不小的作用呢。

參考鏈接：

[1]https://x.com/Kimi_Moonshot/status/1965785427530629243

[2]https://github.com/MoonshotAI/checkpoint-engine[3]https://arxiv.org/abs/2507.20534