[rollout] feat: support over sampling rollout in SGLang Rollout (volcengine#2929)

This PR introduces an **over-sample strategy** for verl's SGLang
multi-turn rollout to address the long-tail problem, where a few slow
requests disproportionately increase the overall rollout time. The core
idea is to over-sample the number of requests at the start of the
rollout and then aggressively cancel any requests that haven't finished
once a target number of completions is met.

- **Improves rollout efficiency** for multi-turn conversations by
reducing total time spent waiting for slow requests.
- **Implements a new request monitoring and cancellation mechanism** to
cut off unnecessary computation.

wandb results is as follow:

https://wandb.ai/zhaochenyang20/benchmark_over_sample_2/workspace?nw=nwuserzhaochenyang20

-----

Of course, this strategy has its share of issues. For example, many
might question why the over-long requests that are dropped aren't simply
saved and continued in the next round. This is certainly possible—it's a
partial rollout strategy—but it would require verl to have a data
buffer, which is beyond the scope of this PR. Furthermore, saving and
continuing these requests would introduce an off-policy problem.

There is also a valid concern that this rather "brutal" dropping
strategy could unfairly ignore very long requests. I agree this is a
very reasonable point, but currently, we don't have a lossless solution.
However, our dropping strategy is very flexible and could even change
with our curriculum learning. For instance, in the example I gave, I
just directly dropped the last 20% of requests. **In practice, we can
dynamically adjust this drop rate and even set different dropping
methods. For example, we could record the return time (t) for the 80% of
requests and then drop any requests that haven't returned after 1.5t.**

We've provided an initial, validated idea and have completed its
implementation. We welcome everyone to join the discussion on how to
accelerate multi-turn rollouts with acceptable losses.

The new over-sample strategy was tested with an 8-GPU setup on the
**gsm8k** dataset, yielding the following results:

- **Rollout Time:** Significant reduction in overall rollout time per
step.
  - **Training Rewards:**
- The reward metric for training steps shows a positive bias. This is
because we exclude the aborted requests (which are typically more
difficult and have lower rewards) from the reward calculation.
- The reward metric for validation steps remains accurate and aligns
with the baseline. This is because the cancellation logic is not
triggered during validation, ensuring a fair and complete evaluation.

This feature modifies `sglang_rollout.py` and `metric_utils.py`. To use
it, follow the standard setup and then run the training script with the
over-sample parameters.

https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/release_log/over_sample.md

The design is centered on three main functions that orchestrate the
over-sampling logic: `run_with_cancellation`,
`process_request_with_monitoring`, and `monitor_and_cancel`. These
functions rely on global variables, such as `all_tasks` and
`completion_lock`, to manage state.

- **`run_with_cancellation`:** This is the entry point. It launches all
requests as `process_request_with_monitoring` tasks concurrently with a
single `monitor_and_cancel` task. It uses `asyncio.gather` to wait for
all tasks to complete (or be canceled) and converts any exceptions from
canceled tasks into padding requests before returning the final output.

- **`process_request_with_monitoring`:** This async function handles a
single request. It waits for the request to complete using
`_async_rollout_a_request` and then checks a shared counter,
`completed_count`, using a `completion_lock` for thread safety. If the
target completion count has not been reached, it returns the real
result. If the target has been met, it returns padding data instead,
effectively "discarding" the late result.

- **`monitor_and_cancel`:** This is a separate async task that polls the
`completed_count`. Once the count reaches the `target_completion`
threshold, it immediately cancels all remaining tasks and sends an
`abort_requests` signal to the SGLang engine, halting any ongoing GPU
computation for those requests.

Key code changes:

  - **`sglang_rollout.py`**:
- Adds the three core asynchronous functions for the over-sample
strategy.
- The `AsyncEngine` class now includes a new `abort_request` method that
calls the synchronous `abort_request` in the `tokenizer_manager`.
  - **`metric_utils.py`**:
- The `compute_data_metrics` function is updated to exclude the aborted
requests (identified by padding) from the denominator when calculating
average rewards during training. This prevents the training reward from
being artificially lowered by the zero-reward aborted requests.

This implementation is designed to be a straightforward and effective
solution for the long-tail problem, though some aspects of the
asynchronous design and the impact on training variance require further
investigation.

---------

Co-authored-by: zhaochenyang <[email protected]>
Co-authored-by: PopSoda2002 <[email protected]>
Co-authored-by: ChangyiYang <[email protected]>
Co-authored-by: PrinsYin <[email protected]>
Co-authored-by: WindowsXp-Beta <[email protected]>
Co-authored-by: zhaochenyang20 <[email protected]>

Author: 6 people

examples/sglang_multiturn/run_qwen2_3b_dapo_multiturn.sh renamed to examples/sglang_multiturn/run_qwen3_4b_dapo_multiturn.sh

@@ -17,6 +17,10 @@ hf download \
     --repo-type dataset \
     --local-dir $HOME/data/Maxwell-Jia/AIME_2024
 
+# Note that this script is using AgentLoop instead of SGLang Multi-Turn
+# We are concerned that the reward is not actually converge, since the
+# reward of retool is encouraging the model to generate more turns to
+# call more tools. The answers are not actually correct.
 
 python3 -m verl.trainer.main_ppo \
     algorithm.adv_estimator=grpo \
@@ -34,7 +38,7 @@ python3 -m verl.trainer.main_ppo \
     data.custom_cls.name=CustomRLHFDataset \
     custom_reward_function.path=$PROJECT_DIR/recipe/retool/retool.py \
     custom_reward_function.name=compute_score \
-    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \
+    actor_rollout_ref.model.path=Qwen/Qwen3-4B-Instruct-2507 \
     actor_rollout_ref.model.use_remove_padding=True \
     actor_rollout_ref.model.enable_gradient_checkpointing=True \
     actor_rollout_ref.actor.use_kl_loss=False \
@@ -43,13 +47,16 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.actor.clip_ratio_high=0.28 \
     actor_rollout_ref.actor.clip_ratio_c=10.0 \
     actor_rollout_ref.actor.optim.lr=1e-6 \
-    actor_rollout_ref.actor.use_dynamic_bsz=True \
-    actor_rollout_ref.actor.ppo_mini_batch_size=8 \
-    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=1024 \
+    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=32768 \
+    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
     actor_rollout_ref.rollout.name=sglang \
     actor_rollout_ref.rollout.mode=async \
     actor_rollout_ref.rollout.tensor_model_parallel_size=1 \
-    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
+    actor_rollout_ref.rollout.update_weights_bucket_megabytes=512 \
     actor_rollout_ref.rollout.multi_stage_wake_up=True \
     actor_rollout_ref.rollout.multi_turn.enable=True \
     actor_rollout_ref.rollout.multi_turn.max_user_turns=16 \
@@ -62,8 +69,8 @@ python3 -m verl.trainer.main_ppo \
     actor_rollout_ref.rollout.val_kwargs.n=30 \
     trainer.logger=['console','wandb'] \
     trainer.project_name=sglang-dapo-multiturn \
-    trainer.experiment_name=qwen2_5-3b_dapo_multiturn \
-    trainer.n_gpus_per_node=4 \
+    trainer.experiment_name=qwen3-4b_dapo_multiturn \
+    trainer.n_gpus_per_node=8 \
     trainer.log_val_generations=20 \
     trainer.val_before_train=True \
     trainer.nnodes=1 \

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -163,6 +163,7 @@ actor_rollout_ref:
     disable_log_stats: true
     do_sample: true
     'n': 1
+    over_sample_rate: 0
     multi_stage_wake_up: false
     engine_kwargs:
       vllm:

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -138,6 +138,7 @@ actor_rollout_ref:
     disable_log_stats: true
     do_sample: true
     'n': 1
+    over_sample_rate: 0
     multi_stage_wake_up: false
     engine_kwargs:
       vllm:

verl/trainer/config/rollout/rollout.yaml

@@ -82,7 +82,13 @@ do_sample: True
 # number of responses (i.e. num sample times). > 1 for grpo
 n: 1
 
-# Whether to wake up inference engine in multi-stage to reduce peak memory during training-rollout transition.
+# The over_sample_rate parameter controls the early termination threshold for training rollouts,
+# where the system will abort remaining requests when (1 - over_sample_rate) * total_requests completions are reached.
+over_sample_rate: 0
+
+# Whether to wake up inference engine in multi-stage for SGLang
+# to reduce peak memory during training-rollout transition.
+# This is only effective for SGLang rollout.
 multi_stage_wake_up: false
 
 # Extra inference engine arguments (vllm, sglang).

verl/trainer/ppo/metric_utils.py

@@ -118,6 +118,20 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str,
     prompt_length = response_info["prompt_length"]
     response_length = response_info["response_length"]
 
+    aborted_mask = (response_length == 0).bool()
+    non_aborted_mask = ~aborted_mask
+
+    non_aborted_sequence_score = sequence_score[non_aborted_mask]
+    non_aborted_sequence_reward = sequence_reward[non_aborted_mask]
+
+    score_mean = torch.mean(non_aborted_sequence_score).detach().item()
+    score_max = torch.max(non_aborted_sequence_score).detach().item()
+    score_min = torch.min(non_aborted_sequence_score).detach().item()
+
+    reward_mean = torch.mean(non_aborted_sequence_reward).detach().item()
+    reward_max = torch.max(non_aborted_sequence_reward).detach().item()
+    reward_min = torch.min(non_aborted_sequence_reward).detach().item()
+
     valid_adv = torch.masked_select(advantages, response_mask)
     valid_returns = torch.masked_select(returns, response_mask)
 
@@ -127,15 +141,30 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str,
         return_diff_var = torch.var(valid_returns - valid_values)
         return_var = torch.var(valid_returns)
 
+    # Aborted samples and non-aborted response length statistics
+    # response_length_non_aborted/*: statistics computed on non-aborted samples only
+    aborted_ratio = torch.mean(aborted_mask.float()).detach().item()
+
+    non_aborted_response_length = response_length[non_aborted_mask]
+    if non_aborted_response_length.numel() > 0:
+        non_aborted_response_length_mean = torch.mean(non_aborted_response_length).detach().item()
+        non_aborted_response_length_max = torch.max(non_aborted_response_length).detach().item()
+        non_aborted_response_length_min = torch.min(non_aborted_response_length).detach().item()
+        non_aborted_response_length_clip_ratio = (
+            torch.mean(torch.eq(non_aborted_response_length, max_response_length).float()).detach().item()
+        )
+    else:
+        raise ValueError("All samples are aborted, this should not happen.")
+
     metrics = {
         # score
-        "critic/score/mean": torch.mean(sequence_score).detach().item(),
-        "critic/score/max": torch.max(sequence_score).detach().item(),
-        "critic/score/min": torch.min(sequence_score).detach().item(),
+        "critic/score/mean": score_mean,
+        "critic/score/max": score_max,
+        "critic/score/min": score_min,
         # reward
-        "critic/rewards/mean": torch.mean(sequence_reward).detach().item(),
-        "critic/rewards/max": torch.max(sequence_reward).detach().item(),
-        "critic/rewards/min": torch.min(sequence_reward).detach().item(),
+        "critic/rewards/mean": reward_mean,
+        "critic/rewards/max": reward_max,
+        "critic/rewards/min": reward_min,
         # adv
         "critic/advantages/mean": torch.mean(valid_adv).detach().item(),
         "critic/advantages/max": torch.max(valid_adv).detach().item(),
@@ -163,6 +192,15 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str,
         "response_length/clip_ratio": torch.mean(torch.eq(response_length, max_response_length).float())
         .detach()
         .item(),
+        # response length (non-aborted only)
+        # These statistics exclude aborted samples to avoid skew from zeros
+        "response_length_non_aborted/mean": non_aborted_response_length_mean,
+        "response_length_non_aborted/max": non_aborted_response_length_max,
+        "response_length_non_aborted/min": non_aborted_response_length_min,
+        "response_length_non_aborted/clip_ratio": non_aborted_response_length_clip_ratio,
+        # aborted ratio
+        # Fraction of samples whose response length is zero
+        "response/aborted_ratio": aborted_ratio,
         # prompt length
         "prompt_length/mean": torch.mean(prompt_length).detach().item(),
         "prompt_length/max": torch.max(prompt_length).detach().item(),

verl/workers/config/rollout.py

@@ -86,6 +86,10 @@ class RolloutConfig(BaseConfig):
     do_sample: bool = True
     n: int = 1
 
+    # Early termination threshold for multi-turn rollout in sglang.
+    # Abort remaining requests when (1 - over_sample_rate) * total_requests are completed.
+    over_sample_rate: float = 0.0
+
     prompt_length: int = 512
     response_length: int = 512