[feature] Support SGLang disk weight recovery#1950
Open
PengchengShi00 wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
目标
在 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。WeightIterator产生的 tensor batch,只构造update_weights_from_disk请求。model_path=disk_weight_path,权重加载由 SGLang server 内部完成。测试覆盖
tests/rl/test_update_weight_disk.py覆盖 SGLang recovery 后从 disk 加载权重的端到端流程:restart_inactive_workers()恢复 inactive worker group。DiskWeightTransport调用 SGLangupdate_weights_from_disk,target_endpoint_ranks={failed_endpoint_rank},只更新刚重启的 rollout endpoint。测试中使用 test-only worker 覆盖
offload()为 no-op。原因是当前主线 recovery 的offload()逻辑主要服务于共卡场景,而该测试验证的是训推分离 recovery 后立即 disk update。正式使用还需要补充
当前实现和测试验证了 SGLang disk recovery 的基本链路,但要进入正式主线使用,还需要补齐以下内容:
active/inactive/recovering之外,引入一个表达“engine 已重启但权重尚未恢复”的状态或者等价标记。health_generate可能会与“engine 已重启但权重尚未恢复”的状态冲突;真实权重通过 disk update 加载完成后,再执行 generate 验证更合理。