大家好，我是 Ai 學(xué)習(xí)的老章

推薦一篇Kimi工程師的硬核文章

在今年 7 月發(fā)布的 Kimi K2[1] 模型中，我們實現(xiàn)了 RL 期間 1T 參數(shù)模型僅用約 20s 完成整個參數(shù)更新流程，顯著優(yōu)化了的 RL E2E 耗時的關(guān)鍵效率瓶頸。在實現(xiàn)高效參數(shù)更新的路上我們也踩了很多坑，希望寫一篇文章也來聊聊我們遇到的一些問題和嘗試的解法。

我們的場景是怎樣的

目前在 LLM 的強(qiáng)化學(xué)習(xí)訓(xùn)練主要分為 colocate 和 disaggregation 兩種架構(gòu)，colocate 即訓(xùn)推共享 GPU 資源，會交替占用 GPU；disaggregation 即訓(xùn)推分離，各自占用不同的 GPU device。在這兩種架構(gòu)里面，一個很重要的階段是參數(shù)更新，即每一輪訓(xùn)練結(jié)束之后需要把參數(shù)同步給推理框架然后開始 rollout，這個階段里面如果占用時間過長，會導(dǎo)致 GPU 一直在空轉(zhuǎn)，無法提高整體的 GPU 利用率，訓(xùn)練端到端的性能會有比較大的損失，因此在 RL 訓(xùn)練中，高效的參數(shù)更新是一個值得優(yōu)化的方向。

我們內(nèi)部的實驗中，colocate 場景會多一些，因此我們也著重針對了 colocate 做了相關(guān)優(yōu)化。我們的 colocate RL 和社區(qū)大部分開源框架的不太一樣，我們并沒有采用 Ray 這種通用的方案，而是讓訓(xùn)練和推理用兩個不同的容器共享一臺機(jī)器來部署，這樣的好處是能夠讓訓(xùn)練和推理完全解耦，各自的功能迭代互相不影響，同時它們的環(huán)境和鏡像也是完全解耦的，這套方案讓我們內(nèi)部訓(xùn)練和推理的迭代非常順暢，同時這種部署方式和我們線上 Kimi 的部署方式比較類似，因此也能較好的復(fù)用一些線上服務(wù)的基礎(chǔ)設(shè)施。當(dāng)然它也帶來了一個問題，即訓(xùn)練和推理較難感知各自的 Process Group，因此給參數(shù)更新帶來了一定的難度。因此我們當(dāng)時在設(shè)計的時候，期望能提供一層輕量的中間層，把訓(xùn)推能夠連起來，盡可能的不去動訓(xùn)練和推理 Process Group 的邏輯，于是我們開始著手設(shè)計checkpoint-engine，希望在訓(xùn)推之間建立一個 bridge 去實現(xiàn)參數(shù)更新。

checkpoint-engine

需要注意的是，我們當(dāng)時選擇這個方案也是經(jīng)過了一些考量和權(quán)衡的。從理論上來說，性能最優(yōu)的方案當(dāng)然是讓 RL 框架能感知訓(xùn)推的并行度，然后有一個統(tǒng)一的控制器去做 reshard 實現(xiàn)參數(shù)更新，這樣肯定是能實現(xiàn)最快速的參數(shù)更新，而且也不會存在冗余的數(shù)據(jù)發(fā)送。但是從工程上來說，我們不想侵入太多邏輯到訓(xùn)練和推理引擎中，因此做了一個權(quán)衡，把 checkpoint-engine 設(shè)計成了每張卡都會收到全量的權(quán)重，這樣在推理已經(jīng)拆分 TP or EP 的情況下肯定會存在數(shù)據(jù)的冗余，但是從網(wǎng)卡和 NVLink 帶寬上來說，H800 這種機(jī)型能提供機(jī)間 or 機(jī)內(nèi)至少 100GiB/s 的性能，因此全量 1TiB 模型權(quán)重在所有機(jī)器通訊一次使用 10s 在我們看來是一個可以接受的范圍，因此我們采用了這個從工程上最解耦的一個方案。

最初的邏輯

