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 サービスとして動作し、主に以下の 2 つの更新方式を提供します。

  • Broadcast: 多数のインスタンス間で同期的に更新を行う最速の方法です。3 段階のプロセス（Host‑to‑Device 転送、CUDA IPC バッファを介したワーカー間ブロードキャスト、推論エンジンによるリロード）をパイプライン化し、通信とコピーを重ね合わせます。
  • P2P: 動的スケーリングや再起動時に使用されます。mooncake-transfer-engine を利用して既存インスタンスの CPU から新インスタンスの GPU へ重みを転送し、バケット割り当てを最適化してネットワーク帯域幅を最大化します。

エンジンはメタデータを収集し、バケットサイズを計画し、ZeroMQ ソケットを介して推論エンジンを制御することで、これらの転送をオーケストレーションします。

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 向けの即時利用可能なサポートと統合パスを提供。