Skip to content

support checkpoint engine[P2P] recovery in colocate mode - #2027

Open
PengchengShi00 wants to merge 13 commits into
InternLM:mainfrom
PengchengShi00:checkpoint-engine-recovery
Open

support checkpoint engine[P2P] recovery in colocate mode#2027
PengchengShi00 wants to merge 13 commits into
InternLM:mainfrom
PengchengShi00:checkpoint-engine-recovery

Conversation

@PengchengShi00

Copy link
Copy Markdown
Collaborator

No description provided.

@PengchengShi00

Copy link
Copy Markdown
Collaborator Author

Checkpoint-Engine P2P 推理引擎故障恢复

在共卡 RL 训练中,rollout worker 负责启动推理引擎并处理生成请求。若某个推理引擎进程异常退出,原有流程会由 RolloutHealthManager 检测失败并重启对应 worker group。
本次故障恢复的目标是:当失败的推理引擎被重启后,不等待本轮 rollout 完全收尾,而是先把该 worker group 标记为 pending_weights,再由训练侧后台线程立即使用 checkpoint-engine 中已注册的权重,通过 P2P 方式把权重推送到恢复后的推理引擎,最后将其重新标记为 active

WorkerLifecycleState新增状态PENDING_WEIGHTS

恢复后的 rollout worker group 不会立刻进入 active,而是先进入 pending_weights。该状态表示:

  • 推理服务已经重启并通过健康检查;
  • 当前服务还没有加载到训练侧最新权重。

checkpoint-engine P2P 更新 pending 目标

RLColocateTrainer 新增了 _update_pending_rollout_weights_from_checkpoint_engine() 路径,用于:

1.. 从 RolloutController 获取 pending_weights 的更新目标;
2.. 调用 bind_rollout_weight_update() 绑定这些 pending target;
3. 调用 weight_update(..., update_pending_only=True) 只更新 pending worker;
4. 调用 pending worker 的 onload_weightsonload_kvcache
5. 更新完成后把对应 group 标记回 active
这样可以避免对所有 active rollout worker 做重复广播,只对刚恢复的失败引擎执行 P2P 权重恢复。

后台监控线程

在 checkpoint-engine 模式下,RLColocateTrainer 会在 rollout 阶段启动后台线程,周期性检查是否存在 pending_weights worker group。

该线程会在满足以下条件时触发恢复:

  • rollout 权重更新锁没有被主流程占用;
  • checkpoint-engine 中已经注册过训练权重;
  • 存在 pending rollout weight-update targets。

权重同步语义

  • update_pending_only=False:正常训练同步,更新所有可更新 rollout target;
  • update_pending_only=True:故障恢复同步,只更新 pending_weights target。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant