[rollout] feat: implement partial rollout feature on rollout engine side by guapisolo · Pull Request #2 · THUDM/slime

guapisolo · 2025-06-19T22:34:34Z

Design overview

This PR completed the first part of the partial rollout feature

Support rollout sample collection and reusage in data buffer.
Support sending rollout abortion to SGLang router.
Support custom buffer read and write operation in buffer_hub.
Update rollout process to better integrate filters.
Update sample status and metadata management.
Add run scripts for deepseek-r1-distill-qwen2.5-1.5B.
Update essential input arguments.

Details

After rollout engine collects enough samples, it will send request to SGLang router to abort running samples.
If partial rollout is enabled, rollout engine will collect redundant samples (aborted, or excess completed samples).
sampling_batch_size is the granularity of the sampling batch in the rollout function. When the number of available samples falls below the target, a sampling operation of size sampling_batch_size will be triggered.
over_sampling_filter_input_size is the input size for the over sampling filter. It will replace the rollout_batch_size as target batch size (total number of valid samples to be collected)
Now support FIFO buffer operation for normal rollout and partial rollout, see fifo.py and partial_fifo.py

Future work

Update gradient calculation to tackle off-policy samples.

Evaluation

[TODO]

Speed up rollout
raw reward converge
eval reward converge

zhuzilin · 2025-06-21T10:49:41Z

maybe changing the sample.status into sample.finish_reason and do:

sample.finish_reason = Sample.FinishReason[ output["meta_info"]["finish_reason"]["type"].upper() ]

would be cleaner.

maybe changing the sample.status into sample.finish_reason and do:

sample.finish_reason = Sample.FinishReason[ output["meta_info"]["finish_reason"]["type"].upper() ]

would be cleaner.

But Pending samples haven't start rollout yet. Use finish_reason may be confusing, setting to None may cause bugs.

make sense. let's keep the origin implementation.

zhuzilin · 2025-06-21T10:56:24Z

we can always add the samples back to the buffer and remove them using the buffer_filter_path,

we can always add the samples back to the buffer and remove them using the buffer_filter_path,

Got it. When I write the code, I never consider the situation when train engine and rollout engine are fully async.
So the staleness info should be maintained in the buffer_filter?

staleness info should be maintained in the buffer_filter

I think we can same the state of the sample into Sample.metadata, and pass the currently rollout id into buffer_filter.

Co-authored-by: Jiajun Li <guapisolo@gmail.com> Co-authored-by: Yuzhen Zhou <492129152@qq.com>

…utdated doc, fix reward return value, cancel filter rename, delete debug output

[rollout] feat: implement partial rollout feature on rollout engine side

code-walk-through eng version

Add Dockerfiles for SGLang and TE FP8 with necessary patches

[rollout] feat: implement partial rollout feature on rollout engine side

Upstream sync 2026 02 22

Keeps sample.rollout_log_probs, loss_masks, response_token_ids, and response length-aligned across turns. Four fixes: 1. Reset stale sample state at entry so a retried (previously aborted) sample doesn't concat new tokens onto old log-probs. 2. Clamp per-turn max_new_tokens to the remaining context budget so total_length can't blow past max_context_length, producing samples larger than the training-side per-partition cap. 3. Abort when sglang returns text without output_token_logprobs. The old fallback retokenized, which grows response_token_ids without matching log-probs — abort lets the rollout manager re-queue cleanly. 4. Trim post-tool-output overflow. Fix #2 clamps the model's generation, but tool output is appended unconstrained. Trim tokens / loss_masks / log-probs together and mark TRUNCATED. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

