feat(rollout): 为异步 RS 精确补齐 prompt-group batch

[inmnin]

摘要

感谢维护 slime。这个 PR 为 disaggregated train_async.py 增加一个默认关闭、fail-closed 的
--rs-batch-refill 路径，用于在 sequence-level / geometric rejection sampling（RS）之后，按完整
prompt group 精确补齐 optimizer step 的 effective batch。

核心目标不是把异步数据变成 on-policy，也不是承诺无条件提速，而是保证：启用该开关后，每个
optimizer step 要么拿到配置数量的完整、通过 RS admission 的 prompt groups，要么在更新前明确失败；
不会静默用 underfilled batch 更新。

背景与问题边界

序列级 / geometric RS 会拒绝训推概率差异过大的 sequence。原始 rollout 的结构 batch 仍然存在，
但 RS 后真正参与梯度计算的有效 sequence 数会波动；按本 PR 采用的 group-atomic 语义，完整且组内所有
sample 均通过 gate 的 prompt group 数也可能不足。当前 actor-side RS 路径不会自动补齐这些有效 group。

这里需要特别说明与 #657 的边界：#657 将 importance sampling
weight 与 rejection sampling mask 解耦；当前 main reducer 的 numerator 使用 modified mask，但 denominator
仍显式使用原始 rollout_mask_sums。本 PR 不依赖重新归一化，而是最终只接纳 gate 全部通过、mask 保持
不变的完整 group，并恢复每个 optimizer step 固定的 prompt-group cardinality。

本 PR 选择严格的 group-atomic 语义：组内任一 sample 未通过 admission，则整个 prompt group 都不进入
本次训练。这样不会通过只补单个 sample 改变 GRPO group 的结构。

最直接的补法是在 underfill 后重新生成完整 rollout_batch_size。但 long-horizon / tool-use rollout 的
尾部延迟通常很高，整批重跑会生成大量最终不会参与训练的轨迹。这里改为由 coordinator 计算真实缺口，
只为 DP/VPP 可调度性做最小向上取整，然后请求对应数量的 replacement groups。

一个最小例子：rollout_batch_size=2，每个 prompt 采 2 个 samples。group A 的两个 samples 都通过，
group B 只有一个 sample 通过。启用本 PR 后，B 会按组拒绝，并只请求 1 个 replacement group；replacement
通过后才以 2 个完整 groups 更新。若达到补采轮次上限仍不足，则在 optimizer.step 前失败。

实现方案

该功能目前只接入 disaggregated train_async.py，执行顺序如下：

1. 首轮只生成配置要求的 rollout_batch_size，本层不做投机式超采样。
2. actor 使用当前 proximal policy 做一次 preflight forward，重新计算 logprob，并执行 sequence/geometric
   RS admission。
3. coordinator 以完整 prompt group 为原子接纳或拒绝候选。
4. 根据已接纳 group 数计算缺口，只请求缺少的 replacement target；target 仅向上取整到 DP/VPP 所需的
   最小 group multiple。
5. replacement 使用当前 rollout policy 生成，再执行同样的 actor preflight 和 group admission。
6. 最多执行 --rs-refill-max-rounds 轮；达到上限仍不足时，在 selected proximal-logprob cache transfer
   和 optimizer step 之前失败。
7. 最终 training forward 继续应用 TIS，并重新执行、校验与 preflight 一致的 RS gate。

preflight 产生的 proximal logprob 先保存在各 actor rank 的 pinned CPU memory 中。group decision 完成后，
只把最终接纳项通过 Ray 传给 RolloutManager；未接纳缓存立即清理。正式 training forward 仍然正常执行，
最终接纳缓存作为 PPO proximal old policy 使用。该设计减少的是跨 actor/manager 的无用传输和最终驻留，
并没有省掉候选 preflight 或 GPU-to-CPU copy。

除主流程外，本 PR 还包含两个由 refill 调用方式暴露出来、各自有回归测试的兼容性修复：

