24-CheckpointEngine训练权重如何同步给推理引擎

Checkpoint Engine：训练权重如何同步给推理引擎

RL 训练里的权重同步不是“顺手保存一次 checkpoint”，而是把刚更新的 actor 重新接到下一轮 rollout。 第24篇为第四组收束：actor更新后，新权重如何真正回到rollout replicas？文章从verl源码追踪CheckpointEngineManager、naive与非naive同步、 NCCL/NIXL/Mooncake/Kimi等后端拓扑，以及vLLM/SGLang/TRT-LLM加载边界，解释权重同步为何是RL step的正式数据面

第 23 篇把训练引擎推到 token 维度：长上下文让 Sequence/Context Parallel 进入 forward、attention 和输出恢复。到这里，第四组还剩最后一个闭环问题：actor/critic 在训练后端里完成 update 之后，新 actor 权重怎样回到 rollout replicas？如果这一步慢、错、或者和推理状态冲突，下一轮样本就可能用旧策略生成，RL step 的训练和推理会脱节。

本文的核心判断是：verl 的 Checkpoint Engine 不是普通存盘模块，而是训练权重到推理引擎的同步抽象。它把 FSDP/Megatron/TorchTitan 等训练后端的 get_per_tensor_param()统一成 named tensor stream，再用 naive、NCCL/HCCL、NIXL、Mooncake、Kimi 等后端传给 rollout worker，最后由 vLLM/SGLang/TRT-LLM adapter 装进推理进程。这个边界把“训练如何切权重”和“推理如何加载权重”隔开，同时也暴露了同步带宽、KV cache 生命周期和未完成请求处理的成本。

先看第 24 篇在一轮 RL step 里的位置。读图时注意：update_actor结束以后，step 还没有真正闭环；update_weights把训练后的 actor 重新变成下一轮 rollout 的 policy。

Checkpoint Engine 把训练后的 actor 接回 rollout

源码上，RayPPOTrainer在初始化 worker 后构造 CheckpointEngineManager，并先让 rollout replicas sleep；训练开始前加载 checkpoint 后会调用一次 checkpoint_manager.update_weights()；主循环里 actor update 之后，也把 update_weights放进计时区间（verl/trainer/ppo/ray_trainer.py:870-884，1294-1297，1581-1583）。这说明权重同步是 RL step 的正式阶段，不是 rollout 层的附属优化。

1. Checkpoint Engine 的合同是 named tensor stream

Checkpoint Engine 的第一层价值，是定义一个训练侧和推理侧都能接受的中间形态：带名字的 tensor 流。训练 engine 负责把自己的 shard、LoRA、QAT、TP/PP/CP 等内部状态导出成这条流；checkpoint engine 负责传输；rollout adapter 负责按推理后端的方式加载。

下面这张图画的是合同边界。看图时注意：checkpoint engine 不需要理解 FSDP 或 Megatron 的内部并行细节，它只要求训练后端给出 name, tensor，再把同样的 stream 交给 rollout。

Checkpoint Engine 的三段式合同

BaseEngine.get_per_tensor_param()定义了训练 engine 的导出接口，返回 named tensor generator 和可选 PEFT config（verl/workers/engine/base.py:150-158）。FSDP 实现会加载 FSDP module 到 GPU，处理 LoRA merge 或 adapter-only，同步 state dict key，并把 DTensor full tensor 转成可传输 tensor（verl/workers/engine/fsdp/transformer_impl.py:785-833）。Megatron/MCore 实现则通过 bridge 导出 full HF weights、adapter weights 或 vanilla bridge weights（verl/workers/engine/megatron/transformer_impl.py:715-744）。

Checkpoint Engine 抽象层只保留 prepare()、build_topology()、init_process_group()、send_weights()、receive_weights()和 finalize()等方法；注释里的示例也是 trainer 调 send_weights(trainer.get_per_tensor_param())，rollout 调 server_adapter.update_weights(engine.receive_weights())（verl/checkpoint_engine/base.py:96-189）。rollout 侧的 BaseRollout.update_weights()也只要求接收 named tensor generator（verl/workers/rollout/base.py:44-64）。这就是统一合同的意义：训练后端可以很复杂，传输层和推理加载层不必反向猜它。

2. naive 路径适合 colocated 训练和推理

verl 的默认配置是 checkpoint_engine.backend: naive，也就是训练 worker 和 rollout 在同一组 actor/rollout worker 内直接同步。它不是没有 checkpoint engine，而是使用 ColocatedCheckpointEngine这条最短路径：send_weights()暂存 generator，receive_weights()再把它 yield 出来（verl/trainer/config/rollout/rollout.yaml:263-292，verl/checkpoint_engine/base.py:209-253）。

下面这张图展示 naive 路径。看图时注意：真正处理内存状态的是 ActorRolloutRefWorker.update_weights()，它要先恢复 rollout weights，再更新权重，最后恢复 KV cache。

naive 路径中的 colocated 权重更新

