Fix(opd): keep teacher weights across sleep/wake in colocate

[ZiyiTsang]

问题

本PR为实现PR#237发现并修复的bug。欢迎交流~

colocate 模式下 managed OPD/SDPO teacher 开启了 enable_memory_saver 但未开启
enable_weights_cpu_backup。teacher 与 actor 共享 GPU，每个训练 step 都经历
release_memory_occupation / resume_memory_occupation；无 CPU backup 时 resume
重新分配的权重页不恢复内容，teacher 自第一次 wake 起输出均匀分布——所有 token 的
logprob 恒等于 -ln(vocab_size)。

蒸馏信号因此完全失效，且方向系统性错误：尖峰 student 对均匀 teacher 的 JSD loss
初始即达理论上界（实测 step-0 train/loss = 0.688 ≈ ln 2），随后把 student 推向
均匀分布（entropy_loss 0.41 → 3.56），表现为多语言乱码、EOS 概率消失、回复全部
截断、eval accuracy 崩到 0、loss 伪收敛到 1e-4 量级。

触发条件

同时满足以下四项即触发（对应 Relax-managed teacher 的 colocate 形态）：

1. --use-opd --opd-type=sglang，且 teacher 走 Relax 托管路径：设置了
   --teacher-hf-checkpoint（单 teacher）或 --opd-teacher-routes（MOPD），
   并且 --resource 含 teacher 条目；
2. --colocate（非 --hybrid），且 --resource 同时含 actor 与 rollout——
   即 is_managed_opd_teacher_colocate 为真，teacher 与 actor 共享 placement group；
3. teacher 随训练步做 sleep/wake：--offload-rollout（colocate 默认开启）使
   teacher 与 actor 锁步 offload/onload，每个 step 经历
   release_memory_occupation → resume_memory_occupation。此时
   build_teacher_overrides(colocate_sync=True) 把 enable_memory_saver 置 True；
4. 用户未显式传 --teacher-sglang-enable-weights-cpu-backup（sglang 默认关闭）。

不触发的形态：独立 PG 的 teacher（enable_memory_saver=False）、外部
--opd-teacher-url 自建 teacher（不经 build_teacher_overrides）、非 OPD 训练。
命中时所有 colocate OPD/OPSD/SDPO 配置在第一个训练步后 teacher 即失效，
蒸馏信号系统性错误且无任何报错。

修复

overrides.setdefault("enable_memory_saver", colocate_sync)
if overrides.get("enable_memory_saver"):
    # Release/resume without a CPU backup leaves weight pages uninitialized (uniform output).
    overrides.setdefault("enable_weights_cpu_backup", True)

• setdefault：显式 --teacher-sglang-enable-weights-cpu-backup 仍可覆盖
• 非 colocate（独立 PG）teacher 不受影响（enable_memory_saver 为 False，guard 不触发）
• 非 OPD 算法不受影响：build_teacher_overrides 唯一调用方是
  relax/distributed/ray/teacher_manager.py 的 TeacherManager，仅 OPD/OPSD/SDPO 路径创建
• 代价：colocate teacher 节点多占一份权重 CPU backup（8B bf16 ≈ 16GB 内存）及
  release/resume 各一次拷贝
• GenRM 路径（sglang_engine.py）已使用同一 flag，本 PR 把 managed teacher 对齐

验证

1. 确定性复现（Qwen3-0.6B + --enable-memory-saver，同 venv sglang 0.5.12.post1）：
   一次 release→resume 后全部 logprob = -11.9312 = -ln(151936)；开启
   enable_weights_cpu_backup 后 3 个 sleep/wake 循环 logprob 逐位不变（max|Δ|=0）。

2. 真实训练（SciKnowEval Biology, Qwen3-8B, 4×H100 colocate）前后对比（wandb）：

   run	状态	步骤	step-0 train/loss	末尾 train/loss	eval acc 轨迹	rollout acc
   修复前 sdpo-sciknoweval-biology-2026-08-20-02:59:55	失败	~189	0.688（≈ln 2 上界）	1e-4（伪收敛）	0.361 → 0.011 → 0.0	0.406 → 骤降
   修复后 sdpo-sciknoweval-biology-2026-08-21-06:45:37	成功	~400（取前 400）	0.026	0.014	0.347 → 0.51（稳步上升）	最高 0.77

   修复前 step-0 train/loss = 0.688 ≈ ln 2 即 JSD 理论上界，说明 student 与
   （均匀化）teacher 分布几乎不相交；随后 loss 单调降到 1e-4 是 student 被均匀化的
   伪收敛。修复后 loss 一开始就落在 ~0.03 以下（真实 hint-only JSD 量级），eval
   accuracy 从 0.35 一路升到 0.51，rollout accuracy 最高 0.77，全程无崩塌。

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[kkyyxhll]

感谢交互
首先说明，这个对目前opd的recipe 并没有任何影响

1. --sglang-enable-weights-cpu-backup Relax默认teacher继承student的sglang args，因此opd recipe中都配置了--sglang-enable-weights-cpu-backup，老师会自动继承，具体逻辑在 https://github.com/redai-infra/Relax/blob/4899b8f3a90489840a736897b4c341d87c6267cf/relax/utils/opd/opd_utils.py#L150 。并且Relax提供的opd_recipe的eval都是正常的，不仅仅是loss的收敛,具体可以参见这个PR里的曲线，chore(opd): update Math OPD training recipe #84 。
2. 我检查了你的pr里面的脚本，都没有配置--sglang-enable-weights-cpu-backup，所以踩到了这个坑，我们会额外增加一些鲁棒性说明。
3. 当前pr中改动可以合，但需要：(a) 重写问题陈述与触发条件、降级严重性定性；(b) 保留一个关闭 CPU backup 的途径（或明确 warning）。
   @NINGBENZHE @Yangruipis cc。

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.