• SGLang GenerateState singleton 提交任务时使用本次 generate_rollout_async 的 per-call args，而不是首次
  初始化时缓存的 args；
• global rollout data source 在单次取样跨越 epoch 时可继续补足请求数量。

正确性与保护条件

补齐数量不会消除 off-policy。本实现有意把第一版范围收窄到 one-step async：

• 强制 --update-weights-interval 1、每个 rollout batch 恰好一个 optimizer step，并继续强制内置 TIS；
• 初始候选只接受 policy staleness [0, 1]；
• replacement 的 rollout policy version 必须与本轮 actor preflight version 完全一致；
• sample ID、group ID、rollout ID、policy version、loss mask、response shape 和 proximal cache 都做严格校验；
• cache 与最终训练样本必须双射，缺失、重复或多余项都会在更新前报错；
• preflight 与正式训练会重复执行同一 RS gate，并验证最终 mask 未被中间 hook 修改；
• refill coordination 的 RPC 有显式超时；异常路径会尝试清理 actor cache 和 manager pending state，同时保留
  原始异常；candidate generation 继续使用 rollout backend 自身的 timeout / health-monitor 语义；
• checkpoint 保存边界不会把已经 prefetch、但尚未对应模型更新的数据源状态写入 checkpoint；
• refill 跨越数据集 epoch 时，data source 不会在单次补采中途错误结束。

针对审计中发现的可变参数与无效传输风险，当前版本还增加了以下保护：

• 每次 candidate generation 都向 custom rollout function 及其 sample hooks 传入 deepcopy(args)，防止其对
  scalar 或 nested config 的修改泄漏到后续 refill round，进而改变 batching、topology 或采样行为；
• custom rollout logger 同样接收独立的 deepcopy(args)；logger/debug saver 前后重新转换 train data，并用
  fingerprint 验证 policy inputs 未变化；
• coordinator 一旦确认 retry budget exhausted，立即报错并进入清理，不再先提取 actor cache 或向 manager
  传输、存储注定不会进入 optimizer step 的 proximal logprob。

为了 fail-closed，参数校验会拒绝尚未证明安全的组合，包括 checkpoint resume、persistent fully-async
queue、token-level RS、external rollout engine、dynamic sampling filter、partial rollout、offload/release、
量化 Megatron、dropout、context parallel、自定义 model provider / Megatron before hooks / TIS / advantage /
loss reducer / train-data converter / reward postprocessor，以及非完整 NCCL 权重同步等。显式
fp8=false / fp4=false 不会被误拒绝。

使用方式

python train_async.py \
  ... \
  --rs-batch-refill \
  --rs-refill-max-rounds 2 \
  --rs-refill-rpc-timeout-seconds 1800 \
  --rs-refill-max-candidate-cache-bytes 1073741824 \
  --update-weights-interval 1 \
  --use-tis \
  --custom-config-path examples/train_infer_mismatch_helper/mis_refill.yaml

examples/train_infer_mismatch_helper/README.md 和 mis_refill.yaml 已补充完整约束与示例；train.py
会明确拒绝该参数，避免误以为 colocated path 已支持。

--rs-refill-max-candidate-cache-bytes 限制的是 proximal-logprob tensor payload，不包含 Python container、
Ray object store / transport buffer，也不会预留 CUDA allocator headroom。最终 batch 的 DP-local proximal
logprob shard 仍可能在 training rank 上 materialize，因此 long-response 任务需要额外预留 GPU 显存；文档已
明确这一点。

请求量对比（非 GPU benchmark）

下面是直接调用本 PR planner 得到的 coordinator target，对比对象是“underfill 后重新请求完整
64-group batch”的朴素补齐策略。它不是 wall-clock benchmark，也不用于声称端到端必然提速。

配置：rollout_batch_size=64 groups、n_samples_per_prompt=8、DP=8、VPP=2、每个 VPP stage 的
microbatch group size=4，因此最小 refill quantum 为 4 groups。

