rose-online-k1024t4096

Online ROSE — K=1024 student prefix + 4096 teacher continuation (Qwen3-32B teacher)

复现"在线 ROSE"蒸馏的一个自包含实验包:Qwen3-1.7B 学生写 1024 token 前缀, Qwen3-32B teacher 接手续写 4096 token,只在 teacher 写的那段上算 CE。

配套完整的评测工具链(含把 teacher 放到推理时用的交错评测),以及一份把这套方法 在 RLVE 上跑了两周得到的全部对照数据。

---

这个配置从哪来

推理时交错实验(eval/interleave_eval.py,学生未经训练、teacher 现场在场)在 RLVE 180 题、16k 口径、n=8 下扫了一个 K1×K2 网格:

keep 率 = avg@8,即"学生最终答对"的轨迹占比。这个数同时是 VERIFY_MODE=keep_correct 下能存活的样本比例 —— batch 256 在 1024+4096 下只剩约 35 条有梯度。

两个方向都很清楚:K1 越小越好(学生自己写的前缀是薄弱环节,越早交给 teacher 越好),K2 是主变量(teacher 写得越多越好,1024→4096 让 pass@8 从 0.2667 跳到 0.4111)。

而训练这边一直固定 K=4096:

也就是说训练侧长期停在推理侧证明最差的那一列上。这个包补的就是推理侧最优的 那一格(1024+4096)在训练侧成不成立。

⚠ 一个必须先说的反向证据

训练与推理的最优方向在已有数据上是相反的:

† 这一格是 4B-Thinking teacher 下测的,两边 teacher 不同,不是严格单变量。

所以这条线有可能得到负结果。那本身也是结论:说明 ROSE 的训练目标(CE 只加在 teacher 的 token 上,学生前缀被 response_mask 置 0、一次梯度都收不到)和 "推理时让 teacher 多写"根本不是一回事 —— 学生学到的是"接着 teacher 往下写", 而推理时它要从 token 0 自己走完全程。

---

方法

prompt ──[student 写 K 个 token]──[teacher 续写 T 个 token]──> 序列结束
          mask = 0 (只当上下文)     mask = 1 (CE 的目标)

response_mask 就是全部机制。verl 的 tool agent loop 本来就会给"模型没生成的 token"写 0,这里只是把角色对调:学生的 token 是上下文,teacher 的是目标。

teacher 是集群外一个普通的 vLLM OpenAI server,只做生成不做打分,所以 reward_model.enable=False,整条 compute_rm_score 路径在这个 recipe 里是死代码。 启动 teacher 必须带 --return-tokens-as-token-ids:续写以 exact token id 返回, id 接 id 拼接,接缝不会被重新分词挪位 —— 这是这个设计唯一的静默失效点。

---

目录

run.sh                              一键: 起 teacher -> 训练 -> 16k pass@8 评测
scripts/
  run_rose_1p7b_32b_k1024t4096.sh   ★ 本实验的配置
  run_rose_online.sh                底层 driver(拼 hydra override)
  bootstrap.sh                      teacher server / train 入口
recipe/rose/                        需要拷进 verl/recipe/ 的文件
  rose_agent_loop.py                ★ 核心: 前缀+续写+mask, 以及 verifier 筛选
  rlve_reward.py                    RLVE verifier 的 verl custom_reward_function 包装
  config/rose_agent.yaml
eval/
  rlve_eval.py                      主评测(16k pass@k)
  interleave_eval.py                ★ 交错推理评测: 学生-teacher-学生 三段
  retrunc_eval.py                   把 32k rollouts 精确截断到小 budget 重判(纯CPU)
  merge_eval.py                     合并分片评测
data/
  rlve_train.parquet                9000 条训练 prompt
  rlve_test.parquet                 180 题测试集
RLVE-Eval/Gym/                      verifier 环境(评测与 verify_mode 都要用)
requirements.txt

不含 verl 本体(134MB)。需要一个装了 recipe/rose 的 verl;把 recipe/rose/ 整个拷进你的 verl/recipe/ 即可。

