高效更新LLM權重：深入探究MoonshotAI的Checkpoint引擎

想象一下，試圖翻新一座摩天大樓的地基，卻不要求其內部人員離開或暫停工作。這正是 MoonshotAI 的 Checkpoint 引擎為 AI 模型所做的。它允許大規模語言模型在執行時更新其“大腦”——權重，因此無需停機。這項突破性技術讓開發者能夠快速高效地改進他們的 AI 模型，即使是在數千個 GPU 上執行超過一萬億個引數的模型也是如此。它快速、可靠，旨在確保 AI 系統在即時演進的同時保持平穩執行，使其成為尖端 AI 應用的重要工具。本文將介紹 Checkpoint Engine 的概念、工作原理以及它對大規模 AI 系統未來的重要性。

什麼是Moonshot AI的Checkpoint Engine？

Moonshot AI 的 Checkpoint Engine 是一款專用中介軟體，旨在在推理過程中即時更新大型語言模型 (LLM) 的權重，而不會中斷正在進行的操作。此功能在需要頻繁更新模型權重的強化學習場景中至關重要。 Checkpoint Engine 目前已與 vLLM 推理框架無縫整合，並透過流水線和記憶體管理技術提供最佳化的效能。它還提供諸如重用現有例項的權重等功能，以減少擴充套件場景中的開銷。

Source: Github

架構

Checkpoint 的核心是 ParameterServer 類，它負責處理權重更新邏輯並協調資料流。

1. H2D（主機到裝置）：使用最佳化的傳輸流水線，將更新後的權重從 CPU 記憶體或儲存移動到 GPU 記憶體。
2. 廣播：利用 CUDA IPC 緩衝區進行共享記憶體通訊，高效地將權重分配到所有推理引擎例項。
3. 重新載入：每個推理引擎隨後根據其分片模式，選擇性地從廣播資料中重新載入相關的權重分片。

這三階段流水線確保了高效的重疊通訊和複製，從而提高了速度。

當 GPU 記憶體有限時，系統可以回退到序列執行以保持可靠性。

使用的方法

Checkpoint 引擎在推理過程中使用兩種主要方法來更新模型權重。

1. 廣播方法：這是最快的預設方法。當需要同時更新大量推理例項時，這是理想的選擇。它將更新後的權重從 CPU 記憶體同步廣播到所有推理 GPU，確保所有例項保持完美同步，並將延遲降至最低。
2. P2P（點對點）方法：用於在執行時動態新增或刪除推理例項的情況。它透過點對點傳輸系統將權重直接從現有例項的 CPU 傳送到新例項的 GPU，從而避免干擾現有的推理工作負載，從而實現流暢靈活的更新。

工作原理

檢查點引擎負責協調整個傳輸過程。它首先收集必要的後設資料來制定計劃，包括確定資料傳輸的合適儲存桶大小。然後，它執行傳輸，並透過 ZeroMQ 套接字控制推理引擎以最大化效能。它將資料傳輸組織到具有重疊通訊和副本的管道中，即使在高負載下也能實現快速高效的權重更新。

Source: Github

透過實現上述方法和架構，檢查點引擎能夠跨數千個 GPU 即時更新 LLM 權重，同時最大程度地降低延遲和服務中斷。

安裝和使用

安裝

要使用最快的廣播，

使用程式碼：

pip install checkpoint-engine

要使用靈活的 P2P 實現：

使用程式碼：

pip install 'checkpoint-engine[p2p]'

這將安裝 mooncake-transfer-engine，以支援不同等級之間的 RDMA 傳輸。

示例用例

步驟 1：

準備一臺配備 8 個 GPU 的 H800 或 H20 機器，並安裝最新的 vLLM。請務必包含 /collective_rpc API 端點提交（可在主分支中找到），因為 checkpoint-engine 將使用此端點來更新權重。

步驟 2：

install checkpoint-engine

程式碼：

uv pip install 'checkpoint-engine[p2p]'

步驟 3：

在我們的用例中，我們將使用 Qwen/Qwen3-235B-A22B-Instruct-2507 作為測試模型。

程式碼：

hf download Qwen/Qwen3-235B-A22B-Instruct-2507 --local-dir /opt/models/Qwen/Qwen3-235B-A22B-Instruct-2507/