实际缺口（groups）	本 planner target	整批重采 target	target 减少
1	4	64	93.75%
5	8	64	87.50%
16	16	64	75.00%
缺口 1..64 的均值	34	64	46.88%

当前 upstream actor-side RS 的行为是“不恢复有效 group 数”，并不是“整批重采”；上表只比较两种可能的
补齐策略。target 是本层传给 rollout function 的请求量，不等同于底层实际 generation task 数。实际耗时
取决于 RS 拒绝率、长尾 rollout、preflight 成本和训推配比；低拒绝率下，本实现仍增加一次 actor preflight，
因此它首先是 effective-batch 正确性选项，而不是无条件性能优化。

可观测性

新增指标覆盖：

• initial / replacement generation time 与 coordinator / preflight time；
• 每轮 candidate、累计 rejected / surplus、accepted groups、gate acceptance rate 和 selection utilization；
• actor candidate cache、aggregate/peak cache、selected transfer、manager retained cache 的字节数与耗时；
• effective trainable tokens、refill path time 和 effective tokens per refill-path second；
• initial policy staleness 与 refill round 数。

这些指标用于真实训练中判断补齐是否值得，而不是在没有 GPU A/B 时替代性能实验。

验证

仓库 CPU 回归

在最终候选 commit 上，按 .github/workflows/pr-test.yml.j2 的 cpu-unittest 与 agent-test 列表，
对 62 个测试文件逐文件启动独立 pytest 进程，避免文件级 module stub 污染：

PYTHONNOUSERSITE=1 \
PYTHONDONTWRITEBYTECODE=1 \
CUDA_VISIBLE_DEVICES='' \
PYTHONPATH="$REPO" \
python -m pytest -q -p no:cacheprovider <test-file>

结果：

项目	结果
workflow CPU 测试文件	62 / 62 通过
pytest cases	703 passed, 1 skipped, 0 failed/error
唯一 skip	lean CPU 环境未安装 SGLang，tests/test_agent/test_adapters.py 按设计 pytest.importorskip("sglang")
test_rs_exact_refill_rollout_manager.py	23 passed

另外完成：

• pre-commit run --all-files：全部 hooks 通过；
• 变更 Python 文件 py_compile、Ruff、Black 与 git diff --check：通过；
• 上表跳过的 adapter 文件另在已安装 SGLang 0.5.12.post1、隐藏 CUDA 的环境中复跑，
  结果为 15 passed；新增的两个 SGLang refill regression 也在该环境中复跑，结果为 2 passed；
• 未随 PR 提交、用于真实 A/B 的 runner CPU harness：113 passed, 3 warnings。

真实 4-GPU 配对 A/B

GPU A/B 在祖先 commit a420a7c8cd8ccf000540447276aa57e21eef029e 上完成一次 clean CAL → A → B。
当前 PR HEAD eba72613258164c5fb6ad54770efe4908a74e77d 相比该实验 commit 只移动了一个测试文件的
__main__ guard，以兼容 workflow 的直接脚本入口；生产代码没有变化。
此前受到非实验 GPU 进程干扰的尝试全部作废，没有混入下面的数字。

项目	固定配置
模型	Qwen3-4B，同一 torch-dist checkpoint，实验结束再次逐文件校验 SHA-256
数据	BytedTsinghua-SIA/DAPO-Math-17k 的固定 12-row 小型 workload：4 个 initial rows（2 个 distinct prompts 各重复一次，用于两个 rollout IDs）+ 8 个 distinct refill reserve rows，共 10 个 distinct prompts；输入 SHA-256 3d1be7cf...5e8de2
评测	BytedTsinghua-SIA/AIME-2024 的 30 个 distinct prompts，仅作 integration smoke
GPU 拓扑	2 张 actor GPU（TP=2）+ 2 张 rollout GPU（两个 TP=1 SGLang engines）
batch	2 prompt groups × 4 samples，GBS=8
训练	CAL → A（feature off）→ B（exact refill），各 2 steps，response length=512
异步约束	TIS，--update-weights-interval 1，固定 seed 与 deterministic inference

