Skip to content

[feature] Support SGLang disk weight recovery#1950

Open
PengchengShi00 wants to merge 1 commit into
InternLM:mainfrom
PengchengShi00:sglang_disk_update_weight
Open

[feature] Support SGLang disk weight recovery#1950
PengchengShi00 wants to merge 1 commit into
InternLM:mainfrom
PengchengShi00:sglang_disk_update_weight

Conversation

@PengchengShi00

Copy link
Copy Markdown
Contributor

目标

在 RL 训练过程中,当 SGLang rollout server 异常退出并重启后,支持从磁盘上的 HuggingFace 格式 checkpoint 恢复模型参数,避免恢复后的 server 只能依赖训练侧重新通过完整的 IPC/NCCL tensor 流同步权重。

当前 PR 的范围收敛为 SGLang server recovery:XTuner 保存 HF checkpoint,并把 checkpoint 目录发送给恢复后的 SGLang server,由 SGLang 通过 update_weights_from_disk 自行加载权重。

LMDeploy 暂不在当前 PR 支持范围内。后续如果需要支持 LMDeploy disk 更新,需要另行设计从 HF checkpoint 读取、bucket 化并复用 IPC/NCCL 发送的链路,或者由LMDeploy提供从disk加载权重的接口。

当前方案

新增 disk 作为 IPC/NCCL 之外的第三种权重更新方式。

  • weight_transport_type="disk":表示从磁盘恢复 rollout 权重。
  • disk_weight_path:待加载的 HF checkpoint 路径,由调用方在 bind rollout weight update 时传入。
  • DiskWeightTransport:当前仅支持 SGLang。
  • SGLang adapter 不消费 WeightIterator 产生的 tensor batch,只构造 update_weights_from_disk 请求。
  • 请求携带 model_path=disk_weight_path,权重加载由 SGLang server 内部完成。

测试覆盖

tests/rl/test_update_weight_disk.py 覆盖 SGLang recovery 后从 disk 加载权重的端到端流程:

  1. 启动训练 workers 和多个 SGLang rollout engines。
  2. 先通过 NCCL 做一次正常权重同步。
  3. 对每个 rollout endpoint 生成一次,记录 baseline。
  4. 训练侧保存 HF checkpoint 到临时目录。
  5. 手动关闭一个 rollout weight-update endpoint。
  6. 调用 restart_inactive_workers() 恢复 inactive worker group。
  7. 使用 DiskWeightTransport 调用 SGLang update_weights_from_disktarget_endpoint_ranks={failed_endpoint_rank},只更新刚重启的 rollout endpoint。
  8. 对恢复后的 endpoint 再 generate 一次,验证输出和 baseline 一致。

测试中使用 test-only worker 覆盖 offload() 为 no-op。原因是当前主线 recovery 的 offload() 逻辑主要服务于共卡场景,而该测试验证的是训推分离 recovery 后立即 disk update。

正式使用还需要补充

当前实现和测试验证了 SGLang disk recovery 的基本链路,但要进入正式主线使用,还需要补齐以下内容:

  • 训推分离场景的 recovery 逻辑
  • 需要在现有 active / inactive / recovering 之外,引入一个表达“engine 已重启但权重尚未恢复”的状态或者等价标记。
  • 需要明确 SGLang dummy load 后的 readiness 语义:当前 health_generate 可能会与“engine 已重启但权重尚未恢复”的状态冲突;真实权重通过 disk update 加载完成后,再执行 generate 验证更合理。

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