在最初開發(fā) k1.5[2] 的時候，我們就實現(xiàn)了 checkpoint-engine 參數(shù)更新的邏輯，思路就是通過 CUDA IPC 來傳輸 tensor 數(shù)據(jù)，在 inference 里提供一個接口接受 CUDA IPC 的 tensor，然后 train 的每個 rank 會對應(yīng)一個 checkpoint-engine，每次 train 完會把 weights 傳給對應(yīng)的 checkpoint-engine，checkpoint-engine 會先通過 broadcast 把參數(shù)傳給每一個 rank，然后再通過 CUDA IPC 把每一個 tensor 打包成 IPC handle 共享給 inference 實現(xiàn)參數(shù)更新，大致流程如下圖所示。在傳輸?shù)臅r候，因為顯存無法裝下所有 weights，因此我們最初采取的分組是 per layer per ep，這個組合當(dāng)時基本能滿足我們的需求。

init

最初的 checkpoint-engine 參數(shù)更新方案

踩坑和填坑

隨著模型尺寸的進(jìn)一步增加，我們發(fā)現(xiàn)這種 per layer per ep 和把每一個 tensor 都打包成 IPC handle 的方式會造成很大的開銷。我們在千卡 H800 上跑 Kimi K2 RL 的時候遇到了很大的性能瓶頸，參數(shù)更新時間甚至能達(dá)到 10min，因此我們不得不著手做深度優(yōu)化。

共用 buffer 并實現(xiàn)兩階段流水線

在做 profile 之后，我們發(fā)現(xiàn)了幾個問題

• per layer per ep 這種模式會導(dǎo)致顯存占用很不穩(wěn)定，有時候會導(dǎo)致參數(shù)更新的時候直接 CUDA OOM

• per layer per ep 是每個 tensor 各自發(fā) async broadcast，最終等發(fā)完再去做 update weights，可能會存在很多細(xì)碎的小通訊

• 每個 tensor 一個 IPC handle 這種模式會導(dǎo)致 vLLM 那邊做序列化和反序列化時間非常長，初版實現(xiàn)會把所有 rank 的 ipc handle 都 gather 到一起傳給 vLLM，vLLM 自己傳給各個 TP 或 EP 的時候會傳全量數(shù)據(jù)

• Checkpoint-engine 里的通訊操作和 vLLM 里 update weights 的操作是串行的

因此基于上面的問題，我們的解決方案也比較清晰了，就是采用 bucket 的方式攢一下 batch，把細(xì)碎的 tensor 放入一個固定 bucket size 的 buffer 里

1. 每次只 broadcast 這個 buffer，這樣就能干掉小通訊帶來的 overhead，也能讓顯存占用非常穩(wěn)定

2. 把這個 buffer 一開始就 share 給 vLLM，二者把 buffer 當(dāng)做 channel 來做數(shù)據(jù)面的傳輸，這樣能避免每次請求傳輸 ipc handle 帶來的開銷，僅第一次請求有傳 ipc handle 的開銷，后面?zhèn)鬏數(shù)亩际?tensor meta 信息，能極大的減輕 vLLM 序列化和反序列化開銷

3. 同時我們針對 ipc handle 的 gather 邏輯做了優(yōu)化

   1. 不用每次請求都 gather 一次 ipc handle，僅第一個請求之前 gather 就行，減少了通訊操作

   2. 不用所有 rank all_gather ipc handles，因為每個 vLLM 實例只會有一個 rank 上的 checkpoint-engine 會發(fā)起請求，因此把 ipc_handle 都 gather 給它就行，不用做全局 all_gather

4. 為了把 broadcast 和 vLLM 里的 update weights 操作 overlap 起來，我們需要使用雙 buffer 實現(xiàn)兩階段流水線

總體的實現(xiàn)思路如下

overlap

雙 buffer 實現(xiàn)兩階段流水線

通過上面的 4 個優(yōu)化，我們成功的把千卡 H800 上參數(shù)更新的時間從 10min 降低到了 2min，基本能滿足內(nèi)部 RL 訓(xùn)練的需求了。

提高 H2D 速度并優(yōu)化推理引擎參數(shù)更新性能

但是從理論上，這個速度應(yīng)該是有提升空間的，在 H800 或者 H20 這類 GPU 機(jī)器之間互聯(lián)的網(wǎng)絡(luò) broadcast 帶寬至少是能達(dá)到 100GiB/s，所以 Kimi K2 1TiB 的權(quán)重文件同步一次的時間理論最優(yōu)應(yīng)該是可以達(dá)到 10s 以內(nèi)的，那這里一定是有一些優(yōu)化空間的。

我們發(fā)現(xiàn)在上面的方案里有幾個問題

