fix(rollout): make dynamic refills granular

.github/workflows/pr-test.yml

@@ -629,6 +629,10 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_rollout_validation.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_sglang_rollout_abort.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_reloadable_process_group_world.py"

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

@@ -82,6 +82,7 @@
         {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0},
         {'test_file': 'test_sample.py', 'num_gpus': 0},
         {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
+        {'test_file': 'test_sglang_rollout_abort.py', 'num_gpus': 0},
         {'test_file': 'test_reloadable_process_group_world.py', 'num_gpus': 0},
         {'test_file': 'test_placement_group.py', 'num_gpus': 0},
         {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0},

docs/en/examples/glm4-9B.md

@@ -243,7 +243,7 @@ slime supports more complex sampling schemes, such as the dynamic sampling in [D
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \
 ```
 
-Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example:
+`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example:
 
 ```bash
    --rollout-batch-size 32 \

docs/en/examples/qwen3-4B.md

@@ -214,7 +214,7 @@ slime supports more complex sampling schemes, such as the dynamic sampling in [D
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \
 ```
 
-Here, `over_sampling_batch_size` needs to be greater than `rollout_batch_size`. For example, you can configure it as:
+`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example:
 
 ```bash
    --rollout-batch-size 32 \

docs/en/get_started/quick_start.md

@@ -345,7 +345,7 @@ slime supports more complex sampling strategies, such as dynamic sampling used i
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
 ```
 
-Here `over_sampling_batch_size` needs to be greater than `rollout_batch_size`, for example, configured as:
+`over_sampling_batch_size` controls the number of prompt groups requested at a time and must be positive. For example:
 
 ```bash
    --rollout-batch-size 32 \
@@ -367,7 +367,7 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs):
     )
 ```
 
-If the filtering function is very strict, causing a large number of prompt groups to be discarded, the system will monitor the number of pending tasks in `remaining_batch_size`. Once the number of pending tasks drops below the target number (32) due to too many being discarded, the system will automatically trigger a new round of oversampling, requesting `over_sampling_batch_size` (64) new prompts again to repeat the above process.
+`remaining_batch_size` tracks accepted and pending prompt groups. Rejected groups decrement it; when it falls below the target (32), the system requests batches of `over_sampling_batch_size` (64) until the accepted and pending total reaches the target again.
 
 ### Partial Rollout
 

docs/zh/examples/glm4-9B.md

@@ -243,7 +243,7 @@ slime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.git
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \
 ```
 
-这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为：
+`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如：
 
 ```bash
    --rollout-batch-size 32 \

docs/zh/examples/qwen3-4B.md

