[Feature] Support weight update from disk

[PengchengShi00]

从 Disk 更新 Rollout 权重方案

一、目标

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

该能力作为现有 IPC、NCCL 权重更新之外的第三种更新方式：disk。它主要服务于 server recovery，也可以作为后续权重同步链路的备选方案。

二、前置条件

1. rollout server 重启时可以先不加载真实权重。

   • LMDeploy 使用 empty_init=True 启动。
   • SGLang 使用 load_format="dummy" 启动。

2. 上游需要传递待加载的 HF checkpoint 路径。

   • checkpoint 必须已经完整写入，避免 rollout server 读取到半写状态。

3. LMDeploy 使用 disk 模式时，还需要知道底层复用 IPC 还是 NCCL。

   • update_weight 当前由 train controller 触发。
   • 如果 train worker 和 rollout worker 在同一台机器上，可以复用 IPC 路径。
   • 如果 train worker 和 rollout worker 在不同机器上，需要复用 NCCL/disaggregated 路径。
   • LMDeploy 根据 IPC/NCCL 选择已有的 tensor 更新接口。

4. SGLang 使用 disk 模式时不需要区分 IPC/NCCL。

   • XTuner 只调用 update_weights_from_disk。
   • SGLang server 自行从 HF checkpoint 路径加载权重。

5. recovery 场景需要明确本次哪些 rollout server 需要重新加载权重，避免影响仍在正常工作的 server。

   • 可以在调用 update_rollout_info() 时，让 server_url_dict 只包含重启后的 server URL（待讨论）

三、设计原则

• disk 是顶层权重更新模式，不等同于 IPC 或 NCCL。
• IPC/NCCL 是 disk 模式下 LMDeploy 可能复用的底层传输路径。
• SGLang 优先走 server-side disk reload，即 XTuner 只传递 HF checkpoint 路径，不显式加载权重。
• LMDeploy 由 XTuner 读取 HF safetensors 权重并 bucket 化，随后复用已有 IPC 或 NCCL 更新接口。
• 尽量保持现有 IPC/NCCL 行为不变，避免影响默认训练流程。
• server recovery 场景下，应支持只恢复失效的 rollout server group。

四、方案

1. 配置与元信息

在 RolloutWeightUpdateInfo 中增加 disk 相关字段：

• transport_type="disk"：表示使用从磁盘更新权重的顶层模式。
• hf_weight_path：待加载的 HF checkpoint 路径，从 rollout_config 或其 extra_rollout_config 获取。
• disk_update_upstream_transport：仅用于 LMDeploy，表示 disk 更新最终复用 ipc 还是 nccl 路径。

2. UpdateWeighter

update_rollout_info 不再接收 train_rollout_mode，而是直接接收 weight_update_mode：

• ipc：现有共卡 IPC 更新。
• nccl：现有训推分离 NCCL 更新。
• disk：新增从 HF checkpoint 路径恢复权重。

weight_update_mode 决定构造哪个 WeightTransport。

3. DiskWeightTransport

DiskWeightTransport 按 backend 委托给不同的 disk adapter：

• SGLang adapter 不消费 WeightIterator 产生的 tensor batch，而是直接使用 hf_weight_path。
• LMDeploy adapter 消费 WeightIterator.iter_disk_hf_batches() 产生的 tensor batch，并复用已有 IPC/NCCL transport。

SGLang adapter：

• 构造请求到 rollout server 的 update_weights_from_disk endpoint。
• 请求中携带 model_path=hf_weight_path。
• 由 train rank 0 给 rollout_server_url_dict 中的目标 rollout server 发送请求。
• 权重加载由 SGLang server 内部完成。

LMDeploy adapter：

• 根据 disk_update_upstream_transport 判断复用 IPC 还是 NCCL。
• IPC 场景下，适用于 train worker 和 rollout server 同机。
• NCCL 场景下，适用于 train worker 和 rollout server 跨机。
• WeightIterator.iter_batch_groups() 在 transport_type="disk" 且 backend 不是 SGLang 时，切换到 iter_disk_hf_batches()。
• iter_disk_hf_batches() 支持读取 HF model.safetensors.index.json 或单文件 model.safetensors，按 update_weight_bucket_size_in_gb 产出 WeightUpdateBatch。
• 产出的 batch 交给现有 IPCWeightTransport 或 NCCLWeightTransport，继续使用 LMDeploy 现有 update_weights / update_weights_from_distributed 接口。
• 对 LMDeploy/pytorch backend，最后发送空的 finished=True batch，用于触发 rollout 侧 finalize。

五、待讨论事项

• 从 disk 恢复 server 时，SGLang 路径本质上是向目标 worker 发送一次 update_weights_from_disk 请求，由 SGLang server 自行从 HF checkpoint 路径加载权重。该路径不会占用无关 train rank 的显存和传输带宽。
• 对 LMDeploy 来说，disk 模式只改变 WeightIterator 的权重来源：从训练侧模型参数导出变为从 HF checkpoint 读取；后续 IPC/NCCL 发送、server 接收和 finalize 流程与原有权重更新路径保持一致。
• 在训推分离场景下，LMDeploy disk 路径会由训练侧节点先读取并打包权重，再发送给推理节点。读取 HF checkpoint 的节点不一定是最终消费权重的节点，这里存在额外的读盘、打包和传输开销。
• disk + NCCL 会复用现有 NCCL 权重更新链路，需要关注它与常规 NCCL 权重更新通信组之间的冲突风险，例如 group name 复用、rollout 侧通信组残留、恢复流程与常规更新并发触发等问题。

lmdeploy 可做优化：

• 引入双 buffer，实现从 disk 读取下一批权重的同时向 server 发送当前批权重，减少磁盘读取与网络/IPC/NCCL 发送之间的串行等待。
• 为 disk + NCCL 使用独立的通信组命名或显式 teardown 策略，避免与常规 NCCL 更新链路共享未清理的通信状态。