• checkpoint-engine 在每次 Broadcast 之前都需要等待單個 rank 做完一次 H2D (Host To Device)，然而機(jī)器上的 H2D 速度基本上只能打到 40-50GiB/s，因此會導(dǎo)致帶寬會被 bound 到單次的 H2D 上，無法發(fā)揮網(wǎng)卡的全部性能

• vLLM 在 update weights 的時候，有一些開銷，具體體現(xiàn)在

• • 每個請求都需要算一遍 dict(self.named\_parameters()) [3] ，這里每次都需要讓 Python 做一些 CPU bound 操作

  • 在 expert 權(quán)重更新的時候 使用.item 會頻繁觸發(fā) GPU -> CPU 的同步 [4] ，會導(dǎo)致 update weights 的速度非常不穩(wěn)定，時快時慢。

因此我們又做了3個優(yōu)化

1. 嘗試把 H2D 和 Broadcast 做 Overlap

2. 緩存dict(self.named_parameters())

3. 在 cpu 上緩存expert_map，避免頻繁 CUDA synchronize

優(yōu)化 2 和 3 實現(xiàn)起來很簡單，優(yōu)化 1 我們想的很完美，認(rèn)為能實現(xiàn)類似下圖一樣的非常完美的 overlap

three-stage-pipeline

但是實測發(fā)現(xiàn)在 H800 和 H20 這種機(jī)型里 Broadcast 和 H2D 會互相搶占 PCIE 帶寬，讓二者都出現(xiàn)了降速，結(jié)果變成了這樣

three-stage-pipeline-pcie

那是否存在能不被單個 PCIE 速度 bound 的方案呢？我們發(fā)現(xiàn)可以讓每個節(jié)點先同時做一遍 H2D，這樣就能達(dá)到比較大的聚合 H2D 帶寬而不是大家都在等一個 rank 做 H2D，到時候需要 broadcast 的時候直接從已經(jīng) H2D 的數(shù)據(jù)里走一次快速的 D2D 就可以把數(shù)據(jù)放到 broadcast 的 buffer 里做 broadcast 了，因為 D2D 的速度非常快，這里基本可以忽略它的開銷。使用這個方案我們就能在 H2D 的時候把所有機(jī)器的 PCIE 用起來，提高整體的吞吐，最終實現(xiàn)的流水線如下圖所示

fixed-two-stage-pipeline

經(jīng)過這些優(yōu)化，我們內(nèi)部測試在千卡 H800 下僅用 20s即可以實現(xiàn) Kimi K2 模型的參數(shù)更新，而且速度穩(wěn)定，基本沒有發(fā)生過某幾次參數(shù)更新很慢的情況。

推理引擎的故障自愈

當(dāng)我們實現(xiàn)了高效的參數(shù)更新邏輯之后，我們發(fā)現(xiàn)在強(qiáng)化學(xué)習(xí)中還存在另一個問題，就是推理引擎會時不時出現(xiàn)一些故障導(dǎo)致 RL 訓(xùn)練 Crash。當(dāng)然這個問題能想到一個簡單的解決方案，即一個推理引擎故障了，我們不要讓整個任務(wù)掛掉，直接重啟它即可。但是需要注意的是，在 RL 流程中，我們是直接通過不落盤的方式把 weights 從訓(xùn)練傳給推理的，這個時候如果推理想重啟，看起來只能從 train 的 checkpoint 中重新 convert 一份重啟，這樣會有比較長的 IO 等待，會導(dǎo)致重啟速度比較慢。最好的方式是希望通過 checkpoint-engine 實現(xiàn)重啟的推理實例在線更新 weights。

在當(dāng)時的設(shè)計里我們其實是無法實現(xiàn)這個功能的，因為所有的 inference engine 的參數(shù)更新邏輯是同步的，無法實現(xiàn)只更新部分 inference 實例的權(quán)重。如果無腦的觸發(fā)一次所有 inference engine 的參數(shù)更新流程，雖然我們可以嘗試讓那些運行中的實例不要 update weights 只做 broadcast，但是這個會導(dǎo)致它們需要額外申請 GPU 顯存，這對本來對顯存敏感的 RL 任務(wù)而言是不可接受的，因此我們需要一種能直接從運行中實例的 CPU RDMA 把 weights 讀到故障實例 GPU 中的傳輸框架。剛好 mooncake-transfer-engine[5] 能完美滿足這個要求。