A/B 共享同一个运行时冻结的 RS 配置，SHA-256 都是 1d4162ee...a8c8。两者除 B 的
--rs-batch-refill 与 A 侧 measurement-only MIS scalar wrapper 外保持一致；归一化命令校验会忽略这两个
预期差异并拒绝其他差异。两步 initial candidate 的 prompt、response、reward、loss mask、rollout logprob、
policy version 和 normalized training content 均在同一 HMAC key 下逐项配对，CAL/A/B 完全一致。
实验 Megatron 环境固定在 revision 3714d81d...2862，其已知 working-tree patch 也单独归档并绑定
SHA-256 0f7295e5...1048；它不是本 PR 的一部分，且 A/B 共用完全相同的环境。

CAL 从同一批自然产生的 stale log-ratio 中选择一个 float32 安全边界：fresh step 拒绝 0/8，stale step
拒绝 1/8，并通过预设的至少 16 个 float32 ULP 间距检查。冻结配置回放证明 A 与 B
得到逐元素相同的 admission mask，因此不是为 B 单独选择更有利的 gate。

stale step 指标	A：feature off	B：exact refill
initial candidate groups / samples	2 / 8	2 / 8
RS 拒绝	1 / 8 sequences（12.5%），影响 1 / 2 groups	1 / 2 groups
replacement 请求	无	1 group / 4 samples
refill rounds	无	1
累计 scored / rejected / accepted groups	不补齐	3 / 1 / 2
surplus groups	无	0
最终 train forward 的 low/high RS mask	有 1 个 sequence 被 mask	0 / 0
stale-step grad norm	0.05251	0.04387

B 的 initial candidate policy version 为 1，actor preflight/scoring version 为 2，replacement version 也
严格为 2；这验证了 replacement 来自当前 policy，而不是重新使用 stale policy。最终 accepted batch 恢复到
2 个完整 groups，且正式 training forward 没有再次拒绝任何 selected sample。A/B fresh step 的 grad norm
也完全一致，均为 0.0628146231，进一步验证配对起点。

真实分项计时如下。它们覆盖 generation + coordinator/preflight/refill path，不等同于完整 optimizer step
吞吐：

B 路径	fresh step	stale + one-refill step
initial generation	15.536 s	14.364 s
preflight	2.036 s	0.464 s
replacement generation	0	14.246 s
coordinator	2.062 s	14.733 s
refill path total	17.598 s	29.097 s
effective trainable tokens	4096	4096
scored trainable tokens	4096	6144

其中 coordinator 是包含 preflight、replacement generation 与 selected-cache transfer 的包容性计时，不能与
这些子项再次相加；refill path total = initial generation + coordinator。coordinator 扣除这些已单列子项后，
两步 residual 分别只有 0.017 s 和 0.018 s。stale step 的 actor/manager candidate cache 峰值为 16 KiB、累计候选 payload 为 24 KiB，
远低于配置的 256 MiB 上限。

本次真实缺口为 1 group。exact refill 返回 4 个 replacement sample slots；若在同一轮按完整 2-group batch
盲目补采，则会返回 8 个 slots。因此这里能报告的是 replacement request slots 减少 50%（8 → 4）。
这是固定观测轮数下的结构性请求量对比，不是实测的 wall-clock speedup，也不推断底层 dynamic-sampling
attempts 或 generated tokens。相对于 feature-off A，B 在该 stale step 实测增加了 14.246 s replacement
generation，以恢复 exact batch；A 不请求 replacement，仍以包含 1 条被 mask sequence 的 nominal GBS=8
batch 更新。二者训练语义不同，因此不将该耗时差解释为通用性能结论。

phase 级计时中，A 的 compute / teardown 为 107.776 / 6.475 s，B 为 124.420 / 6.603 s。为避免人工选取
GPU 交接窗口污染结果，CAL/A/B 的 resident HOLD 分别为 86.936 / 57.078 / 204.059 s，均单独记录并从
elapsed_seconds = compute + teardown 中排除。