CheckpointEngineManager.update_weights()如果发现 backend 是 naive，直接调用 trainer worker group 的 update_weights()并返回（verl/checkpoint_engine/base.py:448-458）。ActorRolloutRefWorker.update_weights()里，非 naive 时只取 self.actor.engine.get_per_tensor_param()并交给 checkpoint engine send_weights()；naive 时则会在 free_cache_engine开启时先 rollout.resume(tags=["weights"])，再从 actor engine 导出权重，处理 LoRA base sync/adapter sync，调用 self.rollout.update_weights()，必要时把 actor 参数 offload 回 CPU，最后 rollout.resume(tags=["kv_cache"])（verl/workers/engine_workers.py:663-743）。

这条路径的工程含义很明确：当训练和推理共置时，权重同步可以绕过额外通信拓扑，但必须和 rollout 的 sleep/resume、KV cache 释放、LoRA base sync 绑定在一起。第 15 篇讲过的 sleep/resume，在这里重新出现为权重同步的内存前置条件。

3. 非 naive 路径先让 rollout 安全停下来

当 trainer 和 rollout 解耦，或者需要独立 rollout replicas 时，Checkpoint Engine Manager 会走非 naive 路径。这里的关键不是“用哪个通信库”，而是更新前后要维护推理服务状态：未完成请求不能继续使用旧权重，KV cache 也不能挡住权重写入。

下面这张图是非 naive 同步生命周期。读图时注意顺序：abort、release KV、build process group、trainer/rollout 双侧 update、finalize、resume KV、resume generation。

非 naive Checkpoint Engine 的同步生命周期

CheckpointEngineManager.update_weights()在非 naive 路径里先 abort_replicas()，再把所有 replica worker 拼成临时 RayWorkerGroup；随后调用 release_kv_cache_replicas()，注释说明它只释放 KV cache，保留 weight buffer，让传输可以直接写入已有权重；之后 build_process_group()调所有 trainer/rollout worker 的 prepare()，由后端类 build_topology()生成双侧 init 参数，再双侧 init_process_group()（verl/checkpoint_engine/base.py:365-390，420-445，448-475）。

真正同步时，manager 同时触发 trainer 的 update_weights(..., mode=backend)和 rollout worker 的 update_weights()；trainer 侧会调用 checkpoint engine 的 send_weights()，rollout 侧的 CheckpointEngineWorker.update_weights()会从 receive_weights()得到 stream，再交给 server adapter。同步完成后，manager 对两侧调用 finalize()，恢复 KV cache，并让 rollout replicas resume_generation()（verl/checkpoint_engine/base.py:300-307，476-492）。这条顺序解释了为什么权重同步会影响推理服务：它必须短暂停止或接管 generation 状态。

4. 后端差异主要体现在拓扑和 bucket

Checkpoint Engine 的后端差异，不应该理解成“谁保存 checkpoint”。它们主要解决两个问题：trainer 和 rollout worker 怎样组成通信拓扑，以及大模型权重怎样按 bucket 流式传输，避免一次性 materialize 太多数据。

下面这张图把几类后端放在同一张图里。看图时注意：无论底层是 collective broadcast、ring P2P、TransferEngine 还是 parameter server，外层仍然是同一组 prepare/build_topology/send/receive/finalize方法。

Checkpoint Engine 后端的拓扑与 bucket 差异

配置文件把 checkpoint engine 放在 rollout 配置下，默认 backend 是 naive，并提供 update_weights_bucket_megabytes和 engine_kwargs；README 列出 naive、NCCL、HCCL、NIXL、Kimi checkpoint engine、Mooncake 等后端及其适用场景（verl/trainer/config/rollout/rollout.yaml:263-292，verl/checkpoint_engine/README.md:13-22）。源码里，__init__.py会尝试导入这些后端模块，后端通过 CheckpointEngineRegistry.register()加入 registry（verl/checkpoint_engine/__init__.py:33-66，verl/checkpoint_engine/base.py:49-93）。

NCCL 后端的拓扑是 rank 0 trainer 加 rollout workers，trainer 其他 rank 只消费 generator 不发送；发送端按 bucket 填充 buffer 并 broadcast metadata，接收端再 merge_weight_chunks()还原 tensor（verl/checkpoint_engine/nccl_checkpoint_engine.py:102-171，229-305）。NIXL 后端也复用 split_weight_chunks()/merge_weight_chunks()，但拓扑是 trainer rank 0 到 rollout workers 的 ring P2P，prepare()会注册 send/recv buffer 的 memory descriptor（verl/checkpoint_engine/nixl_checkpoint_engine.py:238-305，370-443）。Mooncake 后端用 TransferEngine 初始化 session、注册 buffer，并通过 store 传递 bucket metadata（verl/checkpoint_engine/mooncake_checkpoint_engine.py:34-120，149-223）。Kimi checkpoint engine 则把 trainer 侧权重 offload 到 CPU，并注册到 parameter server，再由 rollout 侧读取（verl/checkpoint_engine/kimi_checkpoint_engine.py:222-283，321-368）。