因此我們和 Mooncake 的同學(xué)一起把mooncake-transfer-engine集成到了我們的系統(tǒng)里，我們實現(xiàn)了一個簡單的方案，即讓故障機(jī)器的 rank0 把數(shù)據(jù)按照 bucket_size 從 remote CPU RDMA P2P 的讀到 rank0 的 GPU 里，然后 broadcast 給故障實例的其他 ranks，接著再觸發(fā)參數(shù)更新，這樣就能優(yōu)雅的實現(xiàn)僅更新部分實例的參數(shù)，實現(xiàn)高效的故障自愈。這個方案能實現(xiàn) 40s 更新故障實例的權(quán)重，這個速度對于單點故障恢復(fù)而言完全夠用了。

recover

推理引擎故障自愈

推理啟動加速

在我們內(nèi)部的非 RL 場景通常也會有需求啟動一批推理服務(wù)，針對這個場景我們內(nèi)部已經(jīng)做了一些優(yōu)化，能把權(quán)重先預(yù)熱到/dev/shm里讓推理引擎讀取，這樣能比直接從分布式文件系統(tǒng)讀快很多。但是代價是會占用較多的內(nèi)存空間，同時等待預(yù)熱也需要時間。

但當(dāng)我們把 Kimi K2 參數(shù)同步的開銷優(yōu)化到 20s 這個量級之后，我們發(fā)現(xiàn)這個速度比推理引擎直接從磁盤甚至/dev/shm中讀 weights 都快了很多。而且我們發(fā)現(xiàn) checkpoint-engine 從磁盤里注冊 checkpoint 的操作可以和 vLLM 啟動完全 overlap 起來，vLLM 啟動的時候也有一些類似torch.compile和 Capture CUDA Graph 的操作，我們沒必要串行的讓 vLLM 讀完 weights 在做這些操作。

于是我們讓 vLLM 先 dummy 的啟動，同時開一個 sidecar 啟動 checkpoint engine 注冊 checkpoint。等待 vLLM 就緒了之后直接觸發(fā)所有實例 update weights。實踐下來我們可以做到在和 vLLM dummy 啟動比較接近的時間內(nèi)啟動所有 vLLM 實例，極大的提高了啟動速度。我們內(nèi)部相當(dāng)一部分的推理服務(wù)已經(jīng)在使用這個功能，極大的提高了用戶的使用體驗。

開源

在之后的兩個月里，這套高性能參數(shù)更新的方案穩(wěn)定的支撐了我們的 RL 訓(xùn)練，我們逐漸意識到它具有較好的可擴(kuò)展性和靈活性，因此萌生了把它開源給社區(qū)使用的想法。但是我們內(nèi)部的 checkpoint-engine 有兩層，一層耦合了一些我們的 RL 業(yè)務(wù)邏輯，承擔(dān)了 Convert checkpoint、管理 vLLM 生命周期、自動故障恢復(fù)等邏輯，另一層是核心的參數(shù)更新邏輯，即ParameterServer。我們希望能把ParameterServer解耦出來，提供方便靈活的接口，讓大家能更好更快的用上我們的優(yōu)化。另一方面我們也希望和 vLLM 社區(qū)討論一個高性能的 update weights 接口。

于是我們將我們內(nèi)部 vLLM update weights 的思路 提給了 vLLM 官方[6] ，在和 @游凱超[7] 的討論中也凱超也給了我們一些 idea，最終讓我們敲定了一個比較優(yōu)雅的接口，將和 vLLM 的控制面交互從 HTTP 請求改成了 zmq[8] 隊列，最終 merge 到了 vLLM 官方的 examples[9] 里。

最后我們把ParameterServer單獨拆出來，開源了 checkpoint-engine[10] ，在實現(xiàn)高效參數(shù)更新的情況下也提供了靈活好用的接口，測試下來在各大模型的參數(shù)更新速度都比較高效

Model

設(shè)備信息

收集元數(shù)據(jù)

更新（廣播式）

更新（點對點）

GLM-4.5-Air（BF16）

8xH800 TP8

0.17秒

3.94秒（1.42吉字節(jié)）

8.83秒（4.77吉字節(jié)）

Qwen3-235B-A22B-Instruct-2507 (BF16)

8xH800 TP8

0.46秒

6.75秒（2.69吉字節(jié)）

16.47秒（4.05吉字節(jié)）

深度求索-V3.1（FP8）

16 倍 H20 TP16

1.44秒

12.22秒（2.38吉字節(jié)）

25.77秒（3.61吉字節(jié)）

Kimi-K2-Instruct（FP8）

16xH20 TP16

1.81秒

