检查点引擎

Checkpoint-engine 是一种简单的中间件，用于在 LLM 推理引擎中更新模型权重——这是强化学习中的关键步骤。 我们提供了一种高效轻量级的原地权重更新实现： 在数千张 GPU 上更新我们的 Kimi-K2 模型（1 万亿参数）大约需要 20 秒。

架构

核心权重更新逻辑位于 ParameterServer 类中，该服务与推理引擎部署在一起。它提供了两种权重更新实现：广播和 P2P。

• 广播：当大量推理实例需要同步更新权重时使用。这是最快的实现方式，应作为默认的更新方法。参见 _update_per_bucket 中 ranks == None or [] 的情况。
• P2P：当新的推理实例动态加入（由于重启或动态可用性）而现有实例已经在处理请求时使用。在这种情况下，为避免影响现有实例的工作负载，我们使用 mooncake-transfer-engine 将现有实例 CPU 上的权重通过 P2P 方式传输到新实例的 GPU 上。参见 _update_per_bucket 中指定了 ranks 的情况。

优化的权重广播

在 广播 实现中，checkpoint-engine 会持有 CPU 内存中分片权重的引用，并需要将它们高效地广播到一组推理实例上，这些实例通常采用不同的分片模式。 我们将数据传输分为三个阶段：

1. H2D：将权重移动到 GPU 内存。这些权重可能来自磁盘或训练引擎。
2. 广播：在 checkpoint-engine 工作节点之间进行广播；数据结果是一个与推理引擎共享的 CUDA IPC 缓冲区。
3. 重载：推理引擎决定从广播的数据中复制哪些权重子集。 Checkpoint-engine 协调整个传输过程。它首先收集必要的元数据以制定计划，包括确定合适的数据传输桶大小。 然后执行传输，期间通过 ZeroMQ 套接字控制推理引擎。为了最大化性能，它将数据传输组织成一个通信与拷贝重叠的流水线，如下图所示。详细信息可在 Kimi-K2 技术报告 中找到。

流水线操作自然需要更多的 GPU 内存。当内存不足时，checkpoint-engine 会回退到串行执行。

优化的 P2P 分桶分配

在 P2P 实现中，checkpoint-engine 需要将权重从现有实例发送到新实例。 为了最小化整体传输时间，checkpoint-engine 会为每对发送方和接收方优化分桶分配。 优化目标是充分利用每个发送方和接收方的可用网络带宽。 详情请参阅 issue #25。

基准测试

模型	设备信息	GatherMetas	更新（广播）	更新（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)

以上所有结果均由 examples/update.py 测试，并使用 vLLM v0.10.2rc1 作为推理引擎。几点说明：

• FP8 测试需要额外的 vLLM 补丁，详见 FP8 量化。
• 设备信息：我们测试了多种设备和并行化设置的组合。例如，256-GPU TP16 设置意味着我们部署了 16 个 vLLM 实例，每个实例都采用 16 路张量并行。
• 由于更新时长与 IPC 桶大小有关，我们在表中提供了桶大小。
• P2P 时间是在不超过两个节点（16 张 GPU）(ParameterServer.update(ranks=range(0, 16))) 的范围内测试的。
• 我们将每张 GPU 绑定到其对应的 NUMA 节点，以确保稳定的 H2D 传输速度。

安装

使用最快的广播实现

pip install checkpoint-engine

使用灵活的 P2P 实现，请注意这将安装 mooncake-transfer-engine 以支持不同 Rank 之间的 RDMA 传输。

pip install 'checkpoint-engine[p2p]'

入门指南

准备一台配备 8 张 GPU 和 vLLM 的 H800 或 H20 机器。务必包含 /collective_rpc API 端点 提交（已在主分支中可用），因为 checkpoint-engine 将使用此端点来更新权重。vLLM 版本 v0.10.2 已经过全面测试并推荐使用。

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

安装 checkpoint-engine

uv pip install 'checkpoint-engine[p2p]'

我们使用 Qwen/Qwen3-235B-A22B-Instruct-2507 (BF16) 作为测试模型

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

以开发模式启动 vLLM，并设置 --load-format dummy。请注意，我们还设置了 --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

与此同时，使用以下命令由 checkpoint-engine 更新权重。无需等待 vLLM 准备就绪。

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

重用现有实例的权重

新的 checkpoint-engine 实例可以加入现有的实例，并重用它们的权重。实现起来非常简单。

首先，使用 --save-metas-file global_metas.pkl 启动现有实例，将全局元数据保存到文件中，并通过 --sleep-time 300 确保它们保持运行状态。

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

在检查点注册完成后，新实例可以通过设置 --load-metas-file global_metas.pkl 来获取检查点的副本。

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

FP8 量化

目前，在更新权重时，FP8 量化在 vLLM 中无法原生工作。我们提供了一个简单的补丁 patches/vllm_fp8.patch，用于正确处理权重更新。请注意，此补丁仅在 DeepSeek-V3.1 和 Kimi-K2 上进行了测试，其他模型可能会遇到兼容性问题。

我们已向 vLLM 项目提交了一个 PR，目前正在等待讨论和审查。