THUDM#1 Strip nested chat tokens / chat-template the selector prompt The selector's "### Problem" section was being filled with sample.prompt — which slime had already passed through apply_chat_template, so it contained literal <|im_start|>user / <|im_end|> / <|im_start|>assistant tokens. Embedding those inside the selector's own (unwrapped) prompt produced a nested chat structure the model couldn't parse, leading to malformed output like "Judgment: IDX.AUTHOR / Index ID: 4". Fix: - Add _strip_chat_tokens() to remove Qwen/sglang chat-control tokens from any text before embedding it inside another prompt body. - Add _wrap_user_turn() that uses the tokenizer's chat_template to wrap the FULL selector / rewriter prompt as a proper user turn (so the model receives an unambiguous chat message instead of bare text). - run_agent_system now keeps two versions of the problem: solver_prompt (chat-formatted, sent straight to sglang) and raw_problem (stripped, used inside selector / rewriter templates). - Selector / rewriter solutions are also stripped of stray <|im_end|> tokens before being interpolated into their templates. Applied to both examples (two_agent and multi_agent). THUDM#2 Don't anti-train solvers when the selector fails to parse run_agent_system used the mean of ALL selector samples' rewards (including parse failures with reward 0) to drive the global reward shaping. When the selector regex fails (selector emits malformed "Judgment:" output), every selector sample gets reward=0, mean drops below 0.5, weight = incorrect_reward_weight (0.8) → all solvers get their rewards multiplied by 0.8 → correct solvers are penalized for the selector's parse bug. Fix: track parsed_selector_rewards separately. If the list is empty (every selector failed to parse), skip the global reward shaping entirely and return raw solver / rewriter rewards. This matches the intent: "we have no judgment signal, don't pretend we do." Same fix in both examples.

Combines SPEC §10.1 PRs #2-#5 into one atomic rewrite because all four touch the same single file (M1 decision: no sub-module split). Each PR's intent preserved as a distinct §section banner. PR #2 - §0 TOC docstring + §2 DATACLASSES (Turn, SubSession, SubSnapshot, Session) - Turn loses branch_kind/parent_id/parent_prefix_len/full_ids (0521 legacy) - Turn gains tito_masked: bool (U3) - SubSession gains num_aborts / last_finish_reason / tito_masked_turn_count - SubSnapshot adds finish_reason / num_aborts / tito_masked_turn_count - Session: active_subagent (flat) + completed_subagents + _emit_order (I6) - Session: num_aborts / num_aborts_this_turn / tito_*_turn_count - §3 PRIMITIVES & §4 STORE banners + Store.open_session(record_raw_dump=) PR #3 - §5 TRANSLATE adds verify_tito_for_turn (D2) - pure function: retokenize(decode(output_ids)) == output_ids - §6 ENGINE keeps AbortCoordinator + generate_with_abort_resume (P2 verbatim) PR #4 - §7 SEGMENTS rewrite - delete classifier triple: _classify_branch / _new_turn / 2-pass parent - delete _COMPACT_RESUME_MARKER text sniff + _SUMMARIZATION_MARKERS - delete _SubSession.subagent_stack nested stack -> flat active_subagent - new pick_target with nested-dispatch fail-safe (R2 CC3) - new classify_and_apply (4-condition is_append + pre_wipe snapshot) - new snapshot_subagent + pop_session_split (chronological replay of _emit_order) - segment meta key renamed kind -> segment_kind (U3 / R3) - U3: per-segment finish_reason + num_aborts + tito_masked_turns PR #5 - §8 HANDLER 15-step rewrite + §9 SHELL - _handle_messages: 380 lines -> ~140 lines, numbered 1-15 - step 12: I8 (n>0 guard) + I9 (abort skip TITO) + per-turn mask (only [-n:]) - Session.lock is sole sync primitive (P6); helpers never re-acquire (I3) - open_session(record_tree=) renamed to open_session(record_raw_dump=) - pop_session() single-segment API removed (list mode always on) - _export_raw_dump emits v4 schema (drops full_ids, adds tito_masked/etc) Tests (SPEC §7.1): 22 cases across 5 files, all passing - test_segments_classify.py - 7 cases (pre_wipe, nested fail-safe, etc.) - test_translate_tito.py - 6 cases incl. test_empty_turn_skip (I8 fix) and test_abort_skip_tito (I9) - test_engine_abort_resume.py - 3 cases (concatenate / max-attempts / budget) - test_session_lock.py - 2 cases (same-sid serialize, different sid) - test_pop_session_split.py - 4 cases (chronological, drain, U3 fields) - fixtures/README.md - schema doc

