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

skills/dev/SKILL.md

@@ -25,6 +25,9 @@ The cookbook is the reference implementation of the Fireworks Training SDK. Fork
 | "Custom loss for RL" | [`references/rl/custom-loss.md`](references/rl/custom-loss.md) |
 | "RL hotload / weight sync cadence, on-policy vs off-policy, `weight_sync_timeout`" | [`references/rl/hotload.md`](references/rl/hotload.md) |
 | "Concurrency control for RL rollouts — adaptive vs fixed?" | [`references/rl/concurrency.md`](references/rl/concurrency.md) |
+| "Async RL — overlap rollout with training, off-policy budget, PPO inner minibatches, `rollout_fn(sample_prompt)`" | [`references/rl/async-rl.md`](references/rl/async-rl.md) |
+| "Why is `perf/wait_time_ratio` high?" / `perf/sampler_wait_for_trainer_time` / `perf/trainer_wait_for_sampler_time` | [`references/rl/async-rl.md`](references/rl/async-rl.md#diagnosing-waits) |
+| "How do I size `max_head_offpolicy_versions` and `max_concurrency_rollout_sample`?" | [`references/rl/async-rl.md`](references/rl/async-rl.md#the-off-policy-gate-sample-level) |
 | "How do I promote a checkpoint?" | [`references/tools.md`](references/tools.md#promote_checkpointpy) |
 | "Which checkpoints does the server know about / are promotable?" | [`references/tools.md`](references/tools.md#listing-checkpoints-fireworksclientlist_checkpoints) — `FireworksClient.list_checkpoints(job_id)` |
 | "How do I reconnect a training **client** to a running trainer?" | [`references/tools.md`](references/tools.md#reconnect_and_adjust_lrpy) |

skills/dev/references/recipes.md

@@ -9,6 +9,7 @@ Each recipe is a single Python file in `training/recipes/` that wires the Traini
 | ORPO | `training/recipes/orpo_loop.py` |
 | Importance-weighted GRPO | `training/recipes/igpo_loop.py` |
 | Generic RL loop (GRPO scaffold) | `training/recipes/rl_loop.py` |
+| Async RL loop (rollout/train overlap, PPO inner minibatches) | `training/recipes/async_rl_loop.py` — see [`rl/async-rl.md`](rl/async-rl.md) |
 
 ## "Reference loop" means these files
 
@@ -51,6 +52,7 @@ RL has its own skill folder. Open [`rl/`](rl/) when working with `rl_loop.py`:
 - [`rl/dynamic-filter.md`](rl/dynamic-filter.md) — `should_accept`, why zero-variance groups get dropped
 - [`rl/custom-loss.md`](rl/custom-loss.md) — interface + reference implementation + RL `Config` fields
 - [`rl/hotload.md`](rl/hotload.md) — weight-sync cadence, `weight_sync_timeout`, on-policy vs off-policy, base/delta chain
-- [`rl/concurrency.md`](rl/concurrency.md) — rollout concurrency control; adaptive is the default and the recommendation
+- [`rl/concurrency.md`](rl/concurrency.md) — rollout concurrency control for the **sync** `rl_loop.py` (adaptive is the default)
+- [`rl/async-rl.md`](rl/async-rl.md) — `async_rl_loop.py` overlap recipe: sample-level cap, off-policy budget, PPO inner minibatches
 
 SFT / DPO / ORPO users do not need these.

skills/dev/references/rl/async-rl.md

@@ -0,0 +1,199 @@
+# RL: async recipe (`async_rl_loop.py`)
+
+> ⚠️ **EXPERIMENTAL — under active development.** API surface and config
+> field names may change without backward-compat shims.  Pin to a specific
+> commit if you depend on the current shape.  The recipe also emits a
+> runtime `WARNING` at `main()` start.
+
+The async recipe overlaps rollout sampling with training: while the trainer runs
+`fwd_bwd + optim_step` on batch *N*, the sampler is producing batch *N+1* against
+the previous policy version.  This is distinct from the synchronous
+`rl_loop.py`, which trains and samples in lockstep.
+
+## What you customize (and what you don't)
+
+The recipe is intentionally minimal-surface: **the only thing most users need
+to write is the rollout function**.  Everything else — gate, advantage
+computation, reference-model forwards, weight sync, KL/TIS metrics, PPO inner
+loop, checkpoint plumbing — is handled by `recipes/async_rl_loop.py::main`.
+
+| You write | The recipe handles |
+|---|---|
+| `rollout_fn_factory(setup) -> rollout_fn` (one trajectory per call) | Async fan-out / GroupAssembler / off-policy gate |
+| (optional) `dynamic_filter_fn(pg)` for batch-level filtering | Reference model forwards, KL, TIS, drift metrics |
+| Dataset rows (or pass `rows=` to `main()`) | Weight sync cadence, checkpoint save/promote, WandB axes |
+| `Config(...)` knobs (LR, gate sizes, deployment shape) | PPO inner minibatching, gradient accumulation, advantage z-score |
+
+This is the design intent — extend the rollout, not the loop.  If you find
+yourself forking the recipe, file an issue first; it usually means a knob
+should exist on `Config`.
+
+| | `rl_loop.py` (sync) | `async_rl_loop.py` (async) |
+|---|---|---|
+| Sampler/trainer overlap | None — drains rollouts before each step | Overlapped via an off-policy budget |
+| Rollout API | `make_rollout_fn(rl_cfg, deploy_mgr) -> rollout_fn` | `rollout_fn_factory(setup) -> rollout_fn` |
+| Per-call signature | `rollout_fn(prompt, n)` returns N samples | `rollout_fn(sample_prompt) -> RolloutSample \| None` (one trajectory per call) |
+| Concurrency | `ConcurrencyConfig` (adaptive AIMD) | Sample-level cap on the async runner (`max_concurrency_rollout_sample`) |
+| Off-policy | Strict on-policy | Configurable budget in weight-sync versions |
+| PPO inner steps | One `fwd_bwd + optim_step` per rollout | `ppo_n_minibatches` × inner steps per rollout |
+
+Use the async recipe when (a) you need the rollout/train overlap to fit in your
+GPU budget, or (b) you want PPO inner-loop minibatching.  Use the sync recipe
+otherwise — it's simpler and avoids the off-policy reasoning below.
+
+## The rollout API
+
+The user supplies one trajectory per call:
+
+```python
+async def rollout_fn(sample_prompt: dict) -> RolloutSample | None: ...
+```
+
+`sample_prompt` is the dataset row's dict, renamed once it crosses the
+dataset/sampling seam (the dataset emits a "row"; the sampler treats it as a
+prompt to draw a sample against).  Return `None` to drop the sample (counts as
+one lost sample within the row's group).
+
+`RolloutSample` is three parallel lists plus a scalar reward:
+
+```python
+@dataclass
+class RolloutSample:
+    tokens: list[int]
+    logprobs: list[float]   # 0.0 on non-generated positions
+    loss_mask: list[int]    # 1 on assistant tokens, 0 elsewhere
+    reward: float
+    routing_matrices: list[str] | None = None  # MoE R3 replay
+    finish_reason: str = "stop"
+    text: str = ""
+```
+
+Multi-turn rollouts flatten into the same shape: turn boundaries are implicit
+in `loss_mask` transitions (0 on prompts/user/tool, 1 on assistant).  Per-token
+mask alignment is the contract — the trainer relies on it to mask non-generated
+positions out of the loss, TIS weight, and entropy metric.
+
+The factory pattern keeps per-rollout dependencies (sampler, tokenizer, sample
+kwargs, custom state) out of the per-call signature:
+
+```python
+def make_rollout_fn(setup: RolloutSetup) -> RolloutFn:
+    sampler = DeploymentSampler(setup.inference_base_url, setup.model, setup.api_key, setup.tokenizer)
+    async def rollout_fn(sample_prompt: dict) -> RolloutSample | None:
+        ...
+    return rollout_fn
+```
+
+`RolloutSetup.extras` carries arbitrary caller state (replaces the older
+`RolloutContext.ctx_extras` setattr injection).
+
+## The off-policy gate (sample-level)
+
+The runner gates rollout submission so a sample is admitted iff:
+
+```
+samples_in_flight + samples_accepted < (max_head_offpolicy_versions + version + 1) * batch_size_samples
+                                  AND
+samples_in_flight < max_concurrency_rollout_sample   (when set)
+```
+
+All bookkeeping is in samples (LLM calls), the same unit the deployment's
+`max_batch_size` gates on.
+
+| Knob | Unit | Meaning |
+|---|---|---|
+| `prompt_groups_per_step` | prompts | `B_p`: how many prompts make one optimizer step |
+| `completions_per_prompt` | samples/prompt | `cpp`: GRPO group size per prompt |
+| `max_head_offpolicy_versions` | weight-sync versions | `O`: how many sync boundaries past submit a sample may land at the trainer.  `0` is strict on-policy |
+| `max_concurrency_rollout_sample` | samples | `C_s`: hard cap on in-flight LLM calls; map to `deployment.max_batch_size`.  Must be `>= cpp` or the gate deadlocks |
+| `ppo_n_minibatches` | minibatches | `K`: inner PPO steps per rollout batch (each with `old_policy_logprobs` snapshot reused) |
+
+Derived: `B_s = B_p × cpp` (samples per outer batch), `R = C_s / B_s` (active set
+relative to one batch).
+
+**Version semantics.** `version` increments once per `weight_sync_fn` call
+(once per outer rollout batch when `weight_sync_interval=1`, which the recipe
+pins).  It is **not** an optimizer-step counter — with `K>1` the same version
+spans `K` optim steps.
+
+**Sizing rule of thumb.**  For sustained overlap you need `O >= R - 1`.  AReaL's
+GSM8K example uses `R=1` with `O=2`; we typically run `R=4` with `O=4` (one
+margin step over the minimum).  See WandB `perf/wait_time_ratio` and
+`perf/overlap_ratio` to confirm the rollout side never starves.
+
+**Diagnosing waits.**
+
+- `perf/trainer_wait_for_sampler_time > 0` → the trainer is waiting on rollouts;
+  rollouts are the bottleneck (raise replicas or `C_s`, or accept the wait as
+  Amdahl-bound).
+- `perf/sampler_wait_for_trainer_time > 0` → the rollout side is throttled by
+  the staleness budget; raise `O` or accept that the trainer is the bottleneck.
+- Healthy async has the first metric > 0 and the second `~0`: concurrency cap
+  binds before staleness.
+
+## Configuration cheatsheet
+
+```python
+from training.recipes.async_rl_loop import Config, main
+
+cfg = Config(
+    log_path="./logs",
+    base_model="accounts/fireworks/models/qwen3-1p5b-instruct",
+    learning_rate=1.7e-5,
+    completions_per_prompt=8,
+    prompt_groups_per_step=8,            # B_s = 64 samples per batch
+    max_head_offpolicy_versions=4,       # O = 4 weight-sync versions of slack
+    max_concurrency_rollout_sample=256,  # C_s = 256 -> R = 4x batch in flight
+    ppo_n_minibatches=2,                 # K = 2 inner PPO steps per rollout
+    max_completion_tokens=16384,
+    deployment=DeployConfig(tokenizer_model="Qwen/Qwen2.5-1.5B-Instruct"),
+    infra=InfraConfig(training_shape_id="..."),
+)
+
+main(cfg, rollout_fn_factory=make_rollout_fn, rows=rows)
+```
+
+`synchronous_training=True` forces the loop to drain rollouts before every
+train step; useful as a baseline for measuring the async overlap savings (it
+makes `perf/sampler_wait_for_trainer_time` ≈ train+sync wall time).
+
+## Loss path
+
+Async is **client-side only**.  `loss_path` is fixed; the server-side built-in
+path forbids `kl_beta>0` and `pipeline_parallelism>1`, both of which the async
+loop relies on.  Use `rl_loop.py` if you need the server-side fast path.
+
+## TIS / drift metrics
+
+`utils/rl/common.py::run_loss_loop` computes TIS weights and the
+`inference_diff` / `inference_kld` / `ppo_kl` drift metrics over **active
+positions only** (`loss_mask>0`).  Including masked bridge/user/tool tokens
+biases the geometric-mean TIS toward 1 and the drift metrics toward 0.  This
+matches slime's `tis_level="geometric"` `masked_mean` and AReaL's
+`masked_mean(log_diff, mask, expand=True)` semantics.
+
+The `rollout/entropy` metric in `metrics.py::compute_step_metrics` is also
+masked (`-logprob` averaged over `loss_mask>0`).
+
+## Examples
+
+Two minimal examples ship under `training/examples/rl/`:
+
+- `single_turn_token_in/` — pre-tokenized rows; `rollout_fn` calls the
+  `/v1/completions` token-in/token-out path once per row.
+- `multi_turn_message_in/` — OpenAI-style messages; `rollout_fn` runs a retry
+  loop using `MessageTrajectoryAssembler` to keep prior assistant tokens
+  exact across turns.  Ports AReaL's `examples/multi_turn_math/` to this recipe.
+
+Each example exposes a `rollout_fn_factory(setup)` and a `train.py` that wires
+the dataset and factory into `recipes.async_rl_loop.main`.
+
+## Related
+
+- [`recipes.md`](../recipes.md) — sync `rl_loop.py` overview
+- [`hotload.md`](hotload.md) — weight sync internals (the version counter that
+  `max_head_offpolicy_versions` budgets against)
+- [`gradient-accumulation.md`](gradient-accumulation.md) — PPO minibatch
+  gradient normalization
+- [`dynamic-filter.md`](dynamic-filter.md) — async runner accepts the same
+  `dynamic_filter_fn` signature

training/examples/rl/README.md

@@ -0,0 +1,26 @@
+# Async RL examples
+
+> ⚠️ **EXPERIMENTAL — `async_rl_loop` is under active development.**
+> Config/API may change without backward-compat shims.  See
+> [`/skills/dev/references/rl/async-rl.md`](/skills/dev/references/rl/async-rl.md).
+
+The recipe is intentionally minimal-surface: **the only thing you customize
+is the rollout function** (`rollout_fn(sample_prompt) -> RolloutSample`).
+Gate, advantage, ref forward, weight sync, KL/TIS, PPO inner loop, and
+checkpoints are all handled by `recipes.async_rl_loop.main`.
+
+Two minimal rollouts for the async recipe (`training/recipes/async_rl_loop.py`):
+
+- `single_turn_token_in/` — pre-tokenized rows; one `/v1/completions` call per row.
+- `multi_turn_message_in/` — OpenAI-style messages; ports AReaL's
+  `examples/multi_turn_math/` retry loop.
+
+Each example exposes `rollout_fn_factory(setup) -> rollout_fn` (signature
+`async def rollout_fn(sample_prompt) -> RolloutSample | None`) and a `train.py`
+that wires the dataset and factory into `recipes.async_rl_loop.main`.
+
+For the API contract and recipe knobs, see
+[`/skills/dev/references/rl/async-rl.md`](/skills/dev/references/rl/async-rl.md).
+
+The `train.py` files use placeholder model and tokenizer names — replace with
+your own model and tokenizer identifiers before running.

training/examples/rl/multi_turn_message_in/.gitignore

@@ -0,0 +1,3 @@
+# Materialized GSM8K splits — regenerate with `python prepare_data.py`.
+train.jsonl
+test.jsonl

training/examples/rl/multi_turn_message_in/README.md

@@ -0,0 +1,61 @@
+# Multi-Turn GSM8K (message-in) RL example
+
+A multi-turn math agent that ports
+[AReaL's `examples/multi_turn_math/`](https://github.com/inclusionAI/AReaL/tree/main/examples/multi_turn_math)
+to the cookbook's async RL recipe.
+
+For the recipe API and gate sizing, see
+[`/skills/dev/references/rl/async-rl.md`](/skills/dev/references/rl/async-rl.md).
+
+The model is asked a GSM8K problem and must put its final answer in
+`\boxed{...}`.  If the boxed answer is wrong (verified by `math_verify`),
+the rollout appends a fixed user-feedback message and lets the model retry
+once more (configurable via `--max-turns`, default 2).
+
+The full trajectory — prompt + first attempt + feedback + second attempt —
+is packed into a single `RolloutSample`.  `MessageTrajectoryAssembler`
+keeps the per-token loss mask aligned: assistant tokens are trained on
+across both turns; the original prompt and the user-feedback bridge tokens
+are masked out.
+
+## Files
+
+- `prepare_data.py` — downloads `openai/gsm8k` (config `main`) from HuggingFace
+  via ``load_dataset("openai/gsm8k", "main")`` and writes ``train.jsonl``
+  (7473 rows) and ``test.jsonl`` (1319 rows) with ``{messages, answer}`` rows.
+- `reward.py` — extracts `\boxed{...}` from the completion and verifies
+  against the GSM8K ground-truth (`#### N`) via numeric match plus a
+  `math_verify` fallback.
+- `rollout.py` — per-sample `make_rollout_fn(setup) -> RolloutFn`; runs the
+  retry loop and returns one `RolloutSample` per trajectory.
+- `train.py` — wires the dataset and the rollout factory into
+  `recipes/async_rl_loop.main`.
+- `run.sh` — one-shot end-to-end (auto-runs `prepare_data.py` if needed).
+
+## Usage
+
+```bash
+# 1. Download GSM8K (writes train.jsonl + test.jsonl in this directory).
+python prepare_data.py
+
+# 2. Train.
+python train.py \
+    --base-model accounts/fireworks/models/qwen3-1p5b-instruct \
+    --tokenizer-model Qwen/Qwen2.5-1.5B-Instruct \
+    --max-rows 512 \
+    --completions-per-prompt 4 \
+    --max-turns 2 \
+    --output-model-id accounts/<acct>/models/gsm8k-mt
+```
+
+Or `bash run.sh` for the canned configuration.
+
+## Reference
+
+- AReaL `examples/multi_turn_math/gsm8k_rl_mt.py` — the original
+  `MultiTurnMathAgent` whose retry loop, feedback prompt, and reward
+  function this example mirrors.
+- AReaL `areal/experimental/openai/types.py::to_tensor_dict` — the
+  bridge-masking semantics (prompt and feedback positions get
+  `loss_mask=0`, only assistant tokens are trained on) that
+  `MessageTrajectoryAssembler.to_flat()` reproduces in this recipe.

training/examples/rl/multi_turn_message_in/__init__.py

@@ -0,0 +1 @@
+"""Multi-turn message-in RL example."""