Checkpoint Engine

Checkpoint-engine is a simple middleware to update model weights in LLM inference engines -- a critical step in reinforcement learning. We provide an efficient and lightweight implementation for inplace weight update: updating our Kimi-K2 model (1 Trillion parameters) across thousands of GPUs takes about 20s.

  <img src="https://github.com/MoonshotAI/checkpoint-engine/raw/v0.4.2/figures/checkpoint-engine.png" width="80%" alt="ckpt-engine">

Architecture

The core weight update logic is in ParameterServer class, a service colocated with inference engines. It provides two implementations of weight update: Broadcast and P2P.

• Broadcast: Used when a large number of inference instances need to update weights in synchronous. This is the fastest implementation and should be used as the default update method. See _update_per_bucket with ranks == None or [].
• P2P: Used when new inference instances are dynamically added (due to restarts or dynamic availability) while the existing instances are already serving requests. Under this scenario, to avoid affecting the workloads on existing instances, we use the mooncake-transfer-engine to P2P send weights from CPUs in existing instances to GPUs in new instances. See _update_per_bucket with ranks specified.

Optimized Weight Broadcast

In the Broadcast implementation, the checkpoint-engine holds references to sharded weights in CPU memory, and need to efficiently broadcast them to a cluster of inference instances, often under a different sharding pattern. We arrange the data transfer into 3 stages: 1. H2D: moving weights to GPU memory. These weights may come from disk or the training engine. 2. broadcast: broadcast among checkpoint engine workers; the data results in a CUDA IPC buffer shared with inference engine. 3. reload: inference engine decides what subset of weights to copy from the broadcasted data.

Checkpoint-engine orchestrates the entire transfer process. It first gathers necessary metadata to create a plan, including deciding the proper bucket size for data transfer. It then executes the transfer, where it controls the inference engine through a ZeroMQ socket. To maximize performance, it organizes the data transfers into a pipeline with overlapped communication and copy, illustrated below. The details can be found in Kimi-K2 Technical Report.

  <img src="https://github.com/MoonshotAI/checkpoint-engine/raw/v0.4.2/figures/pipeline.png" width="80%" alt="pipeline">

Pipelining naturally requires more GPU memory. When memory is not enough, checkpoint-engine will fallback to serial execution.

Optimized P2P Bucket Assignment

In the P2P implementation, checkpoint-engine needs to send weights from existing instances to new instances. To minimize the overall transfer time, checkpoint-engine optimizes the bucket assignment for each sender-receiver pair. The optimization goal is to make full use of the available network bandwidth for each sender and receiver. See issue #25

Benchmark

Model	Device Info	GatherMetas	Update (Broadcast)	Update (P2P)
GLM-4.5-Air (BF16)	8xH800 TP8	0.12s	3.47s (3.02GiB)	4.12s (3.02GiB)
Qwen3-235B-A22B-Instruct-2507 (BF16)	8xH800 TP8	0.33s	6.22s (2.67GiB)	7.10s (2.68GiB)
DeepSeek-V3.1 (FP8)	16xH20 TP16	1.17s	10.19s (5.39GiB)	11.80s (5.41GiB)
Kimi-K2-Instruct (FP8)	16xH20 TP16	1.33s	14.36s (5.89GiB)	17.49s (5.91GiB)
DeepSeek-V3.1 (FP8)	256xH20 TP16	0.80s	11.33s (8.00GiB)	11.81s (8.00GiB)
Kimi-K2-Instruct (FP8)	256xH20 TP16	1.22s	16.04s (8.00GiB)	16.75s (8.00GiB)

All results above are tested by examples/update.py and use vLLM v0.10.2rc1 as inference engine. Some notes:

• FP8 test needs additional vLLM patches, see FP8 quantization.
• Device Info: we tested various combination of devices and parallelism setups. For example, a 256-GPU TP16 setup means that we deploy 16 vLLM instances, each with 16-way tensor parallelism.
• Since update duration is related to IPC bucket size, we provide the bucket size in the table.
• The P2P time were tested for updating no more than two nodes (16 GPUs) (ParameterServer.update(ranks=range(0, 16))) out of the entire cluster.
• We bind each GPU to its corresponding NUMA node to ensure stable H2D transfer speeds.

