高效更新LLM权重：深入探究MoonshotAI的Checkpoint引擎

想象一下，试图翻新一座摩天大楼的地基，却不要求其内部人员离开或暂停工作。这正是 MoonshotAI 的 Checkpoint 引擎为 AI 模型所做的。它允许大规模语言模型在运行时更新其“大脑”——权重，因此无需停机。这项突破性技术让开发者能够快速高效地改进他们的 AI 模型，即使是在数千个 GPU 上运行超过一万亿个参数的模型也是如此。它快速、可靠，旨在确保 AI 系统在实时演进的同时保持平稳运行，使其成为尖端 AI 应用的重要工具。本文将介绍 Checkpoint Engine 的概念、工作原理以及它对大规模 AI 系统未来的重要性。

什么是Moonshot AI的Checkpoint Engine？

Moonshot AI 的 Checkpoint Engine 是一款专用中间件，旨在在推理过程中实时更新大型语言模型 (LLM) 的权重，而不会中断正在进行的操作。此功能在需要频繁更新模型权重的强化学习场景中至关重要。 Checkpoint Engine 目前已与 vLLM 推理框架无缝集成，并通过流水线和内存管理技术提供优化的性能。它还提供诸如重用现有实例的权重等功能，以减少扩展场景中的开销。

Source: Github

架构

Checkpoint 的核心是 ParameterServer 类，它负责处理权重更新逻辑并协调数据流。

1. H2D（主机到设备）：使用优化的传输流水线，将更新后的权重从 CPU 内存或存储移动到 GPU 内存。
2. 广播：利用 CUDA IPC 缓冲区进行共享内存通信，高效地将权重分配到所有推理引擎实例。
3. 重新加载：每个推理引擎随后根据其分片模式，选择性地从广播数据中重新加载相关的权重分片。

这三阶段流水线确保了高效的重叠通信和复制，从而提高了速度。

当 GPU 内存有限时，系统可以回退到串行执行以保持可靠性。

使用的方法

Checkpoint 引擎在推理过程中使用两种主要方法来更新模型权重。

1. 广播方法：这是最快的默认方法。当需要同时更新大量推理实例时，这是理想的选择。它将更新后的权重从 CPU 内存同步广播到所有推理 GPU，确保所有实例保持完美同步，并将延迟降至最低。
2. P2P（点对点）方法：用于在运行时动态添加或删除推理实例的情况。它通过点对点传输系统将权重直接从现有实例的 CPU 发送到新实例的 GPU，从而避免干扰现有的推理工作负载，从而实现流畅灵活的更新。

工作原理

检查点引擎负责协调整个传输过程。它首先收集必要的元数据来制定计划，包括确定数据传输的合适存储桶大小。然后，它执行传输，并通过 ZeroMQ 套接字控制推理引擎以最大化性能。它将数据传输组织到具有重叠通信和副本的管道中，即使在高负载下也能实现快速高效的权重更新。

Source: Github

通过实现上述方法和架构，检查点引擎能够跨数千个 GPU 实时更新 LLM 权重，同时最大程度地降低延迟和服务中断。

安装和使用

安装

要使用最快的广播，

使用代码：

pip install checkpoint-engine

要使用灵活的 P2P 实现：

使用代码：

pip install 'checkpoint-engine[p2p]'

这将安装 mooncake-transfer-engine，以支持不同等级之间的 RDMA 传输。

示例用例

步骤 1：

准备一台配备 8 个 GPU 的 H800 或 H20 机器，并安装最新的 vLLM。请务必包含 /collective_rpc API 端点提交（可在主分支中找到），因为 checkpoint-engine 将使用此端点来更新权重。

步骤 2：

install checkpoint-engine

代码：

uv pip install 'checkpoint-engine[p2p]'

步骤 3：

在我们的用例中，我们将使用 Qwen/Qwen3-235B-A22B-Instruct-2507 作为测试模型。

代码：

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

步骤 4：

在开发模式下启动 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

使用检查点引擎更新权重。无需等待 vLLM 准备就绪。使用以下代码。

代码：

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

重用现有实例的权重

新的检查点引擎实例可以加入现有实例并重用其权重。

使用以下方法：

步骤 1：使用 –save-metas-file global_metas.pkl 启动现有实例，将全局元数据保存到文件中。

步骤 2：使用 –sleep-time 300 确保它们保持活动状态。

代码：

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

步骤 3：注册检查点后，新实例可以通过设置 –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 中测试过。因此，与其他模型可能会存在一些兼容性问题。

测试

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

代码

torchrun --nproc-per-node 8 tests/test_update.py

基准测试

见解

以下是我的一些观察结果：

1. 广播方法通常提供最快的更新时间，并针对跨多个推理实例的同步权重更新进行了优化。
2. P2P 方法耗时较长，但可在运行时实例加入或离开时进行动态更新。
3. 这些基准测试展示了 Checkpoint Engine 的可扩展性，能够在 8 到 256 个 GPU 的集群上高效处理数万亿个参数的模型。

Checkpoint Engine的局限性

虽然 Checkpoint Engine 是 LLM 中实时更新权重的强大解决方案，但目前仍存在一些局限性。

• 目前最适合与 vLLM 配合使用：该引擎主要使用 vLLM 框架进行测试。如果您希望将其与其他 AI 框架或自定义设置一起使用，则可能需要进行一些额外的工作才能使其平稳运行。
• 管道仍在改进：理想的无缝管道（能够完美重叠数据移动）尚未完全完成。这意味着仍有潜力进一步加快更新速度。
• P2P 更新可以更顺畅：点对点方法会先将数据发送到一个主节点的瓶颈处，然后再与其他节点共享，这在您拥有大量 GPU 时可能会降低速度。
• 需要额外的 GPU 内存：智能广播系统会使用更多 GPU 内存来加速运行。在内存较少的机器上，它会切换到速度更慢、效率更低的进程。
• 对 FP8 模型的支持有限：如果您使用的是较新的 FP8 量化模型，则需要一些实验性补丁。即便如此，并非所有模型都能良好运行，除了一些经过测试的模型。

小结

Moonshot AI 的 Checkpoint Engine 是一款颠覆性的工具，可以在不停止大型 AI 模型运行的情况下进行更新。即使模型的“大脑”实时变得更加智能，它也能保持一切平稳运行。虽然它仍有一些需要改进的地方，但潜力巨大。如果您正在使用大型 AI 系统，这款工具绝对值得关注。它正在帮助 AI 的未来更快、更高效，并且不会造成任何停机。