---

跑起来

# 0. 环境
pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt --no-build-isolation
cp -r recipe/rose /path/to/verl/recipe/

# 1. teacher server(占 GPU 0-3,常驻)
MODEL_DIR=/path/to/models TEACHER_REPO=Qwen/Qwen3-32B \
TEACHER_GPUS=0,1,2,3 TEACHER_TP=4 TEACHER_MAX_MODEL_LEN=12288 TEACHER_MAX_SEQS=64 \
  bash scripts/bootstrap.sh teacher

# 2. 训练(占 GPU 4-7,70 步)
bash scripts/run_rose_1p7b_32b_k1024t4096.sh

# 或者一键
bash run.sh

形状怎么推导的

MAX_RESP_LENGTH = PREFIX_MAX + TEACHER_MAX_TOKENS = 1024 + 4096 = 5120
MAX_MODEL_LEN   = MAX_RESP_LENGTH + MAX_PROMPT_LENGTH(2048) = 7168
teacher 侧输入   = prompt(≤2048) + 前缀(1024) = 3072,加 max_tokens 4096 → 7168

PREFIX_MAX 必须跟着 CUT_FIXED 一起改。 只改 CUT_FIXED 的话 MAX_RESP_LENGTH 仍按 PREFIX_MAX=4096 推导成 8192,响应张量按 8192 padding, 白占一倍 KV cache。

vLLM 的 context 坑

teacher server 的 --max-model-len 必须 ≥ prompt + 前缀 + TEACHER_MAX_TOKENS。 vLLM 的账是 allowed_input = max_model_len - max_tokens,所以 TEACHER_MAX_TOKENS 一涨,max_model_len 要跟着涨,否则请求直接 400。K=4096 且 t=4096 时默认的 8192 不够,需要 12288。

---

可选:inference-time verifier 筛选

rose_agent_loop.py 支持在 teacher 续写之后把控制权交还给学生,让它写到有 答案,用 RLVE verifier 判分,只保留答对的行:

prompt ─[student K]─[teacher T]─[student 写到有答案]─→ verifier
         \_________ 训练 ______/  \___ 只用于判定, 丢弃 ___/

VERIFY_MODE=keep_correct VERIFY_STUDENT_TOKENS=14336 bash scripts/run_...sh

答错的行 response_mask 全置 0,不产生梯度。第三段不参与训练,所以 MAX_RESP_LENGTH 不变,但 MAX_MODEL_LEN 必须装下 prompt+K+T+续写 (脚本里这两个长度已经解耦)。

代价,不小:

• 每条 rollout 多一次约 14k token 的生成,实测每步从 48 秒涨到 400-500 秒
• 存活比例就是学生自己的准确率。交错实验测得 1024+4096 是 13.82%, 也就是 batch 256 只剩约 35 条有梯度。盯 rose/verify_kept。

---

全部对照数据(RLVE 180 题,16k 口径,CTX=18432 MAX_NEW=16384 N=8 temp0.7 top_p0.9)

teacher = Qwen3-32B,student = Qwen3-1.7B(Instruct)

OPD 领先在线 ROSE 7.2 个点(pass@8)。 差距在 avg@8 上更悬殊(0.0847 vs 0.0521,高 63%),说明 OPD 提高的是单次可靠性,ROSE 更依赖多采样。这与两者的梯度 覆盖一致:OPD 每条 rollout 4096 个 token 全进 loss,ROSE 只有 teacher 写的那 1024 个。

teacher = Qwen3-4B-Thinking-2507

1.7B-Base 学生的 K 消融(32B teacher)

倒 U 型,K=1024 最优;3 epoch 反而更差。注意整组绝对值都极低 —— Base 模型基本 不会 instruction following,OPD 在 Base 上更是直接崩掉(pass@8 0.0056,fmt 0.037)。

蒸馏损失:同配方,推理时用 vs 蒸馏进模型

