時(shí)令 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI
Kimi開(kāi)源又雙叒放大招了！
一個(gè)中間件，就能讓Kimi K2的萬(wàn)億模型參數(shù)進(jìn)入“秒更時(shí)代”。
不僅支持一次性把更新完的權(quán)重從一個(gè)節(jié)點(diǎn)同時(shí)發(fā)送給所有節(jié)點(diǎn)，還能實(shí)現(xiàn)點(diǎn)對(duì)點(diǎn)動(dòng)態(tài)更新。
網(wǎng)友也算是大開(kāi)眼界了，可謂頻頻驚嘆。
下面讓我們一起看看這個(gè)中間件到底是如何發(fā)揮大作用的。
20秒更新萬(wàn)億參數(shù)
此中間件名為checkpoint-engine（檢查點(diǎn)引擎），主要用于強(qiáng)化學(xué)習(xí)中的一個(gè)關(guān)鍵步驟——在大語(yǔ)言模型推理過(guò)程中更新模型權(quán)重。
借助此組件，Kimi-K2只需約20秒就可實(shí)現(xiàn)在數(shù)千個(gè)GPU上更新1萬(wàn)億參數(shù)。
與Kimi k1. 5類似，K2在同步強(qiáng)化學(xué)習(xí)訓(xùn)練中采用了混合共置架構(gòu)，即訓(xùn)練引擎和推理引擎部署在同一組工作節(jié)點(diǎn)上。
當(dāng)一個(gè)引擎處于活躍工作狀態(tài)時(shí)，另一個(gè)引擎會(huì)釋放或卸載其GPU資源以以配合資源調(diào)配。
在每一次強(qiáng)化學(xué)習(xí)訓(xùn)練迭代中，集中式控制器會(huì)先調(diào)用推理引擎生成新的訓(xùn)練數(shù)據(jù)，隨后通知訓(xùn)練引擎基于這些數(shù)據(jù)進(jìn)行訓(xùn)練，并將更新后的參數(shù)發(fā)送至推理引擎，供下一輪迭代使用。
因此，每個(gè)引擎都針對(duì)高吞吐量進(jìn)行了深度優(yōu)化。
然而，隨著模型規(guī)模擴(kuò)展至K2級(jí)別，引擎切換與故障恢復(fù)的延遲變得尤為顯著。
所以，研究團(tuán)隊(duì)思考如何才能更高效地更新模型參數(shù)？
首先，在rollout階段，訓(xùn)練引擎的參數(shù)會(huì)被卸載至DRAM（動(dòng)態(tài)隨機(jī)存取存儲(chǔ)器）中，因此啟動(dòng)訓(xùn)練引擎僅需執(zhí)行一次簡(jiǎn)單的H2D數(shù)據(jù)傳輸。
但在這個(gè)過(guò)程中，啟動(dòng)推理引擎會(huì)面臨更大挑戰(zhàn)，因?yàn)樗仨殢挠?xùn)練引擎獲取更新后的參數(shù)，而兩者的分片范式并不相同。
研究團(tuán)隊(duì)又考慮到K2的規(guī)模及龐大的設(shè)備數(shù)量，使用網(wǎng)絡(luò)文件系統(tǒng)來(lái)重新切分并廣播參數(shù)是不現(xiàn)實(shí)的。因?yàn)閷⑾到y(tǒng)開(kāi)銷(xiāo)保持在較低水平，所需的總帶寬高達(dá)每秒數(shù)PB（千萬(wàn)億字節(jié)）。
在上述背景下，檢查點(diǎn)引擎應(yīng)運(yùn)而生。
研究團(tuán)隊(duì)選擇在訓(xùn)練節(jié)點(diǎn)上部署分布式檢查點(diǎn)引擎來(lái)管理參數(shù)狀態(tài)。
執(zhí)行參數(shù)更新時(shí)，每個(gè)檢查點(diǎn)工作節(jié)點(diǎn)先從訓(xùn)練引擎獲取本地參數(shù)副本，然后將完整參數(shù)集廣播到所有檢查點(diǎn)節(jié)點(diǎn)。
隨后，推理引擎僅從檢查點(diǎn)引擎中獲取自己所需的參數(shù)分片即可。
為了支持1萬(wàn)億參數(shù)的模型更新，他們還選擇采用參數(shù)逐條更新的流水線方式，將內(nèi)存占用降至最低。
理論上的3階段流水線如下所示：
H2D階段：將最新權(quán)重的一個(gè)分片異步復(fù)制到 H2D 緩沖區(qū)；
廣播階段：一旦復(fù)制完成，該分片會(huì)被復(fù)制到其中一個(gè)IPC緩沖區(qū)，并廣播到所有GPU；
重載階段：推理引擎同時(shí)從另一個(gè)IPC緩沖區(qū)加載參數(shù)。
但需注意的是，這種理想的3階段流水線目前尚未實(shí)現(xiàn)，K2應(yīng)用的更簡(jiǎn)單的兩階段方案。
所有設(shè)備先進(jìn)行一次同步的H2D傳輸；
廣播和重載操作隨后并行進(jìn)行。
他們選擇將完整參數(shù)集廣播到整個(gè)集群，而不考慮每個(gè)推理工作節(jié)點(diǎn)的具體切分方式。
雖然這種方式傳輸?shù)臄?shù)據(jù)量會(huì)比理論最優(yōu)方案更多，但它可以簡(jiǎn)化系統(tǒng)設(shè)計(jì)，對(duì)訓(xùn)練和推理引擎的侵入性更低。
研究團(tuán)隊(duì)認(rèn)為，通過(guò)犧牲這一點(diǎn)微小的開(kāi)銷(xiāo)，實(shí)現(xiàn)訓(xùn)練引擎與推理引擎的完全解耦，大大簡(jiǎn)化了維護(hù)和測(cè)試流程。
除了上述問(wèn)題外，像Kimi K2這樣的大模型，優(yōu)化啟動(dòng)時(shí)間也至關(guān)重要。
啟動(dòng)訓(xùn)練引擎時(shí)，他們讓每個(gè)訓(xùn)練工作節(jié)點(diǎn)選擇性地從磁盤(pán)讀取部分或不讀取任何參數(shù)，并將必要參數(shù)廣播至其他對(duì)等節(jié)點(diǎn)。
這么做的目的是確保所有工作節(jié)點(diǎn)只需集體讀取一次檢查點(diǎn)，從而最大限度地減少昂貴的磁盤(pán)IO開(kāi)銷(xiāo)。
除此之外，由于推理引擎是獨(dú)立副本，研究團(tuán)隊(duì)希望避免在它們之間引入額外的同步屏障。
因此，他們選擇在啟動(dòng)階段復(fù)用檢查點(diǎn)引擎。
讓檢查點(diǎn)引擎先像訓(xùn)練引擎啟動(dòng)時(shí)一樣，集體從磁盤(pán)讀取檢查點(diǎn)，然后更新尚未初始化的推理引擎狀態(tài)。
值得一提的是，通過(guò)利用專門(mén)的檢查點(diǎn)引擎，系統(tǒng)還可以抵御單點(diǎn)故障，因?yàn)槟硞€(gè)推理副本可以獨(dú)立重啟，而無(wú)需與其他副本通信。
這么一看，這一中間件真在Kimi K2中起了不小的作用呢。
參考鏈接：
[1]https://x.com/Kimi_Moonshot/status/1965785427530629243
[2]https://github.com/MoonshotAI/checkpoint-engine
[3]https://arxiv.org/abs/2507.20534
一鍵三連「點(diǎn)贊」「轉(zhuǎn)發(fā)」「小心心」
歡迎在評(píng)論區(qū)留下你的想法！
—?完?—
專屬AI產(chǎn)品從業(yè)者的實(shí)名社群，只聊AI產(chǎn)品最落地的真問(wèn)題
掃碼添加小助手，發(fā)送「姓名+公司+職位」申請(qǐng)入群～
進(jìn)群后，你將直接獲得：
最新最專業(yè)的AI產(chǎn)品信息及分析
不定期發(fā)放的熱門(mén)產(chǎn)品內(nèi)測(cè)碼
內(nèi)部專屬內(nèi)容與專業(yè)討論
點(diǎn)亮星標(biāo)
科技前沿進(jìn)展每日見(jiàn)
