[Feat] Support immediate rollout recovery by export hf asynchronously

[YanhuiDua]

背景

当 rollout backend 在两次权重同步之间发生故障时，原有流程会先将故障 worker group 标记为 inactive，等到下一次权重同步前再以skip_load_weights=True 的方式重启，并通过本轮 update_weights() 补齐模型权重。这意味着故障 worker 在下一次权重同步之前无法恢复服务。

本 PR 新增一条可选的即时恢复路径：训练侧提前异步导出可用于恢复的 HF 模型；当 RolloutHealthManager 检测到 rollout worker故障时，可以直接加载最近一次已经就绪的 HF 模型并恢复 worker group，不必等待下一次权重同步。

核心改动

• 在 RL Trainer 配置中新增 enable_immediate_recovery，默认关闭。
• 在 TrainingController 和 TrainingWorker 中新增异步 HF 导出接口：
  • start_hf_export()
  • is_hf_export_done()
  • wait_hf_export()
• 支持 colocate 和 disaggregated 两种 RL Trainer。
• 将 ready recovery HF 的管理和故障恢复统一放在 RolloutHealthManager 中。
• 扩展 RolloutWorker.reinit()，支持在重启时覆盖：
  • model_path
  • tokenizer_path
  • skip_load_weights
• 增加受环境变量保护的 backend crash 注入接口，用于恢复单测和 E2E。
• 增加即时恢复、慢导出回退、多 worker group 恢复等单测。
• 增加独立的 8-GPU Qwen3.5 VLMoE 异步 HF 恢复 E2E。

当前恢复流程

整体流程如下：

完成一次 train -> rollout 权重同步
    |
    +-- 当前 step 已执行常规 HF 保存
    |       |
    |       +-- 直接复用 hf-step-N 作为 recovery HF
    |
    +-- 当前 step 没有常规 HF 保存
            |
            +-- 后台异步导出 hf-step-N
                    |
                    +-- 导出完成后发布为 ready recovery HF
                            |
                            +-- 唤醒 RolloutHealthManager

RolloutHealthManager 检测到 worker group 故障
    |
    +-- 存在 ready recovery HF
    |       |
    |       +-- 关闭故障 group
    |       +-- 使用 recovery HF 重新启动
    |       +-- 执行健康检查
    |       +-- 恢复为 active
    |
    +-- 不存在 ready recovery HF
            |
            +-- 暂不自动重启，保持 inactive
            +-- 等待 recovery HF 发布，或者等待下一次权重同步

另外，在每次开始下一轮权重同步前，Trainer 都会检查上一次异步 HF 导出：

• 如果导出已经完成，则继续使用该 HF 作为 ready recovery HF。
• 如果导出仍未完成：
  1. 等待该导出任务结束；
  2. 关闭本次训练余下阶段的 immediate recovery；
  3. 清除 ready recovery HF；
  4. 后续恢复回退到“空权重启动 + 下一次权重更新”的原有路径。

E2E 测试

新增的 backend crash 接口用于故障注入来跑测试，本e2e测试共运行三步，三步的global batch size分别为8，256，8，第二步需要测试恢复时间，所以将global_batch_size调大
XTUNER_TEST_IMMEDIATE_RECOVERY=1 python .dev_scripts/test_qwen35_vl_moe_async_hf_recovery_e2e.py

运行该脚本的异步HF保存与恢复性能如下：

指标	Step 1 首次导出	Step 2 后续导出
Export launch	39.70 s	2.50 s
Writer 写盘	约 16 s	约 17 s
Finalize + publish	约 2 s	约 4 s
启动到 ready	约 58 s	约 23 s
GPU显存增加	10.4G	10.4 G

该配置下，恢复完需要110s左右，GPU显存增加相对于普通的 save hf 峰值再高 5.43G，增加的额外显存来自于 HF snapshot 阶段的 FSDP/EP all-gather、flatten output

TODO

• 当前自动恢复同步运行在 HealthManager 的健康检查线程上，在一个 worker group 执行 server shutdown、HF reload 和健康检查期间，后续周期性健康检查会被阻塞。后续需要将耗时的 group recovery 移到独立线程或异步任务中
• 没考虑rollout马上就结束时就不需要恢复挂掉的rollout worker，反而会拖慢训练

[YanhuiDua]

@claude review

[claude]

Claude: ## Summary

This PR adds an "immediate recovery" mode for rollout workers: instead of restarting crashed workers with empty weights and waiting for the next weight-sync, it exports an HF checkpoint asynchronously after each weight update and lets the health manager reload workers from that checkpoint immediately.

The refactoring of _skip_load_weights_during_restart into reinit(**kwargs) is a nice simplification. The overall design (background export → publish to health manager → use on restart) is sound.

Issues

Critical

• xtuner/v1/train/rl_trainer.py:1746-1748: close_trace() removed from both fit() methods (and its import dropped). The trace runtime won't flush pending spans at end of training. This appears accidental — the executor shutdown replaced it rather than supplementing it.

Warning

• rl_trainer.py:895-898: Race between rmtree of old recovery path and in-flight health manager restarts that may still be loading from that path.
• rl_trainer.py:935-938: self._ready_recovery_hf_path written from background thread — correctness depends on CPython GIL + Future synchronization. Fragile for maintainers.
• rl_trainer.py:944-951: Feature permanently disabled on first slow export with no re-enablement path.
• rl_trainer.py:2121: New restart_inactive_workers call in disaggregated path — was previously a TODO indicating it needed design evaluation.

Nit

• trainer/worker.py:957-959: cast instead of explicit None guard in is_hf_export_done/wait_hf_export.

Verdict

REQUEST_CHANGES

ProduceBatchResult impact: not directly affected — the recovery path rejoins the normal flow before batch accounting.

RoutedExperts impact: not affected — no changes to routed-experts plumbing or object-ref ownership.

[YanhuiDua]

@claude review