MoonshotAI/checkpoint-engine

Checkpoint-engine is a simple middleware to update model weights in LLM inference engines

What it solves

Checkpoint Engine 是一個中介層，旨在解決 LLM 推論引擎中模型權重更新緩慢的問題。這在需要頻繁且高效地在大型 GPU 叢集上更新權重的強化學習工作流程中尤為關鍵，且不會造成顯著的停機時間或效能退化。

How it works

它作為與推論引擎同址的 ParameterServer 服務運行，使用兩種主要的更新方式：

  • Broadcast：在多個實例之間同步更新的最快方法。它使用三階段流程（Host-to-Device 傳輸、透過 CUDA IPC 緩衝區在工作者之間廣播、以及由推論引擎重新載入），並以管線化方式重疊通信與複製。
  • P2P：用於動態擴縮或重啟。它利用 mooncake-transfer-engine 從現有實例的 CPU 將權重傳送至新實例的 GPU，並優化 bucket 分配以最大化網路頻寬。

引擎透過收集中繼資料、規劃 bucket 大小，並透過 ZeroMQ socket 控制推論引擎，來協調這些傳輸。

Who it’s for

它是為了使用大規模 LLM（例如 1 兆參數的 Kimi‑K2）的開發者與研究者而建，這些使用者需要在數千顆 GPU 上快速更新權重，特別是使用 vLLM 或 SGLang 作為推論後端的情境。

Highlights

  • Extreme Scale：能在約 20 秒內於數千顆 GPU 上更新 1 兆參數模型。
  • Flexible Update Modes：同時支援高速同步廣播與彈性點對點傳輸，以因應動態實例可用性。
  • Pipelined Execution：重疊資料傳輸與記憶體複製，以最大化 GPU 吞吐量。
  • Framework Integration：提供即插即用的支援與 vLLM、SGLang 的整合路徑。