Implements SPEC §10.3 (v0.3 round 3, U5 decision). - mv slime/utils/aiohttp_threaded.py -> examples/coding_agent_rl/aiohttp_threaded.py - middleware.py:69 import: from slime.utils.aiohttp_threaded -> from aiohttp_threaded (bare import to match sub-agent's existing sys.path + bare-import style; examples/coding_agent_rl/ is not a package so SPEC's relative-import form `from .aiohttp_threaded` doesn't apply here) Notes: - Logically part of PR #2 (dataclass/docstring cleanup) per SPEC §10.3, but filed as a standalone commit because the original PR #2-#5 commit (70911892) is no longer HEAD (HEAD is PR THUDM#6 / 159b2b0c); amending HEAD would semantically corrupt PR THUDM#6 scope. User may interactively rebase to fold this into 70911892 if desired. - No 0521 legacy test/script files in this worktree, so only 1 import needed updating (vs SPEC §10.3 listing 5 files — those don't exist here). - All 30 smoke tests still pass.

Combines SPEC §10.1 PRs #2-#5 into one atomic rewrite because all four touch the same single file (M1 decision: no sub-module split). Each PR's intent preserved as a distinct §section banner. PR #2 - §0 TOC docstring + §2 DATACLASSES (Turn, SubSession, SubSnapshot, Session) - Turn loses branch_kind/parent_id/parent_prefix_len/full_ids (0521 legacy) - Turn gains tito_masked: bool (U3) - SubSession gains num_aborts / last_finish_reason / tito_masked_turn_count - SubSnapshot adds finish_reason / num_aborts / tito_masked_turn_count - Session: active_subagent (flat) + completed_subagents + _emit_order (I6) - Session: num_aborts / num_aborts_this_turn / tito_*_turn_count - §3 PRIMITIVES & §4 STORE banners + Store.open_session(record_raw_dump=) PR #3 - §5 TRANSLATE adds verify_tito_for_turn (D2) - pure function: retokenize(decode(output_ids)) == output_ids - §6 ENGINE keeps AbortCoordinator + generate_with_abort_resume (P2 verbatim) PR #4 - §7 SEGMENTS rewrite - delete classifier triple: _classify_branch / _new_turn / 2-pass parent - delete _COMPACT_RESUME_MARKER text sniff + _SUMMARIZATION_MARKERS - delete _SubSession.subagent_stack nested stack -> flat active_subagent - new pick_target with nested-dispatch fail-safe (R2 CC3) - new classify_and_apply (4-condition is_append + pre_wipe snapshot) - new snapshot_subagent + pop_session_split (chronological replay of _emit_order) - segment meta key renamed kind -> segment_kind (U3 / R3) - U3: per-segment finish_reason + num_aborts + tito_masked_turns PR #5 - §8 HANDLER 15-step rewrite + §9 SHELL - _handle_messages: 380 lines -> ~140 lines, numbered 1-15 - step 12: I8 (n>0 guard) + I9 (abort skip TITO) + per-turn mask (only [-n:]) - Session.lock is sole sync primitive (P6); helpers never re-acquire (I3) - open_session(record_tree=) renamed to open_session(record_raw_dump=) - pop_session() single-segment API removed (list mode always on) - _export_raw_dump emits v4 schema (drops full_ids, adds tito_masked/etc) Tests (SPEC §7.1): 22 cases across 5 files, all passing - test_segments_classify.py - 7 cases (pre_wipe, nested fail-safe, etc.) - test_translate_tito.py - 6 cases incl. test_empty_turn_skip (I8 fix) and test_abort_skip_tito (I9) - test_engine_abort_resume.py - 3 cases (concatenate / max-attempts / budget) - test_session_lock.py - 2 cases (same-sid serialize, different sid) - test_pop_session_split.py - 4 cases (chronological, drain, U3 fields) - fixtures/README.md - schema doc

Implements SPEC §10.3 (v0.3 round 3, U5 decision). - mv slime/utils/aiohttp_threaded.py -> examples/coding_agent_rl/aiohttp_threaded.py - middleware.py:69 import: from slime.utils.aiohttp_threaded -> from aiohttp_threaded (bare import to match sub-agent's existing sys.path + bare-import style; examples/coding_agent_rl/ is not a package so SPEC's relative-import form `from .aiohttp_threaded` doesn't apply here) Notes: - Logically part of PR #2 (dataclass/docstring cleanup) per SPEC §10.3, but filed as a standalone commit because the original PR #2-#5 commit (70911892) is no longer HEAD (HEAD is PR THUDM#6 / 159b2b0c); amending HEAD would semantically corrupt PR THUDM#6 scope. User may interactively rebase to fold this into 70911892 if desired. - No 0521 legacy test/script files in this worktree, so only 1 import needed updating (vs SPEC §10.3 listing 5 files — those don't exist here). - All 30 smoke tests still pass.

