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 버퍼를 통한 워커 간 브로드캐스트, 추론 엔진에 의한 재로드의 3단계 프로세스를 파이프라인화하여 통신과 복사를 겹치게 합니다.
  • P2P: 동적 스케일링이나 재시작 시 사용됩니다. mooncake-transfer-engine을 활용해 기존 인스턴스의 CPU에서 새로운 인스턴스의 GPU로 가중치를 전송하고, 버킷 할당을 최적화해 네트워크 대역폭을 최대화합니다.

엔진은 메타데이터를 수집하고, 버킷 크기를 계획하며, ZeroMQ 소켓을 통해 추론 엔진을 제어함으로써 이러한 전송을 오케스트레이션합니다.

Who it’s for

1조 파라미터 Kimi‑K2와 같은 대규모 LLM을 다루는 개발자 및 연구자를 위해 구축되었으며, 수천 대의 GPU에 걸쳐 빠른 가중치 업데이트가 필요한 경우, 특히 vLLM이나 SGLang을 추론 백엔드로 사용하는 경우에 적합합니다.

Highlights

  • Extreme Scale: 약 20초 안에 수천 대의 GPU에 걸쳐 1조 파라미터 모델을 업데이트할 수 있습니다.
  • Flexible Update Modes: 고속 동기식 브로드캐스트와 유연한 피어‑투‑피어 전송을 모두 지원해 인스턴스 가용성 변동에 대응합니다.
  • Pipelined Execution: 데이터 전송과 메모리 복사를 겹쳐 GPU 처리량을 극대화합니다.
  • Framework Integration: vLLM 및 SGLang에 대한 즉시 사용 가능한 지원 및 통합 경로를 제공합니다.