其他验收结果：

• CAL/A/B GPU telemetry 分别覆盖 75 / 89 / 102 个四卡 snapshots，error fraction 均为 0；
• AIME smoke：A 0/30、B 1/30，两者 truncated ratio 都是 96.67%。因为只有一次采样且 response
  上限为 512，这些数字只证明 eval 管线跑通，不用于声称质量提升或不掉点。

这次短程受控实验验证的是 underfill 复现、exact group refill、staleness/version 链、工程开销与失败边界。
它不构成收敛实验，也不支持“端到端训练一定加速”或“off-policy 质量完全无损”的结论。

与现有 PR 的关系

• #657 将 IS weight 与 RS mask 解耦；当前 main reducer 的
  numerator 使用 modified mask，而 denominator 仍显式使用原始 rollout_mask_sums。本 PR 不依赖改变
  reducer 的归一化语义，而是只保留 gate 全部通过且 mask 不变的完整 group，并补回固定 group cardinality。
• #2084 是开放中的通用 off-policy IS hook，让 correction 可使用
  current policy 与 rollout policy 的比值。本 PR 明确限制为每 rollout 一个 optimizer step、使用内置 TIS，
  并拒绝 custom TIS；两者目标不同，本 PR 不依赖 feat(rl): add off-policy IS correction hook (current policy vs rollout) #2084，也不声称覆盖多步更新的通用 off-policy correction。
• #2196 是开放中的 per-call args 修复，避免
  GenerateState singleton 固定使用首次调用的参数。本 PR 因 replacement 必须传独立参数而触及同一问题，
  采用在 task submission 时显式传入本次 args 的方式，并在 RolloutManager 边界 deep-copy custom
  rollout/sample-hook args。如果 fix(rollout): honor per-call args in generate_rollout_async (GenerateState caches first args) #2196 先合入，我会 rebase 并尽量去除或缩小重复改动。
• #2218 是开放中的 generation-stage dynamic-filter granular
  refill / abort 生命周期修复；本 PR 处理的是 rollout 完成后、actor-side RS 才能观察到的 effective-batch
  缺口。两者阶段不同，但当前版本不能直接叠加启用：本 PR 为保证 exact candidate set 会拒绝
  --dynamic-sampling-filter-path。上面的 planner 请求量也没有计入 fix(rollout): make dynamic refills granular #2218 的潜在收益。
• Draft #2278 面向 persistent fully-async queue，增加确定性的
  policy-version lag window 与 stale requeue；#2279 将 fully-async
  policy-version lifecycle 作为较小的基础改动单独提交，本身不加入 staleness rejection。本 PR 第一版显式
  拒绝 persistent fully-async queue，只支持 one-step async 的硬版本约束，因此不替代这两个 PR。

Checklist

• 主路径默认关闭，仅在显式 --rs-batch-refill 时启用
• underfill / retry exhaustion 时不执行 optimizer step
• replacement 按完整 prompt group 原子补齐，并只做 topology 最小向上取整
• 硬 staleness / policy-version / ID / mask / cache 校验
• 最终训练继续应用 TIS；不把 refill 描述为 on-policy
• custom rollout/sample/logger 参数隔离与 observability mutation 检查
• actor / manager cache 上限、RPC 超时和异常清理
• CPU 单测、相关回归、pre-commit 与私有实验 harness 自检通过
• 文档、示例配置、CUDA headroom 说明和运行指标已补充
• 完成目标四卡无外部 GPU 进程混入的真实配对 A/B，并报告真实结果与适用边界

这个改动横跨 actor、RolloutManager、data source 和 async scheduling，代码量较大。我尽量通过 opt-in、
严格参数校验、更新前失败和逐边界回归测试来限制风险。如果维护者希望拆分独立兼容性修复，或调整接口、
指标和命名，我愿意配合修改。感谢审阅。