- #1 add --opsd-offload-teacher-logits to offload full-vocab teacher logits to CPU between forwards (moved back to device per micro-batch in the loss); keep the on-GPU view by default. Chunked JSD remains future work (THUDM#4, deferred). - THUDM#2 warn when privileged_info is empty/None (teacher==student context -> ~0 signal). - THUDM#3 repack teacher micro-batches by teacher lengths on the actor side, keeping the student's sample-to-rank assignment so response positions stay aligned (repack_micro_batches_by_length in dp_schedule.py; forward-only, per-rank). - THUDM#5 validate OPSD is not combined with --disable-compute-advantages-and-returns. - THUDM#6 skip loading the ref model under OPSD (never forwarded). - THUDM#7 descriptive error when a dataset row lacks the privileged-info field. Tests: add repack unit tests (coverage, token budget, oversized-sample-alone).

#1 (critical): vocab-parallel log-softmax normalizer was all-reduced with an identity backward, under-counting the student-logit gradient by ~1/TP when tensor-parallel size > 1 (the global normalizer couples all ranks' log-probs, so its cotangent must be all-reduced). Add _VocabParallelAllReduceSumGradAllReduce (all-reduce forward AND backward) for the normalizer; keep identity-backward _VocabParallelAllReduceSum for the final replicated jsd reduction. Add a distributed (TP=2, gloo) test that shards the vocab and checks the JSD value and student gradient against the dense single-process reference (would fail under the old identity backward; a TP=1 test cannot catch it). THUDM#2 restore the actor as the live model via try/finally around the teacher forward so an OOM there can't leave teacher weights live for backup("actor"). THUDM#3 clone the [R, V] teacher response slice instead of keeping a view into the full [1, T_padded, V] microbatch buffer (frees the padded buffer; cheap, strictly better than the view). --opsd-offload-teacher-logits still offloads to CPU. THUDM#4 reject --opd-type=self with kl_coef!=0 or --use-kl-loss (OPSD is pure distillation; the ref model is intentionally not loaded). THUDM#5 log when auto-setting --loss-type=opsd instead of overwriting silently.