公共的 bucket 逻辑在 split_weight_chunks()和 merge_weight_chunks()：前者把任意 named tensor 切成 TensorMeta + uint8 chunk，后者按 chunk_offset拼回原 tensor（verl/checkpoint_engine/base.py:33-47，495-563）。这解释了为什么 bucket size 是一个性能和显存权衡：大 bucket 可能提高吞吐，但也会增加传输 buffer 的显存占用。

5. rollout adapter 才知道如何装进推理引擎

最后一段边界在 rollout adapter。Checkpoint Engine 把 named tensor stream 送到 rollout worker，但 vLLM、SGLang、TRT-LLM 的加载方式完全不同。verl 把这种差异留给 BaseRollout.update_weights()的具体实现。

vLLM adapter 会先触发 server 侧 update_weights_from_ipc，再用 BucketedWeightSender发送权重，server 侧 BucketedWeightReceiver收到 bucket 后调用 _update_weights()；更新后会清 KV cache 并写入 global step（verl/workers/rollout/vllm_rollout/vllm_rollout.py:164-195，verl/workers/rollout/vllm_rollout/utils.py:163-226）。SGLang adapter 会按 update_weights_bucket_megabytes分批，把 params batch 交给 SGLang 的 weight sync；如果是 LoRA adapter，也会走 load LoRA from tensors（verl/workers/rollout/sglang_rollout/sglang_rollout.py:287-370）。TRT-LLM adapter 则把 tensor reduce 成 IPC handles，leader 聚合 handles 后交给 adapter update，并在最后调用 finalize（verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:423-530）。

所以 Checkpoint Engine 的真实作用不是替代 rollout backend，而是把训练权重搬到 rollout backend 能处理的位置。它停止在 named tensor stream 和通信拓扑处；至于 vLLM/SGLang/TRT-LLM 怎样 patch weight loader、flush cache、处理 LoRA 或量化权重，仍然属于推理后端 adapter。

小结：第四组在权重同步处闭环

到第 24 篇，第四组从训练引擎内部的 shard、并行、batch 和长上下文，走到了训练和推理之间的权重数据面。CheckpointEngineManager.update_weights()把这件事写成一个明确流程：naive 路径服务 colocated 同步；非 naive 路径服务 trainer/rollout 解耦，并显式处理 abort、KV cache、通信拓扑、双侧 update 和 resume generation。

放回系列地图，第四组补完了 workers/resources -> training engine/distributed core -> weight sync -> rollout/serving engine这一段。下一组会进入性能和生产化：既然一轮 RL step 已经拆成 gen、reward、logprob、update、sync 等阶段，第 25 篇就该回答如何 profiling 一次 RL step，判断慢点到底在 rollout、训练还是权重同步。

本文源码索引

• verl/trainer/ppo/ray_trainer.py:870-884、1294-1297、1581-1583：trainer 如何创建 CheckpointEngineManager，并在训练开始和 actor update 后同步权重。
• verl/workers/engine/base.py:150-158、verl/workers/engine/fsdp/transformer_impl.py:785-833、verl/workers/engine/megatron/transformer_impl.py:715-744：训练后端如何导出 named tensor stream。
• verl/checkpoint_engine/base.py:49-189、209-253、323-492、495-563：registry、抽象接口、naive engine、manager 生命周期和 bucket split/merge。
• verl/workers/engine_workers.py:663-743：ActorRolloutRefWorker.update_weights()如何分出 naive 与 checkpoint engine 路径，并处理 LoRA、weights/KV cache resume 和 offload。
• verl/trainer/config/rollout/rollout.yaml:263-292、verl/workers/config/rollout.py:143-151：checkpoint engine 配置入口、backend 和 bucket size。
• verl/checkpoint_engine/__init__.py:33-66、verl/checkpoint_engine/README.md:13-22：后端模块导入与支持后端概览。
• verl/checkpoint_engine/nccl_checkpoint_engine.py:102-171、229-305：NCCL 拓扑、bucket broadcast 和 receive/merge。
• verl/checkpoint_engine/nixl_checkpoint_engine.py:238-305、370-443：NIXL 的 agent metadata、ring P2P 拓扑和 bucket 传输。
• verl/checkpoint_engine/mooncake_checkpoint_engine.py:34-120、149-223：Mooncake TransferEngine 初始化、buffer 注册和权重接收。
• verl/checkpoint_engine/kimi_checkpoint_engine.py:222-283、321-368：Kimi checkpoint engine 的 trainer/rollout topology 与 parameter server 路径。
• verl/workers/rollout/base.py:44-64、verl/workers/rollout/replica.py:265-291：rollout 抽象和 replica sleep/abort/KV cache 控制。
• verl/workers/rollout/vllm_rollout/vllm_rollout.py:164-195、verl/workers/rollout/vllm_rollout/utils.py:163-226、verl/workers/rollout/sglang_rollout/sglang_rollout.py:287-370、verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:423-530：不同推理后端如何消费权重 stream。