Moonshot AI 开源 Checkpoint-Engine：千亿参数模型权重更新仅需20秒

Moonshot AI 近日在 GitHub 上公开了 checkpoint-engine 项目，定位为大语言模型推理引擎的权重更新中间件。在强化学习训练流程中，训练引擎完成参数更新后需要将新权重同步至推理引擎，这一步骤往往成为整体吞吐量的瓶颈。checkpoint-engine 正是为解决这一问题而生，其核心目标是以最低延迟、最小资源占用完成跨 GPU 集群的权重同步。

该工具的架构核心是 ParameterServer 类，以服务形式与推理引擎共同部署。它在 CPU 内存中持有分片权重的引用，并通过精心设计的三阶段流水线完成数据传输：首先将权重从磁盘或训练引擎移入 GPU 显存（H2D），随后在 checkpoint-engine 工作节点间进行广播并写入 CUDA IPC 缓冲区，最后由推理引擎从广播数据中按需加载所需权重子集。整个过程通过 ZeroMQ 套接字协调控制，通信与拷贝操作充分重叠以最大化性能。

广播模式是默认推荐方案，适合需要同步更新大量推理实例的场景。P2P 模式则专为动态扩容设计：当新推理实例因重启或弹性调度加入集群时，为避免影响现有实例的在线服务，系统会利用 mooncake-transfer-engine 将权重从现有实例的 CPU 内存直接 P2P 传输至新实例的 GPU 显存，支持 RDMA 高速互联。P2P 模式还针对发送方与接收方的桶分配进行了优化，目标是充分利用每对节点的可用网络带宽。

基准测试覆盖了多种主流模型和硬件配置。在16张 H20 GPU（TP16）上，Kimi-K2 Instruct（FP8）的广播更新耗时约14.4秒，DeepSeek-V3.1（FP8）约10.2秒；扩展至256张 H20 GPU 后，两者分别约为16秒和11.3秒，体现出良好的规模扩展性。较小规模的 GLM-4.5-Air 和 Qwen3-235B 在8张 H800 上的广播更新则分别仅需3.5秒和6.2秒。所有测试均基于 vLLM v0.10.2rc1 推理引擎。

在工程实践层面，checkpoint-engine 支持通过 pip 直接安装，基础版本提供广播功能，附加 p2p 选项则会同步安装 mooncake-transfer-engine 以支持 RDMA 传输。项目推荐配合 vLLM v0.10.2 使用，并需要启用 collective_rpc API 端点。此外，新加入的 checkpoint-engine 实例可以复用已有实例的权重，无需重新从磁盘加载，进一步降低了动态扩容的时间成本。该项目的技术细节已在 Kimi-K2 技术报告中详细记录。

要点

• checkpoint-engine 将 Kimi-K2 万亿参数模型在数千张 GPU 上的权重更新时间压缩至约20秒，显著降低强化学习训练中的同步开销。
• 广播与P2P双模式设计兼顾了大规模同步更新和动态扩容两种典型场景，P2P模式借助RDMA传输避免影响在线推理服务。
• 三阶段流水线设计（H2D传输、广播、推理引擎加载）通过通信与计算重叠最大化吞吐，内存不足时自动降级为串行执行。
• 基准测试覆盖 GLM-4.5-Air、Qwen3-235B、DeepSeek-V3.1、Kimi-K2 等主流大模型，验证了工具在不同规模下的实用性。
• 新实例可复用已有实例的权重缓存，结合动态桶分配优化，进一步提升了弹性集群场景下的扩容效率。