feat(rl): gate-native async RL loop + flat Rollout contract

[Hecate0821]

What this PR adds

An async 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 (per-sample, matches AReaL/slime)

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. AReaL's arun_episode and slime's generate use the same shape.

Major components

• recipes/async_rl_loop.py — async loop driving rollout fan-out, group assembly, off-policy gating (row-granular staleness budget), the train step, and the inner PPO loop. Uses main's LossArgs Protocol (build_loss_fn(args), validate_loss_path(args) from refactor(losses): collapse build_loss_fn 9-kwarg call site to LossArgs Protocol #412) 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 primitives (RowRequest, run_async_rl_loop, _StalenessController).
• utils/rl/rollout/ — rollout package: 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.

AReaL-parity knobs

• Config.ppo_n_minibatches (default 1, mirrors rl_loop.py) — each rollout batch snapshots old_policy_logprobs once and runs K forward_backward + optim_step minibatches against that snapshot. K=1 reproduces legacy 1:1 behavior; K>1 makes the PPO ratio measure genuine inner-loop drift and the clip do real work. Collapses AReaL's gradient_accumulation_steps and ppo_n_minibatches into one knob.
• Config.max_head_offpolicy_versions ≡ AReaL's max_head_offpolicyness (staleness budget, row-granular).
• weight_sync_interval is pinned to 1 inside the recipe (constant _WEIGHT_SYNC_INTERVAL). Configurable in WeightSyncConfig for sync recipes; the async recipe ignores it because raising it just trades rollout staleness for sync wall-time, which is almost never worth it.

New WandB metrics

• train/ppo_kl — intra-step KL between the current policy and the step-start old_policy_logprobs snapshot, averaged across the K inner minibatches. Slime/AReaL/OpenRLHF naming. ~0 when ppo_n_minibatches=1; >0 when the inner loop drifts.
• perf/sample_wait_time — per-step seconds the trainer waited for the next batch to fill (was hardcoded 0.0).
• perf/wait_time_ratio — wait / (wait + train), the overall step ratio. With ppo_n_minibatches > 1 the denominator covers all K minibatches because they run sequentially before train_step returns.
• perf/overlap_ratio — 1 - wait_time_ratio.

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.

Loss subsystem changes

• PromptGroup gains a per-sample prompt_lens: List[int] | None 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

Unit tests

~600 unit tests pass on this branch. New coverage for this PR:

• tests/unit/test_async_rl_train.py — full async loop behavior (staleness controller, batch assembly, dynamic filter, weight sync cadence, ppo_n_minibatches path).
• tests/unit/test_group_assembler.py — row fan-in semantics.
• tests/unit/test_gsm8k_reward.py — boxed-answer extraction + math_verify.
• tests/unit/test_rollout_{types,assembler,message,helpers,trace}.py — rollout package.
• tests/unit/test_cursor_dataloader.py, tests/unit/test_data_utils.py.

Integration test plan (GSM8K multi-turn)

End-to-end verification on the GSM8K multi-turn example in two phases, run sequentially. Phase 2 only proceeds after Phase 1 finishes cleanly. Both phases share the same cohort and group shape:

• 200 rows of GSM8K (--max-rows 200)
• 8 samples per prompt (--completions-per-prompt 8)
• 8 prompts per optimizer step (--prompt-groups-per-step 8)
• Dynamic filter on (--filter-constant-reward drops zero-advantage groups so the optimizer doesn't no-op on flat-reward batches)
• WandB enabled

Setup (once):

cd training/examples/rl/multi_turn_message_in
python prepare_data.py            # downloads openai/gsm8k -> train.jsonl
export FIREWORKS_API_KEY=<your-key>
export WANDB_ENTITY=<your-entity>
export WANDB_PROJECT=gsm8k-mt-async

Phase 1 — ppo_n_minibatches=1, 1-version off-policy:

python train.py \
  --base-model accounts/fireworks/models/qwen3-1p5b-instruct \
  --tokenizer-model Qwen/Qwen2.5-1.5B-Instruct \
  --max-rows 200 \
  --completions-per-prompt 8 \
  --prompt-groups-per-step 8 \
  --max-head-offpolicy-versions 1 \
  --ppo-n-minibatches 1 \
  --filter-constant-reward

WandB run should show:

• train/ppo_kl ≈ 0 (no inner-loop drift when K=1)
• async/version_offset_max ≤ 1
• perf/sample_wait_time and perf/wait_time_ratio populated and finite
• rollout/reward trending upward across steps
• train/mean_kl finite, no NaN/inf
• Run terminates cleanly after the dataset is exhausted; final checkpoint promotable

Phase 1 success criteria: run finishes without error, reward improves, wait-ratio metric is non-zero (proves the new instrumentation works).

Phase 2 — ppo_n_minibatches=2, 2-version off-policy:

Only after Phase 1 succeeds:

python train.py \
  --base-model accounts/fireworks/models/qwen3-1p5b-instruct \
  --tokenizer-model Qwen/Qwen2.5-1.5B-Instruct \
  --max-rows 200 \
  --completions-per-prompt 8 \
  --prompt-groups-per-step 8 \
  --max-head-offpolicy-versions 2 \
  --ppo-n-minibatches 2 \
  --filter-constant-reward

WandB run should show:

• train/ppo_kl > 0 (the second minibatch sees real drift from the snapshot)
• train/ppo_clip_frac > 0 (PPO clipping fires)
• async/version_offset_max ≤ 2 (staleness budget honored)
• perf/wait_time_ratio ≤ Phase 1 (more concurrency headroom from the larger budget)
• Optimizer-step count is ~2× the accepted-batch count (2 inner steps per batch)
• rollout/reward trending upward
• DCP cadence unchanged vs. Phase 1 (saves still on rollout-batch granularity, not optim-step granularity)

Phase 2 success criteria: run finishes without error, ppo_kl > 0 (proves the inner loop is actually running), version offset bounded by the budget.

References

• AReaL: https://github.com/inclusionAI/AReaL/tree/main/examples/multi_turn_math
• slime: https://github.com/THUDM/slime