15.45 秒（2.93GiB）

36.24 秒（4.46GiB）

DeepSeek-V3.1（FP8）

256xH20 TP16

1.40秒

13.88秒（2.54吉字節(jié)）

33.30秒（3.86吉字節(jié)）

Kimi-K2-Instruct (FP8)

256xH20 TP16

1.88秒

21.50 秒（2.99GiB）

34.49秒（4.57吉字節(jié)）

開源版的 checkpoint-engine 并沒有和推理框架做耦合，而是可以讓用戶在 update weights 的時候提供一個自定義的req_func來決定要怎么和推理引擎交互，這樣能非常方便的去對接各種不同的推理引擎。推理引擎也可以自己按照需求去決定一些量化的邏輯。

通過我們的接口，大家使用下面的代碼就可以比較容易的實現(xiàn)權(quán)重更新邏輯

ps = ParameterServer(auto_pg=True) ps.register_checkpoint(name, files=files, named_tensors=named_tensors) ps.gather_metas() ps.update(name, req_func)

上面的代碼會幫你托管 NCCL Group 的創(chuàng)建和銷毀，如果你想自己管理 NCCL，也可以不配置auto_pg=True，這樣可以自己來管理 NCCL Group。

另一方面 checkpoint-engine 也支持多種不同的用法

• 支持Fully Broadcast和P2P兩種參數(shù)更新方式，前者是我們最快的實現(xiàn)，能夠一把同時更新上千張卡的推理引擎權(quán)重，后者更靈活，適合用于故障自愈、disaggregation這些場景，能夠動態(tài)的從已有權(quán)重的 checkpoint-engine RDMA 的拉取權(quán)重，無需額外占用運行中推理實例的顯存即可快速實現(xiàn)新實例的參數(shù)更新。

• 支持注冊多個 checkpoint，你可以靈活的在多個 checkpoint 中間切換，只需調(diào)用多次 register_checkpoint 把不同 name 的 checkpoint 注冊到 ps 里，在 update 的時候指定 checkpoint_name 即可使用這個 checkpoint 去更新推理引擎，這樣 checkpoint-engine 就成為了一個參數(shù)服務(wù)。在線和離線的推理引擎都能非常方便的去切換模型的不同版本做測試。

需要注意的是，P2P 目前的實現(xiàn)方案還比較 naive，每次都是rank0串行的讀取權(quán)重再廣播給其他 rank，這里其實可以做 overlap，例如讓別的 rank 也同時讀取權(quán)重，實現(xiàn)類似 Fixed two-stage pipeline 的效果，這塊未來我們也會持續(xù)優(yōu)化，也歡迎社區(qū)的朋友們感興趣來一起實現(xiàn)。

最后感謝 vLLM 社區(qū) 游凱超[11] 提供了一個科學(xué)的推理引擎接口。本文的一些相關(guān)技術(shù)方案也能在 Kimi K2 Technical Report[12] 里找到。希望大家來試用一下 checkpoint-engine[13] ，如果發(fā)現(xiàn)有什么問題，求提 Issue or PR，期望我們能和社區(qū)一起持續(xù)迭代優(yōu)化，提高參數(shù)更新的速度和體驗！

參考資料

Kimi K2: https://github.com/MoonshotAI/Kimi-K2

k1.5: https://arxiv.org/abs/2501.12599

每個請求都需要算一遍 dict(self.named_parameters()): https://github.com/vllm-project/vllm/blob/v0.10.2rc1/vllm/model_executor/models/deepseek_v2.py

使用.item 會頻繁觸發(fā) GPU -> CPU 的同步: https://github.com/vllm-project/vllm/blob/v0.10.2rc1/vllm/model_executor/layers/fused_moe/layer.py

[5]

mooncake-transfer-engine: https://github.com/kvcache-ai/Mooncake

[6]

提給了 vLLM 官方: https://github.com/vllm-project/vllm/issues/24163

[7]

@游凱超: https://www.zhihu.com/people/176cf88046a1cae595b55e12d58c95e9

[8]

zmq: https://zeromq.org/

[9]

vLLM 官方的 examples: https://github.com/vllm-project/vllm/pull/24295

[10]

checkpoint-engine: https://github.com/MoonshotAI/checkpoint-engine

[11]

游凱超: https://github.com/youkaichao

[12]

Kimi K2 Technical Report: https://arxiv.org/abs/2507.20534

[13]

checkpoint-engine: https://github.com/MoonshotAI/checkpoint-engine