teacher = 4B-Thinking, K=4096 + t=1024
  推理时交错(teacher 在场)   pass@8 0.2278
  蒸馏后的 1.7B(单模型)      pass@8 0.1833   <- 保住 62% 的增益
  未训练的 1.7B              pass@8 0.1111

---

关于成本

ROSE 贵 1.95 倍的 GPU-hours、8 倍的单位监督成本,pass@8 还低 7 个点。

根因是架构而非算法:4 张卡钉给常驻 teacher,而 teacher 与 student 严格串行、 零时间重叠,44% 的算力在等待。

在线 ROSE 的 29.2 GPU-h
  teacher (g0-3)   占用 14.6   实际在算 ~9.4   空转 5.2 (36%)
  student (g4-7)   占用 14.6   实际在算 ~6.9   空转 7.7 (53%)

代码本身已经是全异步的(run() 是 async def,teacher 走 aiohttp,批内 256 条 rollout 由 asyncio.gather 全并发)。零重叠是参数造成的相位同步: CUT_FIXED 让每条 rollout 的 K 完全相同,于是它们同时生成、同时结束、同时打 teacher。把 CUT_FIXED 去掉、回落到默认的 K~U[PREFIX_MIN,PREFIX_MAX] 就能打散 相位 —— 这也是唯一一个还没在 RLVE 在线版上试过的配置(历次实验全部 pin 死了 K)。

---

已知的坑

1. bootstrap.sh 的 teacher 默认不是 32B。 TEACHER_REPO 默认值是 Qwen/Qwen3-4B-Thinking-2507。不显式 export 就会静默用 4B。 核对某条线用了什么 teacher,不要只看脚本(脚本可能事后改过),要查 teacher*.out 里的 teacher: /path/to/XXX 那行实际加载记录。

2. 在线 ROSE 长期没有 LR schedule,这是遗漏不是设计。 run_rose_online.sh 以前只传 optim.lr,verl 默认 lr_scheduler_type=constant、 lr_warmup_steps_ratio=0.0,训练日志实测 actor/lr 从 step0 到 step69 恒为 1e-5。 而离线迭代版一直带 cosine + 10% warmup ——"在线 vs 离线"的历次对比里这是个未受控 变量。现已补上开关(SCHEDULER / WARMUP / MIN_LR_RATIO),默认保持旧行为。

调度粒度:fsdp_workers.py 在 update_actor 末尾每个 training step 走一格, num_training_steps 从 trainer.total_training_steps 注入。 ppo_mini_batch_size == train_batch_size 时一步一更新一格,70 步 = 70 格,余弦在 最后一步走完。把 mini_batch 调小则一个 tick 内会有多次 optimizer step 而 lr 不变。

3. 评测口径统一用 16k。 32k 与 16k 下各方案的排序不一致(例:iter ROSE t1024 r1 在 32k 下 pass@8 0.3389 居首,16k 下只有 0.1556 落到 OPD 之后)。已有 32k rollouts 的话用 eval/retrunc_eval.py 16384 <SRC_TAG> <DST_TAG> 精确截断重判即可 —— vLLM 的 max_tokens 不影响前面 token 的采样分布,这是精确重测不是近似,纯 CPU 每条约 10 秒。 实测同一模型原生 16k 与截断重判差约 0.011,在采样噪声内。

4. checkpoint 会撑爆磁盘。 run_rose_online.sh / run_opd_baseline.sh 都没设 max_ckpt_to_keep, 70 步 × 7 个 checkpoint × 11G ≈ 77G。建议加 trainer.max_ckpt_to_keep=2, 或用一个守护脚本只保留最新一个 step 的 optimizer 分片。

5. teacher 几乎从不自然收尾。 所有档位的 teacher_finish 里 stop 都只有 1.5% 左右(K2=1024 时是 0%)。 32B 在 RLVE 上就是写不完,给到 4096 也一样。所以交错的增益与"teacher 有没有写完" 无关,只与"注入了多少 teacher 内容"有关 —— 别指望靠加大 TEACHER_MAX_TOKENS 让它收尾。