feat: async partial rollout trainer with sample supplementation and caching - #58
feat: async partial rollout trainer with sample supplementation and caching#58mamazi0131 wants to merge 1 commit into
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a comprehensive asynchronous partial rollout (APR) system to the verl framework, designed to dramatically enhance the efficiency of reinforcement learning training, particularly when dealing with datasets containing samples of highly varying lengths. By intelligently managing inference tasks, dynamically supplementing samples, and caching partial results, the system minimizes idle GPU time and accelerates the overall training process, leading to significant performance gains without compromising algorithmic correctness. Highlights
Changelog
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an Async Partial Rollout (APR) mechanism to enhance training efficiency, particularly for datasets with long-tail samples. The implementation is comprehensive, adding new components like PRv3AgentLoopManager, RolloutPromptManager, and specialized agent loops for partial generation. The overall design is solid and effectively addresses the stated problem. My review focuses on improving code clarity, maintainability, and fixing a few minor issues. I've identified opportunities for improvement regarding magic numbers, a potential performance concern with busy-waiting, and some inconsistencies in documentation and script files.
|
Hello, thank you very much for your work. I understand that an asynchronous training architecture with colocation similar to Kimi has been implemented now, which is also a missing part of the current Verl. In terms of design, Verl 0.7.1 supports an auto-resume mechanism, decoupling the complex state storage logic between the server and the agent. Meanwhile, parameter synchronization uniformly adopts the checkpoint engine approach, vLLM supports a multi-process mode, and the training engine is integrated through the Model Engine interface in a unified manner. All these modifications facilitate subsequent development and iteration. It is suggested to refer to the following PRs and the current code to refactor this PR: the rollout module shall leverage the auto-resume capability, the training module shall adopt the Model Engine, and parameter synchronization shall use the checkpoint engine, so as to align with the current code and future planning. [Completed] vLLM multi-process: verl-project/verl#4280 |
|
hello @mamazi0131 , i'm a new beginer of verl, i'm glad to help you to refactor this code, could i help you to do that? |
|
@ArronHZG do you still need this feature, in v0.8.0 ? |
I’d be happy to, of course. I’ve been so busy with work lately that I haven’t had time to take care of this. |
thank you! |
## 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.
3 participants
What does this PR do?
Checklist Before Starting
[{modules}] {type}: {description}(This will be checked by the CI){modules}includefsdp,megatron,veomni,sglang,vllm,rollout,trainer,ci,training_utils,recipe,hardware,deployment,ray,worker,single_controller,misc,perf,model,algo,env,tool,ckpt,doc,data,cfg,reward,like[megatron, fsdp, doc]{type}is infeat,fix,refactor,chore,test[BREAKING]to the beginning of the title.[BREAKING][fsdp, megatron] feat: dynamic batchingKey Accomplishments:
Implemented Sample Supplementation and Interruption Mechanisms (SSIM) for dynamic sample replenishment.
Introduced Rollout Caching via a state-aware PromptsManager to resume partial generations, effectively managing sample staleness.
Ensured Off-Policy Correctness for PPO-style algorithms (GRPO/DAPO) using decoupled importance sampling.
Achieved up to 51.1% reduction in end-to-end training time on complex reasoning datasets.
Test
We validated the APR mechanism on two benchmarks using 2 nodes with 8 H20 GPUs and the Qwen3-4B model:
Under consistent convergence, training time was reduced by 11.7% with a 5.93% boost in GPU utilization.
In the presence of 160k-token long-tail samples, the APR achieved a 51.1% reduction in total training time while maintaining superior final performance.
API and Usage Example
Users can now trigger the partial rollout mode by using the specific recipes provided in the recipe/partial_rollout/ directory.
# Run DAPO-MATH17k with Partial Rollout on 2 nodes bash recipe/partial_rollout/run_dapo_math17k_pr_4b_2node.shDesign & Code Changes
Sample Supplementation and Interruption Mechanisms:
Introducing sample supplementation and interruption mechanisms to enable dynamic sample replenishment and automated scheduling of inference tasks.
Rollout Caching:
Using a prompt manager to resume partial rollouts, managing complete and partial samples in the buffer based on sample staleness.
Checklist Before Submitting
Important
Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=alwaysci-requestchannel in theverlSlack workspace. (If not accessible, please try the Feishu group (飞书群).)recipesubmodule, please also update the reference to the submodule commit viagit submodule update --remoteorcd recipe && git pull origin main.