Skip to content

[feature] Support SGLang disk weight recovery#1950

Open
PengchengShi00 wants to merge 2 commits into
InternLM:mainfrom
PengchengShi00:sglang_disk_update_weight
Open

[feature] Support SGLang disk weight recovery#1950
PengchengShi00 wants to merge 2 commits 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 验证更合理。

@PengchengShi00

Copy link
Copy Markdown
Contributor Author

支持在colocate模式下使用disk update weight方式,更新时可在save hf后立即offload train model,原有IPC路径流程不变。但存在以下问题:

  • 速度慢:更新时间随推理引擎增加。xtuner中一个推理引擎就是一个dp,disk更新会向每个推理引擎发一次请求,读hf的次数与推理引擎数量一致。(sglang当设置推理引擎当dp>1时,只会读hf一次)

可通过在xtuner中支持dp (但这会与健康监测的恢复机制冲突),或者新增 从disk导出batch再 复用IPC或者NCCL路径

# Optional port used by NCCL external weight update groups.
weight_update_port: int | None = None
# Optional disk weight path used by disk weight transport.
disk_weight_path: str | None = None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

有办法能让不同的传输方式下需要的信息跟传输方式绑定吗?这样即使后面改了传输方式具体的实现(例如接入其他的权重更新引擎),也不同频繁修改接口,这样可能代码设计会更好
weight_update_host: str | None = None weight_update_port: int | None = None disk_weight_path: str | None = None

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.

2 participants