MoonshotAI/checkpoint-engine

What it solves

Checkpoint Engine 은 LLM 추론 엔진의 모델 가중치를 효율적으로 업데이트하도록 설계된 미들웨어입니다. 이는 모델 가중치를 자주 그리고 빠르게 대규모 GPU 클러스터 전체에 적용해야 하며, 서비스 중단 없이 진행해야 하는 강화 학습 워크플로우에서 특히 중요합니다.

How it works

ParameterServer 클래스를 추론 엔진과 같은 위치에 배치하여 두 가지 주요 방법으로 가중치 업데이트를 관리합니다.

• Broadcast: 가장 빠른 방법으로, 다수 인스턴스 간 동기식 업데이트에 사용됩니다. 세 단계 프로세스(Host‑to‑Device 전송, CUDA IPC 버퍼를 통한 워커 간 브로드캐스트, 추론 엔진으로 재로드)를 파이프라인화하여 통신과 복사를 겹치게 합니다.
• P2P: 동적 스케일링에 사용되며, 새로운 인스턴스가 추가되거나 재시작될 때 활용됩니다. mooncake-transfer-engine 을 이용해 기존 인스턴스의 CPU에서 새로운 인스턴스의 GPU로 가중치를 전송하고, 버킷 할당을 최적화해 네트워크 대역폭을 최대화합니다.

Who it’s for

1조 파라미터 규모의 초대형 LLM 및 vLLM, SGLang 같은 분산 추론 프레임워크를 다루는 개발자와 연구자를 대상으로 하며, 고성능 인플레이스 가중치 업데이트가 필요한 경우에 적합합니다.

Highlights

• Extreme Scale: 약 20초 만에 수천 대의 GPU에서 1조 파라미터 모델을 업데이트할 수 있습니다.
• Flexible Update Paths: 동기식 브로드캐스트와 비동기 P2P 전송을 모두 지원해 동적 가용성에 대응합니다.
• Hardware Support: NVIDIA GPU(CUDA)와 Intel XPU(SYCL IPC 메모리)를 지원합니다.
• Framework Integration: vLLM 및 SGLang에 통합되어 모델 로드 및 업데이트 시간을 단축합니다.