Installation

Use the fastest broadcast implementation

pip install checkpoint-engine

Use the flexible P2P implementation, notice this will install mooncake-transfer-engine to support RDMA transfer between different ranks.

pip install 'checkpoint-engine[p2p]'

Getting Started

Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include /collective_rpc API endpoint commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version v0.10.2 is fully tested and recommended.

mkdir -p /opt/vLLM && cd /opt/vLLM
uv venv --python 3.12 --seed
source .venv/bin/activate
uv pip install vllm==0.10.2

Install checkpoint-engine

uv pip install 'checkpoint-engine[p2p]'

We use Qwen/Qwen3-235B-A22B-Instruct-2507 (BF16) as the test model

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

Start vLLM in dev mode and set --load-format dummy. Notice that we also set --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

Meanwhile, use this command to update weights by checkpoint-engine. No need to wait for vLLM to get ready.

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

Reuse weights from existing instances

New checkpoint-engine instances can join existing instances and reuse their weights. This is simple to achieve.

First, start the existing instances with --save-metas-file global_metas.pkl to save global metas to a file and use --sleep-time 300 to make sure they stay alive.

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

After a checkpoint is registered, new instances can obtain a copy of the checkpoint by setting --load-metas-file global_metas.pkl.

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

FP8 quantization

FP8 quantization currently do not natively work in vLLM when updating weights. We provide a simple patch in patches/vllm_fp8.patch to handle the correct weight update. Notice this patch is only tested in DeepSeek-V3.1 and Kimi-K2. Other models may meet some compatible issues.

A PR is opened to the vLLM project and waiting to discuss and review.

Test

Run a simple correctness test for checkpoint_engine

pytest tests/test_update.py

test_update.py are only designed to run with pytest. Please don't run it directly with torchrun.

Other unit tests can also be done with pytest. Only test_update.py requires GPUs, other tests can be run on CPUs. Only to run CPU tests, use:

pytest tests/ -m "not gpu"

Environment Variables

• PS_MAX_BUCKET_SIZE_GB: An integer is used to set the maximum bucket size for checkpoint-engine. If not set, 8GB is used as default.
• PS_P2P_STORE_RDMA_DEVICES: Comma-separated RDMA devices' names for P2P transfer. If not set, checkpoint-engine will fall back to use NCCL_IB_HCA to detect RDMA devices.
• NCCL_IB_HCA: Available patterns can be found from NCCL documentation. If also not set, all RDMA devices will be used and divided evenly among the ranks.

SGLang Integration

Checkpoint Engine provides efficient distributed checkpoint loading for SGLang inference servers, significantly reducing model loading time for large models and multi-node setups.

Quick Start

1. Install checkpoint-engine:

pip install 'checkpoint-engine[p2p]'

2. Launch SGLang server:

python -m sglang.launch_server \
    --model-path $MODEL_PATH \
    --tp 8 \
    --load-format dummy \
    --wait-for-initial-weights

3. Run checkpoint engine:

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path $MODEL_PATH \
    --inference-parallel-size 8

Multi-Node Setup

For 2-node setup, run the same commands on both nodes with appropriate --host and distributed training parameters.

Key Options

SGLang Server: - --wait-for-initial-weights: Wait for checkpoint engine before becoming ready - --load-format dummy: Enable overlapping initialization tasks

Checkpoint Engine: - --update-method: Choose broadcast, p2p, or all - --inference-parallel-size: Number of parallel processes - --checkpoint-path: Model checkpoint directory

Limitations and Future Work

• This project is currently tested with vLLM and SGLang. Integration with other frameworks is planned for future releases.
• The perfect three-stage pipeline mentioned in our paper is currently not implemented. This could be useful for architectures where H2D and broadcast do not conflict in PCIE.

Acknowledgments

This open source project uses the same vLLM interface in https://github.com/vllm-project/vllm/pull/24295 . Thanks for the comments and insights from youkaichao.