Skip to content

Commit e0f4dc2

Browse files
startjustartju
authored
feat: add partial_rollout recipe (#96)
## Summary Adds `partial_rollout/` to the recipe submodule: APRIL-style ([paper](https://arxiv.org/pdf/2509.18521)) synchronous RL with cross-step rollout interruption + resume to reclaim long-tail GPU bubbles. Aborted gens carry their conversation state across the step boundary and resume on the next step while their KV cache may still be live on the rollout server. Based on upstream `verl-project/verl@8ebccd44` (full pin in [`recipe/partial_rollout/REQUIRED_VERL.txt`](https://github.com/startju/verl-recipe/blob/partial_rollout/partial_rollout/REQUIRED_VERL.txt)). ## Relationship to #58 Open PR #58 (`mamazi0131:main`, 2026-03-01) lands the same recipe directory and was the starting point for this work. This PR is materially different on four architectural axes: 1. **`LLMServerManager` / `AgentLoopManager` split** (verl#6117). Current upstream separates the rollout server manager from the agent-loop manager. This recipe ships `llm_server.py` (`PartialRolloutLLMServerManager`) so cancel/resume fan-out lives on the new server-manager surface; a small symbol-swap in `ray_trainer.init_workers` injects it because upstream `RayPPOTrainer.init_workers` hardcodes `LLMServerManager` with no FQN config knob (unlike the parallel `agent_loop_manager_class` knob it does have). #58's tree doesn't import on current `main`. 2. **Cancel/retry path absorbed inside `FullyLLMServerClient`** (upstream verl#5631). Current upstream's [`FullyLLMServerClient.generate()`](https://github.com/verl-project/verl/blob/main/verl/workers/rollout/llm_server.py#L160) already has an abort-then-retry loop — when a generate is aborted mid-flight by `cancel()`, it parks and resumes against the next weight version's accumulated context without ever returning to AgentLoop. This recipe **gates** that retry branch with `+async_training.partial_rollout=True` and **forces** `get_client(fully_async=True)` for every caller. Net effect: the recipe's `AgentLoop` only handles pull/push and trajectory-grained pull pacing (see axis 3); validation goes through upstream's untouched `AgentLoopWorker.generate_sequences`. #58 instead returns an ABORT sentinel to the agent loop, which then re-enqueues the prompt into `pending_queue` — a structurally heavier path that requires a forked `tool_agent_loop` for state snapshot/restore. 3. **Trajectory-grained pull pacing**. `PartialRolloutAgentLoopWorker.generate_for_prompt` replaces upstream's trailing `outputs = await asyncio.gather(*tasks)` with an `asyncio.wait(FIRST_COMPLETED)` loop that decrements `self.inflight_traj` and signals `self._slot_event` after every per-trajectory completion. The `run_continuous` outer loop then pulls the next prompt as soon as `inflight_traj + n <= max_inflight_prompts * n` — long-tail trajectories inside one prompt don't block new prompts from entering across the budget freed by other in-flight prompts' completions. Pull RPC, `_run_one` tasks, and the slot-wait sentinel share a single `asyncio.wait(running)` set; identity checks dispatch the three task kinds. Validation keeps the simpler upstream gather path (we add `generate_for_prompt` as a new method rather than overriding `generate_sequences`). #58 reaches a similar trajectory-level effect via `pending_queue` re-enqueue + `last_agent_loop_output` snapshot/restore — heavier mechanism, requires forked agent loops. 4. **Engine-level cancel via Python `_resume_event` + `abort_all_requests` drain** (vLLM 0.11 stopgap). vLLM <0.12 doesn't expose `pause_generation`, so `PartialRolloutvLLMHttpServer` adds a Python-side `_resume_event` gate around `generate()` plus an `inflight`-counted `abort_all_requests(reset_prefix_cache=False)` drain loop in `cancel()`. Deletable once verl moves to vLLM ≥0.12. #58 uses per-request `asyncio.Event` + `Lock` — every in-flight generate awaits its own cancel handle; PartialRollout substitutes one engine-core batch call. In addition this PR adds: - `gsm8k_tool_config.yaml` (recovered after upstream #6126 deletion of `examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml`). - `run_qwen3-0.6b_gsm8k_grpo_tool{,_baseline}.sh` plus non-tool 0.6B variants for laptop-scale repro, under `recipe/partial_rollout/run/`. - `README_zh.md`, `REFERENCE.md`, `REQUIRED_VERL.txt`. Happy to fold these into #58 if @mamazi0131 prefers — opening separately because #58 does not apply to current `verl-project/verl@main` and the rebase is nontrivial. ## Test plan Full 1-epoch chain on 2× RTX 3090, Qwen3-0.6B, gsm8k, GRPO + token-level rollout-IS, `max_response_length=4096`, `max_model_len=4608`, batch=8, TP=1. ### Single-turn, PR vs baseline (completed, 934 steps each) ```bash bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo.sh bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo_baseline.sh ``` - [x] 934 steps each, no OOM, no hang - [x] PR `timing_s/gen` avg **18.5s** vs baseline **26.1s** — **−29% gen time** - [x] PR `perf/throughput` avg **1060** tok/s/GPU vs baseline **851** — **+24%** - [x] Learning curves overlap — no learning regression - [x] `pre-commit run --all-files` (ruff, ruff-format) clean. ## Test Result <img width="2780" height="1230" alt="image" src="https://github.com/user-attachments/assets/e719d604-9272-41ac-8c4d-e183867c75bd" /> Single-turn 934-step run on 2× RTX 3090, Qwen3-0.6B, gsm8k, GRPO + token-level rollout-IS, `max_response_length=4096`, batch=8, 1 epoch. Six panels — **green = baseline, pink = partial_rollout**. ### Learning-quality panels (top-left, bottom-left, bottom-right) These three panels exist to falsify "PR breaks the algorithm." If PR shifted training dynamics, one of these would diverge. - **`critic/rewards/mean`** — both runs climb from ~0 to ~0.8 by step ~200, then track together at 0.7–0.9 for the rest of training. **No reward divergence**; PR's cross-step interrupt + resume does not introduce bias into the policy gradient. - **`response_length/mean`** — both rise from ~600 to ~1100–1200 over the run, near-overlapping. PR is marginally higher (matches the 1142 vs 1068 averages reported in the comparison comment), within noise. - **`actor/entropy`** — both decay from ~0.4 to ~0.15 along the same trajectory. **Same exploration / collapse rate**. → PR is policy-correctness-neutral. The cancel-resume mechanism doesn't perturb the optimization. ### Performance panels (top-middle, top-right, bottom-middle) These show the actual speedup. - **`timing_s/gen`** ⭐ — the panel that matters most. **baseline sits at ~25–35s, PR sits at ~15–20s, consistently and across the entire 934 steps**. The two curves almost never cross. Spikes at multiples of 50 are `test_freq=50` validation steps (validation goes through upstream's no-PR path, so both runs pay the same validation cost — those spikes overlap). Excluding warmup, PR averages 18.5s, baseline 26.1s — **−29%**. - **`perf/throughput`** (tok/s/GPU) — mirror of gen timing. **PR ~1100–1400, baseline ~800–1000**, sustained gap, **+30% throughput**. Both curves are noisy step-to-step (batch=8 means high per-step variance — a single long-tail prompt dominates), but the bands clearly separate. - **`timing_s/step`** — total step wall time. **PR ~30–40s, baseline ~40–50s**. Same direction as gen but smaller relative gap (≈ −16%) because non-gen phases (`update_actor` ~12s, `ref` + `old_log_prob` ~5s, etc.) are unchanged by PR. PR's win is concentrated entirely inside the gen phase; the rest is identical work. ### Why the chart is convincing 1. **Sustained, not warmup-bounded**: the gap shows up by step 5 and stays for 900 more steps. Not an outlier of a particular batch. 2. **Two curves never cross on `timing_s/gen`**: any single step PR ≤ baseline (modulo the shared validation spikes). System-level effect, not statistical noise. 3. **Learning curves overlap pixel-for-pixel**: the speedup is **not** "PR took shortcuts and generated less / worse." Reward, length, entropy match. ### Headline **29% faster gen, 30% higher throughput, no learning-curve regression.** At this scale (`max_response=4096`, batch=8, long-tail driven) PR is in its design sweet spot. ## AI-assistance disclosure This PR was drafted with AI assistance (Claude Opus 4.7, 1M context window). The commit carries a `Co-authored-by: Claude` trailer. The submitting human (@startju) reviewed every changed line, ran the test above, and is the accountable owner of this change end-to-end.
1 parent a809c5c commit e0f4dc2

24 files changed

Lines changed: 2370 additions & 0 deletions

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ The script requires only `bash`, `git`, `awk`, and `pip`/`pip3` on `PATH`. It do
8080
| minicpmo | [`recipe/minicpmo/REQUIRED_VERL.txt`](minicpmo/REQUIRED_VERL.txt) |
8181
| nemo_gym | [`recipe/nemo_gym/REQUIRED_VERL.txt`](nemo_gym/REQUIRED_VERL.txt) |
8282
| open_math_reasoning | [`recipe/open_math_reasoning/REQUIRED_VERL.txt`](open_math_reasoning/REQUIRED_VERL.txt) |
83+
| partial_rollout | [`recipe/partial_rollout/REQUIRED_VERL.txt`](partial_rollout/REQUIRED_VERL.txt) |
8384
| prime | [`recipe/prime/REQUIRED_VERL.txt`](prime/REQUIRED_VERL.txt) |
8485
| qat | [`recipe/qat/REQUIRED_VERL.txt`](qat/REQUIRED_VERL.txt) |
8586
| r1 | [`recipe/r1/REQUIRED_VERL.txt`](r1/REQUIRED_VERL.txt) |
@@ -97,6 +98,7 @@ The script requires only `bash`, `git`, `awk`, and `pip`/`pip3` on `PATH`. It do
9798
- [retool](https://github.com/verl-project/verl-recipe/tree/main/retool): Reinforcement Learning for Strategic Tool Use in LLMs
9899
- [langgraph_agent](https://github.com/verl-project/verl-recipe/tree/main/langgraph_agent): A tiny example to demonstrate multi-turn rollout with [LangGraph ReactAgent](https://langchain-ai.github.io/langgraph/agents/overview/) to solve math expression.
99100
- [spo](https://github.com/verl-project/verl-recipe/tree/main/spo): [Single-stream Policy Optimization](https://arxiv.org/abs/2509.13232).
101+
- [partial_rollout](./partial_rollout/): synchronous RL with cross-step rollout interruption + resume to reclaim long-tail GPU bubbles ([APRIL](https://arxiv.org/pdf/2509.18521)-style).
100102
- TBA...
101103

102104
## Contribution

‎partial_rollout/README.md‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
# Recipe: Partial Rollout
2+
3+
English | [简体中文](README_zh.md)
4+
5+
A partial-rollout pipeline for **synchronous RL training**, designed to reclaim GPU bubbles caused by **long-tail response lengths** via sample supplementation and mid-generation interruption with cross-step resume.
6+
7+
> ⚠️ **Don't confuse this with the fully-async framework's partial rollout.** This pipeline still runs the synchronous loop "rollout → wait for batch → one train step"; the only twist is that long-tail samples can be interrupted mid-rollout and resumed in a later step. Trainer and rollout phases remain serial. If you want the trainer/rollout fully decoupled and advancing concurrently, see `verl/experimental/fully_async_policy/` — not here.
8+
9+
> 🔗 **Verl dependency.** Pinned against `verl-project/verl@60546ef2` ([on GitHub](https://github.com/verl-project/verl/commit/60546ef2a7464a158cd170f58f852a62a4e552ba)). Exact `pip install` / `git checkout` recipe in [REQUIRED_VERL.txt](REQUIRED_VERL.txt). Rolling against `main`; bump the pin when refreshing.
10+
11+
## Background
12+
13+
Synchronous PPO/GRPO training in verl waits for every prompt in a batch to finish generating before stepping. RL training datasets are highly **right-skewed in response length**: empirically a small fraction of samples (~3%) emit dramatically longer responses than the median, and these long-tail samples are often the harder, more informative ones — you can't drop them without hurting accuracy. The whole batch stalls on the slowest few; GPUs idle.
14+
15+
![Response Length Distribution across the RL Training Dataset](https://raw.githubusercontent.com/mamazi0131/verl_doc/fca7a6d3acbeca12d69c5de6f85c312c1c9e47b6/Response_Length_Distribution_across_the_RL_Training_Dataset.png)
16+
17+
Partial rollout closes this bubble with two ideas (see the [APRIL paper](https://arxiv.org/pdf/2509.18521) for the academic treatment):
18+
19+
- **Sample supplementation**: when faster workers run out of work, immediately top them up with the next prompt — don't sit idle waiting for the long tail.
20+
- **Mid-generation interruption + cross-step resume**: when the batch's "done" target is met but a few samples are still mid-generation, interrupt them at the step boundary, cache their KV state, and resume them in the next training step. The recipe pays one weight-version drift per resumed sample, corrected by token-level rollout importance sampling (off-policy correction).
21+
22+
![Comparison of GPU Execution Timelines between Standard Synchronous Training and the Proposed Async Partial Rollout](https://raw.githubusercontent.com/mamazi0131/verl_doc/fca7a6d3acbeca12d69c5de6f85c312c1c9e47b6/Comparison_of_GPU_Execution_Timelines_between_Standard_Synchronous_Training_and_the_Proposed_Async_Partial_Rollout.png)
23+
24+
Net effect: convert long-tail GPU bubbles into useful work on the next batch's prompts, at the cost of mild off-policy drift on resumed samples — which token-level IS handles cleanly.
25+
26+
---
27+
28+
## When to use
29+
30+
Use it when:
31+
32+
- The dataset has a **long-tailed response length distribution** (a small fraction of very long samples drags down each step).
33+
- Synchronous PPO/GRPO shows a visible GPU bubble waiting on those long-tail samples.
34+
- Training tolerates **mild off-policy drift** — partial rollout inherently spans multiple weight versions; pair it with IS correction.
35+
- Multi-turn / tool-call workloads — upstream vLLM ≥ 0.12's `pause_generation` + abort is enough; no recipe-side server fork.
36+
37+
Skip it when:
38+
39+
- Response lengths are uniform with no long-tail bubble — the plain synchronous trainer is simpler.
40+
- Strict on-policy is required (every trajectory must come from the current weights).
41+
- You need to mix this pipeline with the upstream sync trainer's batch-shape assumptions — PartialRollout-specific fields (dummy `gen_batch`, continuous-worker semantics) would break them.
42+
43+
---
44+
45+
## Architecture
46+
47+
```
48+
trainer (PartialRolloutRayPPOTrainer)
49+
│
50+
push_batch │ pull_batch
51+
▼
52+
┌──────────────────────────────────┐
53+
│ RolloutPromptManager (Ray) │
54+
│ │
55+
│ pending ─pull─► ongoing ─push─► done
56+
└──────────────────────────────────┘
57+
│
58+
pull_prompts│push_prompts
59+
▼
60+
PartialRolloutAgentLoopWorker ×N (run_continuous loop)
61+
│
62+
llm_client│ generate (FullyLLMServerClient retries aborted)
63+
▼
64+
PartialRolloutvLLMReplica ×replicas
65+
↳ PartialRolloutvLLMHttpServer (Python `paused` gate + abort drain)
66+
▲
67+
cancel/ │
68+
resume │
69+
│
70+
PartialRolloutLLMServerManager
71+
```
72+
73+
| Component | File | Role |
74+
|---|---|---|
75+
| `PartialRolloutRayPPOTrainer` | `ray_trainer.py` | Main trainer loop. `_fit_generate` pushes prompts into the manager, awaits one full batch via `async_rollout_manager.generate_sequences`, runs log_prob / advantage / policy update, then `update_weights`. |
76+
| `RolloutPromptManager` | `prompt_manager.py` | Single-threaded Ray actor holding the three queues (pending / ongoing / done). `pull_batch` and `pull_prompts` both block on `asyncio.Event`s — no busy polling, no per-empty-pull RPCs. |
77+
| `PartialRolloutAgentLoopManager` / `PartialRolloutAgentLoopWorker` | `agent_loop/agent_loop.py` | Worker runs a persistent loop (`run_continuous`) sharing one `asyncio.wait` across the pull RPC, the `_run_one` rollout tasks, and a slot-wait sentinel. Instead of overriding upstream `generate_sequences` (left intact for the validation path), the worker adds `generate_for_prompt`, which is upstream's `generate_sequences` with the trailing `outputs = await asyncio.gather(*tasks)` replaced by an `asyncio.wait(FIRST_COMPLETED)` loop that decrements `self.inflight_traj` and signals `self._slot_event` after each trajectory completion. The outer loop pulls the next prompt as soon as `inflight_traj + n <= max_inflight_prompts * n` — so a long-tail trajectory doesn't block other in-flight prompts from making room for the next pull. Manager exposes `cancel()` / `resume()` (delegated to `PartialRolloutLLMServerManager`) for trainer-side bracketing of `update_weights`. |
78+
| `PartialRolloutvLLMHttpServer` / `PartialRolloutvLLMReplica` | `vllm_rollout/vllm_async_server.py` | vLLM HTTP server with a Python-side `_resume_event` gate (vLLM <0.12 lacks `pause_generation`). `cancel()` clears the gate so new `generate()` calls hang at the wrapper layer, then loops `abort_all_requests(reset_prefix_cache=False)` until in-flight drains. `resume()` sets the gate, releasing queued callers. Drop this whole layer once vLLM ≥0.12 is the floor. |
79+
| `PartialRolloutLLMServerManager` | `llm_server.py` | Thin override of upstream `LLMServerManager`: swaps `rollout_replica_class` to `PartialRolloutvLLMReplica`, forces `get_client(fully_async=True)` so callers receive the retry-on-abort `FullyLLMServerClient`, and fans `cancel` / `resume` out to each replica. Installed via a monkey-patch in `PartialRolloutRayPPOTrainer.init_workers` because upstream has no FQN config knob for `LLMServerManager`. |
80+
81+
---
82+
83+
## Key invariants
84+
85+
1. **Prompt ownership during scheduling**: at any instant a prompt lives in exactly one of pending / ongoing / done. `pull_prompts` moves pending→ongoing; `push_prompts` is ongoing→done — terminal only. There is no aborted-back-to-pending re-queue path, because `FullyLLMServerClient.generate()` absorbs the abort/retry cycle inside one logical generate call.
86+
2. **Cross-step abort/resume bracket**: `PartialRolloutAgentLoopManager.generate_sequences` runs `await self.resume()` at entry and `await self.cancel()` after `pull_batch` returns. The naive `checkpoint_engine` backend (PartialRollout's default) short-circuits before its own abort, so the recipe wires this itself; workers' aborted `client.generate(...)` calls wait inside `FullyLLMServerClient`'s retry loop until the next step's `resume()`.
87+
3. **Per-sample weight-version tracking** lives in `gen_batch.meta_info["global_steps"]` (set by the trainer) and `FullyLLMServerClient.generate()` (records the actual versions each retry submitted against). The worker does no per-call version tracking.
88+
4. **Continuous worker loop + trajectory-grained pull pacing**: each `PartialRolloutAgentLoopWorker` runs `run_continuous` for the actor's lifetime. Budget is counted in trajectories (`max_inflight_prompts * n`), not prompts. Inside each in-flight prompt, `generate_for_prompt` awaits the n trajectories via `asyncio.wait(FIRST_COMPLETED)` — every completion decrements `self.inflight_traj` and sets `self._slot_event`, waking the outer loop to pull the next prompt as soon as a prompt's worth of trajectory slots have freed up across all in-flight prompts (no need to wait for any one prompt to fully complete). Pull RPC, `_run_one` tasks, and the slot-wait sentinel all live in the same `asyncio.wait(running)` set; identity checks dispatch the three task kinds.
89+
5. **The dummy `gen_batch` must carry `uid`**: when the dataloader is exhausted at the end of an epoch, the placeholder batch built to drain in-flight prompts still needs `non_tensor_batch["uid"]`, otherwise the manager can't compute the row count.
90+
6. **Stateful dataloader resume**: `PartialRolloutRayPPOTrainer.fit()` resumes via the stateful dataloader automatically — **don't** add manual skip-on-resume logic.
91+
7. **No graceful shutdown**: `run_continuous` runs for the actor's lifetime; Ray terminates workers at process exit (after `fit()` returns). Any rollouts in flight at exit time are abandoned — acceptable because the trainer isn't going to consume them anyway.
92+
93+
---
94+
95+
## Quick start
96+
97+
### Partial-rollout runs
98+
99+
Single-turn:
100+
```bash
101+
bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo.sh
102+
```
103+
104+
Multi-turn with tool calls — first generate the tool-agent dataset:
105+
```bash
106+
python3 examples/data_preprocess/gsm8k_multiturn_w_tool.py \
107+
--local_save_dir $HOME/data/gsm8k_tool
108+
```
109+
then launch:
110+
```bash
111+
bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo_tool.sh
112+
```
113+
114+
### Baseline runs (vanilla GRPO, for A/B against partial rollout)
115+
116+
Same model / data / batch / hyperparameters; the only differences from the PR variants are the upstream entry, the upstream agent loop, and no IS correction:
117+
```bash
118+
bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo_baseline.sh
119+
bash recipe/partial_rollout/run/run_qwen3-0.6b_gsm8k_grpo_tool_baseline.sh
120+
```
121+
122+
### PartialRollout-specific Hydra overrides
123+
124+
| Key | Value | Notes |
125+
|---|---|---|
126+
| `actor_rollout_ref.rollout.agent.default_agent_loop` | `single_turn_agent` / `tool_agent` | Use upstream agent loops directly — `FullyLLMServerClient` handles abort/retry, no PartialRollout wrapper needed. |
127+
| `+async_training.partial_rollout` | `True` | Gates the retry-on-abort path inside `FullyLLMServerClient`. Required for any PartialRollout run; set via Hydra `+` (the key isn't pre-declared in the trainer config). |
128+
| `algorithm.rollout_correction.rollout_is` | `token` (recommended) / `sequence` | Sequence-level ratios easily hit the `exp(±20)` safety clamp once rollouts span several weight versions. |
129+
| `algorithm.rollout_correction.rollout_is_threshold` | `2.0` | TIS upper bound; for IcePop pass `"0.5_5.0"`. |
130+
| `actor_rollout_ref.rollout.multi_turn.enable` | `True` *(tool variant only)* | Enables multi-turn. |
131+
| `actor_rollout_ref.rollout.multi_turn.tool_config_path` | YAML path | Tool registry; shared across backends. |
132+
133+
---
134+
135+
## File layout
136+
137+
```
138+
partial_rollout/
139+
├── README.md / README_zh.md / REQUIRED_VERL.txt
140+
├── main_ppo.py # @hydra entry; wraps PartialRolloutTaskRunner
141+
├── ray_trainer.py # PartialRolloutRayPPOTrainer
142+
├── prompt_manager.py # RolloutPromptManager (Ray actor)
143+
├── llm_server.py # PartialRolloutLLMServerManager (force fully_async + cancel/resume)
144+
├── agent_loop/
145+
│ └── agent_loop.py # PartialRolloutAgentLoopManager / PartialRolloutAgentLoopWorker
146+
├── vllm_rollout/
147+
│ └── vllm_async_server.py # PartialRolloutvLLMHttpServer / PartialRolloutvLLMReplica (paused gate)
148+
└── run/
149+
├── gsm8k_tool_config.yaml # tool registry for the tool variants
150+
├── run_qwen3-0.6b_gsm8k_grpo.sh # PR, single-turn
151+
├── run_qwen3-0.6b_gsm8k_grpo_tool.sh # PR, tool-call
152+
├── run_qwen3-0.6b_gsm8k_grpo_baseline.sh # baseline, single-turn
153+
├── run_qwen3-0.6b_gsm8k_grpo_tool_baseline.sh # baseline, tool-call
154+
└── run_{dapomath,gsm8k}_{nopr,pr}_grpo_4b_*.sh # 4B Qwen3 ports of the recipe PR
155+
```

0 commit comments

Comments
 (0)