feat(rl): gate-native async RL loop + flat Rollout contract#382
Merged
Conversation
Hecate0821
force-pushed
the
chengxi/rl-async-gate-v2
branch
2 times, most recently
from
April 30, 2026 04:12
6b05eb6 to
b66f45c
Compare
3 tasks
Hecate0821
force-pushed
the
chengxi/rl-async-gate-v2
branch
2 times, most recently
from
April 30, 2026 23:10
fbb69a1 to
ef8998e
Compare
Hecate0821
force-pushed
the
chengxi/rl-async-gate-v2
branch
4 times, most recently
from
May 1, 2026 20:31
caab904 to
9aff7ae
Compare
Adds an async-pipelined RL training recipe (recipes/async_rl_loop.py)
that decouples rollout sampling from optimizer steps, plus the
supporting rollout package and a real GSM8K multi-turn example modeled
after AReaL's examples/multi_turn_math/.
User-facing rollout contract is per-sample (one trajectory per call),
matching AReaL's ``arun_episode`` and slime's ``generate``:
async def rollout_fn(row) -> RolloutSample | None: ...
The framework fans each row out to ``completions_per_prompt`` parallel
calls and joins them by row id via ``GroupAssembler`` before handing
the assembled ``PromptGroup`` to the trainer.
Major components
- recipes/async_rl_loop.py: async loop driving rollout fan-out, group
assembly, off-policy gating (row-granular staleness budget), and the
train step. Uses the new ``LossArgs`` Protocol from main's #412
(build_loss_fn(args), validate_loss_path(args)) and is fixed to
``loss_path="client"`` -- this recipe runs the loss closure in Python
via forward_backward_custom, no server-side builtin path.
- utils/rl/async_train.py: low-level async loop primitives (RowRequest,
run_async_rl_loop, _StalenessController).
- utils/rl/rollout/: rollout package -- types (RolloutSample, Rollout,
rollout_to_prompt_group), GroupAssembler, MessageTrajectoryAssembler
(TITO-style multi-turn), TrajectoryAssembler (token-native), service
Protocol + remote adapter, native trajectory analysis.
- utils/dataloader.py: cursor-based dataloader with epoch shuffle.
- utils/data.py::replicate_rows_for_epochs: deepcopy-per-epoch helper
so rollout fns mutating their input row don't poison later epochs.
- utils/timer.py::elapsed_timer: span-based timer for nested timing.
Examples
- examples/rl/multi_turn_message_in/: GSM8K multi-turn agent ported from
AReaL. Up to N turns; on a wrong boxed answer, append AReaL's
verbatim retry-prompt and let the model try again. Reward via
math_verify with numeric fallback. prepare_data.py downloads
openai/gsm8k from HuggingFace.
- examples/rl/single_turn_token_in/: token-in single-turn baseline.
- examples/rl/vanilla_sampler.py: shared deployment sampler helper
(renamed from sampler.py for clarity).
PromptGroup gains a per-sample ``prompt_lens`` field for heterogeneous
rollouts (multi-turn, tool branches) where each sample has a different
prefix length. ``combine_prompt_groups`` prefers it when set.
``_get_loss_mask`` reads ``loss_fn_inputs["weights"]`` first (new SDK
contract) with legacy fallback to ``"loss_mask"``.
SDK pin: fireworks-ai[training]>=1.2.0a65,<2 (includes the chunk
transient backoff/retry from a64+).
Tests: 596 unit tests pass. GSM8K reward, group assembler, async loop,
and rollout adapter all covered.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Hecate0821
force-pushed
the
chengxi/rl-async-gate-v2
branch
from
May 1, 2026 20:34
9aff7ae to
b6996ae
Compare
…lag, ConcurrencyConfig Cleans up the async RL recipe surface following review. - Delete utils/rl/rollout/concurrency.py (FixedRequestGate, RequestGate Protocol, DEFAULT_REQUEST_GATE_CONCURRENCY) and its unit tests. The per-HTTP gate was redundant: row-level scheduling in _StalenessController already enforces capacity (cfg.sample_max_concurrency) and staleness ((max_offpolicy + version + 1) * batch_size - inflight) via min(), so a third lower-level gate added complexity without buying anything. Examples now use DeploymentSampler directly with no concurrency_controller wired in. - Drop Config.provision_inference: deployment provisioning is implicit in RL (DeployConfig.deployment_id None vs set already encodes create-vs-reuse). No external caller flipped this to False. Removes 8 sites of dead branching including ``or ""`` fallbacks on inference_base_url and inference_model. - Drop Config.concurrency: ConcurrencyConfig was imported, defaulted, and never read in the async recipe. sample_max_concurrency on Config is the actually-wired knob. - Tighten tokenizer load: require deployment.tokenizer_model upfront, drop the conditional ``tokenizer = None; if cfg.deployment.tokenizer_model:`` guard since the tokenizer is mandatory. - pyproject: bump fireworks-ai[training] floor to 1.2.0a65 to match the SDK pin documented in the prior commit message. - smoke imports: add new single_turn_token_in/multi_turn_message_in rollout modules. Tests: 636 unit tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…nc interval to 1 Brings the async RL recipe to AReaL parity for the inner-loop knob and surfaces two new wandb metrics that were missing or dead. - Add ``Config.ppo_n_minibatches`` (default 1) mirroring ``rl_loop.py``. The per-step path now snapshots ``old_policy_logprobs`` once and runs K ``forward_backward + optim_step`` minibatches against that snapshot, so the PPO ratio measures genuine inner-loop drift. DCP cadence stays in rollout-batch units so ``dcp_save_interval`` is unchanged. - Pin ``weight_sync_interval`` to a module-level constant ``_WEIGHT_SYNC_INTERVAL = 1``. Raising it trades rollout staleness for sync wall-time, which is almost never worth it in fully-async RL. ``WeightSyncConfig.weight_sync_interval`` is still honored by the sync recipes; the async recipe ignores it. - Emit ``train/ppo_kl`` from ``run_loss_loop`` (mean of ``exp(diff) - diff - 1`` between current policy logprobs and ``old_policy_logprobs``). Averaged across minibatches by the existing ``compute_step_metrics`` reducer. - Populate ``perf/sample_wait_time`` and fix ``perf/wait_time_ratio`` for the async path: ``async_train.py`` tracks ``last_step_end`` so the trainer's per-step wait gap is recorded, and the recipe sets ``step_wall_time = wait + train`` before ``compute_step_metrics`` fires -- so the ratio is wait/(wait+train), the overall step ratio. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…versions, dynamic filter on multi_turn_message_in Adds three CLI flags to the GSM8K multi-turn example so the test plan in PR #382 is fully reproducible without source edits. - ``--ppo-n-minibatches`` (default 1) -- inner PPO minibatch count. - ``--max-head-offpolicy-versions`` (default 0) -- staleness budget. - ``--filter-constant-reward`` -- drops prompt groups whose rewards are identical across all samples (GRPO advantage is 0 there, optimizer step is a no-op). Standard async-RL hygiene; matches AReaL's ``adv filter`` in spirit. Also drops the ``WeightSyncConfig(weight_sync_interval=1)`` override since the async recipe pins the interval to 1 internally now. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…fn(row) -> rollout_fn(sample_prompt)
The async gate now counts samples (LLM calls) end-to-end instead of
mixing samples (user-facing knob) with rows (internal state).
``_StalenessController`` tracks ``running_samples`` / ``accepted_samples``
in the same unit as ``deployment.max_batch_size``; the implicit
``// completions_per_prompt`` division at the recipe seam (and its
companion ``>= cpp`` validation) is gone. The loop floor-divides
``capacity() // completions_per_prompt`` once at the only place per-prompt
submission granularity matters.
User-facing rollout signature renamed from ``rollout_fn(row)`` to
``rollout_fn(sample_prompt)``. The dataset-cursor layer
(``CursorDataLoader``, ``--max-rows``, ``RowRequest.row_id``,
``make_row_requests``) keeps "row" terminology because those genuinely
operate on dataset rows; the rollout layer (where one row fans out to
``cpp`` parallel sample draws) is now sample-named consistently.
Other fixes in this commit:
* ``is_staleness_bound`` no longer over-attributes deployment-saturation
wall time as ``sampler_wait_for_trainer`` when both staleness and
concurrency caps are simultaneously zero. Empirically validated on
Qwen3.5-35B-A3B GSM8K: ``perf/sampler_wait_for_trainer_time`` dropped
from ~1000s/step (v13) to 0.0s/step across all 8 batches (v14).
* New ``[batch-ready ...]`` debug log line + ``async/running_samples``,
``async/accepted_samples``, ``async/staleness_capacity_at_step``,
``async/concurrency_capacity_at_step`` telemetry.
* ``perf/sample_wait_time`` split into ``perf/trainer_wait_for_sampler_time``
(gap between successive train_steps, includes weight_sync) and
``perf/sampler_wait_for_trainer_time`` (rollout-side blocked-on-trainer
cumulative). ``perf/wait_time_ratio`` and ``perf/overlap_ratio`` rebased
on the sample-level definitions; docstrings clarify the weight_sync
inclusion (Amdahl floor).
* ``_weight_sync`` logs hotload wall inline (matches the
``[step K] weight_sync (X.Xs)`` format) so the reader can correlate
per-step against the next batch-ready line without the Timer-flush
off-by-one.
* CLI flags added on the multi_turn_message_in example:
``--training-shape-id``, ``--lora-rank``, ``--max-concurrency-rollout-sample``,
``--synchronous-training``, ``--replica-count``, ``--wandb-run-name``.
* Stale docs/comments updated: ``utils/rl/README.md``,
``examples/rl/README.md``, ``rollout/{__init__,types,renderer,remote}.py``
describe the post-PR382 ``rollout_fn(sample_prompt)`` + ``RolloutSetup``
contract (no more ``rollout_fn(row, ctx)`` / ``RolloutContext`` /
``ctx_extras``).
* Test rewires for the new constructor signature
(``batch_size_samples``, ``completions_per_prompt``,
``max_concurrent_samples``, ``accepted_samples``).
Tests: 654 passed, 0 failed (full unit suite).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
In multi-turn rollouts, the response window contains masked bridge / user-feedback / tool-result tokens (``loss_mask = 0``). The previous code averaged drift metrics (``perf/inf_diff``, ``perf/inf_kld``, ``train/ppo_kl``) over the whole response slice, and -- more critically -- handed the unmasked ``resp_prox`` / ``resp_inf`` to ``compute_tis_weight``. With ``TISConfig.level == "sequence"`` the geometric mean of log-ratios includes those masked positions and the broadcast weight then multiplies into the **actual loss** at active positions (since most policy losses use ``weight * resp_mask``, the masked positions zero out, but active positions carry a contaminated weight). Fix at the two call sites (``utils/rl/common.py`` for the client-side loss, ``utils/rl/losses.py`` for the per-token IS server-side path): filter ``resp_prox``, ``resp_inf``, and ``pi_detached`` to active positions before any ``.mean()`` or ``compute_tis_weight()`` call, then expand the resulting weight back to the full response shape with ``1.0`` at masked positions (the identity weight, harmless under both mask-multiplied losses and ``dro`` which doesn't multiply by mask). For ``level="token"`` this is a metrics fix only (per-token weights at active positions are unchanged; masked positions get a clean 1.0 instead of nonsense × 0). For ``level="sequence"`` it's a loss correctness fix. Adds ``TestActiveFilterRegression`` to ``test_decoupled_is.py`` covering both levels with a worst-case mask layout. Tests: 656 passed (+2 from the new regression). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Net -91 comment lines (113 removed, 22 tighter replacements) across async_train.py, async_rl_loop.py, metrics.py, common.py, losses.py. Drops what-restating and elaboration prose, keeps only WHY for non-obvious behavior: - ``_StalenessController`` docstring: removes the unification rationale paragraph (the field names + class docstring already convey it). - ``is_staleness_bound`` docstring: keeps the misattribution warning, drops the "why this matters" elaboration. - ``_refill``: drops three explanatory comments; the helper is small and the wait-window state machine is self-evident from the calls. - ``[batch-ready ...]`` log site: drops the 6-line "Debug: surface the gating attribution" preamble; keeps the log line. - ``_weight_sync``: drops the 4-line off-by-one Timer-flush rationale (the inline ``logger.info`` is self-explanatory). - ``Config.max_concurrency_rollout_sample``: 13-line docstring -> 4. - ``compute_step_metrics`` perf block: 3 multi-line WHY comments collapsed to one line each (Amdahl note, weight_sync inclusion, staleness vs concurrency boundary). - TIS active-mask comments in common.py + losses.py: 7-line WHY -> 2-line "matches slime/AReaL" pointer; 5-line expansion rationale -> 2-line identity-weight note. No code semantics change; no silent-fallback or error-handling adjustments. Tests: 656 passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- max_head_offpolicy_versions docstring/CLI now says "weight-sync (policy)
versions" -- the gate's version increments per weight_sync_fn call, not
per optimizer step (so K-minibatch runs don't undercount staleness).
- rollout/entropy now averages -logprob over loss_mask>0 positions only.
Multi-turn samples emit logprobs=0.0 / loss_mask=0 on bridge/user/tool
tokens, which biased the metric toward 0 under the old window slice.
- RolloutSetup docstring: drop the stale "divides by completions_per_prompt"
text -- the runner takes the sample-level cap directly.
- test_async_rl_train: fix max_concurrent docstring (samples, not rows),
add a cpp>1 sample-level cap test, and pin the cpp>cap deadlock guard.
- Code-style pass: trim multi-line comment blocks in async_rl_loop.py,
async_train.py, metrics.py, rollout/{types,assembler,remote}.py and
drop a defensive try/except in datum_target_len.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- vanilla_sampler.build_deployment_sampler: docstring still claimed the recipe divides max_concurrency_rollout_sample by completions_per_prompt to derive a row-level cap; the runner takes the sample-level cap directly. - test_caps_in_flight_at_sample_level_with_cpp_gt_1: docstring promised a version-offset assertion the test never made; trimmed to describe what the test actually pins (the <=8 sample cap) and dropped the unused ``admitted_per_version`` dict. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Architecture/contract docs for the async RL recipe now live in skills/dev/references/rl/async-rl.md (the canonical place for day-to-day training reference per skills/dev/SKILL.md), not under training/. Covers: - sync vs async recipe comparison table - rollout API (rollout_fn(sample_prompt) -> RolloutSample, RolloutSetup, factory pattern) - sample-level off-policy gate semantics (B_s / O / C_s / R, weight-sync versions, sizing rule O >= R-1) - diagnosing perf/wait_time_ratio - TIS/drift/entropy active-mask convention - example pointers SKILL.md routing table gains three entries (async-rl, perf wait diagnosis, gate sizing). recipes.md lists async_rl_loop.py. In-tree READMEs trimmed to brief pointers + a layout table; the substantive prose moved into the skill so readers find it through the skill index instead of stumbling on it under training/. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Module docstring + Config.main() now log a WARNING that async_rl_loop is under active development and the Config / RolloutSetup surface may change without backward-compat shims. - skills/dev/references/rl/async-rl.md gains an EXPERIMENTAL banner and a "What you customize (and what you don't)" section calling out the recipe's design intent: extend the rollout function; everything else (gate, advantage, ref forward, weight sync, KL/TIS, PPO inner loop, checkpoints) is owned by main(). - In-tree READMEs (training/utils/rl, training/examples/rl, training/examples/rl/multi_turn_message_in) gain the same EXPERIMENTAL note and switch their cross-link to the skill from ../../.. relative paths to repo-root absolute paths. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Hecate0821
marked this pull request as ready for review
May 4, 2026 06:29
xiaoyifan
approved these changes
May 4, 2026
Add a short acknowledgements block to the docstrings of recipes/async_rl_loop.py and utils/rl/async_train.py listing the three RL frameworks whose designs informed this loop, with their GitHub URLs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ng section Compact the in-code docstrings on Config and RolloutSetup -- at most two lines per arg, with full explanations moved to the skill. Replace the short "Diagnosing waits" stub with a dedicated "Metrics: tuning staleness and the trainer/sampler GPU split" section that maps each wait metric to the concrete knob to turn (``max_head_offpolicy_versions``, replica count, ``max_concurrency_rollout_sample``), and frames the goal as the sampler-bound, minimal-trainer-wait regime. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The async recipe supports both strict on-policy (max_head_offpolicy_versions=0) and off-policy with bounded staleness; it is a strict superset of the sync recipe. Reframe the doc to lead with that, recommend async as the default for new RL work, and note that sampler_wait_for_trainer_time is structurally non-zero in O=0 mode (correct behavior, not a misconfiguration). Drop two internal-flavored references: * "we typically run R=4 with O=4" -- replaced with neutral guidance. * "deployment shape mis-configured (devShmSize, base hotload)" -- replaced with a generic "investigate the deployment configuration separately". Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Phrased as "we've tested ... and found it healthy" rather than the original "we typically run", so the line is a validated reference point readers can use, not an internal default. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Gate formula in skill doc now reflects row-atomic admission: a row admits iff both staleness_slots and concurrency_slots have >= cpp free, matching the runner's ``capacity() // completions_per_prompt``. * "single_turn_token_in" example: clarify the call is per ``rollout_fn`` invocation (cpp invocations per dataset row), not per row. * ``async_train.py`` module docstring: replace "gating is per dataset row" with the precise framing -- submission is row-atomic; bookkeeping in samples; one row consumes cpp sample slots. * ``async_rl_loop.py`` dual-axis logging note: clarify ``ctx.current_version`` and checkpoint ids stay as optimizer-step labels while the off-policy budget is weight-sync-version based, so the two axes are independent. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This was referenced May 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR adds
A new async RL recipe that overlaps rollout sampling with training, plus the supporting rollout package and a GSM8K multi-turn example ported from AReaL's
examples/multi_turn_math/.Design intent: the only thing users customize is the rollout function. Gate, advantage, ref forward, weight sync, KL/TIS metrics, PPO inner loop, and checkpoints are owned by the recipe.
One trajectory per call (matches AReaL
arun_episode/ slimegenerate). The framework fans each row out tocompletions_per_promptparallel calls, joins them viaGroupAssembler, and hands the assembledPromptGroupto the trainer.Components
recipes/async_rl_loop.py— async loop: rollout fan-out, off-policy gate (sample-level), train step, PPO inner loop. Client-side losses only (forward_backward_custom).utils/rl/async_train.py—RowRequest,run_async_rl_loop,_StalenessController(sample-unit bookkeeping).utils/rl/rollout/—RolloutSample,Rollout,rollout_to_prompt_group,GroupAssembler,MessageTrajectoryAssembler(TITO multi-turn), serviceProtocol+ remote adapter.utils/dataloader.py,utils/data.py::replicate_rows_for_epochs,utils/timer.py::elapsed_timer— supporting plumbing.Knobs (gate, all sample-unit)
prompt_groups_per_stepB_p: prompts per optimizer stepcompletions_per_promptcpp: GRPO group size per promptmax_head_offpolicy_versionsO: weight-sync versions of staleness slack (= AReaL'smax_head_offpolicyness)max_concurrency_rollout_sampleC_s: hard cap on in-flight LLM calls (maps todeployment.max_batch_size)ppo_n_minibatchesK: inner PPO steps per rollout batch (each reuses oneold_policy_logprobssnapshot)Sizing rule:
O ≥ R−1whereR = C_s / (B_p·cpp).weight_sync_intervalis pinned to 1 (raising it just trades staleness for sync wall-time).New WandB metrics
train/ppo_kl— intra-step KL between current policy and theold_policy_logprobssnapshot.perf/trainer_wait_for_sampler_time,perf/sampler_wait_for_trainer_time— symmetric idle accounting.perf/wait_time_ratio,perf/overlap_ratio— overall step async overlap.rollout/entropy— averaged overloss_mask>0only (matches slime/AReaLmasked_mean).Examples
examples/rl/multi_turn_message_in/— GSM8K agent ported from AReaL (retry on wrong\boxed{...}; reward viamath_verify).examples/rl/single_turn_token_in/— token-in single-turn baseline.SDK pin
fireworks-ai[training]>=1.2.0a65,<2Tests
658 unit tests pass. New coverage:
test_async_rl_train.py,test_group_assembler.py,test_rollout_{types,assembler,message,helpers,trace}.py,test_gsm8k_reward.py,test_cursor_dataloader.py,test_data_utils.py,test_decoupled_is.py::TestActiveFilterRegression(TIS active-mask).E2E (GSM8K multi-turn, Qwen3.5-35B-A3B LoRA-r16, 8×B200)
Run is purely sampler-bound (
perf/sampler_wait_for_trainer_time ≈ 0,perf/trainer_wait_for_sampler_time>> 0) — confirms the gate is healthy and the trainer is fully overlapped behind rollouts:References