测试

运行一个简单的 checkpoint_engine 正确性测试：

pytest tests/test_update.py

test_update.py 仅设计为与 pytest 一起运行，请勿直接使用 torchrun 运行。

其他单元测试也可以使用 pytest 完成。只有 test_update.py 需要 GPU，其他测试可以在 CPU 上运行。如果仅需运行 CPU 测试，可使用以下命令：

pytest tests/ -m "not gpu"

环境变量

• PS_MAX_BUCKET_SIZE_GB：用于设置 checkpoint-engine 的最大分桶大小的整数。若未设置，则默认为 8GB。
• PS_P2P_STORE_RDMA_DEVICES：用于 P2P 传输的 RDMA 设备名称，以逗号分隔。若未设置，checkpoint-engine 将回退到使用 NCCL_IB_HCA 自动检测 RDMA 设备。
• NCCL_IB_HCA：可用模式可参考 NCCL 文档。若也未设置，则会使用所有 RDMA 设备，并均匀分配给各个进程。

SGLang 集成

Checkpoint Engine 为 SGLang 推理服务器提供了高效的分布式检查点加载功能，显著减少了大型模型和多节点部署下的模型加载时间。

快速开始

1. 安装 checkpoint-engine：

pip install 'checkpoint-engine[p2p]'

2. 启动 SGLang 服务器：

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

3. 运行 checkpoint engine：

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

多节点设置

对于 2 节点设置，可在两个节点上分别运行相同的命令，并设置适当的 --host 和分布式训练参数。

关键选项

SGLang 服务器：

• --wait-for-initial-weights：等待 checkpoint engine 完成后再进入就绪状态。
• --load-format dummy：启用初始化任务的重叠执行。

Checkpoint Engine：

• --update-method：可选择 broadcast、p2p 或 all。
• --inference-parallel-size：并行进程数。
• --checkpoint-path：模型检查点目录。

局限性与未来工作

• 本项目目前仅在 vLLM 和 SGLang 上进行了测试。与其他框架的集成计划将在未来的版本中实现。
• 我们论文中提到的完美三阶段流水线目前尚未实现。这可能对那些 H2D 和广播操作在 PCIE 总线上不会产生冲突的架构有所帮助。

致谢

本开源项目使用了与 https://github.com/vllm-project/vllm/pull/24295 相同的 vLLM 接口。感谢 youkaichao 提供的评论和见解。

Checkpoint Engine 快速上手指南

Checkpoint Engine 是一款专为大语言模型（LLM）推理引擎设计的轻量级中间件，用于在强化学习等场景中高效地原地更新模型权重。它支持广播（Broadcast）和点对点（P2P）两种更新模式，能在数千张 GPU 上于秒级完成万亿参数模型的权重更新。

环境准备

系统要求

• GPU: 推荐 NVIDIA H800 或 H20（测试基于 8 卡或多节点集群）。
• Python: 3.12+。
• 推理框架:
  • vLLM: 推荐使用 v0.10.2 版本（需包含 collective_rpc API 支持）。
  • SGLang: 支持分布式检查点加载。
• 网络: 若使用 P2P 模式，建议配置 RDMA 网络以获得最佳性能。

前置依赖

确保已安装基础深度学习环境（如 PyTorch, CUDA）。若使用国内网络，建议在 pip 命令中指定清华或阿里镜像源加速下载。

安装步骤

根据需求选择安装模式：

1. 基础模式（仅使用最快的广播更新）

pip install checkpoint-engine -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 完整模式（支持灵活的 P2P 更新，含 RDMA 传输支持） 推荐安装此版本以兼容动态节点加入场景

pip install 'checkpoint-engine[p2p]' -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下以 vLLM 为例，演示如何启动推理服务并实时更新权重。

第一步：准备模型与环境

创建虚拟环境并安装指定版本的 vLLM：

mkdir -p /opt/vLLM && cd /opt/vLLM
uv venv --python 3.12 --seed
source .venv/bin/activate
uv pip install vllm==0.10.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
uv pip install 'checkpoint-engine[p2p]' -i https://pypi.tuna.tsinghua.edu.cn/simple

# 下载测试模型 (以 Qwen3 为例)
hf download Qwen/Qwen3-235B-A22B-Instruct-2507 --local-dir /opt/models/Qwen/Qwen3-235B-A22B-Instruct-2507/

第二步：启动 vLLM 推理服务

启动时需设置 --load-format dummy 并加载 Checkpoint Engine 的 Worker 扩展类，以便接收权重更新指令：

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

第三步：执行权重更新

在另一个终端运行更新脚本。无需等待 vLLM 完全就绪，该工具会自动协调权重传输：

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

---

进阶：SGLang 集成简述

若使用 SGLang 框架，流程更为简化：

1. 启动 SGLang 服务：

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

2. 运行更新引擎：

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

常用环境变量

• PS_MAX_BUCKET_SIZE_GB: 设置最大数据桶大小（默认 8GB），可根据显存调整。
• NCCL_IB_HCA: 指定 RDMA 网卡设备，优化多节点传输性能。