步驟 4：

在開發模式下啟動 vLLM，並設定 –load-format dummy。確保設定 –worker-extension-cls=checkpoint_engine.worker.VllmColocateWorkerExtension。

程式碼：

VLLM_SERVER_DEV_MODE=1 python3 -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 19730 --trust-remote-code \
    --tensor-parallel-size=8 --max-model-len 4096 --load-format dummy \
    --served-model-name checkpoint-engine-demo --model /opt/models/Qwen/Qwen3-235B-A22B-Instruct-2507/ \
    --worker-extension-cls checkpoint_engine.worker.VllmColocateWorkerExtension

使用檢查點引擎更新權重。無需等待 vLLM 準備就緒。使用以下程式碼。

程式碼：

torchrun --nproc-per-node 8 examples/update.py --update-method all --checkpoint-path /opt/models/Qwen/Qwen3-235B-A22B-Instruct-2507/

重用現有例項的權重

新的檢查點引擎例項可以加入現有例項並重用其權重。

使用以下方法：

步驟 1：使用 –save-metas-file global_metas.pkl 啟動現有例項，將全域性後設資料儲存到檔案中。

步驟 2：使用 –sleep-time 300 確保它們保持活動狀態。

程式碼：

torchrun --nproc-per-node 8 examples/update.py --checkpoint-path $MODEL_PATH \
    --sleep-time 300 --save-metas-file global_metas.pkl

步驟 3：註冊檢查點後，新例項可以透過設定 –load-metas-file global_metas.pkl 獲取檢查點的副本。

程式碼：

torchrun --nproc-per-node 8 examples/update.py --load-metas-file global_metas.pkl

FP8量化

目前，FP8 量化在 vLLM 中更新權重時不起作用。它使用 patches/vllm_fp8.patch 中的一個簡單補丁來處理正確的權重更新。此外，此補丁僅在 DeepSeek-V3.1 和 Kimi-K2 中測試過。因此，與其他模型可能會存在一些相容性問題。

測試

對 checkpoint_engine 執行一個簡單的正確性測試

程式碼

torchrun --nproc-per-node 8 tests/test_update.py

基準測試

見解

以下是我的一些觀察結果：

1. 廣播方法通常提供最快的更新時間，並針對跨多個推理例項的同步權重更新進行了最佳化。
2. P2P 方法耗時較長，但可在執行時例項加入或離開時進行動態更新。
3. 這些基準測試展示了 Checkpoint Engine 的可擴充套件性，能夠在 8 到 256 個 GPU 的叢集上高效處理數萬億個引數的模型。

Checkpoint Engine的侷限性

雖然 Checkpoint Engine 是 LLM 中即時更新權重的強大解決方案，但目前仍存在一些侷限性。

• 目前最適合與 vLLM 配合使用：該引擎主要使用 vLLM 框架進行測試。如果您希望將其與其他 AI 框架或自定義設定一起使用，則可能需要進行一些額外的工作才能使其平穩執行。
• 管道仍在改進：理想的無縫管道（能夠完美重疊資料移動）尚未完全完成。這意味著仍有潛力進一步加快更新速度。
• P2P 更新可以更順暢：點對點方法會先將資料傳送到一個主節點的瓶頸處，然後再與其他節點共享，這在您擁有大量 GPU 時可能會降低速度。
• 需要額外的 GPU 記憶體：智慧廣播系統會使用更多 GPU 記憶體來加速執行。在記憶體較少的機器上，它會切換到速度更慢、效率更低的程序。
• 對 FP8 模型的支援有限：如果您使用的是較新的 FP8 量化模型，則需要一些實驗性補丁。即便如此，並非所有模型都能良好執行，除了一些經過測試的模型。

小結

Moonshot AI 的 Checkpoint Engine 是一款顛覆性的工具，可以在不停止大型 AI 模型執行的情況下進行更新。即使模型的“大腦”即時變得更加智慧，它也能保持一切平穩執行。雖然它仍有一些需要改進的地方，但潛力巨大。如果您正在使用大型 AI 系統，這款工具絕對值得關注。它正在幫助 AI 的未來更快、更高效，並且不會造成任何停機。