[Feature] Add Checkpoint Engine as a transport between train and rollout

[PengchengShi00]

Checkpoint Engine 权重同步

1. 背景

XTuner RL 训练需要周期性把 train engine 中的最新权重同步到 rollout engines。现有权重同步路径包括：

路径	适用场景	说明
IPC	共卡模式	训练侧直接把权重通过 IPC 更新给 rollout backend。
Checkpoint Engine	共卡模式	训练侧把权重注册到 in-process ParameterServer，再由 ParameterServer 推送给 rollout engines。
NCCL	分离模式	train workers 和 rollout workers 通过 NCCL broadcast 同步权重。

Checkpoint Engine 路径的收益：

• 降低共卡权重更新时的显存峰值。训练侧可以先把权重注册到 ParameterServer，再 offload train engine、onload rollout engine，最后由 ParameterServer 推送权重。
• 支持恢复失败的 rollout engine。rollout engine 重启后，可以复用上一次注册成功的 checkpoint 重新推送权重。

当前实现只在共卡训练路径中启用 Checkpoint Engine。分离模式检测到 enable_checkpoint_engine=True 时会打印 warning，并回退到 NCCL weight transport。

2. 架构设计

共卡模式下，在 rollout config 中设置 enable_checkpoint_engine=True，可使用 Checkpoint-Engine 进行权重更新， 分离模式下即使设置 enable_checkpoint_engine=True，也会回退到 NCCL transport。

2.1 初始化

1. 创建ParameterServer

训练 worker 内部创建 ParameterServer，读取 train world size，作为 ParameterServer world size，不额外启动独立 PS 进程。

ParameterServer(
    auto_pg=False,
    rank=self.rank,
    world_size=self.ps_world_size,
)

auto_pg=False 表示 Checkpoint Engine 复用 XTuner 已经初始化好的 torch.distributed 默认进程组，不在 update 后销毁该进程组。

2. 划分参数shard

根据 HF checkpoint 的 model.safetensors.index.json 计算当前 PS-rank 负责的参数 key 集合。每个 PS rank 只负责注册自己分到的参数 shard。这些 key 后续用于从 WeightIterator 中过滤当前 PS rank 需要注册的 tensor，避免每个 rank 都注册全量参数。

2.2 更新

CheckpointEngineWeightTransport.update(...) 被拆成两个可选阶段：

update(weight_iterator, need_register=True, need_update=True)

1. Register 阶段

当 need_register=True 时，transport 会：

• 从 weight_iterator.iter_batch_groups() 收集 train engine 当前权重。
• 根据本 rank 的 local checkpoint keys 过滤 tensor。
• 注销上一轮 checkpoint 名称，限制 pinned host memory 占用。
• 调用 ParameterServer.register_checkpoint(...) 注册新 checkpoint：

如果 HF index 中的某些 key 没有从 train engine 收集到，transport 会记录错误日志。MTP-only key 和非 MTP key 会分开打印，便于排查模型结构差异。

2. Update 阶段

当 need_update=True 时，transport 会把当前 checkpoint 推送到 rollout engines：

• 从 rollout_info.active_update_targets 获取当前活跃 rollout targets。
• 校验每个 target 声明的 update_ranks 非空、不越界、不重复。
• 判断本次更新是否覆盖全部 PS ranks：
  • 覆盖全部 ranks 时，调用 ParameterServer.update(..., ranks=None)，使用 Checkpoint Engine broadcast 路径。
  • 只覆盖部分 ranks 时，调用 ParameterServer.update(..., ranks=active_ranks)，使用 p2p update 路径。

3. need_register 和 need_update

need_register 和 need_update 用于控制 Checkpoint Engine 更新的两个阶段。

参数	True	False
need_register	从 train engine 注册一个新的 checkpoint。	复用上一次注册成功的 checkpoint 名称。
need_update	把当前 checkpoint 更新到 rollout engines。	只完成注册，不推送到 rollout engines。

典型使用方式：

# 常规更新：注册 train 权重并同步到 rollout engines
train_controller.update_weights(need_register=True, need_update=True)

# rollout 恢复：复用上一次 checkpoint，只重新更新 rollout engines
train_controller.update_weights(need_register=False, need_update=True)

# 显存紧张：先注册 checkpoint，稍后再更新 rollout engines
train_controller.update_weights(need_register=True, need_update=False)
offload(train)
onload(rollout)
train_controller.update_weights(need_register=False, need_update=True)

need_update=False 的核心用途是把注册和 rollout 更新拆成两个时刻执行。这样训练侧可以先把权重注册到 ParameterServer，然后释放或切换部分资源，再让 rollout engines 加载 checkpoint，从而降低峰值显存压力。

3. 共卡模式 IPC / Checkpoint Engine 速度对比

使用两种更新路径的端到端权重更新时间，包括了rollout onload_weights、train offload时间。

