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

.github/workflows/pr-test.yml

@@ -573,6 +573,30 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_megatron_argument_validation.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_rs_exact_refill_utils.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_rs_exact_refill_actor.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_rs_exact_refill_rollout_manager.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_rs_exact_refill_sglang_rollout.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_rs_exact_refill_train_async.py"
+            },
+            {
+              "num_gpus": 0,
+              "test_file": "test_data_source_epoch_wrap.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_deep_ep_tms_patch.py"

.github/workflows/pr-test.yml.j2

@@ -69,6 +69,12 @@
       'extra_pip_deps': 'transformers wandb',
       'tests': [
         {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_rs_exact_refill_utils.py', 'num_gpus': 0},
+        {'test_file': 'test_rs_exact_refill_actor.py', 'num_gpus': 0},
+        {'test_file': 'test_rs_exact_refill_rollout_manager.py', 'num_gpus': 0},
+        {'test_file': 'test_rs_exact_refill_sglang_rollout.py', 'num_gpus': 0},
+        {'test_file': 'test_rs_exact_refill_train_async.py', 'num_gpus': 0},
+        {'test_file': 'test_data_source_epoch_wrap.py', 'num_gpus': 0},
         {'test_file': 'test_deep_ep_tms_patch.py', 'num_gpus': 0},
         {'test_file': 'test_stateless_adam.py', 'num_gpus': 0},
         {'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0},

docs/en/get_started/usage.md

@@ -205,6 +205,7 @@ The recommended contract is to put the source identifier in `metadata["source_na
   Note: On-policy distillation (OPD) is now orthogonal to the advantage estimator. Use `--use-opd` and `--opd-kl-coef` to enable OPD on top of any estimator.
 - `--calculate-per-token-loss`: By default, slime calculates loss on a per-sample basis, i.e., `mean(sum(sample_i) / len(sample_i))`. Enable this flag to calculate loss on a per-token basis, i.e., `sum(sum(sample_i)) / sum(len(sample_i))`.
 - `--use-tis`: Enable this setting to use TIS (Truncated Importance Sampling) (https://fengyao.notion.site/off-policy-rl).
+- `--rs-batch-refill`: With sequence/geometric RS on disaggregated `train_async.py`, preflight complete prompt groups and reactively generate only the missing groups before an optimizer step. `--rs-refill-max-rounds` bounds retries; exhaustion aborts the step instead of silently using an underfilled effective batch. This is a correctness option with additional actor preflight and rollout latency, not an unconditional speedup. See [Rollout Correction Methods](../../../examples/train_infer_mismatch_helper/README.md#effective-batch-refill-after-sequence-rs) for the current constraints and a complete command.
 
 #### GRPO Algorithm
 

docs/zh/get_started/usage.md

@@ -209,6 +209,7 @@ sglang 的加载非常简单，只需要：
   注意：在策略蒸馏 (OPD) 现在与 advantage estimator 正交，使用 `--use-opd` 和 `--opd-kl-coef` 可以在任意 estimator 之上启用 OPD。
 - `--calculate-per-token-loss`：slime 中默认的方案是 per sample loss，即 `mean(sum(sample_i) / len(sample_i))`，如果需要计算 per token loss，即 `sum(sum(sample_i)) / sum(len(sample_i))`，可以开启 `--calculate-per-token-loss`；
 - `--use-tis`：如果需要开启 tis（https://fengyao.notion.site/off-policy-rl），可以开启这一设置；
+- `--rs-batch-refill`：在训推分离的 `train_async.py` 中使用 sequence/geometric RS 时，先按完整 prompt group 做 actor preflight，再在 optimizer step 前只生成缺失的 group。`--rs-refill-max-rounds` 限制补齐轮数；超过上限会中止该 step，而不是静默使用缩水后的 effective batch。该选项增加 actor preflight 和补采样延迟，属于 correctness 机制，不保证无条件提速。完整命令与当前限制见 [Rollout Correction Methods](../../../examples/train_infer_mismatch_helper/README.md#effective-batch-refill-after-sequence-rs)。
 
 #### GRPO 算法
 

examples/train_infer_mismatch_helper/README.md

@@ -22,6 +22,61 @@ You may specify the **IS/RS configs** with a config file using `--custom-config-
 `--get-mismatch-metrics`: When you don't want to add TIS/MIS, but still want to monitor the mismatch-related metrics (e.g. rollout-training KL). It will **only return mismatch metrics** but not change the loss in any way.
 
 
+### Effective-batch refill after sequence RS
+
+Sequence-level RS can leave an optimizer step with fewer independent prompt groups than configured. In disaggregated
+`train_async.py`, `--rs-batch-refill` makes that loss of effective batch cardinality fail-safe and explicit:
+
+1. Generate exactly `rollout_batch_size` initial prompt groups; there is no speculative over-generation.
+2. Before `optimizer.step`, recompute proximal log probabilities with the actor and apply the configured sequence or
+   geometric RS gate atomically to each complete prompt group.
+3. Keep only the selected groups and their in-memory proximal log-probability cache. Generate the exact deficit,
+   rounded only to the smallest DP/VPP scheduling multiple, from the current rollout policy.
+4. Repeat for at most `--rs-refill-max-rounds`. If the target batch is still incomplete, fail the job before an
+   optimizer step instead of silently training on an underfilled batch.
+
+After an initial or replacement candidate generation returns, every coordinator and actor wait in the refill loop is
+bounded by `--rs-refill-rpc-timeout-seconds` (30 minutes by default). Candidate generation retains the rollout
+backend's existing timeout and health-monitor behavior. `--rs-refill-max-candidate-cache-bytes` (1 GiB by default)
+bounds the proximal-logprob tensor payload retained by any one process: an actor rank checks its current candidate
+round before allocating pinned CPU memory, and the RolloutManager checks its accumulated accepted cache before pulling
+the selected tensors from Ray. `peak_actor_candidate_logprob_cache_bytes` is the per-actor high-water mark to compare
+with that limit; `aggregate_candidate_logprob_cache_bytes` and
+`peak_aggregate_candidate_logprob_cache_bytes` report cumulative and per-round aggregate payload across reporting
+actors. Selected-transfer and manager-retained metrics cover the coordinator side. Leave headroom for Python
+containers, Ray object-store/transport buffers, and other process memory, which are not included in this tensor-payload
+limit. The limit also does not reserve CUDA allocator headroom: training pipeline ranks that receive a final batch may
+materialize their DP-local proximal-logprob shard on device, so long-response jobs must budget GPU memory for it.
+
+The refill path still applies TIS during training; completing the batch does not make stale initial trajectories
+on-policy. Initial candidates are limited to one policy version of staleness and reactive replacements must match the
+actor version used for preflight.
+
+```bash
+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
+```
+
+This first implementation intentionally supports fresh/finetune runs from rollout 0 only. It requires the
+disaggregated Megatron path, a global rollout dataset, one optimizer step per rollout batch, resident actor and rollout
+engines, full NCCL weight updates, dynamic batching, zero actor dropout, the default model provider, and non-quantized
+Megatron training. Checkpoint resume, fully async rollout, token-level RS, external rollout engines, and custom
+model/loss/data hooks are rejected during argument validation. Dynamic sampling filters and custom rollout logging are
+handled conservatively: dynamic sampling filters are rejected because they may change the exact candidate set, while a
+custom rollout logger or debug saver is allowed only if a final train-data fingerprint proves that it left policy
+inputs unchanged.
+
+Refill adds actor preflight work and replacement latency, so it is a correctness option rather than an unconditional
+wall-clock speedup. Its cost depends on the observed rejection rate and the rollout/training time balance.
+
+
 ## Algorithms
 
 We give examples of the algorithms for solving the training-inference mismatch issue.

examples/train_infer_mismatch_helper/mis_refill.yaml

@@ -0,0 +1,22 @@
+# TIS weights correct the policy-gradient loss but do not mask individual tokens.
+# Keep batch normalization disabled: RS refill cannot yet normalize globally
+# across DP-local dynamic microbatches.
+use_tis: true
+tis_level: "token"
+tis_mode: "truncate"
+tis_lower_bound: 0.1
+tis_upper_bound: 10.0
+tis_batch_normalize: false
+
+# Proximal replay also requires rollout_temperature > 0, 0 < rollout_top_p <= 1,
+# rollout_top_k: -1, attention_dropout: 0, and hidden_dropout: 0 in the CLI or
+# surrounding custom configuration.
+attention_dropout: 0.0
+hidden_dropout: 0.0
+
+# RS is a sequence-level gate so a rejected unit maps to an unambiguous batch deficit.
+use_rs: true
+rs_level: "geometric"
+rs_lower_bound: 0.01
+rs_upper_bound: 100.0
+rs_veto_threshold: null