…smatches Logs model output calls, label calls, matched pairs, and unmatched output/label calls at INFO level during reward computation. Works in both label mode and RM mode. Example output: [tool_rl] Model calls (2): [tool_rl] [1] get_weather(location="Beijing", unit="celsius") [tool_rl] [2] get_time() [tool_rl] Label calls (2): [tool_rl] [1] get_weather(location="Beijing") [tool_rl] [2] get_forecast(city="Shanghai") [tool_rl] Unmatched output (1): [tool_rl] [THUDM#2] get_time() [tool_rl] Match result: name=0.333 param=0.250 (matched 1/2 label calls) Co-Authored-By: Claude <noreply@anthropic.com>

guapisolo force-pushed the partial_rollout branch 5 times, most recently from 21f2226 to d9930ee Compare June 20, 2025 21:22

guapisolo marked this pull request as ready for review June 20, 2025 21:22

guapisolo changed the title ~~Implement partial rollout feature for SGLang~~ rollout: Implement partial rollout feature for SGLang (Step 1) Jun 20, 2025

guapisolo changed the title ~~rollout: Implement partial rollout feature for SGLang (Step 1)~~ [rollout] feat: implement partial rollout feature for SGLang (Step 1) Jun 20, 2025

zhuzilin reviewed Jun 21, 2025

View reviewed changes

guapisolo force-pushed the partial_rollout branch 4 times, most recently from 51514d7 to 02b78af Compare June 23, 2025 06:22

zyzshishui and others added 4 commits June 24, 2025 18:29

Squash partial rollout implementation

8765a0f

Co-authored-by: Jiajun Li <guapisolo@gmail.com> Co-authored-by: Yuzhen Zhou <492129152@qq.com>

fix bug on data metadata initialization

3a2232d

update qwen3-4B and qwen2.5-1.5B config for partial rollout, delete o…

7f3713a

…utdated doc, fix reward return value, cancel filter rename, delete debug output

update buffer read and write operation

cebe34b

guapisolo force-pushed the partial_rollout branch 2 times, most recently from a50d504 to cebe34b Compare June 24, 2025 19:06

guapisolo added 2 commits June 25, 2025 01:22

Update stats calculation

1ae9b16

format code updates using pre-commit

1d1a0ed

guapisolo force-pushed the partial_rollout branch from 97f6682 to 1d1a0ed Compare June 25, 2025 01:40

guapisolo changed the title ~~[rollout] feat: implement partial rollout feature for SGLang (Step 1)~~ [rollout] feat: implement partial rollout feature on rollout generator and update buffer management (Step 1) Jun 25, 2025

guapisolo changed the title ~~[rollout] feat: implement partial rollout feature on rollout generator and update buffer management (Step 1)~~ [rollout] feat: implement partial rollout feature on rollout engine and update buffer management Jun 25, 2025

guapisolo changed the title ~~[rollout] feat: implement partial rollout feature on rollout engine and update buffer management~~ [rollout] feat: implement partial rollout feature on rollout engine side Jun 25, 2025

guapisolo requested a review from zhuzilin June 25, 2025 01:49

Merge branch 'THUDM:main' into partial_rollout

0e8b004

zhuzilin merged commit b627fd6 into THUDM:main Jun 30, 2025

rysaya pushed a commit to rysaya/slime that referenced this pull request Aug 15, 2025

Merge pull request THUDM#2 from guapisolo/partial_rollout

c0aee07

[rollout] feat: implement partial rollout feature on rollout engine side

zyzshishui mentioned this pull request Sep 28, 2025

APRIL and slime / Report Template Issue RLsys-Foundation/APRIL#5

Closed

m1nttee pushed a commit to m1nttee/slime that referenced this pull request Nov 24, 2025

Merge pull request THUDM#2 from shuaills/code-walk-through-eng

01c0412

code-walk-through eng version

sniper35 pushed a commit to sniper35/slime that referenced this pull request Dec 21, 2025

Merge pull request THUDM#2 from xieck13/fp8_docker

8a4d9da

Add Dockerfiles for SGLang and TE FP8 with necessary patches

yueming-yuan pushed a commit to yueming-yuan/slime that referenced this pull request Dec 29, 2025

Add end-to-end accuracy tests to CI (THUDM#2)

7b9ea6d

Yangruipis pushed a commit to redai-infra/slime that referenced this pull request Feb 28, 2026

Merge pull request THUDM#2 from guapisolo/partial_rollout

7a2f14d

[rollout] feat: implement partial rollout feature on rollout engine side

mouad-hpc pushed a commit to mouad-hpc/slime that referenced this pull request Mar 12, 2026

Merge pull request THUDM#2 from Osmosis-AI/upstream-sync-2026-02-22

12c1163

Upstream sync 2026 02 22

WindowsXp-Beta mentioned this pull request Mar 18, 2026

Super tiny fix typo #738

Merged

nanjiangwill mentioned this pull request Apr 18, 2026

[1/n] Disaggregated rollout: delta weight sync (disk + NCCL transports) #1806

Merged

leofan-i mentioned this pull request Apr 24, 2026

fix: harden retool rollout against multi-turn / retry desync #1861

Merged

Navigation Menu

Uh oh!

[rollout] feat: implement partial rollout feature on rollout engine side - #2

[rollout] feat: implement partial rollout feature on rollout engine side#2
zhuzilin merged 7 commits into
THUDM:mainfrom
guapisolo:partial_rollout

guapisolo commented Jun 19, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhuzilin Jun 21, 2025

Uh oh!

guapisolo Jun 21, 2025 •

edited

Loading

Uh oh!

zhuzilin Jun 22, 2025

Uh oh!

Uh oh!

zhuzilin Jun 21, 2025

Uh oh!

guapisolo Jun 21, 2025

Uh oh!

zhuzilin Jun 22, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

guapisolo commented Jun 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Design overview

Details

Future work

Evaluation

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhuzilin Jun 21, 2025

Choose a reason for hiding this comment

Uh oh!

guapisolo Jun 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhuzilin Jun 22, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zhuzilin Jun 21, 2025

Choose a reason for hiding this comment

Uh oh!

guapisolo Jun 21, 2025

Choose a reason for hiding this comment

Uh oh!

zhuzilin Jun 22, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

guapisolo commented Jun 19, 2025 •

edited

Loading

guapisolo Jun 21, 2025 •

edited

Loading