模型	并行配置	Backend	IPC更新时间	Checkpoint Engine更新时间
Qwen3-4B	TP=2	SGLang	0.86s	2.3s
Qwen3-8B	TP=2	SGLang	1.1s	2.8s
Qwen3-30B-A3B	TP=2	SGLang	3.7s	5s
Qwen3.5-35B-A3B	EP=4	SGLang	4.3s	6.5s

[hhaAndroid]

训练侧 ep_size > 1 时，专家权重可能被错误切分并绑定到错误的 HF key。xtuner/v1/rl/weight_update/weight_iterator.py:193 只有 NCCL 会 gather train-EP tensor；新加入的 checkpoint_engine 不会走这个逻辑，应该是不对的

[PengchengShi00]

是的，checkpoint-engine 应该始终要走gather_train_ep_shards

[hhaAndroid]

这个地方有比较大的显存峰值吧？要想想这个地方会不会成为 oom 瓶颈点。ce 默认这个地方就是这么做的吗？

[PengchengShi00]

是的，每张卡会有checkpoint shard的显存增加，但是这部分可以放到CPU上（但不建议这么做，注册时还会有额外的H2H 操作，到pinning memory的copy）。checkpoint shard的大小不会高于rollout weight，因此这里的显存峰值也不会高于IPC的峰值

[YanhuiDua]

@claude review

[claude]

Claude: [Critical] rename 漏改，这个 GPU 测试的断言会失效。

第 347 行取的是新名字 weight_update，但这里仍然把 wrapper 赋给旧名字 update_weights，而 TrainingController 在本 PR 中已经没有 update_weights 了。结果是 wrapper 永远不会被调用，self.update_weight_calls 恒为 0，test_qwen35_vl_moe_async_train_2step.py#L418 的 assertEqual(self.update_weight_calls, TOTAL_TRAIN_STEPS - 1) 必然失败。

Suggested change

	trainer.train_controller.update_weights = update_weights_wrapper
	trainer.train_controller.weight_update = update_weights_wrapper

顺带：同类漏改还有 tests/rl/test_rl_colocate_trainer.py#L166（fake controller 仍暴露 update_weights=，因为该测试 mock 掉了 _sync_weights_and_save 才没暴露问题）和 docs/design/checkpoint-engine.md#L40-L43（新文档写的还是 update_weights(...)）。

[claude]

Claude: [Critical] 权重缺失只打日志、不中断，会造成静默的错误权重。

missing 非空意味着这一轮注册的 checkpoint 缺少参数，ParameterServer 仍会照常注册并推送，rollout engine 上这些参数保持上一轮（或初始化）的值 —— 训练/推理权重不一致，而且只有一行 error log。RL 里这种静默不一致极难定位，建议 fail fast，只对已知可忽略的前缀（MTP）放行：

missing_mtp = {k for k in missing if k.startswith("mtp.")}
if missing - missing_mtp:
    raise RuntimeError(
        f"[checkpoint_engine] rank={self.rank} missing keys: {sorted(missing - missing_mtp)}"
    )
if missing_mtp:
    self.logger.warning(f"[checkpoint_engine] rank={self.rank} skip MTP-only keys: {sorted(missing_mtp)}")

另外两个相关问题：

1. 当前 if missing_non_mtp_keys: ... else: ... 在「同时缺 MTP 和非 MTP」时只打印非 MTP，MTP 那部分信息丢了。
2. 反向缺失完全没有检查：local_keys 来自 rollout_config.model_path 的 HF index，如果 train engine 导出的某个 key 不在 index 里（tied lm_head、fused expert 命名差异等），它会被每个 rank 过滤掉，永远不会推给 rollout，且没有任何日志。建议在 filter 之前先统计一次本 rank 见到的全部 key，与 HF index 全集做差集校验：

seen_all_keys |= set(sd.keys())          # filter 之前
...
extra = seen_all_keys - hf_index_keys     # 任何 rank 都不会注册的 key
if extra:
    raise RuntimeError(f"train engine exports keys absent from HF index: {sorted(extra)}")

[YanhuiDua]

什么情况下会缺失呢？如果缺失的话，是不是下一次rollout的结果就一定不对了

[claude]

Claude: [Warning] ruff lint 会失败：类定义前只有 1 个空行，E302 expected 2 blank lines, found 1。pre-commit 覆盖 xtuner/v1/，请跑一次 pre-commit run --all-files。

Suggested change

	class CheckpointEngineAdapter:
	class CheckpointEngineAdapter:

另外从抽象上看，CheckpointEngineAdapter 目前只有一个 build_update_url() 一行拼串，且没有实现 WeightTransportAdapter Protocol（没有 before_update / after_update_all_groups），构造时又立刻做了 if adapter is None: raise 这种不可能发生的检查（transport.py:1034）。按 Deletion test：删掉它，复杂度不会散落到多个调用者，只是把一行 f-string 内联回 _make_req_func。它现在是一个纯转发的 Shallow Module，一个 Adapter 也不构成真实 Seam；建议内联，等真的出现第二个 backend（LMDeploy）时再抽。