@@ -214,7 +214,7 @@ slime 支持了更复杂的 sampling 方案，例如 [DAPO](https://dapo-sia.git
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std \
 ```
 
-这里 `over_sampling_batch_size` 需要大于 ``rollout_batch_size`，例如配置为：
+`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如：
 
 ```bash
    --rollout-batch-size 32 \
@@ -298,4 +298,4 @@ ray job submit ... \
 
 `train.py` 和 `train_async.py` 的差别只在于 train loop 的同步逻辑，我们通过 ray 的异步（`.remote`, `ray.get`）实现了这点。
 
-⚠️  在异步训练时，sglang 的性能检测日志与训练日志可能会混到一起，不易区分，可以通过 `--sglang-log-level` 来减少 sglang 的日志。
+⚠️  在异步训练时，sglang 的性能检测日志与训练日志可能会混到一起，不易区分，可以通过 `--sglang-log-level` 来减少 sglang 的日志。

docs/zh/get_started/quick_start.md

@@ -347,7 +347,7 @@ slime 支持更复杂的采样策略，例如 [DAPO](https://dapo-sia.github.io/
      slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
 ```
 
-这里 `over_sampling_batch_size` 需要大于 `rollout_batch_size`，例如配置为：
+`over_sampling_batch_size` 控制每次请求的 prompt 组数，且必须为正数。例如：
 
 ```bash
    --rollout-batch-size 32 \
@@ -369,7 +369,7 @@ def check_reward_nonzero_std(args, samples: list[Sample], **kwargs):
     )
 ```
 
-如果过滤函数非常严格，导致大量 prompt 组被丢弃，系统会监控 ` remaining_batch_size` 中待处理的任务数量。一旦待处理的任务数因丢弃过多而降至目标数 (32) 以下，系统会自动触发新一轮的过采样，再次请求  `over_sampling_batch_size` (64) 个新的 prompt 重复上述流程。
+`remaining_batch_size` 记录已接收和待处理的 prompt 组数。每丢弃一组，该值就会减一；当它低于目标值 (32) 时，系统会按 `over_sampling_batch_size` (64) 为一批补充请求，直到已接收和待处理的总数重新达到目标值。
 
 ### Partial Rollout
 

slime/rollout/sglang_rollout.py

@@ -334,40 +334,51 @@ async def generate_and_rm_group(
 
 
 async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]:
-    aborted_samples = []
-
     state = GenerateState(args)
     assert not state.aborted
     state.aborted = True
 
-    if parse(sglang_router.__version__) <= parse("0.2.1"):
-        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers")
-        urls = response["urls"]
-    else:
-        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
-        urls = [worker["url"] for worker in response["workers"]]
-
-    await abort_servers_until_idle(urls)
+    done = set()
+    try:
+        if not args.partial_rollout:
+            for task in state.pendings:
+                task.cancel()
+            if state.pendings:
+                await asyncio.gather(*state.pendings, return_exceptions=True)
+            state.pendings.clear()
+
+        if parse(sglang_router.__version__) <= parse("0.2.1"):
+            response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers")
+            urls = response["urls"]
+        else:
+            response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
+            urls = [worker["url"] for worker in response["workers"]]
 
-    # make sure all the pending tasks are finished
-    count = 0
-    while state.pendings:
-        done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED)
+        await abort_servers_until_idle(urls)
 
         if not args.partial_rollout:
-            continue
-
-        # for partial rollout, collect the partial samples into the data buffer
-        for task in done:
-            group = task.result()
-            for sample in group:
-                if sample.response and "start_rollout_id" not in sample.metadata:
-                    sample.metadata["start_rollout_id"] = rollout_id
-            aborted_samples.append(group)
-            count += len(group)
-
-    if args.partial_rollout:
-        logger.info(f"Collected {count} partial samples into the data buffer")
+            return []
+
+        aborted_samples = []
+        count = 0
+        while state.pendings:
+            done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED)
+
+            for task in done:
+                group = task.result()
+                for sample in group:
+                    if sample.response and "start_rollout_id" not in sample.metadata:
+                        sample.metadata["start_rollout_id"] = rollout_id
+                aborted_samples.append(group)
+                count += len(group)
+    except BaseException:
+        for task in state.pendings:
+            task.cancel()
+        await asyncio.gather(*done, *state.pendings, return_exceptions=True)
+        state.pendings.clear()
+        raise
+
+    logger.info(f"Collected {count} partial samples into the data buffer")
 
     return aborted_samples
 

slime/utils/arguments.py

@@ -1930,10 +1930,9 @@ def slime_validate_args(args):
     if args.over_sampling_batch_size is None:
         args.over_sampling_batch_size = args.rollout_batch_size
 
-    assert args.over_sampling_batch_size >= args.rollout_batch_size, (
-        f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than or equal to "
-        f"rollout_batch_size {args.rollout_batch_size}"
-    )
+    assert (
+        args.over_sampling_batch_size > 0
+    ), f"over_sampling_batch_size {args.over_sampling_batch_size} should be greater than zero"
 
     if args.num_epoch is not None:
         if args.num_rollout is not None:

tests/test_megatron_argument_validation.py

@@ -305,6 +305,24 @@ def test_slime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkey
     assert args.offload_rollout is False
 
 
+@pytest.mark.unit
+def test_slime_validate_args_allows_granular_dynamic_sampling_refill(monkeypatch):
+    module = load_slime_arguments_module(monkeypatch)
+    args = make_slime_validate_args(rollout_batch_size=64, over_sampling_batch_size=1)
+
+    module.slime_validate_args(args)
+
+    assert args.over_sampling_batch_size == 1
+
+
+@pytest.mark.unit
+def test_slime_validate_args_rejects_empty_dynamic_sampling_refill(monkeypatch):
+    module = load_slime_arguments_module(monkeypatch)
+
+    with pytest.raises(AssertionError, match="greater than zero"):
+        module.slime_validate_args(make_slime_validate_args(over_sampling_batch_size=0))
+
+
 @pytest.mark.unit
 def test_update_weight_delta_requires_disk_transport(monkeypatch):
     module = load_slime_arguments_module(monkeypatch)