[claude]

Claude: [Warning] WORLD_SIZE 语义在本仓库里是二义的，这里读环境变量既多余又危险，并且让下面的 assert 变成死代码。

1. train worker 里 WORLD_SIZE 由 ray_accelerator_worker.py#L245 在 init_process_group 之前设为 train world size，所以目前等价于 dist.get_world_size()；但在 driver / launcher 侧 WORLD_SIZE 一律表示 节点数（examples/v1/scripts/run_rl.sh、各 config 里的 NNODE = int(os.environ.get("WORLD_SIZE", "1"))，本 PR 新增的 tests/rl/test_update_weight_colocate.py:75 也是这个含义）。一旦这个 transport 被在非 actor 进程里构造，ps_world_size 会静默变成 nnodes，进而污染 split_tensors_for_rank 的分片、_get_target_update_ranks 的越界校验、broadcast/p2p 的判定以及 ParameterServer(world_size=...)，表现为错分片或 hang 而不是清晰报错。
2. dist.get_world_size() 作为 os.environ.get 的默认值是立即求值的，PG 未初始化时会先在这一行抛 ValueError: Default process group has not been initialized，第 901 行那条友好的 assert 永远走不到。

Suggested change

	self.ps_world_size = int(os.environ.get("WORLD_SIZE", dist.get_world_size()))
	assert dist.is_initialized(), (
	"Checkpoint Engine requires an initialized torch.distributed process group."
	)
	self.ps_world_size = dist.get_world_size()

（对应地把 901-903 的 assert 精简掉 ps_world_size > 0 这半个恒真条件。）

[claude]

Claude: [Warning] 这个函数有几处需要收拾（性能 + 可读性 + 契约）：

1. docstring 与实现相反：写的是 "onto CPU"，实际 named[key] = tensor 保留的是 GPU tensor（下面注释也写了「占显存多，速度快」）。这直接决定了 register 阶段的显存峰值，是调用者必须知道的 Interface 信息，请改成准确描述。
2. 热路径上的 empty_cache()：每个 batch 调 1 次、每个 batch group 再调 1 次。empty_cache() 会同步并归还所有 cache block，后续分配要重新走 cudaMalloc，几十~上百个 bucket 累加下来开销可观（PR 里 CE 比 IPC 慢 2-3x，这里可能是一部分原因）。建议只在 group 边界调用，或按 checkpoint_engine_sync_after_register 一样做成可配置。
3. 注释掉的备选实现（967-969 行）请删掉，两种策略的选择应该由 checkpoint_engine_sync_after_register 这类配置表达，而不是留在代码里让人手改。
4. 缺类型标注（CLAUDE.md 要求新代码必须有 type hints）：def _collect_named_tensors(self, weight_iterator: WeightIterator, local_keys: set[str] | None = None) -> dict[str, torch.Tensor]:。
5. 建议加一次连续性保证：checkpoint-engine 的 _register_checkpoint 内部是 buffer[off:off+nbytes] = tensor.view(-1).view(dtype=torch.uint8)，非 contiguous 会直接抛错。这里加 tensor = tensor.contiguous()（或 assert）可以把失败点提前到 XTuner 侧，报错信息更可定位。
6. 中文注释和文件里其余英文注释风格不一致，建议统一为英文。

[claude]

Claude: [Critical] 库代码里无条件覆盖集群级环境变量，会破坏用户环境并在非 Mellanox / 非 8 卡 RDMA 机器上直接失败。

两点问题：

1. NCCL_IB_HCA 是全局 NCCL 配置，这里在 dist.init_process_group() 之后（transport 是在 bind_rollout_weight_update 里才构造的）无条件覆盖用户设置。对已建立的默认 PG 无效，但会影响 checkpoint-engine 之后新建的通信域，且静默丢弃用户显式配置。
2. checkpoint-engine 的 _get_rdma_devices() 优先读 PS_P2P_STORE_RDMA_DEVICES，取不到设备时 _get_my_rdma_device 会 raise RuntimeError("no rdma devices found")，并且 ParameterServer.__init__ 里只 catch 了 ImportError。所以在 RDMA 设备名不是 mlx5_0..7（或没有 RDMA 网卡）的机器上，这两行会把一个「能 broadcast 跑通」的场景变成构造期崩溃。

这属于把部署环境细节硬编码进 Implementation，调用者无法覆盖。建议交回给环境/配置，代码只做「未设置时不干预」：

# 不要覆盖；需要时由 RolloutConfig 显式提供
if rollout_config.checkpoint_engine_rdma_devices:
    os.environ.setdefault("PS_P2P_STORE_RDMA_DEVICES", rollout_config.checkpoint_engine_rdma_devices)

（同样的硬编码也出现在 tests/rl/test_update_weight_colocate.py:42-43，测试里设置是可以接受的，库里不行。）