diff --git a/skills/dev/SKILL.md b/skills/dev/SKILL.md index 9a9b888f..e488329b 100644 --- a/skills/dev/SKILL.md +++ b/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) | diff --git a/skills/dev/references/recipes.md b/skills/dev/references/recipes.md index 3e999c9d..4f39a880 100644 --- a/skills/dev/references/recipes.md +++ b/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. diff --git a/skills/dev/references/rl/async-rl.md b/skills/dev/references/rl/async-rl.md new file mode 100644 index 00000000..6e4c6a50 --- /dev/null +++ b/skills/dev/references/rl/async-rl.md @@ -0,0 +1,289 @@ +# 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 runs rollout sampling and training as concurrent tasks behind a +gate that bounds how stale a sample may be when it lands at the trainer. It +covers the full spectrum: + +- **Strict on-policy** (`max_head_offpolicy_versions=0`): the gate admits at + most one outer batch worth of samples per policy version; samples that would + arrive after the next weight sync are held until the sync. No off-policy + drift, but rollouts and training serialize at each batch boundary. +- **Off-policy with bounded staleness** (`max_head_offpolicy_versions > 0`): + samples may land up to `O` weight-sync versions past their submit version, + which lets rollout sampling overlap with training in steady state. + +`rl_loop.py` is the older synchronous recipe (always strict on-policy, drains +rollouts before each step). The async recipe is a strict superset of that +behavior and is the recommended starting point for new RL work. + +## 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 | Always concurrent; off-policy budget controls how much overlap | +| 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`) | +| On-/off-policy | Strict on-policy only | `max_head_offpolicy_versions=0` for strict on-policy, `>0` for off-policy with bounded staleness | +| PPO inner steps | One `fwd_bwd + optim_step` per rollout | `ppo_n_minibatches` × inner steps per rollout | + +Prefer the async recipe for new work — its `O=0` mode is equivalent to the sync +recipe, and raising `O` later is a single-knob change. The sync recipe stays +for users who already depend on it. + +## 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) + +Bookkeeping is in samples (LLM calls), the same unit the deployment's +`max_batch_size` gates on, but admission is **row-atomic**: the runner +submits whole rows (`cpp` samples each) so every row in the assembler has +the same submit version. A row is admitted iff *both* caps have at least +`cpp` sample slots free: + +``` +staleness_slots = (max_head_offpolicy_versions + version + 1) * batch_size_samples + - (samples_in_flight + samples_accepted) +concurrency_slots = max_concurrency_rollout_sample - samples_in_flight (∞ if unset) + +admit one row iff min(staleness_slots, concurrency_slots) >= completions_per_prompt +``` + +In the runner this is `slots = capacity() // completions_per_prompt`; that +many rows are submitted in the current admission tick. + +| 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`. At +`O = 0` the loop runs strict on-policy regardless of `R`; raising `O` opens the +overlap window. AReaL's GSM8K example uses `R=1` with `O=2`; we've tested +`R=4` with `O=4` (one margin step over the minimum) and found it healthy. See +`## Metrics` below for tuning from a live run. + +## Metrics: tuning staleness and the trainer/sampler GPU split + +The target regime is **sampler-bound with minimal trainer wait**: the trainer +finishes its step + sync just before the next batch is ready, so it idles only +briefly waiting on the sampler. The four core wall-time metrics tell you +where you sit on that frontier and what to change. + +| Metric | Means | Healthy value | +|---|---|---| +| `perf/trainer_wait_for_sampler_time` | Trainer idle, waiting for the next batch to assemble | **Small but >0** (sampler-bound, minimal slack) | +| `perf/sampler_wait_for_trainer_time` | Sampler blocked on the staleness gate, waiting for a weight sync to release budget | **≈0** for off-policy; expected `>0` at `O=0` | +| `perf/wait_time_ratio` | `(trainer_wait + sampler_wait) / step_time` | < ~0.1 in the off-policy regime | +| `perf/overlap_ratio` | Fraction of step wall-time the two sides ran concurrently | → 1.0 in the off-policy regime | + +In off-policy steady state the two wait metrics are mutually exclusive: exactly +one side waits per step. Read the two ratios first to triage, then drop into +the wait pair to decide what to change. + +> At `max_head_offpolicy_versions=0` (strict on-policy), `sampler_wait_for_trainer_time` +> is structurally non-zero because the gate refuses to admit the next batch +> until the current weight sync. That is correct behavior, not a bug. The +> goal in that mode is just to keep the *trainer* fully utilized; the sampler +> wait is the cost of strict on-policy. + +### Reading a run (off-policy regime, `O > 0`) + +1. **`sampler_wait_for_trainer_time` >> 0** → samplers are blocked on the + staleness budget. The trainer is the slow side; admit more off-policyness + or reduce trainer load. + - First lever: raise `max_head_offpolicy_versions` by 1 (`O += 1`). Cheap + and reversible; KL/PPO-clip will absorb modest extra drift. + - If `O` is already at the sizing rule (`O ≥ R−1`) and the wait persists, + the trainer step itself is too long. Add training replicas / + pipeline-parallel ranks, or lower `ppo_n_minibatches`. +2. **`trainer_wait_for_sampler_time` >> 0 and `sampler_wait` ≈ 0** → the + intended sampler-bound regime. Check whether the wait is *minimized*: + - If the wait is small and `wait_time_ratio < ~0.1`, you're done. + - If the wait is large, the sampler is the slow side. Raise inference + replicas, raise `max_concurrency_rollout_sample` toward the deployment's + `max_batch_size`, or shrink the per-sample work (lower + `max_completion_tokens`, tighter retry budget in the rollout). +3. **Both waits >0** → the gate is mis-sized; one side is starving the other + on alternating steps. Usually `R = C_s / (B_p·cpp)` is high but `O` is too + low: each batch admits fast, but no off-policy slack means the next batch + stalls until the sync. Bump `O` until `sampler_wait` collapses to ~0. +4. **Both waits ≈ 0, low `overlap_ratio`** → step times are dominated by + non-overlapped phases (weight sync, checkpoint save). Re-check + `weight_sync_interval=1` is in effect; if the sync itself is the long pole, + investigate the deployment configuration separately. + +### Choosing the trainer/sampler GPU split + +Same metrics, in aggregate: + +- **Sampler-bound run, large `trainer_wait_for_sampler`** → shift GPUs from + trainer to inference: more replicas, larger TP, or a larger + `max_concurrency_rollout_sample`. +- **Trainer-bound run, large `sampler_wait_for_trainer` after maxing out `O`** + → shift GPUs from inference to trainer: more training replicas, raise + pipeline parallelism, or accept lower `ppo_n_minibatches`. + +In the off-policy regime, tune until `sampler_wait_for_trainer_time ≈ 0` *and* +`trainer_wait_for_sampler_time` is small — the trainer is the marginal-cost +resource and is fully utilized; the sampler always has work queued +just-in-time. In strict on-policy (`O=0`), only the second condition is +achievable; `sampler_wait_for_trainer_time` is bounded below by the train + +sync wall time. + +### Other useful metrics + +- `train/ppo_kl` — intra-step KL between the current policy and the + `old_policy_logprobs` snapshot. Large with `ppo_n_minibatches > 1` means + the inner loop is genuinely doing work; a near-zero value means + `ppo_n_minibatches=1` would have the same effect at lower cost. +- `train/inference_kld`, `train/inference_diff` — drift between the policy + used to sample and the policy at training time. Should track `O`: at `O=0` + these are ~0; at `O=4` they grow but should remain bounded. Spikes that + don't decay across steps often indicate `weight_sync_fn` is silently failing. +- `rollout/entropy` — averaged over `loss_mask>0` only. Sudden collapse is + the usual mode-collapse signal. + +## 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` makes one + `/v1/completions` token-in/token-out call per invocation (the recipe + invokes it `completions_per_prompt` times per dataset 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 diff --git a/training/examples/rl/README.md b/training/examples/rl/README.md new file mode 100644 index 00000000..e9390efe --- /dev/null +++ b/training/examples/rl/README.md @@ -0,0 +1,28 @@ +# 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 + `rollout_fn` invocation (recipe invokes `rollout_fn` `completions_per_prompt` + times 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. diff --git a/training/examples/rl/multi_turn_message_in/.gitignore b/training/examples/rl/multi_turn_message_in/.gitignore new file mode 100644 index 00000000..e268ae2e --- /dev/null +++ b/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 diff --git a/training/examples/rl/multi_turn_message_in/README.md b/training/examples/rl/multi_turn_message_in/README.md new file mode 100644 index 00000000..87c75470 --- /dev/null +++ b/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//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. diff --git a/training/examples/rl/multi_turn_message_in/__init__.py b/training/examples/rl/multi_turn_message_in/__init__.py new file mode 100644 index 00000000..5aaa1e96 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/__init__.py @@ -0,0 +1 @@ +"""Multi-turn message-in RL example.""" diff --git a/training/examples/rl/multi_turn_message_in/prepare_data.py b/training/examples/rl/multi_turn_message_in/prepare_data.py new file mode 100755 index 00000000..9d18d791 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/prepare_data.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Download GSM8K from HuggingFace and convert to JSONL for the recipe. + +Pulls the real ``openai/gsm8k`` dataset (config ``main``) — both ``train`` and +``test`` splits in one call — and writes ``train.jsonl`` and ``test.jsonl`` +next to this script (or wherever ``--output-dir`` points). + +Output format per row:: + + {"id": "gsm8k-train-0", + "messages": [{"role": "user", + "content": "\\nPlease put your final answer within \\boxed{}."}], + "answer": ""} + +Mirrors AReaL's ``examples/multi_turn_math/`` data shape so the reward function +can verify ``\\boxed{...}`` against the GSM8K ground-truth ``#### N`` token. + +Usage:: + + python prepare_data.py # writes train.jsonl + test.jsonl + python prepare_data.py --split train # only train + python prepare_data.py --max-rows 100 # cap each split +""" + +from __future__ import annotations + +import argparse +import json +import os + +from datasets import load_dataset + +PROMPT_SUFFIX = "\nPlease put your final answer within \\boxed{}." + +DEFAULT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def _write_split(split_name: str, split_data, out_path: str, max_rows: int | None) -> int: + n = 0 + with open(out_path, "w") as f: + for idx, row in enumerate(split_data): + if max_rows is not None and n >= max_rows: + break + entry = { + "id": f"gsm8k-{split_name}-{idx}", + "messages": [ + {"role": "user", "content": row["question"] + PROMPT_SUFFIX}, + ], + "answer": row["answer"], + } + f.write(json.dumps(entry, ensure_ascii=False) + "\n") + n += 1 + return n + + +def main(): + parser = argparse.ArgumentParser(description="Prepare GSM8K JSONL") + parser.add_argument("--split", default="all", choices=["all", "train", "test"], + help="Which split to write (default: both)") + parser.add_argument("--output-dir", default=DEFAULT_DIR, + help="Directory for {split}.jsonl files") + parser.add_argument("--max-rows", type=int, default=None, + help="Optional cap on rows per split") + args = parser.parse_args() + + ds = load_dataset("openai/gsm8k", "main") + print(f"Loaded openai/gsm8k(main): " + f"train={len(ds['train'])} rows, test={len(ds['test'])} rows") + + splits = ["train", "test"] if args.split == "all" else [args.split] + os.makedirs(args.output_dir, exist_ok=True) + for s in splits: + out_path = os.path.join(args.output_dir, f"{s}.jsonl") + n = _write_split(s, ds[s], out_path, args.max_rows) + print(f"Wrote {n} rows to {out_path}") + + +if __name__ == "__main__": + main() diff --git a/training/examples/rl/multi_turn_message_in/reward.py b/training/examples/rl/multi_turn_message_in/reward.py new file mode 100644 index 00000000..adb21616 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/reward.py @@ -0,0 +1,87 @@ +"""GSM8K reward function for the multi-turn message-in recipe. + +Mirrors AReaL's ``examples/multi_turn_math/gsm8k_rl_mt.py::gsm8k_reward_fn``: +parse the model's completion (``\\boxed{...}``) and the GSM8K ground-truth +answer string (whose final number follows ``#### ``) with ``math_verify`` and +return ``1.0`` on a verified match, ``0.0`` otherwise. + +A pure-regex numeric fallback handles the common case where ``math_verify``'s +LaTeX parsing rejects an otherwise-correct integer answer. +""" + +from __future__ import annotations + +import logging +import re + +logger = logging.getLogger(__name__) + +_BOXED_RE = re.compile(r"\\boxed\s*\{", re.DOTALL) +_GSM8K_GT_RE = re.compile(r"####\s*([\-\+]?[\d,]*\.?\d+)") +_NUMERIC_TOL = 1e-6 + + +def _extract_boxed(text: str) -> str | None: + """Return the content of the LAST ``\\boxed{...}`` in ``text``. + + Walks brace depth so nested braces (``\\frac{1}{2}``) are preserved. + """ + matches = list(_BOXED_RE.finditer(text)) + if not matches: + return None + last = matches[-1] + start = last.end() + depth = 1 + i = start + while i < len(text) and depth > 0: + c = text[i] + if c == "{": + depth += 1 + elif c == "}": + depth -= 1 + i += 1 + if depth != 0: + return None + return text[start : i - 1].strip() + + +def _gsm8k_ground_truth_number(answer: str) -> str | None: + """Strip the GSM8K chain-of-thought; return the final number after ``#### ``.""" + m = _GSM8K_GT_RE.search(answer) + if not m: + return None + return m.group(1).replace(",", "") + + +def _try_numeric_match(pred: str, gt: str) -> bool: + try: + return abs(float(pred.replace(",", "")) - float(gt.replace(",", ""))) < _NUMERIC_TOL + except (ValueError, OverflowError): + return False + + +def gsm8k_reward(completion: str, answer: str) -> float: + """Return ``1.0`` if the model's boxed answer matches the GSM8K ground truth. + + The cheap numeric path catches the GSM8K-typical case (integer answers). + Falls through to ``math_verify`` (a project dep) for fractions / surds / + LaTeX. Returns ``0.0`` on any failure -- never raises. + """ + pred = _extract_boxed(completion) + if pred is None: + return 0.0 + gt_num = _gsm8k_ground_truth_number(answer) + if gt_num is not None and _try_numeric_match(pred, gt_num): + return 1.0 + + try: + from math_verify import parse as math_parse, verify as math_verify_fn + + pred_parsed = math_parse(f"\\boxed{{{pred}}}") + gt_parsed = math_parse(answer) + if pred_parsed and gt_parsed and math_verify_fn(gt_parsed, pred_parsed): + return 1.0 + except Exception: + logger.debug("math_verify failed; pred=%r gt=%r", pred, answer, exc_info=True) + + return 0.0 diff --git a/training/examples/rl/multi_turn_message_in/rollout.py b/training/examples/rl/multi_turn_message_in/rollout.py new file mode 100644 index 00000000..e9f360d1 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/rollout.py @@ -0,0 +1,111 @@ +"""Multi-turn GSM8K rollout (per-sample) with retry-on-wrong feedback. + +Mirrors AReaL's ``examples/multi_turn_math/gsm8k_rl_mt.py``: ask GSM8K, score +the boxed answer, and -- if the model is wrong on the first try -- append a +fixed user-feedback message and let it try once more (configurable via +``setup.extras["max_turns"]``, default ``2``). + +The whole trajectory (prompt + assistant turn 1 + feedback + assistant turn 2) +is packed into a single ``RolloutSample``. ``MessageTrajectoryAssembler`` +keeps the per-token loss mask aligned: ``1`` on every assistant-generated +token (across all turns), ``0`` everywhere else (original prompt, the +user-feedback bridge between turns). The scalar ``reward`` is the +last-turn verification result (``0.0`` or ``1.0``); rolling-up across turns +isn't useful in concat mode -- the GRPO advantage compares trajectories, so +"wrong then right" naturally beats "wrong then wrong" without an explicit +per-turn discount. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from training.examples.rl.multi_turn_message_in.reward import gsm8k_reward +from training.examples.rl.vanilla_sampler import build_deployment_sampler +from training.utils.rl.rollout import ( + MessageTrajectoryAssembler, + RolloutSample, + TITOTokenizer, +) + +if TYPE_CHECKING: + from training.recipes.async_rl_loop import RolloutFn, RolloutSetup + +logger = logging.getLogger(__name__) + +RETRY_PROMPT = ( + "Your answer is either wrong or not parsable to the reward function. " + "You may misunderstand the original question. Please carefully read the " + "original question, check the previous errors, and try to answer it again." +) + + +def make_rollout_fn(setup: "RolloutSetup") -> "RolloutFn": + sampler = build_deployment_sampler(setup) + sample_kwargs = dict(setup.sample_kwargs) + tokenizer = setup.tokenizer + max_turns = int(setup.extras.get("max_turns", 2)) + if max_turns < 1: + raise ValueError(f"max_turns must be >= 1, got {max_turns}") + + async def rollout_fn(sample_prompt: dict) -> RolloutSample | None: + messages = list(sample_prompt.get("messages") or []) + answer = sample_prompt.get("answer") + if not messages or answer is None: + return None + + assembler = MessageTrajectoryAssembler(TITOTokenizer(tokenizer)) + current_messages = messages + last_reward = 0.0 + + for turn in range(max_turns): + prompt_tokens = assembler.prepare_next_input(current_messages) + completions = await sampler.sample_with_prompt_tokens( + prompt_tokens, n=1, **sample_kwargs, + ) + if not completions: + return None + completion = completions[0] + + prompt_len = int(completion.prompt_len) + output_tokens = list(completion.full_tokens[prompt_len:]) + output_logprobs = list(completion.inference_logprobs or []) + if ( + getattr(completion, "logprobs_echoed", False) + and len(output_logprobs) == len(completion.full_tokens) + ): + output_logprobs = output_logprobs[prompt_len:] + if not output_tokens or len(output_logprobs) != len(output_tokens): + return None + + assistant_text = getattr(completion, "text", "") or tokenizer.decode(output_tokens) + assistant_message = {"role": "assistant", "content": assistant_text} + assembler.add_assistant_response( + request_messages=current_messages, + assistant_message=assistant_message, + prompt_token_ids=prompt_tokens, + completion_token_ids=output_tokens, + completion_logprobs=output_logprobs, + finish_reason=getattr(completion, "finish_reason", "stop"), + ) + + last_reward = gsm8k_reward(assistant_text, str(answer)) + if last_reward >= 1.0: + break + + if turn + 1 < max_turns: + current_messages = current_messages + [ + assistant_message, + {"role": "user", "content": RETRY_PROMPT}, + ] + + tokens, logprobs, loss_mask = assembler.trajectory.to_flat() + return RolloutSample( + tokens=tokens, + logprobs=logprobs, + loss_mask=loss_mask, + reward=last_reward, + ) + + return rollout_fn diff --git a/training/examples/rl/multi_turn_message_in/run.sh b/training/examples/rl/multi_turn_message_in/run.sh new file mode 100755 index 00000000..03f74ae1 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/run.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +REPO_ROOT="$(cd "$HERE/../../../.." && pwd)" +export PYTHONPATH="${REPO_ROOT}:${PYTHONPATH:-}" + +if [[ ! -f "$HERE/train.jsonl" ]]; then + echo "train.jsonl not found; downloading openai/gsm8k from HuggingFace..." + python "$HERE/prepare_data.py" +fi + +python "$HERE/train.py" \ + --base-model accounts/fireworks/models/qwen3-1p5b-instruct \ + --tokenizer-model Qwen/Qwen2.5-1.5B-Instruct \ + --dataset-path "$HERE/train.jsonl" \ + --max-rows 512 \ + --epochs 1 \ + --completions-per-prompt 4 \ + --prompt-groups-per-step 8 \ + --max-completion-tokens 1024 \ + --max-turns 2 \ + --learning-rate 1.7e-5 \ + --kl-beta 0.0 \ + --output-model-id "${OUTPUT_MODEL_ID:-accounts/fireworks/models/gsm8k-mt-$(date +%s)}" diff --git a/training/examples/rl/multi_turn_message_in/train.py b/training/examples/rl/multi_turn_message_in/train.py new file mode 100755 index 00000000..8371f2c7 --- /dev/null +++ b/training/examples/rl/multi_turn_message_in/train.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Multi-turn GSM8K async RL training. + +Run ``prepare_data.py`` first to materialize ``train.jsonl`` (and +``test.jsonl``) from HuggingFace ``openai/gsm8k`` (config ``main``). +Then:: + + python train.py \\ + --base-model accounts/fireworks/models/qwen3-1p5b-instruct \\ + --tokenizer-model Qwen/Qwen2.5-1.5B-Instruct \\ + --output-model-id accounts//models/gsm8k-mt + +The rollout (see ``rollout.py``) does up to ``max_turns`` (default 2) +LLM calls per row: if the boxed answer is wrong, a fixed user-feedback +message is appended and the model retries. Reward is the verified +boxed-answer match for the last turn. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import time +from typing import Iterator + +from training.examples.rl.multi_turn_message_in.rollout import make_rollout_fn +from training.recipes.async_rl_loop import Config, main +from training.utils import DeployConfig, InfraConfig, WandBConfig + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%H:%M:%S", +) +logger = logging.getLogger(__name__) + +DEFAULT_DATASET = os.path.join(os.path.dirname(__file__), "train.jsonl") + + +def _iter_rows(path: str, max_rows: int | None) -> Iterator[dict]: + n = 0 + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + yield json.loads(line) + n += 1 + if max_rows is not None and n >= max_rows: + return + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Multi-turn GSM8K async RL") + p.add_argument("--base-model", default="accounts/fireworks/models/qwen3-1p5b-instruct") + p.add_argument("--tokenizer-model", default="Qwen/Qwen2.5-1.5B-Instruct") + p.add_argument("--dataset-path", default=DEFAULT_DATASET) + p.add_argument("--output-model-id", required=False, default=None) + p.add_argument("--max-rows", type=int, default=512) + p.add_argument("--epochs", type=int, default=1) + p.add_argument("--completions-per-prompt", type=int, default=4) + p.add_argument("--max-completion-tokens", type=int, default=1024) + p.add_argument("--temperature", type=float, default=1.0) + p.add_argument("--prompt-groups-per-step", type=int, default=8) + p.add_argument("--learning-rate", type=float, default=1.7e-5) + p.add_argument("--kl-beta", type=float, default=0.0) + p.add_argument("--max-head-offpolicy-versions", type=int, default=0, + help="Off-policy staleness budget in weight-sync (policy) versions; " + "one weight-sync per outer rollout batch, regardless of ppo-n-minibatches " + "(0 = strict on-policy).") + p.add_argument("--ppo-n-minibatches", type=int, default=1, + help="Inner PPO minibatches per rollout batch (1 = legacy 1:1).") + p.add_argument("--filter-constant-reward", action="store_true", + help="Drop prompt groups whose samples all share the same reward (zero GRPO advantage).") + p.add_argument("--max-turns", type=int, default=2, + help="Max LLM calls per trajectory (AReaL default: 2)") + p.add_argument("--log-path", default="./gsm8k_mt_logs") + p.add_argument("--wandb-entity", default=os.environ.get("WANDB_ENTITY", "")) + p.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "gsm8k-mt")) + p.add_argument("--training-shape-id", default=None, + help="Full training shape resource name (accounts/.../trainingShapes/...).") + p.add_argument("--lora-rank", type=int, default=0, + help="LoRA rank (0 = full-parameter training).") + p.add_argument("--synchronous-training", action="store_true", + help="Force fully-synchronous mode (no rollout/train overlap). " + "Drains in-flight rollouts before each train_step and " + "marks the rollout side blocked-on-trainer; " + "perf/sampler_wait_for_trainer_time then reflects train+sync wall time.") + p.add_argument("--max-concurrency-rollout-sample", type=int, default=None, + help="Cap in-flight LLM calls against the inference deployment " + "(maps to deployment max_batch_size; e.g. 64 for a 64-slot " + "shape with completions_per_prompt=8 = 8 rows in flight). " + "Must be >= completions_per_prompt. Default unbounded.") + p.add_argument("--wandb-run-name", default=None, + help="Override the WandB run name.") + p.add_argument("--replica-count", type=int, default=None, + help="Pin the inference deployment to a fixed replica count " + "(default 1). Higher values fan out rollout sampling.") + return p.parse_args() + + +def run(): + args = parse_args() + if not os.path.exists(args.dataset_path): + raise FileNotFoundError( + f"Dataset not found at {args.dataset_path}. " + "Run `python prepare_data.py` first." + ) + + rows = list(_iter_rows(args.dataset_path, args.max_rows)) + logger.info("Loaded %d GSM8K rows from %s", len(rows), args.dataset_path) + + cfg = Config( + log_path=args.log_path, + base_model=args.base_model, + learning_rate=args.learning_rate, + kl_beta=args.kl_beta, + completions_per_prompt=args.completions_per_prompt, + max_completion_tokens=args.max_completion_tokens, + temperature=args.temperature, + epochs=args.epochs, + max_rows=args.max_rows, + lora_rank=args.lora_rank, + prompt_groups_per_step=args.prompt_groups_per_step, + max_head_offpolicy_versions=args.max_head_offpolicy_versions, + ppo_n_minibatches=args.ppo_n_minibatches, + max_concurrency_rollout_sample=args.max_concurrency_rollout_sample, + synchronous_training=args.synchronous_training, + output_model_id=args.output_model_id, + infra=InfraConfig(training_shape_id=args.training_shape_id), + deployment=DeployConfig( + tokenizer_model=args.tokenizer_model, + replica_count=args.replica_count, + ), + wandb=WandBConfig( + entity=args.wandb_entity, + project=args.wandb_project, + run_name=args.wandb_run_name or f"gsm8k-mt-{int(time.time()) % 100000}", + ), + ) + # Drop prompt groups whose rewards are constant across all samples -- + # GRPO z-score advantage is 0 there, so the optimizer step is a no-op. + dynamic_filter_fn = ( + (lambda pg: len(set(pg.rewards)) > 1) + if args.filter_constant_reward else None + ) + main( + cfg, + rollout_fn_factory=make_rollout_fn, + rows=rows, + rollout_extras={"max_turns": args.max_turns}, + dynamic_filter_fn=dynamic_filter_fn, + ) + + +if __name__ == "__main__": + run() diff --git a/training/examples/rl/single_turn_token_in/__init__.py b/training/examples/rl/single_turn_token_in/__init__.py new file mode 100644 index 00000000..bec3e0d4 --- /dev/null +++ b/training/examples/rl/single_turn_token_in/__init__.py @@ -0,0 +1 @@ +"""Single-turn token-in RL example.""" diff --git a/training/examples/rl/single_turn_token_in/rollout.py b/training/examples/rl/single_turn_token_in/rollout.py new file mode 100644 index 00000000..d22b3907 --- /dev/null +++ b/training/examples/rl/single_turn_token_in/rollout.py @@ -0,0 +1,52 @@ +"""Single-turn token-in rollout (per-sample). + +Input dataset rows carry ``prompt_token_ids``. No messages or chat +templates. The framework fans each dataset row out to +``completions_per_prompt`` parallel calls; each call returns one +:class:`RolloutSample`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from training.examples.rl.vanilla_sampler import build_deployment_sampler +from training.utils.rl.rollout import RolloutSample + +if TYPE_CHECKING: + from training.recipes.async_rl_loop import RolloutFn, RolloutSetup + + +def make_rollout_fn(setup: "RolloutSetup") -> "RolloutFn": + sampler = build_deployment_sampler(setup) + sample_kwargs = dict(setup.sample_kwargs) + + async def rollout_fn(sample_prompt: dict) -> RolloutSample | None: + prompt_token_ids = list(sample_prompt.get("prompt_token_ids") or []) + if not prompt_token_ids: + return None + + completions = await sampler.sample_with_prompt_tokens( + prompt_token_ids, n=1, **sample_kwargs, + ) + if not completions: + return None + completion = completions[0] + prompt_len = int(completion.prompt_len) + tokens = list(completion.full_tokens) + output_tokens = tokens[prompt_len:] + output_logprobs = list(completion.inference_logprobs or []) + if getattr(completion, "logprobs_echoed", False) and len(output_logprobs) == len(tokens): + output_logprobs = output_logprobs[prompt_len:] + if not output_tokens or len(output_logprobs) != len(output_tokens): + return None + return RolloutSample( + tokens=tokens, + logprobs=[0.0] * prompt_len + output_logprobs, + loss_mask=[0] * prompt_len + [1] * len(output_tokens), + reward=float(sample_prompt.get("reward", 0.0)), + finish_reason=getattr(completion, "finish_reason", "stop"), + text=getattr(completion, "text", ""), + ) + + return rollout_fn diff --git a/training/examples/rl/single_turn_token_in/train.py b/training/examples/rl/single_turn_token_in/train.py new file mode 100644 index 00000000..b576e1f1 --- /dev/null +++ b/training/examples/rl/single_turn_token_in/train.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python3 +"""Minimal async RL wiring for the single-turn token-in rollout.""" + +from __future__ import annotations + +from training.examples.rl.single_turn_token_in.rollout import make_rollout_fn +from training.recipes.async_rl_loop import Config, main +from training.utils import DeployConfig + + +if __name__ == "__main__": + cfg = Config( + log_path="/tmp/rl-single-turn-token-in", + base_model="accounts/fireworks/models/example", + prompt_groups_per_step=1, + completions_per_prompt=2, + deployment=DeployConfig(tokenizer_model="example-tokenizer"), + ) + rows = [{"id": "tok-1", "prompt_token_ids": [1, 2, 3], "reward": 1.0}] + main(cfg, rollout_fn_factory=make_rollout_fn, rows=rows) diff --git a/training/examples/rl/vanilla_sampler.py b/training/examples/rl/vanilla_sampler.py new file mode 100644 index 00000000..135dfbbb --- /dev/null +++ b/training/examples/rl/vanilla_sampler.py @@ -0,0 +1,25 @@ +"""Example-owned sampler helper for async RL rollouts.""" + +from __future__ import annotations + +from fireworks.training.sdk.deployment import DeploymentSampler + +from training.recipes.async_rl_loop import RolloutSetup + + +def build_deployment_sampler(setup: RolloutSetup) -> DeploymentSampler: + """Construct a :class:`DeploymentSampler` from a :class:`RolloutSetup`. + + The training recipe assembles the setup once at startup and hands it + to the rollout factory; the factory uses this helper to materialize + a sampler bound to the inference deployment. Concurrency is enforced + by the async runner in sample (LLM-call) units via + ``cfg.max_concurrency_rollout_sample`` -- the same unit the + deployment's ``max_batch_size`` gates on. No HTTP-layer gate. + """ + return DeploymentSampler( + inference_url=setup.inference_base_url, + model=setup.model, + api_key=setup.api_key, + tokenizer=setup.tokenizer, + ) diff --git a/training/pyproject.toml b/training/pyproject.toml index a0b238e7..e7a1d14d 100644 --- a/training/pyproject.toml +++ b/training/pyproject.toml @@ -7,7 +7,7 @@ name = "fireworks-training-cookbook" version = "0.1.0" requires-python = ">=3.11" dependencies = [ - "fireworks-ai[training]>=1.1.0a64,<2", + "fireworks-ai[training]>=1.2.0a65,<2", "tqdm", "torch", "datasets", diff --git a/training/recipes/async_rl_loop.py b/training/recipes/async_rl_loop.py new file mode 100644 index 00000000..a947bbd3 --- /dev/null +++ b/training/recipes/async_rl_loop.py @@ -0,0 +1,689 @@ +#!/usr/bin/env python3 +"""Async RL recipe with per-sample rollouts and recipe-owned training. + +EXPERIMENTAL -- under active development. API surface (``Config`` field +names, ``RolloutSetup`` shape, gate semantics) may change without a +backward-compat shim. The recipe is intentionally minimal-surface: the +only thing most users need to write is the rollout function; everything +else (gate, advantage, ref forward, weight sync, KL/TIS, PPO inner loop, +checkpoints) is handled by ``main()``. See +``skills/dev/references/rl/async-rl.md`` for the full contract. + +Acknowledgements -- prior art referenced while designing this loop: + +* AReaL (https://github.com/inclusionAI/AReaL) +* slime (https://github.com/THUDM/slime) +* Miles (https://github.com/radixark/miles) + +Users write ``rollout_fn(sample_prompt) -> RolloutSample | None`` -- one +trajectory per call. ``sample_prompt`` is the dataset row's dict re-named +once it crosses the dataset/sampling seam. The recipe fans each dataset +row out to ``completions_per_prompt`` parallel calls and assembles the +resulting samples into a PromptGroup inside the loop. + +Rollout dependencies (tokenizer, sampler, request gate, etc.) flow +through :class:`RolloutSetup`. The user supplies a +``rollout_fn_factory(setup) -> rollout_fn`` callable that closes over +the setup; this matches AReaL's workflow construction pattern. +""" + +from __future__ import annotations + +import asyncio +import logging +import math +import os +import signal +import time +from contextlib import ExitStack +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable + +import tinker +import transformers + +from fireworks.training.sdk import DeploymentManager, TrainerJobManager +from training.utils.client import GradAccNormalization +from fireworks.training.sdk.weight_syncer import WeightSyncer +from training.utils import ( + DEFAULT_ADAM, + DeployConfig, + InfraConfig, + ResourceCleanup, + WandBConfig, + WeightSyncConfig, + load_jsonl_dataset, + read_api_extra_headers_env, + setup_wandb, + validate_config, + wandb_finish, + wandb_log, +) +from training.utils.checkpoints import TrainingCheckpoints +from training.utils.dataloader import CursorDataLoader +from training.utils.rl import PromptGroup, setup_infra +from training.utils.rl.async_train import RowRequest, run_async_rl_loop +from training.utils.rl.cispo import CISPOConfig +from training.utils.rl.dapo import DAPOConfig +from training.utils.rl.dro import DROConfig +from training.utils.rl.gspo import GSPOConfig +from training.utils.rl.losses import ( + LossConfig, + PolicyLoss, + build_loss_fn, + combine_prompt_groups, + validate_loss_path, +) +from training.utils.rl.metrics import compute_minibatch_metrics, compute_step_metrics +from training.utils.rl.tis import TISConfig +from training.utils.rl.train import DynamicFilterFn, TrainStepFns +from training.utils.rl.rollout import RolloutSample +from training.utils.timer import elapsed_timer, flush_timing, timer + +logger = logging.getLogger(__name__) + +__all__ = [ + "Config", + "RolloutFn", + "RolloutFnFactory", + "RolloutSetup", + "main", +] + +# Pinned: raising sync interval trades rollout staleness for sync wall-time, +# almost never worth it in fully-async RL. +_WEIGHT_SYNC_INTERVAL = 1 + + +@dataclass +class Config: + log_path: str + base_model: str = "accounts/fireworks/models/qwen3-8b" + dataset: str | None = None + """JSONL path/URL; optional when passing ``rows=`` to ``main()``.""" + + learning_rate: float = 1e-5 + kl_beta: float = 0.001 + completions_per_prompt: int = 4 + max_completion_tokens: int = 1024 + temperature: float = 1.0 + epochs: int = 1 + shuffle: bool = True + seed: int = 0 + max_rows: int = 100 + max_seq_len: int | None = None + lora_rank: int = 0 + + prompt_groups_per_step: int = 1 + max_head_offpolicy_versions: int = 0 + """Staleness budget in weight-sync versions; ``0`` = strict on-policy. + See ``skills/dev/references/rl/async-rl.md`` (gate semantics).""" + max_concurrency_rollout_sample: int | None = None + """In-flight LLM-call cap (same unit as ``deployment.max_batch_size``); + must be ``>= completions_per_prompt`` or the gate stalls.""" + min_group_size: int = 1 + """Minimum surviving samples per row to emit a PromptGroup.""" + grad_accumulation_normalization: GradAccNormalization | str | None = ( + GradAccNormalization.NUM_LOSS_TOKENS + ) + + policy_loss: PolicyLoss = "grpo" + """One of the registered RL policy losses (see :data:`PolicyLoss`). + Client-side only -- ``loss_path`` is not exposed, see skill doc.""" + + dapo: DAPOConfig = field(default_factory=DAPOConfig) + dro: DROConfig = field(default_factory=DROConfig) + gspo: GSPOConfig = field(default_factory=GSPOConfig) + cispo: CISPOConfig = field(default_factory=CISPOConfig) + eps_clip: float = 0.2 + eps_clip_high: float | None = None + ratio_log_cap: float = 20.0 + ppo_n_minibatches: int = 1 + """Inner PPO steps per rollout batch sharing one ``old_policy_logprobs`` + snapshot; ``1`` reproduces the legacy 1:1 behavior.""" + synchronous_training: bool = False + """Drain rollouts before each train step (no overlap); baseline knob + for measuring async savings.""" + tis: TISConfig = field(default_factory=TISConfig) + + infra: InfraConfig = field(default_factory=InfraConfig) + deployment: DeployConfig = field(default_factory=DeployConfig) + weight_sync: WeightSyncConfig = field(default_factory=WeightSyncConfig) + wandb: WandBConfig = field(default_factory=lambda: WandBConfig(project="rl-async")) + + init_from_checkpoint: str | None = None + """Resume from prior checkpoint; bare name = this job, ``"job:name"`` + = cross-job.""" + save_final_checkpoint: bool = True + """Save a resumable+promotable checkpoint at the end of training.""" + output_model_id: str | None = None + """Promote the final checkpoint to this 4-segment model id on clean exit.""" + policy_job_id: str | None = None + """Reuse an existing policy trainer job (e.g. when resuming).""" + + +@dataclass +class RolloutSetup: + """Dependencies the recipe hands the rollout factory once at startup. + + Inference endpoint, tokenizer, sampling kwargs, plus an ``extras`` dict + for caller state. See ``skills/dev/references/rl/async-rl.md``. + """ + + tokenizer: Any + tokenizer_id: str + sample_kwargs: dict[str, Any] + inference_base_url: str + api_key: str + model: str + completions_per_prompt: int + extras: dict[str, Any] = field(default_factory=dict) + + +RolloutFn = Callable[[dict], Awaitable[RolloutSample | None]] +RolloutFnFactory = Callable[[RolloutSetup], RolloutFn] + + +def _save_checkpoint( + ckpt: TrainingCheckpoints, + *, + name: str, + data_consumed: int, + resumable: bool = True, + promotable: bool = False, +) -> None: + logger.info("[%s] dcp_save...", name) + with elapsed_timer("dcp_save") as span: + ckpt.save( + name, + resumable=resumable, + promotable=promotable, + data_consumed=data_consumed, + ) + logger.info("[%s] dcp_save: done (%.1fs)", name, span.elapsed) + + +def main( + config: Config, + *, + rollout_fn_factory: RolloutFnFactory, + dynamic_filter_fn: DynamicFilterFn | None = None, + rows: list[dict] | None = None, + cancel_on_exit: bool = False, + rollout_extras: dict[str, Any] | None = None, +) -> None: + """Run the async RL loop with a user-supplied rollout factory. + + ``rollout_fn_factory(setup) -> rollout_fn`` is called once at startup + with the assembled :class:`RolloutSetup`. The returned + ``rollout_fn(sample_prompt) -> RolloutSample | None`` is invoked + ``completions_per_prompt`` times per dataset row (each invocation is + one sample draw against the inference deployment). + + ``cancel_on_exit=True`` tears down provisioned remote resources on exit. + """ + cfg = config + + logger.warning( + "async_rl_loop is EXPERIMENTAL and under active development; " + "the Config / RolloutSetup API may change without backward-compat " + "shims. See skills/dev/references/rl/async-rl.md.", + ) + + def _signal_handler(signum, _): + name = signal.Signals(signum).name + raise SystemExit(f"Terminated by {name}") + + signal.signal(signal.SIGTERM, _signal_handler) + signal.signal(signal.SIGINT, _signal_handler) + + if rows is None and not cfg.dataset: + raise ValueError("Provide either cfg.dataset or rows= to main().") + if not cfg.deployment.tokenizer_model: + raise ValueError("deployment.tokenizer_model is required.") + validate_config( + cfg.base_model, + cfg.dataset or None, + cfg.weight_sync, + cfg.deployment, + output_model_id=cfg.output_model_id, + require_dataset=(rows is None), + ) + if cfg.completions_per_prompt < 2: + raise ValueError( + "async_rl_loop requires cfg.completions_per_prompt >= 2: the " + "default GRPO-style advantage normalizer (z-score by " + "torch.std(rewards)) is undefined on length-1 reward tensors " + "and would drop every group, silently consuming the dataset " + "without ever training. Set completions_per_prompt >= 2 (the " + f"default is 4); got {cfg.completions_per_prompt}." + ) + if cfg.ppo_n_minibatches < 1: + raise ValueError( + f"ppo_n_minibatches must be >= 1; got {cfg.ppo_n_minibatches}." + ) + setup_wandb( + cfg.wandb, + { + "completions_per_prompt": cfg.completions_per_prompt, + "prompt_groups_per_step": cfg.prompt_groups_per_step, + "max_head_offpolicy_versions": cfg.max_head_offpolicy_versions, + "ppo_n_minibatches": cfg.ppo_n_minibatches, + "tokenizer_id": cfg.deployment.tokenizer_model, + "shuffle": cfg.shuffle, + "seed": cfg.seed, + "kl_beta": cfg.kl_beta, + "lr": cfg.learning_rate, + }, + ) + # Dual axis: train/* per inner PPO minibatch, rollout/* per outer batch. + try: + import wandb as _wandb # noqa: WPS433 (deliberate local import) + + if _wandb.run is not None: + _wandb.define_metric("rollout/step") + _wandb.define_metric("rollout/*", step_metric="rollout/step") + _wandb.define_metric("perf/*", step_metric="rollout/step") + _wandb.define_metric("infra/*", step_metric="rollout/step") + _wandb.define_metric("ctx/*", step_metric="rollout/step") + _wandb.define_metric("batch/*", step_metric="rollout/step") + _wandb.define_metric("async/*", step_metric="rollout/step") + _wandb.define_metric("version/*", step_metric="rollout/step") + except ImportError: + pass + + api_key = os.environ["FIREWORKS_API_KEY"] + base_url = os.environ.get("FIREWORKS_BASE_URL", "https://api.fireworks.ai") + additional_headers = read_api_extra_headers_env() + + rlor_mgr = TrainerJobManager( + api_key=api_key, base_url=base_url, additional_headers=additional_headers, + ) + deploy_mgr = DeploymentManager( + api_key=api_key, base_url=base_url, additional_headers=additional_headers, + ) + + with ResourceCleanup(rlor_mgr, deploy_mgr) as cleanup, ExitStack() as stack: + infra = setup_infra( + rlor_mgr=rlor_mgr, + deploy_mgr=deploy_mgr, + base_model=cfg.base_model, + infra_cfg=cfg.infra, + deploy_cfg=cfg.deployment, + lora_rank=cfg.lora_rank, + max_seq_len=cfg.max_seq_len, + learning_rate=cfg.learning_rate, + policy_job_id=cfg.policy_job_id, + needs_reference=(cfg.kl_beta > 0), + needs_inference=True, + role_prefix="rl-async", + api_key=api_key, + cleanup=cleanup if cancel_on_exit else None, + ) + policy_job_id = infra.policy_job_id + assert infra.inference_model is not None # needs_inference=True invariant + inference_model = infra.inference_model + for closeable in infra.closeables: + stack.callback(closeable.close) + + wandb_log({**infra.boot_metrics, "rollout/step": 0}) + + policy = infra.policy + reference = infra.reference + + tokenizer = transformers.AutoTokenizer.from_pretrained( + cfg.deployment.tokenizer_model, trust_remote_code=True, + ) + weight_syncer = WeightSyncer( + policy_client=policy, + deploy_mgr=deploy_mgr, + deployment_id=infra.deployment_id, + base_model=cfg.base_model, + hotload_timeout=cfg.weight_sync.weight_sync_timeout, + first_checkpoint_type=cfg.weight_sync.first_checkpoint_type, + lora_rank=cfg.lora_rank, + ) + + ckpt = TrainingCheckpoints( + policy, + rlor_mgr, + trainer_id=policy_job_id, + log_path=cfg.log_path, + lora_rank=cfg.lora_rank, + ) + + resume_info = ckpt.resume( + init_from_checkpoint=cfg.init_from_checkpoint, + ) + step_offset = resume_info.step if resume_info else 0 + if step_offset: + logger.info("Resuming from step %d", step_offset) + rollout_offset = step_offset // max(1, cfg.ppo_n_minibatches) + wandb_log( + {"train/step": step_offset, "rollout/step": rollout_offset}, + ) + + if cfg.weight_sync.weight_sync_before_training: + name = f"resume-{step_offset}-base" if step_offset > 0 else "step-0-base" + weight_syncer.save_and_hotload(name, checkpoint_type="base") + + current_version = step_offset + + if rows is None: + rows = load_jsonl_dataset(cfg.dataset, cfg.max_rows) + else: + rows = list(rows) + + prior_rows_consumed = resume_info.data_consumed if resume_info else 0 + row_loader = CursorDataLoader( + rows, + start_cursor=prior_rows_consumed, + epochs=cfg.epochs, + shuffle=cfg.shuffle, + seed=cfg.seed, + ) + + adam_params = tinker.AdamParams(learning_rate=cfg.learning_rate, **DEFAULT_ADAM) + # This recipe is client-side only. ``LossConfig`` adapts the cfg + # fields to the ``LossArgs`` Protocol that ``build_loss_fn`` reads; + # ``loss_path="client"`` is fixed (no builtin server-side path). + loss_args = LossConfig( + policy_loss=cfg.policy_loss, + loss_path="client", + kl_beta=cfg.kl_beta, + eps_clip=cfg.eps_clip, + eps_clip_high=cfg.eps_clip_high, + ratio_log_cap=cfg.ratio_log_cap, + dapo=cfg.dapo, + dro=cfg.dro, + gspo=cfg.gspo, + cispo=cfg.cispo, + tis=cfg.tis, + ) + validate_loss_path(loss_args) + client_loss_builder = build_loss_fn(loss_args) + + sample_kwargs: dict = dict( + max_tokens=cfg.max_completion_tokens, + temperature=cfg.temperature, + max_seq_len=infra.max_seq_len, + http_timeout=cfg.deployment.sample_timeout, + logprobs=True, + ) + + rollout_setup = RolloutSetup( + tokenizer=tokenizer, + tokenizer_id=cfg.deployment.tokenizer_model, + sample_kwargs=sample_kwargs, + inference_base_url=deploy_mgr.inference_url, + api_key=api_key, + model=inference_model, + completions_per_prompt=cfg.completions_per_prompt, + extras=dict(rollout_extras or {}), + ) + rollout_fn = rollout_fn_factory(rollout_setup) + + ctx_metadata: dict[str, Any] = { + "completions_per_prompt": cfg.completions_per_prompt, + "prompt_groups_per_step": cfg.prompt_groups_per_step, + "max_head_offpolicy_versions": cfg.max_head_offpolicy_versions, + "ppo_n_minibatches": cfg.ppo_n_minibatches, + "max_concurrency_rollout_sample": cfg.max_concurrency_rollout_sample, + "weight_sync_interval": _WEIGHT_SYNC_INTERVAL, + "max_completion_tokens": cfg.max_completion_tokens, + "temperature": cfg.temperature, + "shuffle": cfg.shuffle, + "seed": cfg.seed, + "tokenizer_id": cfg.deployment.tokenizer_model, + "model": inference_model, + } + + def make_row_requests(): + for item in row_loader: + row = item.value + idx = item.index + + def factory(_sub_index: int, sample_prompt=row): + return rollout_fn(sample_prompt) + + yield RowRequest( + row_id=idx, + sample_factory=factory, + row_meta={"row_id": row.get("id")}, + on_resolved=lambda _reason, idx=idx: row_loader.mark_resolved(idx), + ) + + def ref_forward(groups: list[PromptGroup]) -> None: + if reference is None: + return + all_ref_data = [d for pg in groups for d in pg.ref_data] + ref_fwd = reference.forward(all_ref_data, "cross_entropy") + idx = 0 + for pg in groups: + n = len(pg.ref_data) + pg.ref_logprobs = [ + ref_fwd.loss_fn_outputs[idx + i]["logprobs"].data for i in range(n) + ] + idx += n + + def fwd_bwd_minibatch( + data, adv, ref_lp, prompt_lens, inf_lp, old_policy_logprobs, + ): + """One inner PPO minibatch. + + Callers pre-compute ``old_policy_logprobs`` once per rollout + batch and pass a slice of the flattened rollout tensors for + this minibatch. + """ + return policy.forward_backward_custom( + data, client_loss_builder(adv, ref_lp, prompt_lens, inf_lp, old_policy_logprobs), + ) + + def train_step( + step: int, prompt_groups: list[PromptGroup], loop_stats: dict | None = None, + ) -> tuple[int, dict]: + """ref_forward + old_policy_logprobs snapshot + num_minibatches x (fwd_bwd + optim_step) + metrics. + + ``num_minibatches = cfg.ppo_n_minibatches``. ``old_policy_logprobs`` + is snapshotted once per rollout batch and reused across every + inner optim step so the PPO ratio measures genuine policy drift. + DCP checkpoints fire only at rollout boundaries (cadence in + rollout batches, not optim steps) so resume accounting is + independent of the minibatch count. + + Slime-aligned dual-axis logging: per-minibatch ``train/*`` lands + on the ``train/step`` axis (one point per inner PPO step); + per-batch ``rollout/*`` / ``perf/*`` / ``async/*`` / ``version/*`` + land on ``rollout/step`` (one point per outer rollout batch). + ``ctx.current_version`` and checkpoint identities remain + optimizer-step labels (resume math is in optim steps); the + off-policy budget is accounted in weight-sync versions inside + ``_StalenessController`` and is independent of those labels. + """ + train_start = time.monotonic() + num_minibatches = max(1, cfg.ppo_n_minibatches) + # 1-indexed outer-batch counter. ``step`` here is the optim-step + # count carried over from prior batches, which is always a + # multiple of num_minibatches at batch boundaries. + rollout_id = step // num_minibatches + 1 + + with elapsed_timer("ref_forward") as span: + ref_forward(prompt_groups) + logger.info( + "[rollout %d] ref_forward (%.1fs)", rollout_id, span.elapsed, + ) + + data, adv, ref_lp, prompt_lens, inf_lp = combine_prompt_groups(prompt_groups) + with elapsed_timer("old_policy_forward") as span: + old_policy_fwd = policy.forward(data, "cross_entropy") + old_policy_logprobs = [ + old_policy_fwd.loss_fn_outputs[i]["logprobs"].data + for i in range(len(data)) + ] + logger.info( + "[rollout %d] old_policy_forward (%.1fs)", + rollout_id, span.elapsed, + ) + + n = len(data) + minibatch_size = max(1, math.ceil(n / num_minibatches)) + fwd_bwd_results: list = [] + optim_result: Any = None + for minibatch_idx in range(num_minibatches): + mb_start = minibatch_idx * minibatch_size + mb_end = min(mb_start + minibatch_size, n) + if mb_start >= mb_end: + break + + with elapsed_timer("fwd_bwd") as span: + fwd_bwd_result = fwd_bwd_minibatch( + data[mb_start:mb_end], + adv[mb_start:mb_end], + ref_lp[mb_start:mb_end], + prompt_lens[mb_start:mb_end], + inf_lp[mb_start:mb_end], + old_policy_logprobs[mb_start:mb_end], + ) + fwd_bwd_results.append(fwd_bwd_result) + logger.info( + "[rollout %d step %d] fwd_bwd (mb %d/%d) (%.1fs)", + rollout_id, step + 1, minibatch_idx + 1, num_minibatches, + span.elapsed, + ) + + with elapsed_timer("optim_step") as span: + optim_result = policy.optim_step( + adam_params, + grad_accumulation_normalization=cfg.grad_accumulation_normalization, + ) + step += 1 + logger.info( + "[rollout %d step %d] optim_step (mb %d/%d) (%.1fs)", + rollout_id, step, minibatch_idx + 1, num_minibatches, + span.elapsed, + ) + + # Per-minibatch train/* on train/step axis (each inner step + # is genuinely distinct data, not an average). + mb_metrics = compute_minibatch_metrics(fwd_bwd_result, optim_result) + if mb_metrics: + mb_metrics["train/step"] = step + mb_metrics["train/minibatch_idx"] = minibatch_idx + 1 + mb_metrics["train/num_minibatches"] = num_minibatches + wandb_log(mb_metrics) + + # ``step_wall_time`` covers the full step (queue wait + all K + # minibatches), so ``perf/wait_time_ratio`` = wait / (wait + train). + if loop_stats is not None: + train_wall = time.monotonic() - train_start + loop_stats["step_wall_time"] = ( + loop_stats.get("trainer_wait_for_sampler_time", 0.0) + train_wall + ) + metrics = compute_step_metrics( + prompt_groups=prompt_groups, + fwd_bwd_results=fwd_bwd_results, + optim_result=optim_result, + n_accum=len(fwd_bwd_results), + timing_metrics=flush_timing(), + loop_stats=loop_stats, + completions_per_prompt=cfg.completions_per_prompt, + ) + # Capture the per-rollout-mean KL for the human summary line + # before the train/* stripping below removes it. + mean_kl = metrics.get("train/mean_kl", 0.0) + # train/* already logged per-minibatch above; strip the averages. + metrics = {k: v for k, v in metrics.items() if not k.startswith("train/")} + metrics["rollout/step"] = rollout_id + metrics["train/step"] = step # monotonic fallback for the wandb global step + metrics["ctx/current_version"] = current_version + for k, v in ctx_metadata.items(): + if isinstance(v, bool): + metrics[f"ctx/{k}"] = int(v) + elif isinstance(v, (int, float)) and v is not None: + metrics[f"ctx/{k}"] = v + + logger.info( + "Rollout %d (step %d) | Reward %.3f | Acc %.1f%% | KL %.4f", + rollout_id, + step, + metrics.get("rollout/reward", 0.0), + metrics.get("rollout/accuracy", 0.0) * 100, + mean_kl, + ) + wandb_log(metrics) + # DCP cadence is in rollout batches, not optim steps, so + # ppo_n_minibatches doesn't change save frequency. + rollouts_completed = (step - step_offset) // num_minibatches + interval = cfg.weight_sync.dcp_save_interval + if ( + loop_stats is not None + and interval > 0 + and rollouts_completed > 0 + and rollouts_completed % interval == 0 + ): + try: + _save_checkpoint( + ckpt, + name=f"step-{step}", + data_consumed=int(loop_stats["resolved_rows"]), + ) + except (OSError, RuntimeError) as e: + # Periodic save: surface the failure but keep training. + # The final save (after the loop) is allowed to propagate + # so orchestration sees terminal save problems. + logger.warning("[step %d] dcp_save failed: %s", step, e) + return step, metrics + + def _weight_sync(step: int) -> None: + nonlocal current_version + with elapsed_timer("weight_sync") as span: + weight_syncer.save_and_hotload(f"step-{step}") + current_version = step + logger.info("[step %d] weight_sync (%.1fs)", step, span.elapsed) + + global_step, final_stats = asyncio.run( + run_async_rl_loop( + rows=make_row_requests(), + train_fns=TrainStepFns(train_step=train_step), + completions_per_prompt=cfg.completions_per_prompt, + prompt_groups_per_step=cfg.prompt_groups_per_step, + max_head_offpolicy_versions=cfg.max_head_offpolicy_versions, + with_reference=(reference is not None), + min_group_size=cfg.min_group_size, + weight_sync_fn=_weight_sync, + weight_sync_interval=_WEIGHT_SYNC_INTERVAL, + max_concurrent=cfg.max_concurrency_rollout_sample, + dynamic_filter_fn=dynamic_filter_fn, + synchronous_training=cfg.synchronous_training, + global_step=step_offset, + resolved_rows_fn=lambda: row_loader.data_consumed, + return_final_stats=True, + ) + ) + + # Save resume progress even if all remaining rows were dropped. + # Promotion still requires at least one optimizer step. + resume_row_cursor = int(final_stats["resolved_rows"]) + has_trained_steps = global_step > step_offset + has_advanced_dataset = resume_row_cursor > prior_rows_consumed + if cfg.save_final_checkpoint and (has_trained_steps or has_advanced_dataset): + cp_name = f"step-{global_step}" + ckpt.save( + cp_name, + resumable=True, + promotable=has_trained_steps, + data_consumed=resume_row_cursor, + ) + if cfg.output_model_id and has_trained_steps: + ckpt.promote_latest(cfg.output_model_id, cfg.base_model) + + logger.info( + "Async RL training complete: %d steps (%d new in this run)", + global_step, global_step - step_offset, + ) + wandb_finish() + return { + "steps": global_step, + "policy_job_id": policy_job_id, + "reference_job_id": infra.reference_job_id, + } diff --git a/training/tests/test_smoke_imports.py b/training/tests/test_smoke_imports.py index c2c6bacc..bce07d4c 100644 --- a/training/tests/test_smoke_imports.py +++ b/training/tests/test_smoke_imports.py @@ -83,6 +83,8 @@ def test_util_imports(module: str): EXAMPLE_MODULES = [ "training.examples.rl.deepmath.prepare_data", + "training.examples.rl.single_turn_token_in.rollout", + "training.examples.rl.multi_turn_message_in.rollout", "training.examples.tools.promote_checkpoint", ] diff --git a/training/tests/unit/test_async_rl_train.py b/training/tests/unit/test_async_rl_train.py new file mode 100644 index 00000000..4c7e6c42 --- /dev/null +++ b/training/tests/unit/test_async_rl_train.py @@ -0,0 +1,647 @@ +"""Unit tests for training.utils.rl.async_train.run_async_rl_loop. + +Exercises the per-sample API: rows fan out to N samples, the +GroupAssembler joins them by row id, and the loop trains on assembled +PromptGroups. No GPU / network / fireworks SDK required. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from training.utils.rl.async_train import ( + RowRequest, + _StalenessController, + run_async_rl_loop, +) +from training.utils.rl.rollout import RolloutSample +from training.utils.rl.train import TrainStepFns + + +def _run(coro): + return asyncio.run(coro) + + +def _sample(reward: float = 0.0) -> RolloutSample: + """Minimal valid RolloutSample for loop-level tests.""" + return RolloutSample( + tokens=[1, 2], + logprobs=[0.0, -0.1], + loss_mask=[0, 1], + reward=reward, + ) + + +def _passthrough_advantages(rewards): + """REINFORCE-style: raw reward as advantage. Safe on N=1 groups.""" + return list(rewards) + + +def _row( + row_id: int, + *, + reward: float = 0.0, + fail_indices: tuple[int, ...] = (), + on_resolved=None, + row_meta=None, +) -> RowRequest: + """Build a RowRequest whose sample_factory returns a fresh sample + per sub_index (or ``None`` for indices in ``fail_indices``).""" + + async def factory(sub_index: int): + if sub_index in fail_indices: + return None + return _sample(reward=reward) + + return RowRequest( + row_id=row_id, + sample_factory=factory, + row_meta=row_meta, + on_resolved=on_resolved, + ) + + +def _noop_train_step(step, groups, extra): + return step + 1, {} + + +# Sensible defaults for tests that don't care about the new knobs. +_DEFAULTS = dict( + completions_per_prompt=1, + advantage_fn=_passthrough_advantages, +) + + +class TestArgValidation: + def test_rejects_invalid_prompt_groups_per_step(self): + with pytest.raises(ValueError, match="prompt_groups_per_step"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=0, + max_head_offpolicy_versions=0, + **_DEFAULTS, + )) + + def test_rejects_invalid_completions_per_prompt(self): + with pytest.raises(ValueError, match="completions_per_prompt"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + completions_per_prompt=0, + advantage_fn=_passthrough_advantages, + )) + + def test_rejects_min_group_size_above_completions(self): + with pytest.raises(ValueError, match="min_group_size"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + completions_per_prompt=2, + min_group_size=3, + advantage_fn=_passthrough_advantages, + )) + + def test_rejects_negative_offpolicy(self): + with pytest.raises(ValueError, match="max_head_offpolicy_versions"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=-1, + **_DEFAULTS, + )) + + def test_rejects_zero_weight_sync_interval(self): + with pytest.raises(ValueError, match="weight_sync_interval"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + weight_sync_interval=0, + **_DEFAULTS, + )) + + def test_rejects_zero_max_concurrent(self): + with pytest.raises(ValueError, match="max_concurrent"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + max_concurrent=0, + **_DEFAULTS, + )) + + def test_rejects_sync_interval_gt_offpolicy_plus_one(self): + """``weight_sync_interval > 1`` + ``max_head_offpolicy_versions == 0`` + is mathematically guaranteed to stall.""" + with pytest.raises(ValueError, match="max_head_offpolicy_versions"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + weight_sync_interval=5, + **_DEFAULTS, + )) + + def test_accepts_balanced_sync_interval_and_offpolicy(self): + result = _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=2, + weight_sync_interval=3, + **_DEFAULTS, + )) + assert isinstance(result, int) + + +class TestStalenessController: + def test_capacity_matches_areal_formula(self): + # cpp=1 collapses sample-level math to batch/group units, so the + # values below double as a reference for the AReaL row-level form. + ctl = _StalenessController( + batch_size_samples=4, + completions_per_prompt=1, + max_staleness=1, + max_concurrent_samples=3, + ) + assert ctl.capacity() == 3 + + ctl.submit() + ctl.submit() + assert ctl.capacity() == 1 + + ctl.accept() + ctl.accept() + assert ctl.capacity() == 3 + + ctl.advance_version() + assert ctl.capacity() == 3 + + def test_reject_frees_running_capacity_without_charging_accepted(self): + ctl = _StalenessController( + batch_size_samples=2, + completions_per_prompt=1, + max_staleness=0, + max_concurrent_samples=2, + ) + ctl.submit() + ctl.submit() + assert ctl.capacity() == 0 + + ctl.reject("none") + assert ctl.accepted_samples == 0 + assert ctl.sample_fails == 1 + assert ctl.capacity() == 1 + + +class TestHappyPath: + def test_strict_onpolicy_runs_all_steps(self): + """budget=0, gpb=2, 4 rows -> 2 training steps.""" + n_rows = 4 + calls = [] + + def train_step(step, groups, extra): + calls.append((step, len(groups), extra.get("async/version_offset_max"))) + return step + 1, {} + + result = _run(run_async_rl_loop( + rows=(_row(i, reward=float(i)) for i in range(n_rows)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + **_DEFAULTS, + )) + assert result == 2 + assert len(calls) == 2 + assert all(offset == 0 for _, _, offset in calls) + + def test_sample_failures_skip_batch_not_charge_gate(self): + """Sample factory returning None drops the row from training.""" + n_rows = 6 + calls = [] + + def train_step(step, groups, extra): + calls.append(extra["sample_fails"]) + return step + 1, {} + + # Even rows fail (single sample == row drops). + rows = ( + _row(i, fail_indices=(0,) if i % 2 == 0 else ()) + for i in range(n_rows) + ) + + _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=2, + **_DEFAULTS, + )) + assert any(fails >= 1 for fails in calls) + + def test_filter_drops_exclude_from_buffer(self): + """dynamic_filter_fn=False drops groups, doesn't train on them.""" + n_rows = 6 + + def filter_fn(pg): + return pg.rewards[0] >= 3.0 + + trained_groups = [] + + def train_step(step, groups, extra): + trained_groups.extend(groups) + return step + 1, {} + + _run(run_async_rl_loop( + rows=(_row(i, reward=float(i)) for i in range(n_rows)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=3, + max_head_offpolicy_versions=2, + dynamic_filter_fn=filter_fn, + **_DEFAULTS, + )) + assert len(trained_groups) == 3 + assert all(g.rewards[0] >= 3.0 for g in trained_groups) + + def test_iterator_exhausted_flushes_partial_final_step(self): + sizes: list[int] = [] + + def train_step(step, groups, extra): + sizes.append(len(groups)) + return step + 1, {} + + final_step = _run(run_async_rl_loop( + rows=(_row(i) for i in range(5)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + **_DEFAULTS, + )) + assert sizes == [2, 2, 1] + assert final_step == 3 + + def test_partial_only_buffer_still_trains(self): + sizes: list[int] = [] + + def train_step(step, groups, extra): + sizes.append(len(groups)) + return step + 1, {} + + _run(run_async_rl_loop( + rows=(_row(i) for i in range(2)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=4, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + **_DEFAULTS, + )) + assert sizes == [2] + + def test_stalled_gate_does_not_consume_remaining_iterator(self): + """A stalled strict-on-policy loop must not drain a lazy row iterator + just to report how many rows remain.""" + + class Rows: + def __init__(self): + self.consumed = 0 + + def __iter__(self): + return self + + def __next__(self): + self.consumed += 1 + if self.consumed > 4: + raise StopIteration + return _row(self.consumed) + + rows = Rows() + _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=None, + **_DEFAULTS, + )) + assert rows.consumed == 2 + + +class TestVersionTracking: + def test_weight_sync_increments_version(self): + sync_calls = [] + + def weight_sync(step): + sync_calls.append(step) + + _run(run_async_rl_loop( + rows=(_row(i) for i in range(6)), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=weight_sync, + weight_sync_interval=1, + **_DEFAULTS, + )) + assert sync_calls == [1, 2, 3] + + def test_weight_sync_interval_skips(self): + sync_calls = [] + + def weight_sync(step): + sync_calls.append(step) + + _run(run_async_rl_loop( + rows=(_row(i) for i in range(8)), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=2, + weight_sync_fn=weight_sync, + weight_sync_interval=2, + **_DEFAULTS, + )) + assert sync_calls == [2, 4] + + def test_version_offset_emitted(self): + metric_snapshots = [] + + def train_step(step, groups, extra): + metric_snapshots.append(dict(extra)) + return step + 1, {} + + _run(run_async_rl_loop( + rows=(_row(i) for i in range(4)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + **_DEFAULTS, + )) + assert len(metric_snapshots) == 2 + for m in metric_snapshots: + assert m["async/version_offset_mean"] == 0 + assert m["async/version_offset_max"] == 0 + assert m["async/version_offset_min"] == 0 + + +class TestMaxConcurrent: + def test_caps_in_flight_even_with_budget(self): + """max_concurrent counts in-flight SAMPLES; with cpp=1 the row + and sample bookkeeping coincide.""" + peak_in_flight = [0] + + def train_step(step, groups, extra): + peak_in_flight[0] = max(peak_in_flight[0], extra["async/in_flight"]) + return step + 1, {} + + async def slow_factory(_sub): + await asyncio.sleep(0.01) + return _sample() + + rows = ( + RowRequest(row_id=i, sample_factory=slow_factory) + for i in range(10) + ) + + _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=5, + max_head_offpolicy_versions=10, + max_concurrent=2, + **_DEFAULTS, + )) + assert peak_in_flight[0] <= 2 + + def test_caps_in_flight_at_sample_level_with_cpp_gt_1(self): + """cpp=4, max_concurrent=8: never more than 8 samples in flight. + + Pins the unit semantic for max_concurrent: it counts samples + (LLM calls), not rows. Each row submits cpp=4 sample tasks + atomically, so under an 8-sample cap at most 2 rows are admitted. + """ + peak_samples = [0] + + async def slow_factory(_sub): + await asyncio.sleep(0.005) + return _sample() + + def train_step(step, groups, extra): + peak_samples[0] = max(peak_samples[0], extra["async/in_flight"]) + return step + 1, {} + + rows = ( + RowRequest(row_id=i, sample_factory=slow_factory) + for i in range(6) + ) + + _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=10, + completions_per_prompt=4, + max_concurrent=8, + weight_sync_fn=lambda _step: None, + advantage_fn=_passthrough_advantages, + )) + assert peak_samples[0] <= 8 + + def test_rejects_max_concurrent_below_cpp(self): + """max_concurrent < cpp would deadlock the gate (can't admit a row).""" + with pytest.raises(ValueError, match="max_concurrent"): + _run(run_async_rl_loop( + rows=iter([]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=2, + completions_per_prompt=4, + max_concurrent=3, + advantage_fn=_passthrough_advantages, + )) + + +class TestTaskExceptions: + def test_exception_in_sample_propagates(self): + """Sample-task exceptions abort the run rather than getting silently + counted as ``sample_fails``.""" + + async def factory(sub_index): + raise RuntimeError("boom") + + bad_row = RowRequest(row_id=99, sample_factory=factory) + + with pytest.raises(RuntimeError, match="boom"): + _run(run_async_rl_loop( + rows=iter([_row(0), bad_row, _row(2), _row(3)]), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=2, + **_DEFAULTS, + )) + + +class TestRowResolutionHooks: + def test_on_resolved_drives_external_cursor(self): + cursor = {"value": 0} + + def make_row(i): + return _row(i, on_resolved=lambda _reason, i=i: cursor.update(value=i + 1)) + + final_step, final = _run(run_async_rl_loop( + rows=(make_row(i) for i in range(3)), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + resolved_rows_fn=lambda: cursor["value"], + return_final_stats=True, + **_DEFAULTS, + )) + assert final_step == 3 + assert final["resolved_rows"] == 3 + + def test_final_stats_include_tail_sample_failures(self): + """Rows where the only sample fails count as sample_fails in stats.""" + + rows = ( + _row(i, fail_indices=(0,) if i >= 3 else ()) + for i in range(5) + ) + + final_step, final = _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + weight_sync_fn=lambda _step: None, + return_final_stats=True, + **_DEFAULTS, + )) + # Rows 0,1 train as step 1; row 2 trains as a partial step 2. + assert final_step == 2 + assert final["sample_fails"] == 2 + assert final["total_accepted"] == 3 + assert final["resolved_rows"] == 5 + + def test_filter_drops_recorded_in_final_stats(self): + def filter_fn(_pg): + return False # drop everything + + final_step, final = _run(run_async_rl_loop( + rows=(_row(i) for i in range(3)), + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=0, + dynamic_filter_fn=filter_fn, + return_final_stats=True, + **_DEFAULTS, + )) + assert final_step == 0 + assert final["filter_drops"] == 3 + assert final["sample_fails"] == 0 + assert final["resolved_rows"] == 3 + + def test_resolved_rows_offset_applies_to_metrics(self): + snapshots = [] + + def train_step(step, groups, extra): + snapshots.append(dict(extra)) + return step + 1, {} + + final_step, final = _run(run_async_rl_loop( + rows=(_row(i) for i in range(2)), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=1, + max_head_offpolicy_versions=1, + resolved_rows_offset=10, + return_final_stats=True, + **_DEFAULTS, + )) + assert final_step == 2 + assert [m["resolved_rows"] for m in snapshots] == [12, 12] + assert final["resolved_rows"] == 12 + + +class TestPerSampleFanout: + def test_grpo_fanout_assembles_n_samples_per_row(self): + """completions_per_prompt=4 with default GRPO advantages: every row + produces a 4-sample group.""" + trained_groups = [] + + def train_step(step, groups, extra): + trained_groups.extend(groups) + return step + 1, {} + + # Two rows -> 1 step at gpb=2. Each row fans out to 4 samples. + _run(run_async_rl_loop( + rows=iter([_row(0, reward=1.0), _row(1, reward=2.0)]), + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + completions_per_prompt=4, + weight_sync_fn=lambda _step: None, + )) + assert len(trained_groups) == 2 + # Each PromptGroup has 4 datums (one per sample). + assert all(len(g.data) == 4 for g in trained_groups) + + def test_partial_group_emits_when_min_group_size_satisfied(self): + """If 1 of 4 samples fails but min_group_size=2, the row still emits.""" + trained_groups = [] + + def train_step(step, groups, extra): + trained_groups.extend(groups) + return step + 1, {} + + # Fail 2 of 4 samples per row -> 2 surviving samples meets min=2. + rows = iter([ + _row(0, reward=0.5, fail_indices=(0, 2)), + _row(1, reward=1.5, fail_indices=(1, 3)), + ]) + _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + completions_per_prompt=4, + min_group_size=2, + weight_sync_fn=lambda _step: None, + )) + assert len(trained_groups) == 2 + assert all(len(g.data) == 2 for g in trained_groups) + + def test_row_fully_dropped_when_below_min_group_size(self): + """When more samples fail than min_group_size allows, the row drops + entirely and counts as sample_fail.""" + rows = iter([ + _row(0, fail_indices=(0, 1, 2, 3)), # all fail + _row(1), # all succeed + _row(2), + ]) + final_step, final = _run(run_async_rl_loop( + rows=rows, + train_fns=TrainStepFns(train_step=_noop_train_step), + prompt_groups_per_step=2, + max_head_offpolicy_versions=0, + completions_per_prompt=4, + min_group_size=2, + weight_sync_fn=lambda _step: None, + return_final_stats=True, + )) + # Only 2 surviving rows -> 1 step at gpb=2. + assert final_step == 1 + assert final["sample_fails"] == 1 + assert final["total_accepted"] == 2 diff --git a/training/tests/unit/test_cursor_dataloader.py b/training/tests/unit/test_cursor_dataloader.py new file mode 100644 index 00000000..8dcd7a66 --- /dev/null +++ b/training/tests/unit/test_cursor_dataloader.py @@ -0,0 +1,70 @@ +from training.utils.dataloader import CursorDataLoader + + +def test_cursor_starts_from_resume_point(): + loader = CursorDataLoader(["a", "b", "c"], start_cursor=1) + + assert next(loader).value == "b" + assert loader.data_consumed == 1 + + +def test_cursor_advances_in_order_only(): + loader = CursorDataLoader(["a", "b", "c"]) + + loader.mark_resolved(1) + assert loader.data_consumed == 0 + + loader.mark_resolved(0) + assert loader.data_consumed == 2 + + loader.mark_resolved(2) + assert loader.data_consumed == 3 + + +def test_cursor_ignores_already_resolved_indices(): + loader = CursorDataLoader(["a", "b"], start_cursor=1) + + loader.mark_resolved(0) + assert loader.data_consumed == 1 + + +def test_cursor_can_resume_past_available_rows(): + loader = CursorDataLoader(["a"], start_cursor=3) + + assert loader.data_consumed == 3 + assert list(loader) == [] + + +def test_cursor_owns_epochs_without_reusing_row_objects(): + loader = CursorDataLoader([{"id": 1}], epochs=2) + + first = next(loader) + first.value["seen"] = True + second = next(loader) + + assert first.index == 0 + assert second.index == 1 + assert second.value == {"id": 1} + + +def test_cursor_shuffle_is_deterministic_per_epoch(): + rows = list(range(10)) + first = [item.value for item in CursorDataLoader(rows, epochs=2, shuffle=True, seed=7)] + second = [item.value for item in CursorDataLoader(rows, epochs=2, shuffle=True, seed=7)] + unshuffled = [item.value for item in CursorDataLoader(rows, epochs=2)] + + assert first == second + assert first != unshuffled + assert sorted(first[:10]) == rows + assert sorted(first[10:]) == rows + + +def test_cursor_resume_reconstructs_epoch_shuffle_position(): + rows = list(range(6)) + full = [item.value for item in CursorDataLoader(rows, epochs=3, shuffle=True, seed=4)] + resumed = [ + item.value + for item in CursorDataLoader(rows, start_cursor=8, epochs=3, shuffle=True, seed=4) + ] + + assert resumed == full[8:] diff --git a/training/tests/unit/test_data_utils.py b/training/tests/unit/test_data_utils.py new file mode 100644 index 00000000..558b921a --- /dev/null +++ b/training/tests/unit/test_data_utils.py @@ -0,0 +1,53 @@ +"""Unit tests for ``training.utils.data`` helper utilities.""" + +from __future__ import annotations + +from training.utils import replicate_rows_for_epochs + + +def test_replicate_rows_for_epochs_each_row_is_independent(): + """The naive ``rows * epochs`` only multiplies list references -- + every epoch shares the same dict instances. Any rollout function + that mutates its input row in place (attaching scratch fields, + normalizing prompt data, caching renders) leaks that mutation + into every later epoch and subsequent passes train on the + already-mutated row instead of the original dataset. + ``replicate_rows_for_epochs`` returns ``epochs * len(rows)`` + INDEPENDENT dict instances so per-epoch mutations cannot leak. + """ + rows = [{"prompt": "a", "scratch": []}, {"prompt": "b", "scratch": []}] + out = replicate_rows_for_epochs(rows, epochs=3) + assert len(out) == 6 # 2 rows x 3 epochs + # Every output dict is a distinct object -- id() differs even for + # the same logical row across epochs. + assert len({id(r) for r in out}) == 6 + # The nested mutable container (``scratch``) is also independent + # -- that's what makes ``deepcopy`` (vs shallow ``dict(r)``) the + # right primitive here. + assert len({id(r["scratch"]) for r in out}) == 6 + # Mutating one copy does not affect any other copy or the originals. + out[0]["scratch"].append("epoch0-mutation") + assert all(r["scratch"] == [] for r in rows), ( + "Mutating a replicated copy leaked back into the originals -- " + "deepcopy expected." + ) + assert all(out[i]["scratch"] == [] for i in range(1, 6)), ( + "Mutating one copy leaked into other copies -- sibling rows " + "must be independent across all epoch slots." + ) + + +def test_replicate_rows_for_epochs_zero_epochs_returns_empty(): + """``epochs=0`` is a degenerate but valid input; the helper + returns an empty list rather than raising.""" + rows = [{"x": 1}] + assert replicate_rows_for_epochs(rows, epochs=0) == [] + + +def test_replicate_rows_for_epochs_preserves_per_epoch_order(): + """Rows are emitted epoch-by-epoch so the resume slicing logic + (which advances a raw-row cursor through ``all_rows``) sees the + same order it did when ``rows * epochs`` was used.""" + rows = [{"i": 0}, {"i": 1}, {"i": 2}] + out = replicate_rows_for_epochs(rows, epochs=2) + assert [r["i"] for r in out] == [0, 1, 2, 0, 1, 2] diff --git a/training/tests/unit/test_decoupled_is.py b/training/tests/unit/test_decoupled_is.py index cf076cac..c0784598 100644 --- a/training/tests/unit/test_decoupled_is.py +++ b/training/tests/unit/test_decoupled_is.py @@ -55,3 +55,58 @@ def test_token_level_gives_different_weights(self): weight, _ = compute_tis_weight(prox, inf, TISConfig(cap=100.0)) assert weight[0] != weight[1] + + +class TestActiveFilterRegression: + """Regression: TIS/drift averaged over the full response slice + (including masked bridge / user-feedback positions in multi-turn + rollouts) silently contaminated the per-sample weight at active + positions when ``level="sequence"``. Both ``utils/rl/common.py`` and + ``utils/rl/losses.py`` now pre-filter to ``loss_mask > 0`` positions + before calling :func:`compute_tis_weight` and expand the result back + to the full response shape with 1.0 at masked positions. + """ + + def test_sequence_weight_ignores_masked_positions(self): + # 4 response tokens: indices 0 and 3 are active assistant tokens; + # indices 1 and 2 are a masked bridge / user-feedback span with + # an extreme log-ratio that, if averaged in, would dominate. + resp_prox = torch.tensor([-0.5, 5.0, 5.0, -0.5]) + resp_inf = torch.tensor([-0.5, 0.0, 0.0, -0.5]) + resp_mask = torch.tensor([1.0, 0.0, 0.0, 1.0]) + active = resp_mask > 0.5 + + tis_weight_active, _ = compute_tis_weight( + resp_prox[active], resp_inf[active], TISConfig(level="sequence", cap=100.0), + ) + tis_weight = torch.ones(resp_prox.shape[0]) + tis_weight[active] = tis_weight_active.to(tis_weight.dtype) + + # Active-only mean of (prox - inf) is 0 -> weight 1.0 at active + # positions. Contaminated full-slice mean would be (5+5)/4 = 2.5 + # -> exp(2.5) ~= 12.18; that's the wrong answer the fix prevents. + torch.testing.assert_close(tis_weight[0], torch.tensor(1.0)) + torch.testing.assert_close(tis_weight[3], torch.tensor(1.0)) + assert tis_weight[1].item() == 1.0 # masked: identity + assert tis_weight[2].item() == 1.0 + + def test_token_weight_isolates_active_positions(self): + # With per-token IS, the active-only filter shouldn't change + # active values (a no-op there), but it lets us assign 1.0 to + # masked positions instead of whatever the per-token formula + # would have produced from masked log-ratios. + resp_prox = torch.tensor([-0.5, 5.0, -0.5]) + resp_inf = torch.tensor([-0.5, 0.0, -0.5]) + resp_mask = torch.tensor([1.0, 0.0, 1.0]) + active = resp_mask > 0.5 + + tis_weight_active, _ = compute_tis_weight( + resp_prox[active], resp_inf[active], TISConfig(level="token", cap=100.0), + ) + tis_weight = torch.ones(resp_prox.shape[0]) + tis_weight[active] = tis_weight_active.to(tis_weight.dtype) + + torch.testing.assert_close(tis_weight[0], torch.tensor(1.0)) + torch.testing.assert_close(tis_weight[2], torch.tensor(1.0)) + # Masked position is the identity weight, not exp(5.0) ≈ 148.4. + assert tis_weight[1].item() == 1.0 diff --git a/training/tests/unit/test_group_assembler.py b/training/tests/unit/test_group_assembler.py new file mode 100644 index 00000000..624e30de --- /dev/null +++ b/training/tests/unit/test_group_assembler.py @@ -0,0 +1,200 @@ +"""Unit tests for training.utils.rl.rollout.group_assembler.GroupAssembler. + +Covers row-level fan-in semantics: full-group emit, partial-group on +quorum, single-sample dropped under default advantages, custom +advantage_fn passthrough, row_meta propagation, and drain-at-shutdown. +""" + +from __future__ import annotations + +from training.utils.rl.rollout import ( + GroupAssembler, + RolloutSample, +) + + +def _sample(reward: float = 0.0) -> RolloutSample: + return RolloutSample( + tokens=[1, 2, 3], + logprobs=[0.0, -0.1, -0.2], + loss_mask=[0, 1, 1], + reward=reward, + ) + + +def _passthrough(rewards): + return list(rewards) + + +class TestSettlement: + def test_full_group_emits_after_n_samples(self): + asm = GroupAssembler(completions_per_prompt=3) + for _ in range(3): + asm.note_started("row-A", submit_version=0) + for r in (0.5, 1.5, 2.5): + res = asm.add_sample("row-A", _sample(reward=r)) + # Last add_sample should return resolution; earlier two return None. + assert res is not None + assert res.pg is not None + assert len(res.pg.data) == 3 + assert res.pg.rewards == [0.5, 1.5, 2.5] + + def test_intermediate_calls_return_none(self): + asm = GroupAssembler(completions_per_prompt=3) + for _ in range(3): + asm.note_started("r", submit_version=0) + assert asm.add_sample("r", _sample(0.0)) is None + assert asm.add_sample("r", _sample(1.0)) is None + # Only the third call settles. + out = asm.add_sample("r", _sample(2.0)) + assert out is not None + + def test_drops_when_all_samples_fail(self): + asm = GroupAssembler(completions_per_prompt=2) + asm.note_started("r", submit_version=0) + asm.note_started("r", submit_version=0) + assert asm.note_dropped("r") is None + out = asm.note_dropped("r") + assert out is not None + assert out.pg is None # settled-but-empty + + def test_min_group_size_drops_below_threshold(self): + asm = GroupAssembler(completions_per_prompt=4, min_group_size=3) + for _ in range(4): + asm.note_started("r", submit_version=0) + asm.add_sample("r", _sample(0.0)) + asm.add_sample("r", _sample(1.0)) + asm.note_dropped("r") + out = asm.note_dropped("r") + # 2 surviving < min=3 -> settled-but-empty. + assert out is not None + assert out.pg is None + + def test_partial_group_emits_when_min_satisfied(self): + asm = GroupAssembler( + completions_per_prompt=4, + min_group_size=2, + advantage_fn=_passthrough, + ) + for _ in range(4): + asm.note_started("r", submit_version=0) + asm.add_sample("r", _sample(0.5)) + asm.add_sample("r", _sample(1.5)) + asm.note_dropped("r") + out = asm.note_dropped("r") + assert out is not None + assert out.pg is not None + assert len(out.pg.data) == 2 + assert out.pg.rewards == [0.5, 1.5] + + +class TestVersionTracking: + def test_min_submit_version_returned(self): + asm = GroupAssembler(completions_per_prompt=2, advantage_fn=_passthrough) + asm.note_started("r", submit_version=5) + asm.note_started("r", submit_version=7) + asm.add_sample("r", _sample(0.0)) + out = asm.add_sample("r", _sample(1.0)) + assert out is not None + # Stale-as-its-stalest-sample: oldest version dominates. + assert out.min_submit_version == 5 + + +class TestAdvantageFn: + def test_default_grpo_drops_singleton(self): + """Default GRPO z-score is NaN on N=1 -- the row drops.""" + asm = GroupAssembler(completions_per_prompt=1) + asm.note_started("r", submit_version=0) + out = asm.add_sample("r", _sample(1.0)) + assert out is not None + assert out.pg is None + + def test_custom_passthrough_keeps_singleton(self): + """REINFORCE-style custom advantage_fn is well-defined on N=1.""" + asm = GroupAssembler(completions_per_prompt=1, advantage_fn=_passthrough) + asm.note_started("r", submit_version=3) + out = asm.add_sample("r", _sample(0.7)) + assert out is not None + assert out.pg is not None + assert out.pg.rewards == [0.7] + assert out.pg.advantages == [0.7] + assert out.min_submit_version == 3 + + +class TestRowMeta: + def test_row_meta_threaded_to_prompt_group(self): + asm = GroupAssembler(completions_per_prompt=2, advantage_fn=_passthrough) + meta = {"row_id": "abc", "extra": 1} + asm.note_started("r", submit_version=0, row_meta=meta) + asm.note_started("r", submit_version=0) # second call ignores row_meta + asm.add_sample("r", _sample(0.0)) + out = asm.add_sample("r", _sample(1.0)) + assert out is not None + assert out.pg is not None + assert out.pg.row_meta == meta + # Defensive copy. + assert out.pg.row_meta is not meta + + +class TestMultipleRows: + def test_independent_row_assembly(self): + asm = GroupAssembler(completions_per_prompt=2, advantage_fn=_passthrough) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.note_started("B", submit_version=0) + asm.note_started("B", submit_version=0) + + # Interleave samples across rows. + assert asm.add_sample("A", _sample(1.0)) is None + assert asm.add_sample("B", _sample(2.0)) is None + out_a = asm.add_sample("A", _sample(3.0)) + out_b = asm.add_sample("B", _sample(4.0)) + + assert out_a is not None and out_a.pg is not None + assert out_b is not None and out_b.pg is not None + assert out_a.pg.rewards == [1.0, 3.0] + assert out_b.pg.rewards == [2.0, 4.0] + + def test_pending_rows_count(self): + asm = GroupAssembler(completions_per_prompt=2) + asm.note_started("A", submit_version=0) + asm.note_started("B", submit_version=0) + assert asm.pending_rows() == 2 + asm.note_started("A", submit_version=0) + assert asm.pending_rows() == 2 # still A and B + asm.add_sample("A", _sample(0.0)) + asm.add_sample("A", _sample(0.0)) + assert asm.pending_rows() == 1 + + +class TestDrain: + def test_drain_emits_partial_groups_meeting_min(self): + asm = GroupAssembler( + completions_per_prompt=4, + min_group_size=2, + advantage_fn=_passthrough, + ) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.add_sample("A", _sample(1.0)) + asm.add_sample("A", _sample(2.0)) + # Two more never landed; drain emits the partial. + drained = asm.drain() + assert len(drained) == 1 + assert drained[0].pg is not None + assert drained[0].pg.rewards == [1.0, 2.0] + + def test_drain_skips_below_min(self): + asm = GroupAssembler( + completions_per_prompt=4, + min_group_size=3, + ) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.note_started("A", submit_version=0) + asm.add_sample("A", _sample(0.0)) + # Only 1 sample; below min=3. + assert asm.drain() == [] diff --git a/training/tests/unit/test_gsm8k_reward.py b/training/tests/unit/test_gsm8k_reward.py new file mode 100644 index 00000000..5cb04115 --- /dev/null +++ b/training/tests/unit/test_gsm8k_reward.py @@ -0,0 +1,62 @@ +"""Unit tests for the GSM8K multi-turn example's reward function.""" + +from __future__ import annotations + +from training.examples.rl.multi_turn_message_in.reward import gsm8k_reward + + +GT_FROM_GSM8K = ( + "Janet sells 16 - 3 - 4 = <<16-3-4=9>>9 duck eggs a day.\n" + "She makes 9 * 2 = $<<9*2=18>>18 every day at the farmer's market.\n" + "#### 18" +) + + +class TestCorrectAnswers: + def test_boxed_integer_matches_gt(self): + assert gsm8k_reward("The answer is \\boxed{18}.", GT_FROM_GSM8K) == 1.0 + + def test_boxed_with_spaces(self): + assert gsm8k_reward("Therefore \\boxed{ 18 }", GT_FROM_GSM8K) == 1.0 + + def test_uses_last_boxed_when_multiple(self): + # First guess wrong, second guess right — last boxed wins. + completion = "First I thought \\boxed{16}, but actually \\boxed{18}." + assert gsm8k_reward(completion, GT_FROM_GSM8K) == 1.0 + + def test_comma_separated_thousands(self): + gt = "#### 1,200" + assert gsm8k_reward("\\boxed{1200}", gt) == 1.0 + assert gsm8k_reward("\\boxed{1,200}", gt) == 1.0 + + +class TestWrongAnswers: + def test_wrong_number(self): + assert gsm8k_reward("\\boxed{17}", GT_FROM_GSM8K) == 0.0 + + def test_no_boxed(self): + assert gsm8k_reward("The answer is 18.", GT_FROM_GSM8K) == 0.0 + + def test_unmatched_braces(self): + assert gsm8k_reward("\\boxed{18", GT_FROM_GSM8K) == 0.0 + + def test_empty_completion(self): + assert gsm8k_reward("", GT_FROM_GSM8K) == 0.0 + + +class TestEdgeCases: + def test_gt_with_negative_number(self): + assert gsm8k_reward("\\boxed{-5}", "Some reasoning.\n#### -5") == 1.0 + + def test_nested_braces_preserved(self): + # Should still extract the (incorrect-here) content, then fail to match. + assert gsm8k_reward("\\boxed{\\frac{1}{2}}", GT_FROM_GSM8K) == 0.0 + + def test_missing_gt_marker(self): + # No '#### N' anchor → ground-truth extraction yields None; + # math_verify fallback will try the raw answer string. + assert gsm8k_reward("\\boxed{18}", "no answer marker here") == 0.0 + + def test_close_but_not_equal_floats_reject(self): + gt = "#### 18" + assert gsm8k_reward("\\boxed{18.0001}", gt) == 0.0 diff --git a/training/tests/unit/test_rl_metrics_aliases.py b/training/tests/unit/test_rl_metrics_aliases.py index d286bb9e..0cf88838 100644 --- a/training/tests/unit/test_rl_metrics_aliases.py +++ b/training/tests/unit/test_rl_metrics_aliases.py @@ -36,7 +36,7 @@ def test_excludes_loop_rollout_count_metrics(self): assert "rollout/tokens_completed" not in loop_metrics assert "rollout/sample_fail_ratio" not in loop_metrics assert "rollout/raw_reward" not in loop_metrics - assert "perf/sample_wait_time" not in loop_metrics + assert "perf/trainer_wait_for_sampler_time" not in loop_metrics assert "perf/wait_time_ratio" not in loop_metrics assert "perf/overlap_ratio" not in loop_metrics assert "perf/step_wall_time" not in loop_metrics @@ -59,7 +59,7 @@ def test_uses_canonical_loop_stats_keys(self): "total_sampled": 7, "filter_drops": 1, "sample_fails": 2, - "sample_wait_time": 3.0, + "trainer_wait_for_sampler_time": 3.0, "step_wall_time": 4.0, "all_raw_rewards": [1.0, 0.0], }, @@ -70,7 +70,7 @@ def test_uses_canonical_loop_stats_keys(self): assert metrics["rollout/samples_completed"] == 1 assert metrics["rollout/filter_reject_ratio"] == 1 / 7 assert metrics["rollout/sample_fail_count"] == 2 - assert metrics["perf/sample_wait_time"] == 3.0 + assert metrics["perf/trainer_wait_for_sampler_time"] == 3.0 assert metrics["perf/wait_time_ratio"] == 0.75 assert metrics["perf/overlap_ratio"] == 0.25 assert metrics["perf/step_wall_time"] == 4.0 diff --git a/training/tests/unit/test_rollout_assembler.py b/training/tests/unit/test_rollout_assembler.py new file mode 100644 index 00000000..33a57661 --- /dev/null +++ b/training/tests/unit/test_rollout_assembler.py @@ -0,0 +1,300 @@ +"""Unit tests for ``TrajectoryAssembler``. + +The assembler carries the AReaL prefix-equality invariant for multi-turn +rollouts. Each engine call's ``input_tokens`` must extend the +already-accumulated sequence verbatim; if the rollout function +re-tokenized text between turns the assembler raises +:class:`PrefixMismatch` instead of training on misaligned tokens. +""" + +from __future__ import annotations + +import pytest + +from training.utils.rl.rollout import remote as tr +from training.utils.rl.rollout import ( + InferenceCall, + PrefixMismatch, + TrajectoryAssembler, +) + + +# --------------------------------------------------------------------------- +# Happy paths +# --------------------------------------------------------------------------- + + +def test_single_call_builds_payload(): + asm = TrajectoryAssembler(tokenizer_id="test-tok") + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3], + output_tokens=[10, 11], + output_logprobs=[-0.1, -0.2], + finish_reason="stop", + ), + ) + payload = asm.to_payload(total_reward=0.5) + + assert payload.tokenizer_id == "test-tok" + assert payload.total_reward == 0.5 + assert payload.finish_reason == "stop" + assert getattr(payload, "_assembled") is True + assert [t.role for t in payload.turns] == ["user", "assistant"] + assert payload.turns[0].token_ids == [1, 2, 3] + assert payload.turns[1].token_ids == [10, 11] + assert payload.turns[1].logprobs == [-0.1, -0.2] + + +def test_multi_turn_extends_prefix(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3], + output_tokens=[10, 11], + output_logprobs=[-0.1, -0.2], + ), + ) + # Next call's input is prior prompt + assistant output + a 2-token suffix. + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3, 10, 11, 50, 51], + output_tokens=[20, 21, 22], + output_logprobs=[-0.3, -0.4, -0.5], + ), + role_before="tool", + ) + payload = asm.to_payload(total_reward=1.0) + + roles = [t.role for t in payload.turns] + assert roles == ["user", "assistant", "tool", "assistant"] + assert payload.turns[2].token_ids == [50, 51] + assert payload.turns[3].token_ids == [20, 21, 22] + assert payload.turns[3].logprobs == [-0.3, -0.4, -0.5] + assert asm.accumulated_tokens == [1, 2, 3, 10, 11, 50, 51, 20, 21, 22] + + +def test_to_flat_returns_aligned_arrays(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2], + output_tokens=[10, 11], + output_logprobs=[-0.1, -0.2], + ), + ) + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 10, 11, 99], + output_tokens=[20], + output_logprobs=[-0.3], + ), + ) + tokens, logprobs, mask = asm.to_flat() + + assert tokens == [1, 2, 10, 11, 99, 20] + assert mask == [0, 0, 1, 1, 0, 1] + assert logprobs == [0.0, 0.0, -0.1, -0.2, 0.0, -0.3] + + +def test_add_call_allows_explicit_boundary_trim(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2], + output_tokens=[777], + output_logprobs=[-0.1], + ), + ) + + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 50], + output_tokens=[20], + output_logprobs=[-0.2], + ), + max_trim_tokens=1, + ) + + tokens, logprobs, mask = asm.to_flat() + assert tokens == [1, 2, 50, 20] + assert logprobs == [0.0, 0.0, 0.0, -0.2] + assert mask == [0, 0, 0, 1] + assert asm.accumulated_tokens == [1, 2, 50, 20] + + +def test_add_environment_tokens_records_turn_but_not_engine_visible(): + """``add_environment_tokens`` is for tokens the ENGINE never sees as + input (e.g. incremental-prompt adapters where the engine carries its + own state). The tokens MUST be recorded in the trajectory the trainer + consumes, but they MUST NOT extend the engine-visible accumulated + sequence — otherwise ``add_call``'s strict-prefix invariant on the + very next engine call would reject any input that doesn't start with + those env-only tokens.""" + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1], + output_tokens=[10], + output_logprobs=[-0.1], + ), + ) + asm.add_environment_tokens([42, 43], role="tool") + + # The engine-visible prefix is unchanged by add_environment_tokens. + assert asm.accumulated_tokens == [1, 10] + + # The next engine call's input_tokens starts at the engine-visible + # prefix [1, 10] — NOT [1, 10, 42, 43]. This is the documented + # contract for env-only tokens. + asm.add_call( + InferenceCall( + input_tokens=[1, 10], + output_tokens=[20], + output_logprobs=[-0.2], + ), + ) + + # Trajectory the trainer consumes: includes the tool turn between + # the two assistant turns. + roles = [t.role for t in asm.to_payload(total_reward=0.0).turns] + assert roles == ["user", "assistant", "tool", "assistant"] + + +def test_add_environment_tokens_does_not_break_next_add_call_prefix(): + """Regression: previously ``add_environment_tokens`` extended the + engine-visible ``_seq``, so the next ``add_call`` REQUIRED the env + tokens to appear in its ``input_tokens`` — which contradicts the + documented purpose of the helper. An incremental-engine adapter + that records a tool reply as env-only and then issues the next + engine call (with a fresh, narrower input prefix) used to hit + ``PrefixMismatch`` immediately.""" + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1], + output_tokens=[10], + output_logprobs=[-0.1], + ), + ) + asm.add_environment_tokens([42, 43], role="tool") + # The next engine call still passes only the engine-visible prefix + # plus its own new tokens. Must not raise PrefixMismatch. + asm.add_call( + InferenceCall( + input_tokens=[1, 10], # env tokens [42, 43] are NOT here + output_tokens=[20], + output_logprobs=[-0.2], + ), + ) + + +# --------------------------------------------------------------------------- +# Error paths -- the whole point of the helper +# --------------------------------------------------------------------------- + + +def test_prefix_mismatch_raises_with_index(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3], + output_tokens=[10, 11], + output_logprobs=[-0.1, -0.2], + ), + ) + # Caller "re-tokenized" and produced a different token at position 2. + with pytest.raises(PrefixMismatch) as exc_info: + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 999, 10, 11, 50], + output_tokens=[20], + output_logprobs=[-0.3], + ), + ) + msg = str(exc_info.value) + assert "index 2" in msg + assert "prior_token=3" in msg + assert "input_token=999" in msg + assert "re-tokenized" in msg + + +def test_prefix_too_short_raises(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3, 4, 5], + output_tokens=[10], + output_logprobs=[-0.1], + ), + ) + # Next input is shorter than what the engine has already seen. + with pytest.raises(PrefixMismatch): + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3], + output_tokens=[20], + output_logprobs=[-0.2], + ), + ) + + +def test_logprob_length_mismatch_raises(): + asm = TrajectoryAssembler() + with pytest.raises(ValueError, match="output_logprobs length"): + asm.add_call( + InferenceCall( + input_tokens=[1], + output_tokens=[10, 11, 12], + output_logprobs=[-0.1, -0.2], + ), + ) + + +def test_empty_to_payload_raises(): + asm = TrajectoryAssembler() + with pytest.raises(RuntimeError, match="empty"): + asm.to_payload(total_reward=0.0) + + +def test_no_assistant_to_payload_raises(): + asm = TrajectoryAssembler() + asm.add_environment_tokens([1, 2, 3], role="user") + with pytest.raises(RuntimeError, match="no assistant"): + asm.to_payload(total_reward=0.0) + + +# --------------------------------------------------------------------------- +# Round-trip through the existing packer +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_assembler_payload_round_trips_through_packer(): + asm = TrajectoryAssembler(tokenizer_id="test-tok") + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3], + output_tokens=[10, 11], + output_logprobs=[-0.1, -0.2], + ), + ) + asm.add_call( + InferenceCall( + input_tokens=[1, 2, 3, 10, 11, 50], + output_tokens=[20, 21], + output_logprobs=[-0.3, -0.4], + ), + role_before="user", + ) + payload = asm.to_payload(total_reward=1.5) + + sample = await tr.pack_payload_to_sample( + payload, + tokenizer_id="test-tok", + ) + # Concatenated tokens / mask / logprobs come straight from the assembler. + assert sample.tokens == [1, 2, 3, 10, 11, 50, 20, 21] + assert sample.loss_mask == [0, 0, 0, 1, 1, 0, 1, 1] + assert sample.logprobs == [0.0, 0.0, 0.0, -0.1, -0.2, 0.0, -0.3, -0.4] + assert sample.reward == pytest.approx(1.5) diff --git a/training/tests/unit/test_rollout_helpers.py b/training/tests/unit/test_rollout_helpers.py new file mode 100644 index 00000000..94f1a918 --- /dev/null +++ b/training/tests/unit/test_rollout_helpers.py @@ -0,0 +1,120 @@ +"""Unit tests for ``training.utils.rl.rollout.extract_completion``. + +Locks in the token / logprob alignment contract: + + * ``token_ids`` is required; missing or empty raises. + * ``token_ids`` and ``token_logprobs`` are filtered IN LOCKSTEP + when the provider emits a null placeholder, so a leading or + middle ``None`` token doesn't shift every remaining logprob + onto the wrong token (which would silently corrupt PPO/GRPO + ratio + KL for the affected completion). +""" + +from __future__ import annotations + +import pytest + +from training.utils.rl.rollout import extract_completion + + +def _choice(token_ids, token_logprobs=None, finish_reason="stop"): + return { + "token_ids": token_ids, + "logprobs": ( + {"token_logprobs": token_logprobs} + if token_logprobs is not None + else None + ), + "finish_reason": finish_reason, + } + + +def test_basic_extract(): + call = extract_completion( + _choice([10, 11, 12], [-0.1, -0.2, -0.3]), + input_tokens=[1, 2, 3], + ) + assert call.input_tokens == [1, 2, 3] + assert call.output_tokens == [10, 11, 12] + assert call.output_logprobs == [-0.1, -0.2, -0.3] + + +def test_missing_token_ids_raises(): + with pytest.raises(ValueError, match="missing 'token_ids'"): + extract_completion(_choice(None), input_tokens=[1]) + + +def test_empty_token_ids_raises(): + with pytest.raises(ValueError, match="missing 'token_ids'"): + extract_completion(_choice([]), input_tokens=[1]) + + +def test_null_token_at_start_filters_logprobs_in_lockstep(): + """A provider that emits ``None`` at index 0 of ``token_ids`` used + to drop the entry from tokens but keep ALL logprobs, then + tail-trim the resulting length mismatch. Every remaining + logprob shifted onto the wrong token, silently corrupting + PPO/GRPO ratio + KL for that completion. Filtering both lists + in lockstep keeps surviving (token, logprob) pairs aligned.""" + call = extract_completion( + _choice( + token_ids=[None, 11, 12, 13], + token_logprobs=[-9.0, -0.1, -0.2, -0.3], + ), + input_tokens=[1], + ) + assert call.output_tokens == [11, 12, 13] + # The logprob paired with the null token (-9.0) is dropped, NOT + # the trailing logprob (-0.3) — which would have shifted every + # remaining logprob one slot to the left. + assert call.output_logprobs == [-0.1, -0.2, -0.3] + + +def test_null_token_in_middle_filters_logprobs_in_lockstep(): + call = extract_completion( + _choice( + token_ids=[10, None, 12, 13], + token_logprobs=[-0.1, -9.0, -0.2, -0.3], + ), + input_tokens=[1], + ) + assert call.output_tokens == [10, 12, 13] + assert call.output_logprobs == [-0.1, -0.2, -0.3] + + +def test_no_null_tokens_unchanged(): + call = extract_completion( + _choice( + token_ids=[10, 11, 12], + token_logprobs=[-0.1, -0.2, -0.3], + ), + input_tokens=[1], + ) + assert call.output_tokens == [10, 11, 12] + assert call.output_logprobs == [-0.1, -0.2, -0.3] + + +def test_logprob_list_padded_by_one_truncates_to_token_count(): + """Some providers emit ``len(logprobs) == len(tokens) + 1``; + the helper trims to ``len(tokens)`` (back-compat with the + pre-lockstep behavior for non-null token cases).""" + call = extract_completion( + _choice( + token_ids=[10, 11, 12], + token_logprobs=[-0.1, -0.2, -0.3, -0.4], # 1 too many + ), + input_tokens=[1], + ) + assert call.output_tokens == [10, 11, 12] + assert call.output_logprobs == [-0.1, -0.2, -0.3] + + +def test_logprob_list_too_short_raises(): + with pytest.raises(ValueError, match="aligned with token_ids"): + extract_completion( + _choice( + token_ids=[10, 11, 12], + token_logprobs=[-0.1], # 2 too few + ), + input_tokens=[1], + ) diff --git a/training/tests/unit/test_rollout_message.py b/training/tests/unit/test_rollout_message.py new file mode 100644 index 00000000..7fdec9ba --- /dev/null +++ b/training/tests/unit/test_rollout_message.py @@ -0,0 +1,157 @@ +import pytest + +from training.utils.rl.rollout import ( + MessageTrajectoryAssembler, + MessageValidationError, + TITOTokenizer, +) + + +class FakeTokenizer: + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, tools=None, **kwargs): + rendered = "".join(f"<{msg['role']}>{msg.get('content') or ''}\n" for msg in messages) + if add_generation_prompt: + rendered += "" + if tokenize: + return self.encode(rendered, add_special_tokens=False) + return rendered + + def encode(self, text, add_special_tokens=False): + return [ord(ch) for ch in text] + + +def test_message_in_preserves_generated_assistant_tokens_across_tool_turn(): + assembler = MessageTrajectoryAssembler(TITOTokenizer(FakeTokenizer())) + first_messages = [{"role": "user", "content": "hi"}] + + first_prompt = assembler.prepare_next_input(first_messages) + generated_tokens = [9001, 9002] + assembler.add_assistant_response( + request_messages=first_messages, + assistant_message={"role": "assistant", "content": "retokenized text would differ"}, + prompt_token_ids=first_prompt, + completion_token_ids=generated_tokens, + completion_logprobs=[-0.1, -0.2], + ) + + second_messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "retokenized text would differ"}, + {"role": "tool", "content": "42", "tool_call_id": "call_1"}, + ] + second_prompt = assembler.prepare_next_input(second_messages) + + assert second_prompt[: len(first_prompt) + len(generated_tokens)] == first_prompt + generated_tokens + + +def test_message_in_user_and_system_appends_are_allowed_by_default(): + assembler = MessageTrajectoryAssembler(TITOTokenizer(FakeTokenizer())) + first_messages = [{"role": "user", "content": "hi"}] + first_prompt = assembler.prepare_next_input(first_messages) + assembler.add_assistant_response( + request_messages=first_messages, + assistant_message={"role": "assistant", "content": "hello"}, + prompt_token_ids=first_prompt, + completion_token_ids=[101], + completion_logprobs=[-0.1], + ) + + next_prompt = assembler.prepare_next_input( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "system", "content": "retry carefully"}, + {"role": "user", "content": "again"}, + ] + ) + + assert next_prompt[: len(first_prompt) + 1] == first_prompt + [101] + + +def test_message_in_rejects_modified_prior_message(): + assembler = MessageTrajectoryAssembler(TITOTokenizer(FakeTokenizer())) + first_messages = [{"role": "user", "content": "hi"}] + first_prompt = assembler.prepare_next_input(first_messages) + assembler.add_assistant_response( + request_messages=first_messages, + assistant_message={"role": "assistant", "content": "hello"}, + prompt_token_ids=first_prompt, + completion_token_ids=[101], + completion_logprobs=[-0.1], + ) + + with pytest.raises(MessageValidationError): + assembler.prepare_next_input([{"role": "user", "content": "changed"}]) + + +def test_message_in_one_step_rollback_to_assistant_checkpoint(): + assembler = MessageTrajectoryAssembler(TITOTokenizer(FakeTokenizer())) + + m1 = [{"role": "user", "content": "hi"}] + p1 = assembler.prepare_next_input(m1) + assembler.add_assistant_response( + request_messages=m1, + assistant_message={"role": "assistant", "content": "call tool"}, + prompt_token_ids=p1, + completion_token_ids=[101], + completion_logprobs=[-0.1], + ) + m2 = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "call tool"}, + {"role": "tool", "content": "first", "tool_call_id": "call_1"}, + ] + p2 = assembler.prepare_next_input(m2) + assembler.add_assistant_response( + request_messages=m2, + assistant_message={"role": "assistant", "content": "final"}, + prompt_token_ids=p2, + completion_token_ids=[202], + completion_logprobs=[-0.2], + ) + + retry_prompt = assembler.prepare_next_input( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "call tool"}, + {"role": "tool", "content": "retry", "tool_call_id": "call_1"}, + ] + ) + + assert assembler.token_ids == p1 + [101] + assert retry_prompt[: len(p1) + 1] == p1 + [101] + + +def test_message_in_multi_step_rollback_raises(): + assembler = MessageTrajectoryAssembler(TITOTokenizer(FakeTokenizer()), max_assistant_rollback_steps=1) + + m1 = [{"role": "user", "content": "hi"}] + p1 = assembler.prepare_next_input(m1) + assembler.add_assistant_response( + request_messages=m1, + assistant_message={"role": "assistant", "content": "a1"}, + prompt_token_ids=p1, + completion_token_ids=[101], + completion_logprobs=[-0.1], + ) + m2 = m1 + [{"role": "assistant", "content": "a1"}, {"role": "tool", "content": "t1"}] + p2 = assembler.prepare_next_input(m2) + assembler.add_assistant_response( + request_messages=m2, + assistant_message={"role": "assistant", "content": "a2"}, + prompt_token_ids=p2, + completion_token_ids=[202], + completion_logprobs=[-0.2], + ) + m3 = m2 + [{"role": "assistant", "content": "a2"}, {"role": "tool", "content": "t2"}] + p3 = assembler.prepare_next_input(m3) + assembler.add_assistant_response( + request_messages=m3, + assistant_message={"role": "assistant", "content": "a3"}, + prompt_token_ids=p3, + completion_token_ids=[303], + completion_logprobs=[-0.3], + ) + + with pytest.raises(MessageValidationError, match="exceeds"): + assembler.prepare_next_input(m1 + [{"role": "assistant", "content": "a1"}, {"role": "tool", "content": "retry"}]) diff --git a/training/tests/unit/test_rollout_trace.py b/training/tests/unit/test_rollout_trace.py new file mode 100644 index 00000000..9fad783c --- /dev/null +++ b/training/tests/unit/test_rollout_trace.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from training.utils.rl.rollout.trace import ( + analyze_flat_sample, + analyze_token_turn_traces, + analyze_turns, +) +from training.utils.rl.rollout.assembler import InferenceCall, TrajectoryAssembler + + +class _Tokenizer: + def decode(self, ids, skip_special_tokens: bool = False) -> str: + del skip_special_tokens + return "".join(f"<{int(i)}>" for i in ids) + + +def test_token_turn_trace_is_native_trajectory_not_artifact(): + trajectory = analyze_token_turn_traces( + [ + {"prompt_ids": [1, 2], "completion_ids": [3], "completion_logprobs": [-0.1]}, + {"prompt_ids": [1, 2, 3, 4], "completion_ids": [5], "completion_logprobs": [-0.2]}, + ], + tokenizer=_Tokenizer(), + ) + + data = trajectory.to_dict() + assert data["kind"] == "rollout_trajectory" + assert data["summary"]["prefix_mismatch_count"] == 0 + assert [row["token_id"] for row in data["tokens"]] == [1, 2, 3, 4, 5] + assert [row["chunk_source"] for row in data["tokens"]] == [ + "prompt_delta", + "prompt_delta", + "completion", + "prompt_delta", + "completion", + ] + + +def test_token_turn_trace_surfaces_prefix_mismatch_on_trajectory_tokens(): + trajectory = analyze_token_turn_traces( + [ + {"prompt_ids": [1], "completion_ids": [2]}, + {"prompt_ids": [9], "completion_ids": [10]}, + ], + tokenizer=_Tokenizer(), + ) + + data = trajectory.to_dict() + assert data["summary"]["prefix_mismatch_count"] == 1 + assert data["issues"][0]["code"] == "prefix_mismatch" + diverged = [row for row in data["tokens"] if "prefix_mismatch" in row["issues"]] + assert [row["token_id"] for row in diverged] == [9] + + +def test_turn_trace_keeps_roles_and_logprob_issue(): + trajectory = analyze_turns( + [ + {"role": "user", "token_ids": [1, 2]}, + {"role": "assistant", "token_ids": [3], "logprobs": []}, + ], + tokenizer=_Tokenizer(), + ) + + data = trajectory.to_dict() + assert data["summary"]["logprob_mismatch_count"] == 1 + assert [row["role"] for row in data["tokens"]] == ["user", "user", "assistant"] + assert [row["renderer_claim_weight"] for row in data["tokens"]] == [0.0, 0.0, 1.0] + + +def test_flat_sample_trace_uses_loss_mask_as_training_boundary(): + trajectory = analyze_flat_sample( + {"tokens": [1, 2, 3, 4], "loss_mask": [0, 1, 0, 1]}, + tokenizer=_Tokenizer(), + ) + + data = trajectory.to_dict() + assert data["summary"]["generated_token_count"] == 2 + assert [row["role"] for row in data["tokens"]] == [ + "context", + "assistant", + "context", + "assistant", + ] + + +def test_trajectory_assembler_emits_native_trajectory(): + asm = TrajectoryAssembler() + asm.add_call( + InferenceCall( + input_tokens=[1, 2], + output_tokens=[3], + output_logprobs=[-0.1], + ) + ) + + data = asm.to_trajectory(tokenizer=_Tokenizer()).to_dict() + assert data["kind"] == "rollout_trajectory" + assert [row["token_id"] for row in data["tokens"]] == [1, 2, 3] + assert [row["role"] for row in data["tokens"]] == ["user", "user", "assistant"] diff --git a/training/tests/unit/test_rollout_types.py b/training/tests/unit/test_rollout_types.py new file mode 100644 index 00000000..a5dcc90c --- /dev/null +++ b/training/tests/unit/test_rollout_types.py @@ -0,0 +1,278 @@ +"""Unit tests for training.utils.rl.rollout.""" + +from __future__ import annotations + +import pytest + +from training.utils.rl.rollout import Rollout, RolloutSample, rollout_to_prompt_group + + +def _sample( + *, + tokens=(1, 2, 3, 4, 5), + loss_mask=(0, 0, 1, 1, 1), + logprobs=None, + reward=0.0, + finish_reason="stop", + text="x", +): + if logprobs is None: + logprobs = [0.0 if m == 0 else -0.1 for m in loss_mask] + return RolloutSample( + tokens=list(tokens), + logprobs=list(logprobs), + loss_mask=list(loss_mask), + reward=float(reward), + finish_reason=finish_reason, + text=text, + ) + + +class TestValidation: + def test_rejects_empty_samples(self): + with pytest.raises(ValueError, match="empty"): + rollout_to_prompt_group(Rollout(samples=[])) + + def test_rejects_length_mismatch(self): + s = RolloutSample( + tokens=[1, 2, 3], + logprobs=[0.0, -0.1], # wrong length + loss_mask=[0, 1, 1], + reward=1.0, + ) + with pytest.raises(ValueError, match="length mismatch"): + rollout_to_prompt_group(Rollout(samples=[s])) + + def test_rejects_tokens_too_short(self): + s = RolloutSample( + tokens=[1], + logprobs=[0.0], + loss_mask=[1], + reward=0.0, + ) + with pytest.raises(ValueError, match="length >= 2"): + rollout_to_prompt_group(Rollout(samples=[s])) + + def test_rejects_all_zero_loss_mask(self): + s = _sample(tokens=[1, 2, 3], loss_mask=[0, 0, 0], logprobs=[0.0, 0.0, 0.0]) + with pytest.raises(ValueError, match="all zeros"): + rollout_to_prompt_group(Rollout(samples=[s])) + + +class TestSingleTurnPacking: + def test_two_samples_build_two_datums(self): + """Basic packing structure check. Uses 2 samples since + ``rollout_to_prompt_group`` now drops singleton groups (the + default ``compute_advantages`` would produce NaN advantages + on length-1 inputs and poison the training step).""" + r = Rollout(samples=[ + _sample(reward=1.0), + _sample(reward=0.0), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert len(pg.data) == 2 + assert pg.rewards == [1.0, 0.0] + # target lengths are tokens[1:] so 4 entries for tokens of length 5 + td = pg.data[0].loss_fn_inputs["target_tokens"] + assert td.shape == [4] + mask_td = pg.data[0].loss_fn_inputs["weights"] + assert mask_td.shape == [4] + + def test_singleton_rollout_with_default_advantages_dropped(self): + """The default GRPO-style ``compute_advantages`` z-score-normalizes + by ``torch.std(rewards)``, which is NaN on a length-1 tensor. + Without a guard the NaN advantage would flow through to the + loss kernel and poison the entire training step. The helper + validates ``advantage_fn`` output and drops the group if any + advantage is non-finite — preserves the NaN-poison protection + without pre-rejecting all N=1 groups (REINFORCE-style runs + below). + """ + r = Rollout(samples=[_sample(reward=1.0)]) + pg = rollout_to_prompt_group(r) + assert pg is None + + def test_singleton_rollout_with_custom_advantage_fn_emitted(self): + """REINFORCE-style async RL is a legitimate single-sample + objective (``completions_per_prompt=1``); a previous coarse + ``len(samples) < 2`` precheck silently dropped every such + rollout and made the recipe make no training progress despite + advertising REINFORCE support. When the caller supplies a + custom ``advantage_fn`` (e.g. ``lambda r: r``) that returns + finite values on N=1, the rollout MUST be packed into a + PromptGroup like any other. + """ + r = Rollout(samples=[_sample(reward=0.7)]) + pg = rollout_to_prompt_group(r, advantage_fn=lambda rewards: list(rewards)) + assert pg is not None + assert pg.rewards == [0.7] + # The custom advantage_fn passes the reward through unchanged. + assert pg.advantages == [0.7] + # Single-sample group still emits a single Datum at the + # same packing shape as multi-sample groups. + assert len(pg.data) == 1 + + def test_non_finite_advantage_drops_even_with_custom_fn(self): + """The post-compute non-finite check must catch NaN/inf + advantages from ANY ``advantage_fn``, not just the default — + the protection is on the output, not the sample count. A + custom advantage_fn that happens to return NaN on a + degenerate input still produces a non-finite advantage that + would poison the loss; drop the group.""" + r = Rollout(samples=[_sample(reward=1.0), _sample(reward=2.0)]) + pg = rollout_to_prompt_group( + r, advantage_fn=lambda rewards: [float("nan")] * len(rewards), + ) + assert pg is None + + def test_n_samples_center_advantages(self): + r = Rollout(samples=[ + _sample(reward=1.0), + _sample(reward=0.0), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.rewards == [1.0, 0.0] + assert sum(pg.advantages) == pytest.approx(0.0, abs=1e-5) + + def test_prompt_len_derived_from_first_mask_one(self): + """prompt_len field reflects where assistant tokens start in sample 0.""" + r = Rollout(samples=[ + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 0, 1, 1]), + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 0, 1, 1]), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.prompt_len == 3 + # Per-sample list mirrors prompt_len when all samples agree. + assert pg.prompt_lens == [3, 3] + + def test_per_sample_prompt_lens_for_heterogeneous_rollout(self): + """Multi-turn / tool branches produce samples with different + prompt prefix lengths in the same rollout group. ``prompt_lens`` + must record each sample's first-assistant-token index so the + downstream loss slices each sample at its own boundary — + replicating a single ``prompt_len`` would drop tokens for + short-prefix samples and leak prompt tokens for long-prefix ones. + """ + r = Rollout(samples=[ + # Sample A: 2-token prefix, then 3 assistant tokens. + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 1, 1, 1]), + # Sample B (different branch): 4-token prefix, 1 assistant token. + _sample(tokens=[1, 2, 9, 9, 7], loss_mask=[0, 0, 0, 0, 1]), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.prompt_lens == [2, 4] + # Back-compat scalar reflects sample 0 only (callers that already + # consume ``prompt_lens`` see the per-sample truth). + assert pg.prompt_len == 2 + + def test_combine_prompt_groups_uses_per_sample_prompt_lens(self): + """``combine_prompt_groups`` must prefer per-sample ``prompt_lens`` + over the scalar ``prompt_len`` so heterogeneous rollouts flow + through to the loss with correct boundaries. + """ + from training.utils.rl.losses import combine_prompt_groups + + r = Rollout(samples=[ + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 1, 1, 1]), + _sample(tokens=[1, 2, 9, 9, 7], loss_mask=[0, 0, 0, 0, 1]), + ]) + pg = rollout_to_prompt_group(r) + _, _, _, prompt_lens, _ = combine_prompt_groups([pg]) + assert prompt_lens == [2, 4], ( + "combine_prompt_groups must extend with per-sample prompt_lens " + f"when set; got {prompt_lens}" + ) + + def test_with_reference_mirrors_policy_datums(self): + r = Rollout(samples=[_sample(), _sample()]) + pg = rollout_to_prompt_group(r, with_reference=True) + assert pg is not None + assert len(pg.ref_data) == 2 + + def test_completion_lens_counts_mask_ones(self): + r = Rollout(samples=[ + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 1, 1, 1]), # 3 mask ones + _sample(tokens=[1, 2, 3, 4, 5], loss_mask=[0, 0, 1, 1, 1]), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.completion_lens == [3, 3] + + def test_truncated_from_finish_reason(self): + r = Rollout(samples=[ + _sample(finish_reason="length"), + _sample(finish_reason="stop"), + ]) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.truncated == [True, False] + + +class TestMultiTurnPacking: + def test_interleaved_mask_preserved_in_datum(self): + """Tool-use shape: assistant turn, then tool result (mask 0), then + assistant turn again. The loss_mask must be preserved verbatim + (shifted by one for the target-alignment).""" + # tokens: [prompt, prompt, asst1, asst1, tool, tool, asst2, asst2] + # mask: [0, 0, 1, 1, 0, 0, 1, 1 ] + def _interleaved_sample(reward: float): + return _sample( + tokens=[10, 11, 20, 21, 30, 31, 40, 41], + loss_mask=[0, 0, 1, 1, 0, 0, 1, 1], + logprobs=[0.0, 0.0, -0.1, -0.2, 0.0, 0.0, -0.3, -0.4], + reward=reward, + ) + # Two samples since singleton groups now drop (NaN-advantage guard). + pg = rollout_to_prompt_group(Rollout(samples=[ + _interleaved_sample(reward=1.0), + _interleaved_sample(reward=0.0), + ])) + assert pg is not None + mask_td = pg.data[0].loss_fn_inputs["weights"] + # Target is tokens[1:] = 7 entries; mask is loss_mask[1:] = 7 entries. + assert list(mask_td.data) == [0, 1, 1, 0, 0, 1, 1] + # Completion-len counts the original mask ones, not the shifted one. + assert pg.completion_lens == [4, 4] + + +class TestInferenceLogprobs: + def test_inf_logprobs_shifted_by_one(self): + """inf_logprobs on PromptGroup is logprobs[1:] (target-aligned).""" + def mk(reward: float): + return _sample( + tokens=[1, 2, 3, 4, 5], + loss_mask=[0, 0, 1, 1, 1], + logprobs=[0.0, 0.0, -0.1, -0.2, -0.3], + reward=reward, + ) + # Two samples — singleton groups now drop (NaN-advantage guard). + pg = rollout_to_prompt_group(Rollout(samples=[mk(1.0), mk(0.0)])) + assert pg is not None + assert pg.inf_logprobs == [ + [0.0, -0.1, -0.2, -0.3], + [0.0, -0.1, -0.2, -0.3], + ] + + +class TestRowMeta: + def test_row_meta_copied_when_present(self): + r = Rollout( + samples=[_sample(), _sample()], + row_meta={"ground_truth": "42"}, + ) + pg = rollout_to_prompt_group(r) + assert pg is not None + assert pg.row_meta == {"ground_truth": "42"} + # Defensive copy. + assert pg.row_meta is not r.row_meta + + def test_row_meta_none_stays_none(self): + pg = rollout_to_prompt_group( + Rollout(samples=[_sample(), _sample()]) + ) + assert pg is not None + assert pg.row_meta is None diff --git a/training/tests/unit/test_train_frozen_lake.py b/training/tests/unit/test_train_frozen_lake.py index 51a45c6a..9cfc6ec2 100644 --- a/training/tests/unit/test_train_frozen_lake.py +++ b/training/tests/unit/test_train_frozen_lake.py @@ -556,7 +556,7 @@ async def fake_run_rl_loop(**kwargs): "total_sampled": 1, "filter_drops": 0, "sample_fails": 0, - "sample_wait_time": 0.1, + "trainer_wait_for_sampler_time": 0.1, "step_wall_time": 0.2, }, ) diff --git a/training/utils/__init__.py b/training/utils/__init__.py index e29be5b1..45f1ec02 100644 --- a/training/utils/__init__.py +++ b/training/utils/__init__.py @@ -56,6 +56,8 @@ "WandBConfig", "compute_advantages", "compute_pass_at_k", + "CursorDataLoader", + "CursorItem", "create_trainer_job", "request_trainer_job", "wait_trainer_job", @@ -68,6 +70,7 @@ "iter_preference_examples", "load_jsonl_dataset", "load_preference_dataset", + "replicate_rows_for_epochs", "log_metrics_json", "make_render_dataloader", "make_orpo_loss_fn", @@ -116,7 +119,9 @@ find_common_prefix_length, normalize_preference_row, prepare_sampling_messages, + replicate_rows_for_epochs, ) +from training.utils.dataloader import CursorDataLoader, CursorItem from training.utils.dataloader_cursor import RawRowCursor from training.utils.infra import ( Infra, diff --git a/training/utils/data.py b/training/utils/data.py index 8b878134..eff75677 100644 --- a/training/utils/data.py +++ b/training/utils/data.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import json import logging import math @@ -15,6 +16,22 @@ logger = logging.getLogger(__name__) +def replicate_rows_for_epochs( + rows: List[Dict[str, Any]], epochs: int, +) -> List[Dict[str, Any]]: + """Duplicate ``rows`` ``epochs`` times with INDEPENDENT row dicts. + + The naive ``rows * epochs`` only multiplies list references -- every + epoch ends up sharing the same dict instances. Any rollout function + that mutates its input row in place (attaching scratch fields, + normalizing prompt data, caching renders) leaks that mutation into + every later epoch, so subsequent passes train on the already-mutated + row instead of the original dataset. ``copy.deepcopy`` per row per + epoch gives each epoch its own independent copies. + """ + return [copy.deepcopy(r) for _ in range(epochs) for r in rows] + + class RLPromptDataset: """Batch-indexed prompt dataset for RL training. diff --git a/training/utils/dataloader.py b/training/utils/dataloader.py new file mode 100644 index 00000000..935ea06e --- /dev/null +++ b/training/utils/dataloader.py @@ -0,0 +1,89 @@ +"""Cursor-based dataloader utilities.""" + +from __future__ import annotations + +import copy +import random +from dataclasses import dataclass +from typing import Generic, TypeVar + +T = TypeVar("T") + + +@dataclass(frozen=True) +class CursorItem(Generic[T]): + index: int + value: T + + +class CursorDataLoader(Generic[T]): + def __init__( + self, + items: list[T], + start_cursor: int = 0, + *, + epochs: int = 1, + shuffle: bool = False, + seed: int = 0, + ): + if start_cursor < 0: + raise ValueError("start_cursor must be >= 0") + if epochs < 0: + raise ValueError("epochs must be >= 0") + self.items = items + self.epochs = epochs + self.shuffle = shuffle + self.seed = seed + self.cursor = start_cursor + self.next_index = start_cursor + self._resolved: set[int] = set() + self._permutations: dict[int, list[int]] = {} + + def __iter__(self): + return self + + def __next__(self) -> CursorItem[T]: + if self.next_index >= self.total_items: + raise StopIteration + idx = self.next_index + self.next_index += 1 + return CursorItem(index=idx, value=copy.deepcopy(self.items[self._row_index(idx)])) + + @property + def data_consumed(self) -> int: + return self.cursor + + @property + def total_items(self) -> int: + return len(self.items) * self.epochs + + @property + def epoch_id(self) -> int: + return self.cursor // len(self.items) if self.items else 0 + + @property + def sample_offset(self) -> int: + return self.cursor % len(self.items) if self.items else 0 + + def mark_resolved(self, index: int) -> None: + if index < self.cursor: + return + if index >= self.total_items: + raise ValueError("resolved index out of range") + self._resolved.add(index) + while self.cursor in self._resolved: + self._resolved.remove(self.cursor) + self.cursor += 1 + + def _row_index(self, index: int) -> int: + if not self.items: + raise IndexError("empty dataloader") + epoch = index // len(self.items) + offset = index % len(self.items) + if not self.shuffle: + return offset + if epoch not in self._permutations: + perm = list(range(len(self.items))) + random.Random(self.seed + epoch).shuffle(perm) + self._permutations[epoch] = perm + return self._permutations[epoch][offset] diff --git a/training/utils/logging.py b/training/utils/logging.py index 746eb77e..c840045b 100644 --- a/training/utils/logging.py +++ b/training/utils/logging.py @@ -39,6 +39,7 @@ def setup_wandb(wb: WandBConfig, config: dict[str, Any]) -> bool: wandb.define_metric("rollout/*", step_metric="train/step") wandb.define_metric("batch/*", step_metric="train/step") wandb.define_metric("infra/*", step_metric="train/step") + wandb.define_metric("ctx/*", step_metric="train/step") logger.info("WandB: %s", wandb.run.url) return True except ImportError: @@ -72,13 +73,23 @@ def compute_pass_at_k( return results -def wandb_log(metrics: dict[str, Any], step: int) -> None: - """Log metrics to WandB if available.""" +def wandb_log(metrics: dict[str, Any], step: int | None = None) -> None: + """Log metrics to WandB if available. + + When ``step`` is ``None`` we let wandb auto-increment its internal + step counter; metrics with a declared ``step_metric`` (via + ``define_metric``) read their x-axis from the dict instead. Passing + an explicit ``step`` is still supported for single-axis recipes that + want to peg the global step manually. + """ try: import wandb if wandb.run is not None: - wandb.log(metrics, step=step, commit=True) + if step is None: + wandb.log(metrics, commit=True) + else: + wandb.log(metrics, step=step, commit=True) except ImportError: pass diff --git a/training/utils/rl/README.md b/training/utils/rl/README.md new file mode 100644 index 00000000..3617f54a --- /dev/null +++ b/training/utils/rl/README.md @@ -0,0 +1,38 @@ +# RL rollout primitives + +> ⚠️ **EXPERIMENTAL — the async RL recipe these primitives serve is under +> active development.** API may change without backward-compat shims. + +Token-native rollout API for the async RL recipe (`training/recipes/async_rl_loop.py`). + +For the user-facing contract (`rollout_fn(sample_prompt) -> RolloutSample`, +`RolloutSetup`, `rollout_fn_factory`, off-policy gate sizing), see +[`/skills/dev/references/rl/async-rl.md`](/skills/dev/references/rl/async-rl.md). + +## Layout + +| File | Purpose | +| --- | --- | +| `rollout/types.py` | `Rollout`, `RolloutSample`, `rollout_to_prompt_group` (trainer packing). | +| `rollout/assembler.py` | Token-native multi-turn assembly with prefix checks. | +| `rollout/message.py` | Generic message-in TITO bridge that preserves prior assistant tokens. | +| `rollout/renderer.py` | Optional renderer-backed single-turn helper. | +| `rollout/remote.py` | Optional service payload packer. | + +The correctness-critical path is `TrajectoryAssembler`: every next model +request must extend the accumulated token sequence, except for an explicit +generic boundary trim passed by the caller. + +`MessageTrajectoryAssembler` wraps it for OpenAI-style message loops: +preserves prior assistant token IDs exactly; tokenizes only appended `tool` / +`user` / `system` messages; appends the next assistant generation prompt; +rejects edits to prior messages; supports bounded rollback to an earlier +assistant checkpoint. No model-specific TITO subclasses live here — keep +tokenizer-specific boundary policy in user code. + +## Tests + +Invariants only: `test_rollout_types.py`, `test_rollout_assembler.py`, +`test_rollout_message.py`, `test_rollout_helpers.py`. No remote-service +mocks or example-specific policy tests unless they guard a real trainer or +token-alignment invariant. diff --git a/training/utils/rl/__init__.py b/training/utils/rl/__init__.py index 4a538af3..ebcdffc2 100644 --- a/training/utils/rl/__init__.py +++ b/training/utils/rl/__init__.py @@ -33,6 +33,30 @@ "expand_turn_advantages_from_spans", "make_igpo_loss_fn", "score_prefix", + # Async loop + rollout contract + "run_async_rl_loop", + "RowRequest", + "Rollout", + "RolloutSample", + "GroupAssembler", + "rollout_to_prompt_group", + # Service-agnostic rollout adapter + "RolloutPayload", + "RolloutService", + "TurnRecord", + "make_remote_rollout_fn", + # Multi-turn assembly + "InferenceCall", + "MessageTrajectoryAssembler", + "MessageTrajectoryError", + "MessageValidationError", + "PrefixMismatch", + "TrajectoryAssembler", + "TITOTokenizer", + "extract_completion", + "get_tito_tokenizer", + "precompute_chat_suffix", + "TokenizationError", ] from training.utils.rl.pp import PPBatchRecommendation, compute_pp_recommendation @@ -63,3 +87,25 @@ make_igpo_loss_fn, score_prefix, ) +from training.utils.rl.async_train import RowRequest, run_async_rl_loop +from training.utils.rl.rollout import ( + GroupAssembler, + InferenceCall, + MessageTrajectoryAssembler, + MessageTrajectoryError, + MessageValidationError, + PrefixMismatch, + Rollout, + RolloutPayload, + RolloutSample, + RolloutService, + TITOTokenizer, + TrajectoryAssembler, + TurnRecord, + TokenizationError, + extract_completion, + get_tito_tokenizer, + make_remote_rollout_fn, + precompute_chat_suffix, + rollout_to_prompt_group, +) diff --git a/training/utils/rl/async_train.py b/training/utils/rl/async_train.py new file mode 100644 index 00000000..9e523b59 --- /dev/null +++ b/training/utils/rl/async_train.py @@ -0,0 +1,535 @@ +"""Per-sample async RL training loop. + +Acknowledgements -- prior art referenced while designing this loop: + +* AReaL (https://github.com/inclusionAI/AReaL) +* slime (https://github.com/THUDM/slime) +* Miles (https://github.com/radixark/miles) + +The user supplies ``rollout_fn(sample_prompt) -> RolloutSample | None`` -- +one trajectory per call. ``sample_prompt`` is a dataset row's dict +re-named once it crosses into the sampling layer. The loop fans each +dataset row out to ``completions_per_prompt`` parallel sample calls, +joins them by row id via :class:`GroupAssembler`, applies the optional +dynamic filter on the assembled :class:`PromptGroup`, and feeds the +trainer when a batch fills. + +Submission is row-atomic; the gate accounts in samples, with one row +consuming ``completions_per_prompt`` sample slots against the staleness +budget. Because rows are submitted whole, every sample in a row carries +the same submit version, so the group's accountable version is unambiguous. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable, Hashable, Iterable, Iterator, List + +from training.utils.data import compute_advantages +from training.utils.rl.losses import PromptGroup +from training.utils.rl.rollout.group_assembler import ( + AdvantageFn, + GroupAssembler, + RowResolution, +) +from training.utils.rl.rollout.types import RolloutSample +from training.utils.rl.train import DynamicFilterFn, TrainStepFns + +logger = logging.getLogger(__name__) + +__all__ = ["RowRequest", "run_async_rl_loop"] + + +SampleFactory = Callable[[int], Awaitable[RolloutSample | None]] + + +@dataclass +class RowRequest: + """One dataset row's fan-out factory. + + ``sample_factory(sub_index)`` is called by the loop for each + ``sub_index in [0, completions_per_prompt)`` and must return a fresh + coroutine resolving to a :class:`RolloutSample` or ``None``. + ``None`` counts as one lost sample within the row's group (the row + can still produce a valid PromptGroup if at least + ``GroupAssembler.min_group_size`` samples land). + + ``on_resolved`` fires exactly once per row, when the loop decides the + row's fate: ``"accepted"``, ``"filter"`` (assembled group rejected + by ``dynamic_filter_fn``), or ``"none"`` (no surviving samples / + advantage_fn produced non-finite values). + """ + + row_id: Hashable + sample_factory: SampleFactory + row_meta: dict | None = None + on_resolved: Callable[[str], None] | None = None + + +@dataclass +class _StalenessController: + """Sample-level (LLM-call) capacity gate. + + Bookkeeping is in samples to match ``deployment.max_batch_size``. + One prompt submission consumes ``completions_per_prompt`` samples; + a prompt resolution releases the same amount (re-credited to + ``accepted_samples`` on accept, freed on reject). + """ + + batch_size_samples: int # = prompt_groups_per_step * completions_per_prompt + completions_per_prompt: int # samples per prompt + max_staleness: int # versions + max_concurrent_samples: int | None + version: int = 0 + accepted_samples: int = 0 + running_samples: int = 0 + sample_fails: int = 0 + filter_drops: int = 0 + rejected_count: int = 0 + # Wall-clock blocked on the staleness budget; ~0 in healthy async runs. + sampler_wait_for_trainer_total: float = 0.0 + _sampler_wait_for_trainer_start: float | None = field(default=None, repr=False) + + def staleness_capacity(self) -> int: + return ( + (self.max_staleness + self.version + 1) * self.batch_size_samples + - (self.accepted_samples + self.running_samples) + ) + + def concurrency_capacity(self) -> int | None: + if self.max_concurrent_samples is None: + return None + return self.max_concurrent_samples - self.running_samples + + def capacity(self) -> int: + """Binding admit budget (min of staleness and concurrency), in samples.""" + s = self.staleness_capacity() + c = self.concurrency_capacity() + cap = s if c is None else min(c, s) + return max(0, cap) + + def is_staleness_bound(self) -> bool: + """True iff one more prompt is blocked by staleness, not concurrency. + + Returning True when both budgets are zero would wrongly attribute + deployment-saturation wall time as ``sampler_wait_for_trainer``. + """ + cpp = self.completions_per_prompt + if self.staleness_capacity() >= cpp: + return False + if self.max_concurrent_samples is None: + return True + return self.concurrency_capacity() >= cpp + + def mark_sampler_wait_for_trainer_start(self, now: float) -> None: + if self._sampler_wait_for_trainer_start is None: + self._sampler_wait_for_trainer_start = now + + def mark_sampler_wait_for_trainer_end(self, now: float) -> None: + if self._sampler_wait_for_trainer_start is not None: + self.sampler_wait_for_trainer_total += now - self._sampler_wait_for_trainer_start + self._sampler_wait_for_trainer_start = None + + def submit(self) -> None: + self.running_samples += self.completions_per_prompt + + def accept(self) -> None: + self.running_samples -= self.completions_per_prompt + self.accepted_samples += self.completions_per_prompt + + def reject(self, reason: str) -> None: + self.running_samples -= self.completions_per_prompt + self.rejected_count += 1 + if reason == "none": + self.sample_fails += 1 + elif reason == "filter": + self.filter_drops += 1 + + def advance_version(self) -> None: + self.version += 1 + + def resolved_count(self, offset: int) -> int: + # Resolved prompts since start = accepted (in prompts) + rejected. + return offset + (self.accepted_samples // self.completions_per_prompt) + self.rejected_count + + +@dataclass +class _RowState: + request: RowRequest + submit_version: int + sample_tasks: List[asyncio.Task] = field(default_factory=list) + + +async def run_async_rl_loop( + rows: Iterable[RowRequest], + *, + train_fns: TrainStepFns, + completions_per_prompt: int, + prompt_groups_per_step: int, + max_head_offpolicy_versions: int, + advantage_fn: AdvantageFn = compute_advantages, + with_reference: bool = False, + router_replay_completion_only: bool = False, + min_group_size: int = 1, + weight_sync_fn: Callable[[int], None] | None = None, + weight_sync_interval: int = 1, + max_concurrent: int | None = None, + dynamic_filter_fn: DynamicFilterFn | None = None, + global_step: int = 0, + resolved_rows_offset: int = 0, + resolved_rows_fn: Callable[[], int] | None = None, + return_final_stats: bool = False, + synchronous_training: bool = False, +) -> int | tuple[int, dict[str, Any]]: + """Run the per-sample async RL loop. + + Args: + rows: Iterable of :class:`RowRequest`, one per dataset row. The + loop calls each row's ``sample_factory`` ``completions_per_prompt`` + times to materialize the per-sample coroutines. + train_fns: Training callbacks (see :class:`TrainStepFns`). + completions_per_prompt: Number of samples drawn per row. + prompt_groups_per_step: Number of accepted rows that form one + optimizer step. + max_head_offpolicy_versions: Off-policy budget in weight-sync + (policy) versions; the gate's version increments once per + ``weight_sync_fn`` call, not per optimizer step. ``0`` is + strict on-policy. + advantage_fn: Group-level advantage computation. Default is + GRPO z-score normalization. Output is validated for + non-finite values; a row producing NaN/inf advantages is + dropped. + with_reference: Pack reference-model datums alongside the + policy datums (needed for KL). + router_replay_completion_only: When a sample carries + ``routing_matrices``, zero out prompt-position routing so + only completion-token routing is replayed. + min_group_size: Minimum surviving samples for a row to emit a + PromptGroup. Rows with fewer surviving samples are dropped. + weight_sync_fn: Called after every ``weight_sync_interval`` + optimizer steps. Must bump the deployment version; the + loop increments its internal version counter on return. + weight_sync_interval: Fire ``weight_sync_fn`` every N steps. + max_concurrent: Hard cap on **samples (LLM calls)** in flight -- + this is the same unit the deployment's ``max_batch_size`` + gates on, so the recipe-side cap and the deployment cap are + directly comparable. Must be ``>= completions_per_prompt`` + (one row's worth) or the gate deadlocks. ``None`` lets the + staleness budget alone bound concurrency. + dynamic_filter_fn: Post-assembly filter on :class:`PromptGroup`. + ``False`` drops the row from the trainer buffer (without + charging the gate beyond the row slot already consumed). + global_step: Initial optimizer-step counter. + resolved_rows_offset: Resume cursor offset. + resolved_rows_fn: Optional durable cursor provider. + return_final_stats: When ``True``, return ``(global_step, stats)``. + + Returns: + Final ``global_step``, optionally with a stats dict. + """ + if completions_per_prompt < 1: + raise ValueError("completions_per_prompt must be >= 1") + if prompt_groups_per_step < 1: + raise ValueError("prompt_groups_per_step must be >= 1") + if max_head_offpolicy_versions < 0: + raise ValueError("max_head_offpolicy_versions must be >= 0") + if weight_sync_interval < 1: + raise ValueError("weight_sync_interval must be >= 1") + if max_concurrent is not None and max_concurrent < completions_per_prompt: + # Sample-level cap: must fit at least one row's worth (cpp samples). + # Anything smaller would deadlock the gate. + raise ValueError( + f"max_concurrent (samples) = {max_concurrent} must be " + f">= completions_per_prompt ({completions_per_prompt}); " + "smaller values would cap row concurrency at <1 (deadlock)." + ) + if min_group_size < 1: + raise ValueError("min_group_size must be >= 1") + if min_group_size > completions_per_prompt: + raise ValueError( + f"min_group_size ({min_group_size}) must be " + f"<= completions_per_prompt ({completions_per_prompt})", + ) + if weight_sync_interval > max_head_offpolicy_versions + 1: + raise ValueError( + f"weight_sync_interval ({weight_sync_interval}) must be " + f"<= max_head_offpolicy_versions + 1 " + f"({max_head_offpolicy_versions + 1}); otherwise the async " + "gate stalls before the next sync because all rollouts at " + "the current version are exhausted. Either lower " + "weight_sync_interval to <= max_head_offpolicy_versions + 1 " + "(so a sync arrives before capacity hits zero), or raise " + "max_head_offpolicy_versions to >= weight_sync_interval - 1 " + "(so the head budget covers every step between syncs)." + ) + + staleness = _StalenessController( + batch_size_samples=prompt_groups_per_step * completions_per_prompt, + completions_per_prompt=completions_per_prompt, + max_staleness=max_head_offpolicy_versions, + max_concurrent_samples=max_concurrent, + ) + assembler = GroupAssembler( + completions_per_prompt=completions_per_prompt, + advantage_fn=advantage_fn, + with_reference=with_reference, + router_replay_completion_only=router_replay_completion_only, + min_group_size=min_group_size, + ) + + # Each in-flight sample task -> (row_id, sub_index). + in_flight: dict[asyncio.Task, tuple[Hashable, int]] = {} + # Row state by row_id; lifetime spans first sample submit -> row resolution. + rows_state: dict[Hashable, _RowState] = {} + buffer: list[tuple[PromptGroup, int, RowRequest]] = [] + + rows_iter: Iterator[RowRequest] = iter(rows) + iterator_exhausted = False + + def _resolved_rows(fallback: int) -> int: + if resolved_rows_fn is not None: + return resolved_rows_fn() + return resolved_rows_offset + fallback + + def _resolve(request: RowRequest, reason: str) -> None: + if request.on_resolved is not None: + request.on_resolved(reason) + + def _submit_row(request: RowRequest) -> None: + rows_state[request.row_id] = _RowState( + request=request, + submit_version=staleness.version, + ) + staleness.submit() + for sub_index in range(completions_per_prompt): + assembler.note_started( + request.row_id, + submit_version=staleness.version, + row_meta=request.row_meta if sub_index == 0 else None, + ) + coro = request.sample_factory(sub_index) + task = asyncio.ensure_future(coro) + in_flight[task] = (request.row_id, sub_index) + rows_state[request.row_id].sample_tasks.append(task) + + def _refill() -> None: + nonlocal iterator_exhausted + now = time.monotonic() + if iterator_exhausted: + # Draining, not blocked on the trainer. + staleness.mark_sampler_wait_for_trainer_end(now) + return + slots = staleness.capacity() // completions_per_prompt + if slots == 0: + if staleness.is_staleness_bound(): + staleness.mark_sampler_wait_for_trainer_start(now) + return + staleness.mark_sampler_wait_for_trainer_end(now) + for _ in range(slots): + try: + request = next(rows_iter) + except StopIteration: + iterator_exhausted = True + return + _submit_row(request) + + def _can_make_batch() -> bool: + if len(buffer) >= prompt_groups_per_step: + return True + if iterator_exhausted and not in_flight and buffer: + return True + return False + + def _has_outstanding_work() -> bool: + return bool(in_flight) or _can_make_batch() or not iterator_exhausted + + def _on_row_resolved(row_id: Hashable, resolution: RowResolution) -> None: + """Drive row-level outcome bookkeeping when the GroupAssembler + settles a row (returned ``RowResolution``).""" + state = rows_state.pop(row_id, None) + if state is None: + return + if resolution.pg is None: + staleness.reject("none") + _resolve(state.request, "none") + return + if dynamic_filter_fn is not None and not dynamic_filter_fn(resolution.pg): + staleness.reject("filter") + _resolve(state.request, "filter") + return + staleness.accept() + buffer.append((resolution.pg, resolution.min_submit_version, state.request)) + + _refill() + # Gap from here to the next ``step_start`` is ``trainer_wait_for_sampler_time``. + last_step_end = time.monotonic() + prev_sampler_wait_for_trainer_total = staleness.sampler_wait_for_trainer_total + + while _has_outstanding_work(): + if not _can_make_batch(): + if not in_flight: + if iterator_exhausted: + break + _refill() + if not in_flight: + logger.warning( + "Async loop stalled: capacity=0 with no in-flight " + "tasks and the row iterator still open. Increase " + "max_head_offpolicy_versions or supply a " + "weight_sync_fn so versions advance.", + ) + iterator_exhausted = True + break + continue + + done, _ = await asyncio.wait( + set(in_flight), return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + row_id, _sub = in_flight.pop(task) + exc = task.exception() + if exc is not None: + for other in in_flight: + other.cancel() + in_flight.clear() + raise exc + sample = task.result() + if sample is None: + resolution = assembler.note_dropped(row_id) + else: + resolution = assembler.add_sample(row_id, sample) + if resolution is not None: + _on_row_resolved(row_id, resolution) + _refill() + continue + + batch_pairs = buffer[:prompt_groups_per_step] + buffer = buffer[prompt_groups_per_step:] + batch = [pg for pg, _, _ in batch_pairs] + versions = [v for _, v, _ in batch_pairs] + offsets = [staleness.version - v for v in versions] + + step_start = time.monotonic() + for _, _, request in batch_pairs: + _resolve(request, "accepted") + resolved_rows = ( + resolved_rows_fn() + if resolved_rows_fn is not None + else staleness.resolved_count(resolved_rows_offset) + ) + + # Sync mode: drain in-flight rollouts before train_step (kills overlap) + # and explicitly open the wait window so the trainer's wall time lands + # in ``perf/sampler_wait_for_trainer_time``. + if synchronous_training and in_flight: + while in_flight: + done, _ = await asyncio.wait( + set(in_flight), return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + row_id, _sub = in_flight.pop(task) + exc = task.exception() + if exc is not None: + for other in in_flight: + other.cancel() + in_flight.clear() + raise exc + sample = task.result() + if sample is None: + resolution = assembler.note_dropped(row_id) + else: + resolution = assembler.add_sample(row_id, sample) + if resolution is not None: + _on_row_resolved(row_id, resolution) + if synchronous_training: + staleness.mark_sampler_wait_for_trainer_start(time.monotonic()) + + sampler_wait_for_trainer_time = staleness.sampler_wait_for_trainer_total - prev_sampler_wait_for_trainer_total + prev_sampler_wait_for_trainer_total = staleness.sampler_wait_for_trainer_total + trainer_wait_for_sampler_time = step_start - last_step_end + cc = staleness.concurrency_capacity() + logger.info( + "[batch-ready v=%d] in_flight=%d running=%d accepted=%d buffer=%d " + "staleness_cap=%d concurrency_cap=%s " + "wait_for_sampler=%.1fs wait_for_trainer=%.1fs " + "(staleness_bound=%s)", + staleness.version, + len(in_flight), + staleness.running_samples, + staleness.accepted_samples, + len(buffer), + staleness.staleness_capacity(), + "unbounded" if cc is None else cc, + trainer_wait_for_sampler_time, + sampler_wait_for_trainer_time, + staleness.is_staleness_bound(), + ) + extra_metrics: dict[str, Any] = { + "async/version_offset_mean": sum(offsets) / len(offsets), + "async/version_offset_max": max(offsets), + "async/version_offset_min": min(offsets), + "async/in_flight": len(in_flight), + "async/sample_fails": staleness.sample_fails, + "async/filter_drops": staleness.filter_drops, + "async/stale_drops": 0, + "all_raw_rewards": [r for pg in batch for r in pg.rewards], + "valid_prompt_groups": len(batch), + "total_sampled": staleness.accepted_samples + staleness.rejected_count * completions_per_prompt, + "filter_drops": staleness.filter_drops, + "sample_fails": staleness.sample_fails, + "stale_drops": 0, + "resolved_rows": resolved_rows, + "trainer_wait_for_sampler_time": trainer_wait_for_sampler_time, + "sampler_wait_for_trainer_time": sampler_wait_for_trainer_time, + "async/running_samples": staleness.running_samples, + "async/accepted_samples": staleness.accepted_samples, + "async/staleness_capacity_at_step": staleness.staleness_capacity(), + "async/concurrency_capacity_at_step": ( + -1 if cc is None else cc + ), + } + + global_step, _step_metrics = await asyncio.to_thread( + train_fns.train_step, global_step, batch, extra_metrics, + ) + last_step_end = time.monotonic() + + if weight_sync_fn is not None and global_step % weight_sync_interval == 0: + await asyncio.to_thread(weight_sync_fn, global_step) + staleness.advance_version() + + _refill() + + pending_tasks = list(in_flight) + for task in pending_tasks: + task.cancel() + results = await asyncio.gather(*pending_tasks, return_exceptions=True) + for task, result in zip(pending_tasks, results): + if isinstance(result, asyncio.CancelledError): + continue + if isinstance(result, BaseException): + logger.warning( + "Async loop drained an in-flight rollout sample with an " + "unhandled exception: %r", result, + ) + in_flight.clear() + rows_state.clear() + + staleness.mark_sampler_wait_for_trainer_end(time.monotonic()) + + accepted_prompts = staleness.accepted_samples // completions_per_prompt + resolved_this_run = accepted_prompts + staleness.rejected_count + final_stats = { + "sample_fails": staleness.sample_fails, + "filter_drops": staleness.filter_drops, + "stale_drops": 0, + "total_accepted": accepted_prompts, + "resolved_rows": _resolved_rows(resolved_this_run), + "sampler_wait_for_trainer_time_total": staleness.sampler_wait_for_trainer_total, + } + if return_final_stats: + return global_step, final_stats + return global_step diff --git a/training/utils/rl/common.py b/training/utils/rl/common.py index b3541d32..f276b7c4 100644 --- a/training/utils/rl/common.py +++ b/training/utils/rl/common.py @@ -29,12 +29,17 @@ def _get_loss_mask( dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: - """Extract per-position loss mask from ``loss_fn_inputs["loss_mask"]``. + """Extract per-position loss mask from ``loss_fn_inputs["weights"]``. Returns a tensor of shape ``[resp_len]`` sliced from ``response_start``. - Falls back to all-ones (no masking) when the datum has no ``loss_mask``. + Falls back to ``loss_fn_inputs["loss_mask"]`` for legacy datums and + finally to all-ones (no masking) when neither is present. + + The trainer SDK rejects any ``loss_fn_inputs`` key other than + ``{"target_tokens", "weights"}`` for ``forward_backward_custom``, so + new producers MUST write the per-token mask under ``"weights"``. """ - mask_td = datum.loss_fn_inputs.get("loss_mask") + mask_td = datum.loss_fn_inputs.get("weights") or datum.loss_fn_inputs.get("loss_mask") if mask_td is not None: mask_vals = mask_td.data[response_start : response_start + resp_len] if len(mask_vals) < resp_len: @@ -142,6 +147,7 @@ def run_loss_loop( total_kl = 0.0 total_inf_diff = 0.0 total_inf_kld = 0.0 + total_ppo_kl = 0.0 inf_num_samples = 0 num_tokens = 0 tis_metrics_agg: Dict[str, float] = {} @@ -196,12 +202,24 @@ def run_loss_loop( device=resp_pi.device, ) - inf_log_diff = pi_detached - resp_inf + # Filter to loss_mask>0 positions: masked bridge/tool tokens otherwise + # contaminate sequence-level TIS weight (matches slime/AReaL behavior). + active_pi = pi_detached[active] + active_inf = resp_inf[active] + active_prox = resp_prox[active] + + inf_log_diff = active_pi - active_inf total_inf_diff += inf_log_diff.abs().mean().item() total_inf_kld += (torch.exp(inf_log_diff) - inf_log_diff - 1.0).mean().item() + ppo_log_diff = active_pi - active_prox + total_ppo_kl += (torch.exp(ppo_log_diff) - ppo_log_diff - 1.0).mean().item() inf_num_samples += 1 - tis_weight, bm = compute_tis_weight(resp_prox, resp_inf, tis_config) + tis_weight_active, bm = compute_tis_weight(active_prox, active_inf, tis_config) + # Identity (1.0) at masked positions: zeroes under ``resp_mask`` for + # masked-multiplied losses, no-op weight for ``dro``. + tis_weight = torch.ones(resp_len, dtype=resp_pi.dtype, device=resp_pi.device) + tis_weight[active] = tis_weight_active.to(resp_pi.dtype) for k, v in bm.items(): tis_metrics_agg[k] = tis_metrics_agg.get(k, 0.0) + v @@ -232,6 +250,7 @@ def run_loss_loop( if inf_num_samples > 0: base_metrics["inference_diff"] = total_inf_diff / inf_num_samples base_metrics["inference_kld"] = total_inf_kld / inf_num_samples + base_metrics["ppo_kl"] = total_ppo_kl / inf_num_samples for k, v in tis_metrics_agg.items(): base_metrics[k] = v / n_samples diff --git a/training/utils/rl/losses.py b/training/utils/rl/losses.py index 96d5b5ff..8912cc67 100644 --- a/training/utils/rl/losses.py +++ b/training/utils/rl/losses.py @@ -218,6 +218,13 @@ class PromptGroup: """Raw completion texts (for trajectory logging).""" row_meta: dict | None = None """Dataset row metadata, e.g. ground_truth (for trajectory logging).""" + prompt_lens: List[int] | None = None + """Per-sample prompt boundaries. Heterogeneous rollouts (multi-turn, + tool branches) have different prefix lengths per sample, so the scalar + ``prompt_len`` cannot represent them faithfully -- ``prompt_lens[i]`` + is the boundary for ``data[i]``. Left ``None`` for legacy single-turn + rollouts where every sample shares the same prefix; ``combine_prompt_groups`` + then falls back to ``[prompt_len] * len(data)``.""" def combine_prompt_groups( @@ -238,7 +245,10 @@ def combine_prompt_groups( advantages.extend(pg.advantages) if pg.ref_logprobs is not None: ref_logprobs.extend(pg.ref_logprobs) - prompt_lens.extend([pg.prompt_len] * len(pg.data)) + if pg.prompt_lens is not None: + prompt_lens.extend(pg.prompt_lens) + else: + prompt_lens.extend([pg.prompt_len] * len(pg.data)) inf_logprobs.extend(pg.inf_logprobs) return data, advantages, ref_logprobs, prompt_lens, inf_logprobs @@ -291,7 +301,14 @@ def build_builtin_loss_datums( ) resp_prox = torch.tensor(prox_lp[response_start:response_start + resp_len], dtype=torch.float32) resp_inf = torch.tensor(inf_lp[response_start:response_start + resp_len], dtype=torch.float32) - tis_weight, _ = compute_tis_weight(resp_prox, resp_inf, tis_config) + # Active-only filter mirrors common.py: keep masked bridge tokens + # out of the sequence-level TIS weight. + active = loss_mask > 0.5 + tis_weight_active, _ = compute_tis_weight( + resp_prox[active], resp_inf[active], tis_config, + ) + tis_weight = torch.ones(resp_len, dtype=torch.float32) + tis_weight[active] = tis_weight_active.to(torch.float32) else: tis_weight = torch.ones(resp_len, dtype=torch.float32) diff --git a/training/utils/rl/metrics.py b/training/utils/rl/metrics.py index 7a7d3ab1..4d849381 100644 --- a/training/utils/rl/metrics.py +++ b/training/utils/rl/metrics.py @@ -25,20 +25,30 @@ def median(values: Sequence[int]) -> float: def datum_target_len(datum: tinker.Datum) -> int: - """Best-effort extraction of target-token length from a training datum.""" - try: - target = datum.loss_fn_inputs.get("target_tokens") - shape = getattr(target, "shape", None) - if isinstance(shape, (list, tuple)) and shape: - return int(shape[0]) - data = getattr(target, "data", None) - if data is not None: - return len(data) - except Exception: - pass + """Length of the target-token tensor on a training datum (0 if missing).""" + target = datum.loss_fn_inputs.get("target_tokens") + shape = getattr(target, "shape", None) + if isinstance(shape, (list, tuple)) and shape: + return int(shape[0]) + data = getattr(target, "data", None) + if data is not None: + return len(data) return 0 +def datum_loss_mask(datum: tinker.Datum) -> list[float] | None: + """Per-position loss mask for a datum (length == target_tokens), or None. + + The adapter writes the mask under ``"weights"`` (legacy datums use + ``"loss_mask"``). Returned values are floats; callers compare against + a positive threshold to find active positions. + """ + td = datum.loss_fn_inputs.get("weights") or datum.loss_fn_inputs.get("loss_mask") + if td is None: + return None + return list(getattr(td, "data", []) or []) + + def total_target_tokens(prompt_groups: Sequence[PromptGroup]) -> int: return sum( datum_target_len(datum) @@ -75,6 +85,29 @@ def add_train_perf_metrics(metrics: dict[str, Any], *, total_model_tokens: int) metrics["perf/weight_sync_ratio"] = weight_sync_time / step_time +def compute_minibatch_metrics( + fwd_bwd_result: Any, + optim_result: Any, +) -> dict[str, Any]: + """Per-minibatch ``train/*`` metrics for one ``fwd_bwd + optim_step``. + + Recipes that log on a per-PPO-minibatch axis (slime convention: each + inner step is its own data point on ``train/step``) call this once + per minibatch instead of letting :func:`compute_step_metrics` + average across the inner loop. + """ + metrics: dict[str, Any] = {} + if fwd_bwd_result is not None and getattr(fwd_bwd_result, "metrics", None): + for k, v in fwd_bwd_result.metrics.items(): + if k not in _SKIP_REMOTE_KEYS: + metrics[f"train/{k}"] = v + if optim_result is not None and getattr(optim_result, "metrics", None): + for k, v in optim_result.metrics.items(): + if k not in _SKIP_REMOTE_KEYS: + metrics[f"train/{k}"] = v + return metrics + + def build_loop_metrics( *, train_step: int, @@ -119,12 +152,8 @@ def compute_step_metrics( if k not in _SKIP_REMOTE_KEYS: metrics[f"train/{k}"] = v - # Average fwd_bwd metrics across inner minibatches. With ppo_n_minibatches=1 - # this reduces to the pre-PR behavior (one result, mean == that result). With - # K>1 the last minibatch alone is misleading — early minibatches haven't - # drifted yet so their ppo_clip_frac is ~0, while later ones clip more; the - # headline claim of this PR is that clipping fires, which the mean reports - # honestly rather than only showing the most-clipped minibatch. + # Mean across inner minibatches (last-only would hide that early + # minibatches don't clip yet while later ones do). if fwd_bwd_results: accum: dict[str, float] = {} for result in fwd_bwd_results: @@ -157,13 +186,29 @@ def compute_step_metrics( if all_truncated: metrics["rollout/truncated_ratio"] = sum(all_truncated) / len(all_truncated) + # Entropy is a mean of -logprob over loss_mask>0 positions only. + # Multi-turn rollouts emit loss_mask=0 / logprobs=0.0 on bridge/user/tool + # tokens; including them biases entropy toward 0. entropy_vals: list[float] = [] for pg in prompt_groups: - for inf_lp in pg.inf_logprobs: - resp_start = max(0, pg.prompt_len - 1) + per_sample = pg.prompt_lens if pg.prompt_lens is not None else None + for i, inf_lp in enumerate(pg.inf_logprobs): + sample_prompt_len = ( + per_sample[i] if per_sample is not None and i < len(per_sample) + else pg.prompt_len + ) + resp_start = max(0, sample_prompt_len - 1) resp_lp = inf_lp[resp_start:] if len(inf_lp) > resp_start else [] - if resp_lp: - entropy_vals.append(-sum(resp_lp) / len(resp_lp)) + if not resp_lp: + continue + mask = datum_loss_mask(pg.data[i]) if i < len(pg.data) else None + if mask is not None: + resp_mask = mask[resp_start : resp_start + len(resp_lp)] + active = [lp for lp, m in zip(resp_lp, resp_mask) if m > 0.5] + else: + active = list(resp_lp) + if active: + entropy_vals.append(-sum(active) / len(active)) if entropy_vals: metrics["rollout/entropy"] = sum(entropy_vals) / len(entropy_vals) @@ -177,13 +222,21 @@ def compute_step_metrics( metrics["rollout/sample_fail_count"] = loop_stats["sample_fails"] metrics["rollout/fwd_bwd_count"] = n_accum - sample_wait_time = float(loop_stats["sample_wait_time"]) - metrics["perf/sample_wait_time"] = sample_wait_time - # The ratio is defined over the same sampling window that drives - # fwd_bwd firing: queue-wait time divided by sampling-loop wall time. - step_wall_time = float(loop_stats["step_wall_time"]) + # Gap between successive train_steps; folds in weight_sync wall time + # (``last_step_end`` is set before ``weight_sync_fn`` runs), so subtract + # ``perf/weight_sync_time`` to back out the pure-starvation portion. + trainer_wait_for_sampler_time = float(loop_stats.get("trainer_wait_for_sampler_time", 0.0)) + metrics["perf/trainer_wait_for_sampler_time"] = trainer_wait_for_sampler_time + # Should be ~0 in healthy async (concurrency cap binds before staleness); + # large values mean the staleness budget is throttling the rollout side. + metrics["perf/sampler_wait_for_trainer_time"] = float( + loop_stats.get("sampler_wait_for_trainer_time", 0.0) + ) + # = wait / (wait + train_wall). When rollouts are structurally slower + # than train+sync the ratio is Amdahl-bound -- not a pipeline bug. + step_wall_time = float(loop_stats.get("step_wall_time", 0.0)) if step_wall_time > 0: - wait_ratio = sample_wait_time / step_wall_time + wait_ratio = trainer_wait_for_sampler_time / step_wall_time metrics["perf/wait_time_ratio"] = wait_ratio metrics["perf/overlap_ratio"] = 1.0 - wait_ratio diff --git a/training/utils/rl/rollout/__init__.py b/training/utils/rl/rollout/__init__.py new file mode 100644 index 00000000..0c87735a --- /dev/null +++ b/training/utils/rl/rollout/__init__.py @@ -0,0 +1,117 @@ +"""Rollout primitives for the RL recipes. + +The trainer's user-facing contract is per-sample (matches AReaL/slime):: + + async def rollout_fn(sample_prompt) -> RolloutSample | None: ... + +``sample_prompt`` is a dataset row's dict, renamed at the recipe seam +to mark that it is now per-sample input rather than dataset-cursor +state. The recipe fans each dataset row out to +``completions_per_prompt`` parallel calls and joins them by row id via +:class:`GroupAssembler` before handing the assembled +:class:`PromptGroup` to the trainer. + +This package layers the supporting types and helpers: + +* :mod:`.types` — :class:`Rollout` (group of samples) and + :class:`RolloutSample` (single trajectory), plus the + :func:`rollout_to_prompt_group` adapter that packs a group into the + trainer's :class:`PromptGroup`. +* :mod:`.group_assembler` — :class:`GroupAssembler` joins per-sample + rollouts into PromptGroups by row id once all samples for a row have + settled. +* :mod:`.service` — service-agnostic Protocol + payload types + (:class:`RolloutService`, :class:`RolloutPayload`, :class:`TurnRecord`). +* :mod:`.remote` — drives a :class:`RolloutService` and packs payloads + (:func:`make_remote_rollout_fn`, :func:`pack_payload_to_sample`). +* :mod:`.renderer` — renderer-backed single-turn helper + (:func:`single_turn_renderer_rollout`, :func:`model_input_to_token_ids`). +* :mod:`.assembler` — multi-turn token-native stitching + (:class:`TrajectoryAssembler`, :func:`extract_completion`, + :func:`precompute_chat_suffix`). +* :mod:`.message` — message-in multi-turn assembly that preserves prior + assistant tokens with TITO-style incremental tokenization. +* :mod:`.trace` — native rollout trajectory analysis for visualization and + diagnostics without a live verifier probe. +""" + +from training.utils.rl.rollout.assembler import ( + InferenceCall, + PrefixMismatch, + TrajectoryAssembler, + extract_completion, + precompute_chat_suffix, +) +from training.utils.rl.rollout.group_assembler import ( + GroupAssembler, + PendingGroup, +) +from training.utils.rl.rollout.message import ( + MessageTrajectoryAssembler, + MessageTrajectoryError, + MessageValidationError, + TITOTokenizer, + TokenizationError, + get_tito_tokenizer, +) +from training.utils.rl.rollout.remote import ( + make_remote_rollout_fn, + pack_payload_to_sample, +) +from training.utils.rl.rollout.renderer import ( + MultimodalRenderingNotSupported, + model_input_to_token_ids, + single_turn_renderer_rollout, +) +from training.utils.rl.rollout.service import ( + RolloutPayload, + RolloutService, + TurnRecord, +) +from training.utils.rl.rollout.trace import ( + RolloutTrajectory, + TrajectoryIssue, + TrajectoryToken, + analyze_flat_sample, + analyze_token_turn_traces, + analyze_turns, +) +from training.utils.rl.rollout.types import ( + Rollout, + RolloutSample, + rollout_to_prompt_group, +) + + +__all__ = [ + "InferenceCall", + "GroupAssembler", + "MessageTrajectoryAssembler", + "MessageTrajectoryError", + "MessageValidationError", + "MultimodalRenderingNotSupported", + "PendingGroup", + "PrefixMismatch", + "Rollout", + "RolloutPayload", + "RolloutSample", + "RolloutService", + "RolloutTrajectory", + "TITOTokenizer", + "TrajectoryAssembler", + "TrajectoryIssue", + "TrajectoryToken", + "TurnRecord", + "TokenizationError", + "analyze_flat_sample", + "analyze_token_turn_traces", + "analyze_turns", + "extract_completion", + "get_tito_tokenizer", + "make_remote_rollout_fn", + "model_input_to_token_ids", + "pack_payload_to_sample", + "precompute_chat_suffix", + "rollout_to_prompt_group", + "single_turn_renderer_rollout", +] diff --git a/training/utils/rl/rollout/assembler.py b/training/utils/rl/rollout/assembler.py new file mode 100644 index 00000000..98916fbd --- /dev/null +++ b/training/utils/rl/rollout/assembler.py @@ -0,0 +1,433 @@ +"""Multi-turn trajectory assembly with prefix-equality invariant. + +The cookbook's :class:`RolloutPayload` is token-native: every turn carries +``token_ids`` straight from the inference call, and assistant turns carry +per-token ``logprobs`` aligned 1:1 with those token IDs. When a rollout +function makes more than one engine call (multi-turn agents, tool-using +loops, retry-with-feedback patterns), it has to stitch the per-call +results into a single payload *without* re-tokenizing any text. Re- +tokenization silently drifts the loss mask off the BPE boundary the +engine actually generated and misaligns the inference logprobs with the +tokens the trainer sees. + +This module provides :class:`TrajectoryAssembler` -- a thin helper that +carries the AReaL ``MultiTurnWorkflow`` invariant +(``areal/workflow/multi_turn.py``): each engine call's ``input_tokens`` +must start with the already-accumulated sequence. When the invariant +holds, the assembler folds the call into the running trajectory +(non-assistant gap + assistant output). When it breaks -- engine saw +something the assembler didn't record, usually because the rollout +function re-rendered messages or skipped a turn -- the assembler raises +:class:`PrefixMismatch` with the first divergence index instead of +training on misaligned tokens. + +Slime relies on author discipline (no assert), AReaL has the assert in +each workflow's ``arun_episode``, tinker-cookbook silently splits into +extra datums. We pick AReaL's "loud crash" mode and centralize it in +one helper so every rollout function gets the same guarantee for free. + +Usage:: + + assembler = TrajectoryAssembler(tokenizer_id=ctx.tokenizer_id) + + # First engine call -- the input is the initial prompt. + call = extract_completion(response.choices[0]) + assembler.add_call(call) + + # ... feed tool result, build next prompt by concatenating engine + # tokens (NOT re-rendering text), call engine again ... + + call2 = extract_completion(response2.choices[0]) + assembler.add_call(call2, role_before="tool") # gap was a tool reply + + return assembler.to_payload(total_reward=reward) +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Iterable, List, Mapping, Optional, Sequence + +from training.utils.rl.rollout.service import Role, RolloutPayload, TurnRecord + + +__all__ = [ + "InferenceCall", + "PrefixMismatch", + "TrajectoryAssembler", + "extract_completion", + "precompute_chat_suffix", +] + + +@dataclass +class InferenceCall: + """One engine round trip, captured token-natively. + + ``input_tokens`` is the full prompt the engine saw (including any + chat template suffix appended after the prior turn's assistant + message). ``output_tokens`` and ``output_logprobs`` come straight + from the engine response and must align 1:1. + """ + + input_tokens: List[int] + output_tokens: List[int] + output_logprobs: List[float] + finish_reason: str = "stop" + + +class PrefixMismatch(RuntimeError): + """Raised when an engine call's ``input_tokens`` doesn't extend the + accumulated trajectory. + + Almost always caused by re-tokenizing text between turns instead of + concatenating the engine's output tokens verbatim. The exception + message includes the first divergence index and the conflicting + token IDs. + """ + + +def _first_divergence(a: List[int], b: List[int]) -> int: + n = min(len(a), len(b)) + for i in range(n): + if a[i] != b[i]: + return i + return n + + +def _is_strict_prefix(prior: List[int], full: List[int]) -> bool: + if len(full) < len(prior): + return False + for i, t in enumerate(prior): + if full[i] != t: + return False + return True + + +@dataclass +class TrajectoryAssembler: + """Stitch multi-turn engine calls into a token-native trajectory. + + Each :meth:`add_call` records one engine round trip. Between calls, + any tokens injected by the environment (tool replies, user + follow-ups, generation-prompt suffixes) are derived from the next + call's ``input_tokens`` -- the assembler asserts those gap tokens + sit cleanly after the previously accumulated sequence. + + Attributes: + tokenizer_id: Identifier for the tokenizer that produced the + token IDs. Propagated to :class:`RolloutPayload` so the + packer can fail loud on a mismatched-vocab integration. + role_for_input: Default role assigned to non-assistant gap + tokens. Override per call with ``add_call(role_before=...)``. + ``"user"`` is the right default for chat-template gaps; + pass ``"tool"`` when the gap is a tool reply. + """ + + tokenizer_id: Optional[str] = None + role_for_input: Role = "user" + _turns: List[TurnRecord] = field(default_factory=list, init=False, repr=False) + _seq: List[int] = field(default_factory=list, init=False, repr=False) + + def add_call( + self, + call: InferenceCall, + *, + role_before: Optional[Role] = None, + max_trim_tokens: int = 0, + ) -> None: + """Record one engine call and advance the trajectory. + + Raises: + PrefixMismatch: if ``call.input_tokens`` doesn't start with + the already-accumulated sequence, after applying the explicit + tokenizer-boundary trim allowed by ``max_trim_tokens``. + ValueError: if ``output_logprobs`` length doesn't match + ``output_tokens``. + """ + if len(call.output_logprobs) != len(call.output_tokens): + raise ValueError( + f"output_logprobs length ({len(call.output_logprobs)}) " + f"!= output_tokens length ({len(call.output_tokens)})", + ) + + if max_trim_tokens < 0: + raise ValueError("max_trim_tokens must be >= 0") + + prior_len = len(self._seq) + if prior_len: + if not _is_strict_prefix(self._seq, call.input_tokens): + keep_len = prior_len - max_trim_tokens + if keep_len < 0 or not _is_strict_prefix(self._seq[:keep_len], call.input_tokens): + idx = _first_divergence(self._seq, call.input_tokens) + prior_tok = self._seq[idx] if idx < prior_len else None + input_tok = call.input_tokens[idx] if idx < len(call.input_tokens) else None + raise PrefixMismatch( + f"engine input_tokens diverges from accumulated sequence " + f"at index {idx} (prior_len={prior_len}, " + f"input_len={len(call.input_tokens)}, " + f"prior_token={prior_tok}, input_token={input_tok}). " + f"This usually means the rollout function re-tokenized text " + f"between turns instead of concatenating engine output tokens " + f"verbatim. Build the next prompt by appending engine tokens + " + f"a precomputed chat-template suffix (see " + f":func:`precompute_chat_suffix` in this module).", + ) + self._trim_engine_visible_tail(max_trim_tokens) + prior_len = len(self._seq) + gap_tokens = list(call.input_tokens[prior_len:]) + else: + gap_tokens = list(call.input_tokens) + + gap_role: Role = role_before or self.role_for_input + if gap_tokens: + self._turns.append(TurnRecord(role=gap_role, token_ids=gap_tokens)) + self._seq.extend(gap_tokens) + + self._turns.append( + TurnRecord( + role="assistant", + token_ids=list(call.output_tokens), + logprobs=list(call.output_logprobs), + finish_reason=call.finish_reason, + ), + ) + self._seq.extend(call.output_tokens) + + def add_environment_tokens( + self, + tokens: List[int], + *, + role: Role = "tool", + ) -> None: + """Record non-assistant tokens that won't appear in the next call's + ``input_tokens``. + + Most users don't need this -- :meth:`add_call` derives gap tokens + from the prefix delta, which works whenever the engine receives + the full conversation each turn. Use this only when your engine + API takes incremental prompts and the trajectory has to record + tokens the engine never sees as input. + + Important: these tokens are NOT added to ``_seq`` (the + engine-visible accumulated sequence) because they will not appear + in the next ``call.input_tokens``. Adding them would make + :meth:`add_call`'s strict-prefix invariant fail on the very next + engine call: ``input_tokens`` would not start with ``_seq + + env_tokens``. We still record them in ``_turns`` so they show up + in the flat trajectory the trainer consumes. + """ + if not tokens: + return + self._turns.append(TurnRecord(role=role, token_ids=list(tokens))) + + def to_payload(self, *, total_reward: Optional[float] = None) -> RolloutPayload: + """Emit a :class:`RolloutPayload` ready for the trainer packer. + + Sets the ``_assembled`` flag so the packer skips its defensive + prefix-consistency check. + """ + if not self._turns: + raise RuntimeError("assembler is empty; call add_call first") + last_assistant = next( + (t for t in reversed(self._turns) if t.role == "assistant"), + None, + ) + if last_assistant is None: + raise RuntimeError("assembler has no assistant turn; nothing to train on") + payload = RolloutPayload( + turns=list(self._turns), + total_reward=total_reward, + tokenizer_id=self.tokenizer_id, + finish_reason=last_assistant.finish_reason or "stop", + ) + payload._assembled = True # type: ignore[attr-defined] + return payload + + def to_flat(self) -> tuple[List[int], List[float], List[int]]: + """Return ``(tokens, logprobs, loss_mask)`` directly. + + Useful when feeding a custom packer that doesn't take + :class:`RolloutPayload`. + """ + tokens: List[int] = [] + logprobs: List[float] = [] + loss_mask: List[int] = [] + for t in self._turns: + assert t.token_ids is not None + n = len(t.token_ids) + tokens.extend(t.token_ids) + if t.role == "assistant": + assert t.logprobs is not None and len(t.logprobs) == n + logprobs.extend(t.logprobs) + loss_mask.extend([1] * n) + else: + logprobs.extend([0.0] * n) + loss_mask.extend([0] * n) + return tokens, logprobs, loss_mask + + def to_trajectory( + self, + *, + tokenizer: Any | None = None, + source: str = "trajectory_assembler", + ): + """Return a native trajectory analysis for verifier visualization.""" + from training.utils.rl.rollout.trace import analyze_turns + + return analyze_turns(self._turns, tokenizer=tokenizer, source=source) + + @property + def accumulated_tokens(self) -> List[int]: + """The current engine-visible token sequence (read-only view).""" + return list(self._seq) + + def _trim_engine_visible_tail(self, n: int) -> None: + """Drop an explicit tokenizer-boundary tail from the flat trajectory.""" + if n == 0: + return + if n > len(self._seq): + raise PrefixMismatch(f"cannot trim {n} tokens from accumulated sequence of length {len(self._seq)}") + self._seq = self._seq[:-n] + remaining = n + while remaining and self._turns: + turn = self._turns[-1] + assert turn.token_ids is not None + if len(turn.token_ids) > remaining: + del turn.token_ids[-remaining:] + if turn.role == "assistant" and turn.logprobs is not None: + del turn.logprobs[-remaining:] + remaining = 0 + else: + remaining -= len(turn.token_ids) + self._turns.pop() + if remaining: + raise PrefixMismatch(f"cannot trim {n} tokens from trajectory turns") + + +# --------------------------------------------------------------------------- +# Token-native helpers used alongside TrajectoryAssembler +# --------------------------------------------------------------------------- + + +def extract_completion( + choice: Mapping[str, Any], + *, + input_tokens: Sequence[int], +) -> InferenceCall: + """Build an :class:`InferenceCall` from a completions-API choice dict. + + Expects Fireworks/OpenAI-shaped fields: + + * ``token_ids``: list of int, the engine's output token IDs. + * ``logprobs.token_logprobs``: list of float aligned 1:1 with + ``token_ids``. ``None`` entries are coerced to ``0.0`` (some + providers omit the logprob for the first token). + * ``finish_reason``: str (default ``"stop"``). + + Raises: + ValueError: if ``token_ids`` is missing/empty or ``logprobs`` + doesn't align with ``token_ids``. Both are required for + token-native training. + """ + raw_token_ids = choice.get("token_ids") + if not raw_token_ids: + raise ValueError( + "completion choice is missing 'token_ids'; this is required for " + "token-native rollout assembly. Enable token-id passthrough on " + "the inference call (e.g. set ``logprobs=True`` and request " + "``token_ids`` from the Fireworks Completions API).", + ) + + raw_logprobs = choice.get("logprobs") + if isinstance(raw_logprobs, Mapping): + token_logprobs: Iterable[Any] = raw_logprobs.get("token_logprobs") or [] + else: + token_logprobs = [] + raw_logprobs_list: List[Any] = list(token_logprobs) + + # Filter token_ids/token_logprobs in lockstep so a None placeholder + # in tokens drops its paired logprob; otherwise the remaining logprobs + # would shift onto the wrong tokens and corrupt the PPO/GRPO ratio. + output_tokens: List[int] = [] + output_logprobs: List[float] = [] + if len(raw_logprobs_list) == len(raw_token_ids): + for tok, lp in zip(raw_token_ids, raw_logprobs_list): + if tok is None: + continue + output_tokens.append(int(tok)) + output_logprobs.append(float(lp) if lp is not None else 0.0) + else: + output_tokens = [int(t) for t in raw_token_ids if t is not None] + output_logprobs = [ + float(lp) if lp is not None else 0.0 for lp in raw_logprobs_list + ] + if len(output_logprobs) > len(output_tokens): + output_logprobs = output_logprobs[: len(output_tokens)] + + if len(output_logprobs) != len(output_tokens): + raise ValueError( + f"completion has {len(output_tokens)} token_ids but " + f"{len(output_logprobs)} logprobs; the inference call must " + f"return per-token logprobs aligned with token_ids " + f"(slime/AReaL convention).", + ) + + finish_reason = choice.get("finish_reason") or "stop" + + return InferenceCall( + input_tokens=list(input_tokens), + output_tokens=output_tokens, + output_logprobs=output_logprobs, + finish_reason=str(finish_reason), + ) + + +def precompute_chat_suffix( + tokenizer: Any, + *, + follow_up_content: str, + follow_up_role: str = "user", + add_generation_prompt: bool = True, +) -> List[int]: + """Tokenize the chat-template scaffolding that follows an assistant turn. + + Returns the token IDs for ``[end-of-assistant-turn + follow_up_role + wrapper + follow_up_content + generation prompt]``, computed by + diffing two ``apply_chat_template`` outputs. Append this to the + engine's assistant output tokens to build the next prompt without + re-tokenizing the assistant content. + + AReaL reference: ``MultiTurnWorkflow.__init__`` + (``areal/workflow/multi_turn.py:41-57``). + + Raises: + RuntimeError: if the tokenizer's chat template doesn't extend + cleanly when a follow-up message is appended -- usually + because the renderer mutates earlier turns (e.g. strips + ```` blocks). In that case there is no stable + suffix and you need a renderer-aware assembly path. + """ + base = [{"role": "assistant", "content": "x"}] + s1 = list(tokenizer.apply_chat_template(base, tokenize=True)) + + extended = base + [{"role": follow_up_role, "content": follow_up_content}] + s2 = list( + tokenizer.apply_chat_template( + extended, + tokenize=True, + add_generation_prompt=add_generation_prompt, + ), + ) + + if len(s2) < len(s1) or s2[: len(s1)] != s1: + raise RuntimeError( + "tokenizer's chat template does not extend cleanly when adding a " + "follow-up message: the prefix re-tokenizes differently after the " + "second turn is appended. This usually means the renderer mutates " + "earlier turns (e.g. Qwen3 strips blocks, KimiK2 injects " + "default system prompts). No stable suffix exists for this model " + "-- use renderer-aware assembly that re-renders each turn through " + "the renderer's own helpers.", + ) + return s2[len(s1) :] diff --git a/training/utils/rl/rollout/group_assembler.py b/training/utils/rl/rollout/group_assembler.py new file mode 100644 index 00000000..468d9609 --- /dev/null +++ b/training/utils/rl/rollout/group_assembler.py @@ -0,0 +1,231 @@ +"""Per-row group assembler for the per-sample rollout API. + +The per-sample ``rollout_fn`` returns one :class:`RolloutSample` per call. +The async loop fans each row out to ``completions_per_prompt`` parallel +samples, then this assembler joins them back into a :class:`PromptGroup` +once all expected samples for a row land (or a partial-emission policy +fires). Group advantages are computed only at the join point, so the +trainer never sees a half-formed group. + +This mirrors AReaL/slime, where the user-facing function produces one +trajectory and the framework owns group assembly. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import Callable, Dict, Hashable, List, Optional + +from training.utils.data import compute_advantages +from training.utils.rl.losses import PromptGroup +from training.utils.rl.rollout.types import ( + Rollout, + RolloutSample, + rollout_to_prompt_group, +) + + +logger = logging.getLogger(__name__) + + +__all__ = [ + "GroupAssembler", + "PendingGroup", + "RowResolution", +] + + +AdvantageFn = Callable[[List[float]], List[float]] + + +@dataclass +class RowResolution: + """Outcome of a row once all its samples have settled. + + ``pg`` is the assembled PromptGroup, or ``None`` when no group + survived (every sample dropped, advantage_fn produced non-finite + values, or the surviving sample count fell below ``min_group_size``). + ``min_submit_version`` is the oldest submit version among the row's + samples. + """ + + pg: Optional[PromptGroup] + min_submit_version: int + + +@dataclass +class PendingGroup: + """In-flight group of samples for one row.""" + + row_id: Hashable + expected_n: int + samples: List[RolloutSample] = field(default_factory=list) + submit_versions: List[int] = field(default_factory=list) + row_meta: Optional[dict] = None + # Number of samples started for this row. ``samples + dropped == started``; + # the group is "settled" when ``started == expected_n``. + started: int = 0 + dropped: int = 0 + + @property + def settled(self) -> bool: + return (len(self.samples) + self.dropped) >= self.expected_n + + @property + def min_submit_version(self) -> int: + return min(self.submit_versions) if self.submit_versions else 0 + + +class GroupAssembler: + """Join per-sample rollouts into PromptGroups by row id. + + Each row is assigned an ``expected_n`` (typically + ``completions_per_prompt``). The async loop calls :meth:`note_started` + when a sample is submitted and one of :meth:`add_sample` / + :meth:`note_dropped` when the task resolves. Once a row's group is + fully settled, the assembler packs the surviving samples through + :func:`rollout_to_prompt_group` and returns the resulting PromptGroup + (or ``None`` if every sample for the row was dropped or the + ``advantage_fn`` produced non-finite values). + + The assembler is single-threaded -- the async loop drives it from + the event-loop thread. + """ + + def __init__( + self, + *, + completions_per_prompt: int, + advantage_fn: AdvantageFn = compute_advantages, + with_reference: bool = False, + router_replay_completion_only: bool = False, + min_group_size: int = 1, + ) -> None: + if completions_per_prompt < 1: + raise ValueError("completions_per_prompt must be >= 1") + if min_group_size < 1: + raise ValueError("min_group_size must be >= 1") + self._n = completions_per_prompt + self._advantage_fn = advantage_fn + self._with_reference = with_reference + self._r3_completion_only = router_replay_completion_only + self._min_group_size = min_group_size + self._pending: Dict[Hashable, PendingGroup] = {} + + def note_started( + self, + row_id: Hashable, + *, + submit_version: int, + row_meta: Optional[dict] = None, + ) -> None: + """Record that one sample for ``row_id`` was submitted. + + The first call for a given ``row_id`` materializes the + :class:`PendingGroup` slot. Subsequent calls just bump ``started`` + and append the submit version. ``row_meta`` is stored on the + first call; later calls are ignored (all samples for a row share + meta). + """ + group = self._pending.get(row_id) + if group is None: + group = PendingGroup( + row_id=row_id, + expected_n=self._n, + row_meta=dict(row_meta) if row_meta else None, + ) + self._pending[row_id] = group + group.started += 1 + group.submit_versions.append(submit_version) + + def add_sample( + self, + row_id: Hashable, + sample: RolloutSample, + ) -> Optional[RowResolution]: + """Record one resolved sample. + + Returns ``None`` if the row still has samples in flight, or a + :class:`RowResolution` once the row has fully settled. The + resolution carries the assembled :class:`PromptGroup` (or + ``None`` if no group survived). + """ + group = self._require(row_id) + group.samples.append(sample) + return self._maybe_emit(row_id, group) + + def note_dropped( + self, + row_id: Hashable, + ) -> Optional[RowResolution]: + """Record that one sample for ``row_id`` failed. + + Same return contract as :meth:`add_sample`. + """ + group = self._require(row_id) + group.dropped += 1 + return self._maybe_emit(row_id, group) + + def _require(self, row_id: Hashable) -> PendingGroup: + group = self._pending.get(row_id) + if group is None: + raise KeyError( + f"GroupAssembler: row {row_id!r} has no pending group " + "(note_started was not called).", + ) + return group + + def _maybe_emit( + self, + row_id: Hashable, + group: PendingGroup, + ) -> Optional[RowResolution]: + if not group.settled: + return None + del self._pending[row_id] + min_version = group.min_submit_version + if len(group.samples) < self._min_group_size: + logger.info( + "GroupAssembler: dropping row %r (got %d/%d samples; min=%d)", + row_id, len(group.samples), group.expected_n, self._min_group_size, + ) + return RowResolution(pg=None, min_submit_version=min_version) + rollout = Rollout(samples=group.samples, row_meta=group.row_meta) + pg = rollout_to_prompt_group( + rollout, + advantage_fn=self._advantage_fn, + with_reference=self._with_reference, + router_replay_completion_only=self._r3_completion_only, + ) + return RowResolution(pg=pg, min_submit_version=min_version) + + def pending_rows(self) -> int: + """Number of rows with at least one in-flight sample.""" + return len(self._pending) + + def drain(self) -> List[RowResolution]: + """Force-emit any pending rows whose surviving samples meet + ``min_group_size``. + + Use only at shutdown / after the sample iterator is exhausted and + all in-flight tasks have resolved. Rows still missing samples + below ``min_group_size`` are dropped silently. + """ + out: List[RowResolution] = [] + for row_id in list(self._pending): + group = self._pending[row_id] + del self._pending[row_id] + min_version = group.min_submit_version + if len(group.samples) < self._min_group_size: + continue + rollout = Rollout(samples=group.samples, row_meta=group.row_meta) + pg = rollout_to_prompt_group( + rollout, + advantage_fn=self._advantage_fn, + with_reference=self._with_reference, + router_replay_completion_only=self._r3_completion_only, + ) + if pg is not None: + out.append(RowResolution(pg=pg, min_submit_version=min_version)) + return out diff --git a/training/utils/rl/rollout/message.py b/training/utils/rl/rollout/message.py new file mode 100644 index 00000000..f68ef65b --- /dev/null +++ b/training/utils/rl/rollout/message.py @@ -0,0 +1,342 @@ +"""Message-in multi-turn assembly with TITO-style token preservation. + +This module bridges OpenAI-style message loops to the token-native rollout +contract. Prior assistant outputs stay as engine token IDs; only newly +appended non-assistant messages are tokenized before the next model call. +""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, List, Optional + +from training.utils.rl.rollout.assembler import InferenceCall, TrajectoryAssembler +from training.utils.rl.rollout.service import RolloutPayload + + +class MessageTrajectoryError(ValueError): + """Base error for message trajectory validation failures.""" + + +class MessageValidationError(MessageTrajectoryError): + """Raised when messages are not append-only or violate role policy.""" + + +class TokenizationError(MessageTrajectoryError): + """Raised when incremental message tokenization cannot be made safe.""" + + +def _assert_append_only_with_allowed_roles( + old_messages: List[dict[str, Any]], + new_messages: List[dict[str, Any]], + allowed_roles: List[str], +) -> None: + if len(new_messages) < len(old_messages): + raise MessageValidationError( + f"new messages shorter than stored prefix: {len(new_messages)} < {len(old_messages)}" + ) + for idx, old_msg in enumerate(old_messages): + if old_msg != new_messages[idx]: + raise MessageValidationError(f"message at index {idx} modifies stored prefix") + for idx, msg in enumerate(new_messages[len(old_messages) :], start=len(old_messages)): + role = msg.get("role") + if role not in allowed_roles: + raise MessageValidationError( + f"appended role={role!r} at index {idx} is not allowed; allowed={allowed_roles}" + ) + + +class TITOTokenizer: + """Incrementally tokenize appended non-assistant turns. + + ``merge_tokens`` preserves the pretokenized assistant prefix exactly and + appends tokenized user/tool/system deltas plus the next assistant opener. + """ + + max_trim_tokens: int = 0 + + def __init__( + self, + tokenizer: Any, + *, + chat_template_kwargs: Optional[dict[str, Any]] = None, + allowed_append_roles: Optional[List[str]] = None, + ) -> None: + self.tokenizer = tokenizer + self.chat_template_kwargs = dict(chat_template_kwargs or {}) + self.allowed_append_roles = list(allowed_append_roles or ["tool", "user", "system"]) + + def render_initial_prompt( + self, + messages: List[dict[str, Any]], + *, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + rendered = self._render_messages(messages, add_generation_prompt=True, tools=tools) + return self._encode_text(rendered) + + def merge_tokens( + self, + *, + old_messages: List[dict[str, Any]], + new_messages: List[dict[str, Any]], + pretokenized_token_ids: List[int], + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + return list(pretokenized_token_ids) + self.tokenize_additional_non_assistant( + old_messages, + new_messages, + tools=tools, + ) + + def tokenize_additional_non_assistant( + self, + old_messages: List[dict[str, Any]], + new_messages: List[dict[str, Any]], + *, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + _assert_append_only_with_allowed_roles(old_messages, new_messages, self.allowed_append_roles) + appended_messages = new_messages[len(old_messages) :] + incremental: List[int] = [] + for segment in self._split_appended_segments(appended_messages): + role = segment[0].get("role") + if role == "tool": + incremental.extend(self._tokenize_tool_segment(segment, tools=tools)) + elif role in {"user", "system"}: + incremental.extend(self._tokenize_user_or_system_segment(segment[0], tools=tools)) + else: + raise TokenizationError(f"unsupported appended role for TITO tokenization: {role!r}") + incremental.extend(self._tokenize_rendered_suffix(new_messages, [], add_generation_prompt=True, tools=tools)) + return incremental + + def _split_appended_segments(self, appended_messages: List[dict[str, Any]]) -> List[List[dict[str, Any]]]: + segments: List[List[dict[str, Any]]] = [] + i = 0 + while i < len(appended_messages): + role = appended_messages[i].get("role") + if role == "tool": + j = i + 1 + while j < len(appended_messages) and appended_messages[j].get("role") == "tool": + j += 1 + segments.append(appended_messages[i:j]) + i = j + elif role in {"user", "system"}: + segments.append([appended_messages[i]]) + i += 1 + else: + raise TokenizationError(f"unsupported appended role for TITO segmentation: {role!r}") + return segments + + def _tokenize_tool_segment( + self, + appended_messages: List[dict[str, Any]], + *, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + dummy_assistant = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": msg.get("tool_call_id") or f"call_{idx}", + "type": "function", + "function": {"name": msg.get("name") or "tool", "arguments": "{}"}, + } + for idx, msg in enumerate(appended_messages) + ], + } + return self._tokenize_rendered_suffix( + [ + {"role": "system", "content": "dummy system"}, + {"role": "user", "content": "dummy user"}, + dummy_assistant, + ], + appended_messages, + tools=tools, + ) + + def _tokenize_user_or_system_segment( + self, + appended_message: dict[str, Any], + *, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + return self._tokenize_rendered_suffix( + [ + {"role": "system", "content": "dummy system"}, + {"role": "user", "content": "dummy user"}, + ], + [appended_message], + tools=tools, + ) + + def _tokenize_rendered_suffix( + self, + base_messages: List[dict[str, Any]], + appended_messages: List[dict[str, Any]], + *, + add_generation_prompt: bool = False, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + without = self._render_messages(base_messages, add_generation_prompt=False, tools=tools) + with_appended = self._render_messages( + base_messages + appended_messages, + add_generation_prompt=add_generation_prompt, + tools=tools, + ) + if not with_appended.startswith(without): + roles = [msg.get("role") for msg in appended_messages] if appended_messages else ["generation_prompt"] + raise TokenizationError(f"rendered suffix diff failed for roles={roles}") + return self._encode_text(with_appended[len(without) :]) + + def _render_messages( + self, + messages: List[dict[str, Any]], + *, + add_generation_prompt: bool, + tools: Optional[List[dict[str, Any]]] = None, + ) -> str: + return self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=add_generation_prompt, + tools=tools, + **self.chat_template_kwargs, + ) + + def _encode_text(self, text: str) -> List[int]: + if hasattr(self.tokenizer, "encode"): + return list(self.tokenizer.encode(text, add_special_tokens=False)) + encoded = self.tokenizer(text, add_special_tokens=False) + return list(encoded["input_ids"]) + + +def get_tito_tokenizer( + tokenizer: Any, + **kwargs: Any, +) -> TITOTokenizer: + return TITOTokenizer(tokenizer, **kwargs) + + +@dataclass +class _Checkpoint: + messages: List[dict[str, Any]] + seq: List[int] + turns: list + + +@dataclass +class MessageTrajectoryAssembler: + """Message-level adapter over the token-native :class:`TrajectoryAssembler`.""" + + tito_tokenizer: TITOTokenizer + trajectory: TrajectoryAssembler = field(default_factory=TrajectoryAssembler) + max_assistant_rollback_steps: int = 1 + messages: List[dict[str, Any]] = field(default_factory=list) + checkpoints: List[_Checkpoint] = field(default_factory=list) + + @property + def token_ids(self) -> List[int]: + return self.trajectory.accumulated_tokens + + def prepare_next_input( + self, + request_messages: List[dict[str, Any]], + *, + tools: Optional[List[dict[str, Any]]] = None, + ) -> List[int]: + if not self.checkpoints: + return self.tito_tokenizer.render_initial_prompt(request_messages, tools=tools) + + self._try_detect_and_rollback_to_assistant_checkpoint(request_messages) + _assert_append_only_with_allowed_roles( + self.messages, + request_messages, + self.tito_tokenizer.allowed_append_roles, + ) + return self.tito_tokenizer.merge_tokens( + old_messages=self.messages, + new_messages=request_messages, + pretokenized_token_ids=self.token_ids, + tools=tools, + ) + + def add_assistant_response( + self, + *, + request_messages: List[dict[str, Any]], + assistant_message: dict[str, Any], + prompt_token_ids: List[int], + completion_token_ids: List[int], + completion_logprobs: List[float], + finish_reason: str = "stop", + ) -> None: + self.trajectory.add_call( + InferenceCall( + input_tokens=list(prompt_token_ids), + output_tokens=list(completion_token_ids), + output_logprobs=list(completion_logprobs), + finish_reason=finish_reason, + ), + max_trim_tokens=self.tito_tokenizer.max_trim_tokens, + ) + self.messages = list(request_messages) + [assistant_message] + self.checkpoints.append(self._make_checkpoint()) + + def to_payload(self, *, total_reward: Optional[float] = None) -> RolloutPayload: + return self.trajectory.to_payload(total_reward=total_reward) + + def to_trajectory( + self, + *, + tokenizer: Any | None = None, + source: str = "message_trajectory_assembler", + ): + """Return a native trajectory analysis for verifier visualization.""" + return self.trajectory.to_trajectory(tokenizer=tokenizer, source=source) + + def _make_checkpoint(self) -> _Checkpoint: + return _Checkpoint( + messages=deepcopy(self.messages), + seq=list(self.trajectory._seq), + turns=deepcopy(self.trajectory._turns), + ) + + def _restore_checkpoint(self, checkpoint: _Checkpoint) -> None: + self.messages = deepcopy(checkpoint.messages) + self.trajectory._seq = list(checkpoint.seq) + self.trajectory._turns = deepcopy(checkpoint.turns) + + def _try_detect_and_rollback_to_assistant_checkpoint(self, request_messages: List[dict[str, Any]]) -> None: + stored = self.messages + if not stored or not self.checkpoints: + return + + match_len = 0 + for idx in range(min(len(request_messages), len(stored))): + if stored[idx] == request_messages[idx]: + match_len = idx + 1 + else: + break + if match_len >= len(stored): + return + + checkpoint_index = -1 + assistant_count = 0 + for idx in range(match_len): + if stored[idx].get("role") == "assistant": + checkpoint_index = assistant_count + assistant_count += 1 + if checkpoint_index < 0: + raise MessageValidationError(f"rollback failed: no assistant message found in first {match_len} messages") + + discard_count = len(self.checkpoints) - (checkpoint_index + 1) + if discard_count > self.max_assistant_rollback_steps: + raise MessageValidationError( + f"rollback failed: discard_count={discard_count} exceeds " + f"max_assistant_rollback_steps={self.max_assistant_rollback_steps}" + ) + self.checkpoints = self.checkpoints[: checkpoint_index + 1] + self._restore_checkpoint(self.checkpoints[-1]) diff --git a/training/utils/rl/rollout/remote.py b/training/utils/rl/rollout/remote.py new file mode 100644 index 00000000..7b604147 --- /dev/null +++ b/training/utils/rl/rollout/remote.py @@ -0,0 +1,280 @@ +"""Generic packer: :class:`RolloutService` output -> :class:`RolloutSample`. + +The recipe needs ``rollout_fn(sample_prompt) -> RolloutSample | None``, +called once per sample (each dataset row fans out to +``completions_per_prompt`` calls). The most common integration pattern +is "some remote service produces completion data; the trainer accepts +token-level data verbatim". This helper collapses the common case to +one call:: + + rollout_fn = make_remote_rollout_fn(service, sample_kwargs=...) + +Token-native only +----------------- + +Every :class:`TurnRecord` MUST carry ``token_ids``, and assistant +turns MUST carry ``logprobs`` aligned with ``token_ids``. The packer +concatenates the per-turn tokens, derives ``loss_mask`` from roles +(``1`` on assistant, ``0`` elsewhere), and trusts the supplied +per-token ``logprobs`` as-is. + +Re-tokenizing assistant text after the fact silently breaks two +things: (a) the loss mask drifts off the BPE boundary the engine +actually generated, and (b) the per-token logprobs no longer align +with the tokens fed to the trainer. AReaL and slime both refuse to +do it; this packer follows the same rule. + +Reward +------ + +``payload.total_reward`` is authoritative when set ("server wins"). +When ``None``, pass ``reward_fn=...`` to grade client-side. +""" + +from __future__ import annotations + +import logging +from typing import Any, Awaitable, Callable, List, Optional, Union + +from training.utils.rl.rollout.types import RolloutSample +from training.utils.rl.rollout.service import ( + RolloutPayload, + RolloutService, + RolloutServiceCallable, +) + + +__all__ = [ + "make_remote_rollout_fn", + "pack_payload_to_sample", +] + + +logger = logging.getLogger(__name__) + + +ServiceLike = Union[RolloutService, RolloutServiceCallable] +RewardFn = Callable[[dict, RolloutPayload], Awaitable[float]] + + +async def _call_service( + service: ServiceLike, + messages: List[dict], + *, + n: int, + sample_kwargs: dict[str, Any], + row: dict, +) -> List[RolloutPayload]: + if hasattr(service, "rollout"): + return await service.rollout( # type: ignore[union-attr] + messages, n=n, sample_kwargs=sample_kwargs, row=row, + ) + return await service(messages, n, sample_kwargs, row) # type: ignore[misc, operator] + + +def make_remote_rollout_fn( + service: ServiceLike, + *, + sample_kwargs: dict[str, Any] | None = None, + tokenizer_id: str | None = None, + reward_fn: Optional[RewardFn] = None, + messages_key: str = "messages", + allow_empty_messages: bool = False, +): + """Build a per-sample ``rollout_fn`` that calls ``service`` once per + sample and packs the resulting payload. + + The framework fans each row out to ``completions_per_prompt`` parallel + samples; this helper requests ``n=1`` from the service per call. When + a payload carries ``total_reward=None``, ``reward_fn`` grades it + client-side. + + Parameters + ---------- + service + The :class:`RolloutService` (or compatible callable) producing + :class:`RolloutPayload` lists. + sample_kwargs + Forwarded to ``service.rollout`` per call. Typically the + ``RolloutSetup.sample_kwargs`` passed by the recipe. + tokenizer_id + Optional policy tokenizer identifier. When set, payloads with + a mismatched ``tokenizer_id`` raise rather than silently train + on misaligned token IDs. + allow_empty_messages + When ``False`` (default), the rollout returns ``None`` whenever + ``sample_prompt[messages_key]`` is empty or missing -- the right + guard for chat-style services. When ``True``, empty messages + are forwarded through (e.g. env-driven domains where the env + emits the seed observation inside the service). + """ + + sk = dict(sample_kwargs or {}) + + async def rollout_fn(sample_prompt: dict) -> RolloutSample | None: + messages = sample_prompt.get(messages_key) or [] + if not messages and not allow_empty_messages: + return None + + # The service-side keyword stays ``row=`` (user-supplied services + # and reward fns expect that name); only the closure param renames. + payloads = await _call_service( + service, + messages, + n=1, + sample_kwargs=dict(sk), + row=sample_prompt, + ) + + if not payloads: + return None + if len(payloads) > 1: + logger.warning( + "service returned %d payloads for n=1; using the first", + len(payloads), + ) + + try: + return await pack_payload_to_sample( + payloads[0], + tokenizer_id=tokenizer_id, + reward_fn=reward_fn, + row=sample_prompt, + ) + except _PackError as exc: + logger.warning("dropping payload: %s", exc) + return None + + return rollout_fn + + +# --------------------------------------------------------------------------- +# Payload -> RolloutSample +# --------------------------------------------------------------------------- + + +class _PackError(RuntimeError): + pass + + +async def pack_payload_to_sample( + payload: RolloutPayload, + *, + tokenizer_id: str | None = None, + reward_fn: Optional[RewardFn] = None, + row: Optional[dict] = None, +) -> RolloutSample: + """Normalise one :class:`RolloutPayload` into one :class:`RolloutSample`. + + Token-native only: every turn must carry ``token_ids``, and every + assistant turn must carry ``logprobs`` aligned with ``token_ids``. + Raises :class:`_PackError` with a user-actionable message on any + structural problem. + + ``tokenizer_id`` is the policy tokenizer identifier; when both it + and ``payload.tokenizer_id`` are set, they must match. + """ + if not payload.turns: + raise _PackError("payload has no turns") + if ( + payload.tokenizer_id + and tokenizer_id + and payload.tokenizer_id != tokenizer_id + ): + raise _PackError( + "payload tokenizer_id " + f"{payload.tokenizer_id!r} does not match policy tokenizer " + f"{tokenizer_id!r}", + ) + + missing = [i for i, t in enumerate(payload.turns) if t.token_ids is None] + if missing: + raise _PackError( + f"turns {missing} missing token_ids; this packer is token-native " + "only. Have the upstream service emit per-turn token_ids and " + "per-token logprobs from the same call that generated them " + "(see slime / AReaL). Re-tokenizing text post-hoc silently " + "misaligns the loss mask and inference logprobs.", + ) + + # Defensive checks for hand-built payloads (``_assembled=False``). + if not getattr(payload, "_assembled", False): + empty_turns = [i for i, t in enumerate(payload.turns) if not t.token_ids] + if empty_turns: + raise _PackError( + f"hand-built payload has empty turn(s) at indices {empty_turns}; " + "every turn must carry at least one token id. An empty " + "intermediate turn typically means the service mis-rendered " + "(re-tokenized) a turn and emitted a stale or empty span. " + "If this is intentional, drop the turn from the payload " + "instead of leaving an empty token_ids list.", + ) + if payload.turns[-1].role != "assistant": + raise _PackError( + "hand-built payload must end with an assistant turn (got " + f"role={payload.turns[-1].role!r}). The trainer's loss is " + "computed over the final assistant span; a non-assistant " + "tail typically means a gap turn was appended after the " + "last engine call by mistake.", + ) + + tokens, logprobs, loss_mask = _pack_token_native(payload) + + reward = payload.total_reward + if reward is None: + if reward_fn is None: + raise _PackError( + "payload.total_reward is None and no reward_fn was provided", + ) + reward = float(await reward_fn(row or {}, payload)) + + last_assistant = next( + (t for t in reversed(payload.turns) if t.role == "assistant"), None, + ) + text = last_assistant.text if last_assistant else "" + finish_reason = ( + (last_assistant.finish_reason if last_assistant else None) + or payload.finish_reason + or "stop" + ) + + return RolloutSample( + tokens=tokens, + logprobs=logprobs, + loss_mask=loss_mask, + reward=float(reward), + finish_reason=finish_reason, + text=text, + ) + + +def _pack_token_native( + payload: RolloutPayload, +) -> tuple[List[int], List[float], List[int]]: + tokens: List[int] = [] + logprobs: List[float] = [] + loss_mask: List[int] = [] + for t in payload.turns: + assert t.token_ids is not None # checked by caller + n = len(t.token_ids) + if t.role == "assistant": + if t.logprobs is None or len(t.logprobs) != n: + raise _PackError( + f"assistant turn needs per-token logprobs aligned with " + f"token_ids (got {0 if t.logprobs is None else len(t.logprobs)} " + f"for {n} tokens)", + ) + lp = [float(x) for x in t.logprobs] + mask = [1] * n + else: + lp = [0.0] * n + mask = [0] * n + tokens.extend(t.token_ids) + logprobs.extend(lp) + loss_mask.extend(mask) + + if not any(m > 0 for m in loss_mask): + raise _PackError("no assistant tokens in payload; nothing to train on") + if len(tokens) < 2: + raise _PackError("payload shorter than 2 tokens") + return tokens, logprobs, loss_mask diff --git a/training/utils/rl/rollout/renderer.py b/training/utils/rl/rollout/renderer.py new file mode 100644 index 00000000..d77d9393 --- /dev/null +++ b/training/utils/rl/rollout/renderer.py @@ -0,0 +1,218 @@ +"""Renderer-backed RL rollout primitives. + +This module exposes the renderer-backed single-turn helper and the +``ModelInput`` flattening adapter: + +* :func:`single_turn_renderer_rollout` — single-turn helper that turns a + renderer + a pre-tokenized sampling primitive into a flat + :class:`~training.utils.rl.rollout.types.Rollout` whose tokens, logprobs, + and loss mask are derived end-to-end from the renderer-built prompt and + the sampler-returned assistant tokens. +* :func:`model_input_to_token_ids` — flatten a Tinker ``ModelInput`` from a + renderer's ``build_generation_prompt(...)`` into ``list[int]``. Multimodal + chunks are rejected with :class:`MultimodalRenderingNotSupported`. + +Multi-turn flows are shown as concrete ``async def +rollout_fn(sample_prompt) -> RolloutSample | None`` examples under +``cookbook/training/examples/rl/``. Per-rollout context (sampler, +tokenizer, sample kwargs, custom state) is closed over via +:class:`RolloutSetup` at factory time -- the framework no longer threads +a ``ctx`` argument through. Keep environment/tool policy in user code; +this helper only covers single-turn renderer packing. + +Boundary +-------- + +The renderer is consumed inside the rollout; it is not the trainer's +data contract. The trainer's contract is per-sample: +``rollout_fn(sample_prompt) -> RolloutSample | None``. This module is +renderer-name-agnostic and never re-renders chat templates client-side. + +Parse-failure / truncation handling +----------------------------------- + +Parse-failure handling is *not* a framework primitive. This helper hands +``(parsed_message, parse_success)`` back to the user-supplied ``reward_fn``; +the caller chooses what to do (drop by returning ``None``, score zero by +returning ``0.0``, or branch on ``parse_success`` for custom behavior). +``finish_reason='length'`` flows through unchanged unless the user branches +on it. +""" + +from __future__ import annotations + +import logging +from typing import Any, Awaitable, Callable, List, Optional, Protocol + +import tinker + +from training.utils.rl.rollout.types import RolloutSample + + +logger = logging.getLogger(__name__) + + +__all__ = [ + "MultimodalRenderingNotSupported", + "model_input_to_token_ids", + "single_turn_renderer_rollout", +] + + +class MultimodalRenderingNotSupported(RuntimeError): + """Raised when a ``ModelInput`` carries non-text chunks. + + Renderer-backed RL rollouts are text-only in this iteration. A multimodal + chunk (image asset pointer, image bytes) reaching this adapter indicates a + multimodal rollout flow that has not yet been scoped for RL training. + """ + + +def model_input_to_token_ids(model_input: tinker.ModelInput) -> List[int]: + """Flatten a renderer ``ModelInput`` to ``list[int]``. + + Accepts only :class:`tinker.EncodedTextChunk` chunks. Any other chunk + type — image asset pointers or raw image bytes — raises + :class:`MultimodalRenderingNotSupported`. + """ + out: List[int] = [] + for chunk in model_input.chunks: + if isinstance(chunk, tinker.EncodedTextChunk): + out.extend(chunk.tokens) + else: + raise MultimodalRenderingNotSupported( + f"chunk type {type(chunk).__name__} is not supported by " + "renderer-backed RL rollouts" + ) + return out + + +# A renderer-shaped protocol. Avoids a hard import of the Tinker base class +# in this module's signature surface so tests can pass simple stubs without +# subclassing the heavy upstream class. We do not export this — the helper +# accepts any object that responds to these methods. +class _RendererLike(Protocol): + def build_generation_prompt(self, messages: List[Any]) -> tinker.ModelInput: ... + def parse_response(self, tokens: List[int]) -> Any: ... + def get_stop_sequences(self) -> List[Any]: ... + + +SampleWithPromptTokens = Callable[..., Awaitable[List[Any]]] +"""Callable matching :meth:`DeploymentSampler.sample_with_prompt_tokens`.""" + + +MessageBuilder = Callable[[Any], Awaitable[List[Any]]] +"""``async (row) -> messages`` — builds the seed conversation.""" + + +RewardFn = Callable[[Any, Any, bool], Awaitable[Optional[float]]] +"""``async (row, parsed_message, parse_success) -> float | None``. + +Return ``None`` to drop the completion (no sample emitted). Return a float +to emit a sample with that reward. ``parse_success`` is the second element +returned by the renderer's ``parse_response``. +""" + + +async def single_turn_renderer_rollout( + row: Any, + *, + renderer: _RendererLike, + sample_with_prompt_tokens: SampleWithPromptTokens, + message_builder: MessageBuilder, + reward_fn: RewardFn, + sample_kwargs: dict[str, Any] | None = None, + tokenizer: Any | None = None, + max_tokens: int | None = None, + stop: List[str] | List[int] | None = None, +) -> RolloutSample | None: + """Single-turn renderer-backed rollout (per-sample). + + Builds messages via ``message_builder(row)``, calls + ``renderer.build_generation_prompt(...)``, flattens the resulting + ``ModelInput`` via :func:`model_input_to_token_ids`, samples ONE + completion through ``sample_with_prompt_tokens`` (the SDK's + pre-tokenized sampling primitive), and packs the completion into a + :class:`RolloutSample` whose tokens / logprobs / loss-mask come straight + from the renderer + sampler. No chat-template re-rendering, no + re-tokenization of decoded assistant text. + + The framework fans each row out to ``completions_per_prompt`` parallel + calls; this helper requests ``n=1`` per call. + + ``stop`` defaults to ``renderer.get_stop_sequences()`` and preserves its + ``list[str] | list[int]`` shape. The user-supplied ``reward_fn`` + receives the renderer's parsed message and parse-success flag and + returns ``None`` (drop) or a float. ``tokenizer`` is required only + when the renderer returns integer stop token IDs that need decoding + to strings for the inference API. Multimodal prompts raise + :class:`MultimodalRenderingNotSupported` via the adapter. + """ + messages = await message_builder(row) + model_input = renderer.build_generation_prompt(messages) + prompt_token_ids = model_input_to_token_ids(model_input) + + if stop is None: + stop = renderer.get_stop_sequences() + + # The Fireworks inference completions API rejects integer stop token + # IDs; it requires string stop sequences. Renderers like Qwen's return + # ``[151645]`` (an int token id) from ``get_stop_sequences()``. Decode + # via the tokenizer when one is supplied. + if stop and all(isinstance(s, int) for s in stop): + if tokenizer is None: + raise ValueError( + "Renderer returned integer stop token IDs but no tokenizer " + "was passed to decode them; the inference API only accepts " + "string stop sequences." + ) + stop = [tokenizer.decode([s], skip_special_tokens=False) for s in stop] + + sk: dict[str, Any] = dict(sample_kwargs or {}) + sk["n"] = 1 + sk["stop"] = stop + if max_tokens is not None: + sk["max_tokens"] = max_tokens + + completions = await sample_with_prompt_tokens(prompt_token_ids, **sk) + if not completions: + return None + + c = completions[0] + prompt_len = int(c.prompt_len) + out_tokens: List[int] = list(c.full_tokens[prompt_len:]) + if not out_tokens: + return None + out_logprobs_raw = getattr(c, "inference_logprobs", None) + if out_logprobs_raw is None: + logger.warning( + "single_turn_renderer_rollout: dropping completion with " + "no inference_logprobs (got None). Configure the sampler " + "with logprobs=True so PPO/GRPO ratio/KL math sees real " + "behavior-policy probabilities." + ) + return None + out_logprobs: List[float] = list(out_logprobs_raw) + if getattr(c, "logprobs_echoed", False) and len(out_logprobs) == prompt_len + len(out_tokens): + out_logprobs = out_logprobs[prompt_len:] + if len(out_logprobs) != len(out_tokens): + logger.warning( + "single_turn_renderer_rollout: dropping completion with " + "misaligned logprobs (got %d, expected %d for assistant tokens).", + len(out_logprobs), len(out_tokens), + ) + return None + + parsed_message, parse_success = renderer.parse_response(out_tokens) + reward = await reward_fn(row, parsed_message, bool(parse_success)) + if reward is None: + return None + + return RolloutSample( + tokens=list(prompt_token_ids) + out_tokens, + logprobs=[0.0] * len(prompt_token_ids) + out_logprobs, + loss_mask=[0] * len(prompt_token_ids) + [1] * len(out_tokens), + reward=float(reward), + finish_reason=getattr(c, "finish_reason", "stop"), + text=getattr(c, "text", ""), + ) diff --git a/training/utils/rl/rollout/service.py b/training/utils/rl/rollout/service.py new file mode 100644 index 00000000..ce558f0c --- /dev/null +++ b/training/utils/rl/rollout/service.py @@ -0,0 +1,117 @@ +"""Service-agnostic rollout protocol. + +The async RL recipe has exactly one extension point -- ``rollout_fn`` -- +and everything inside it is user territory. In practice most users plug +in *some* remote service (an agent framework, a RAG-with-verifier stack, +an LLM-as-judge loop, eval-protocol, ...). This module defines the +shape that lives between the service adapter and the generic packer in +:mod:`training.utils.rl.rollout`, so the integration split is: +*service adapter* on one side, *trainer packing* on the other, with +neither importing the other's dependencies. + +No external deps. In particular, this module does **not** import +``eval_protocol`` or any SDK -- pick it up from a plain service class +and the cookbook has no opinion on which one. + +Token-native contract +--------------------- + +A :class:`RolloutPayload` is **token-native**: every :class:`TurnRecord` +carries ``token_ids`` and every assistant turn carries per-token +``logprobs``, both straight from the same inference call that generated +them (slime/AReaL convention). The trainer never re-tokenizes text; +re-tokenization silently misaligns the loss mask and inference +logprobs. Services that today only emit text (e.g. EP's +``RemoteRolloutProcessor``) must grow a token-native trace before they +can drive RL training; see +``https://github.com/fw-ai/fireworks/issues/23756``. + +Use :class:`training.utils.rl.rollout.TrajectoryAssembler` +to build :class:`RolloutPayload` from multi-turn engine calls. It +carries the AReaL prefix-equality invariant for free and sets +``_assembled=True`` so the packer skips its defensive check. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable, List, Literal, Optional, Protocol + + +__all__ = [ + "TurnRecord", + "RolloutPayload", + "RolloutService", +] + + +Role = Literal["system", "user", "assistant", "tool"] + + +@dataclass +class TurnRecord: + """One conversational turn, token-native. + + ``token_ids`` is required on every turn. ``logprobs`` is required + on assistant turns and must align 1:1 with ``token_ids``; both must + come from the same inference call that generated them. ``text`` is + optional and used only for human-readable logging. + """ + + role: Role + text: str = "" + token_ids: Optional[List[int]] = None + logprobs: Optional[List[float]] = None + finish_reason: Optional[str] = None + + +@dataclass +class RolloutPayload: + """One completion's worth of rollout data, service-emitted. + + ``total_reward`` carries the "server wins" convention: when set, the + trainer trusts it; when ``None``, the packer expects the caller to + attach a reward (e.g. by grading downstream). This deprecates + in-band reward sentinels. + """ + + turns: List[TurnRecord] + total_reward: Optional[float] = None + tokenizer_id: Optional[str] = None + """Identifier for the tokenizer that produced ``token_ids``. When + set, the packer asserts it matches the trainer's tokenizer so a + mismatched-vocab integration fails loud instead of training on the + wrong token IDs.""" + finish_reason: str = "stop" + extras: dict = field(default_factory=dict) + _assembled: bool = False + """Set by :class:`TrajectoryAssembler.to_payload`. Tells the packer + that the per-turn ``token_ids`` were stitched with the prefix-equality + invariant already enforced, so the defensive consistency check can be + skipped. Hand-built payloads default to ``False`` and get the check.""" + + +class RolloutService(Protocol): + """What the cookbook needs from any rollout backend. + + Given a dataset row's prompt messages, return ``n`` completed + rollouts. Everything else -- multi-turn loops, tool calls, + retries, grading -- is the service's business. + """ + + async def rollout( + self, + messages: List[dict], + *, + n: int, + sample_kwargs: dict[str, Any], + row: dict, + ) -> List[RolloutPayload]: ... + + +RolloutServiceCallable = Callable[ + [List[dict], int, dict[str, Any], dict], + Awaitable[List[RolloutPayload]], +] +"""Plain-function form of :class:`RolloutService`. Either shape works +with :func:`training.utils.rl.rollout.make_remote_rollout_fn`.""" diff --git a/training/utils/rl/rollout/trace.py b/training/utils/rl/rollout/trace.py new file mode 100644 index 00000000..11895b67 --- /dev/null +++ b/training/utils/rl/rollout/trace.py @@ -0,0 +1,330 @@ +"""Native rollout trajectory analysis. + +Rollout correctness is a trajectory property, not a renderer artifact +property. This module keeps the primary shape close to the RL data model: +turns, token rows, and issues derived from token-native rollout data. + +The verifier UI can visualize this shape directly and reuse its rule engine +over ``trajectory.tokens``. No live probe or renderer round trip is needed. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any, Iterable, Mapping, Sequence + + +@dataclass +class TrajectoryIssue: + code: str + severity: str + message: str + turn_idx: int | None = None + token_idx: int | None = None + + +@dataclass +class TrajectoryToken: + idx: int + token_id: int + decoded: str = "" + role: str = "" + turn_idx: int | None = None + source: str = "" + train_weight: float = 0.0 + provenance: str = "prompt_hard_append" + issues: list[str] = field(default_factory=list) + + def to_viewer_row(self) -> dict[str, Any]: + return { + "idx": self.idx, + "token_id": self.token_id, + "decoded": self.decoded, + "chunk_source": self.source, + "msg_idx": None, + "turn_idx": self.turn_idx, + "role": self.role, + "renderer_claim_weight": self.train_weight, + "provenance": self.provenance, + "issues": list(self.issues), + } + + +@dataclass +class RolloutTrajectory: + source: str + tokens: list[TrajectoryToken] + issues: list[TrajectoryIssue] = field(default_factory=list) + turns: list[dict[str, Any]] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + + def summary(self) -> dict[str, Any]: + return { + "source": self.source, + "turn_count": len(self.turns), + "token_count": len(self.tokens), + "generated_token_count": sum(1 for t in self.tokens if t.train_weight > 0.0), + "issue_count": len(self.issues), + "prefix_mismatch_count": sum(1 for i in self.issues if i.code == "prefix_mismatch"), + "logprob_mismatch_count": sum(1 for i in self.issues if i.code == "logprob_mismatch"), + } + + def to_dict(self) -> dict[str, Any]: + return { + "kind": "rollout_trajectory", + "source": self.source, + "metadata": dict(self.metadata), + "summary": self.summary(), + "turns": list(self.turns), + "tokens": [t.to_viewer_row() for t in self.tokens], + "issues": [asdict(i) for i in self.issues], + } + + +def analyze_token_turn_traces( + token_turn_traces: Sequence[Mapping[str, Any]], + *, + tokenizer: Any | None = None, + source: str = "token_turn_traces", + metadata: Mapping[str, Any] | None = None, +) -> RolloutTrajectory: + """Analyze per-engine-call traces with the prefix-extension invariant.""" + tokens: list[TrajectoryToken] = [] + issues: list[TrajectoryIssue] = [] + turns: list[dict[str, Any]] = [] + seq: list[int] = [] + + def add_token_rows( + ids: Sequence[int], + *, + turn_idx: int, + role: str, + source_name: str, + train_weight: float, + provenance: str, + row_issues: list[str] | None = None, + ) -> None: + for tok_id in ids: + tokens.append( + TrajectoryToken( + idx=len(tokens), + token_id=int(tok_id), + decoded=_decode_one(tokenizer, int(tok_id)), + role=role, + turn_idx=turn_idx, + source=source_name, + train_weight=train_weight, + provenance=provenance, + issues=list(row_issues or []), + ) + ) + + for turn_idx, trace in enumerate(token_turn_traces): + prompt_ids = [int(x) for x in trace.get("prompt_ids") or []] + completion_ids = [int(x) for x in trace.get("completion_ids") or []] + completion_logprobs = trace.get("completion_logprobs") + turns.append( + { + "turn_idx": turn_idx, + "prompt_len": len(prompt_ids), + "completion_len": len(completion_ids), + "finish_reason": trace.get("finish_reason"), + } + ) + + if completion_logprobs is not None and len(list(completion_logprobs)) != len(completion_ids): + issues.append( + TrajectoryIssue( + code="logprob_mismatch", + severity="error", + message=( + f"turn {turn_idx} has {len(completion_ids)} completion tokens " + f"but {len(list(completion_logprobs))} logprobs" + ), + turn_idx=turn_idx, + ) + ) + + if len(prompt_ids) >= len(seq) and prompt_ids[: len(seq)] == seq: + gap = prompt_ids[len(seq) :] + add_token_rows( + gap, + turn_idx=turn_idx, + role="user", + source_name="prompt_delta", + train_weight=0.0, + provenance="prompt_hard_append", + ) + seq.extend(gap) + else: + idx = _first_divergence(seq, prompt_ids) + issues.append( + TrajectoryIssue( + code="prefix_mismatch", + severity="error", + message=( + f"turn {turn_idx} prompt does not extend accumulated trajectory " + f"at index {idx}" + ), + turn_idx=turn_idx, + token_idx=idx, + ) + ) + add_token_rows( + prompt_ids, + turn_idx=turn_idx, + role="user", + source_name="prompt_full_after_mismatch", + train_weight=0.0, + provenance="tokenization_diverged", + row_issues=["prefix_mismatch"], + ) + seq = list(prompt_ids) + + add_token_rows( + completion_ids, + turn_idx=turn_idx, + role="assistant", + source_name="completion", + train_weight=1.0, + provenance="native_generated", + ) + seq.extend(completion_ids) + + return RolloutTrajectory( + source=source, + tokens=tokens, + issues=issues, + turns=turns, + metadata=dict(metadata or {}), + ) + + +def analyze_turns( + turns: Iterable[Any], + *, + tokenizer: Any | None = None, + source: str = "rollout_payload", + metadata: Mapping[str, Any] | None = None, +) -> RolloutTrajectory: + """Analyze ``RolloutPayload.turns`` or equivalent dictionaries.""" + out_tokens: list[TrajectoryToken] = [] + issues: list[TrajectoryIssue] = [] + out_turns: list[dict[str, Any]] = [] + for turn_idx, turn in enumerate(turns): + role = str(_get(turn, "role") or "") + token_ids = [int(x) for x in (_get(turn, "token_ids") or [])] + logprobs = _get(turn, "logprobs") + out_turns.append({"turn_idx": turn_idx, "role": role, "token_len": len(token_ids)}) + if role == "assistant" and (logprobs is None or len(list(logprobs)) != len(token_ids)): + issues.append( + TrajectoryIssue( + code="logprob_mismatch", + severity="error", + message=( + f"assistant turn {turn_idx} has {len(token_ids)} tokens " + f"but {0 if logprobs is None else len(list(logprobs))} logprobs" + ), + turn_idx=turn_idx, + ) + ) + train_weight = 1.0 if role == "assistant" else 0.0 + for tok_id in token_ids: + out_tokens.append( + TrajectoryToken( + idx=len(out_tokens), + token_id=tok_id, + decoded=_decode_one(tokenizer, tok_id), + role=role, + turn_idx=turn_idx, + source="turn", + train_weight=train_weight, + provenance="native_generated" if train_weight else "prompt_hard_append", + ) + ) + + if out_turns and out_turns[-1]["role"] != "assistant": + issues.append( + TrajectoryIssue( + code="terminal_non_assistant", + severity="warning", + message=f"trajectory ends with role={out_turns[-1]['role']!r}", + turn_idx=int(out_turns[-1]["turn_idx"]), + ) + ) + + return RolloutTrajectory( + source=source, + tokens=out_tokens, + issues=issues, + turns=out_turns, + metadata=dict(metadata or {}), + ) + + +def analyze_flat_sample( + sample: Any, + *, + tokenizer: Any | None = None, + source: str = "rollout_sample", + metadata: Mapping[str, Any] | None = None, +) -> RolloutTrajectory: + """Analyze a flat ``RolloutSample`` or equivalent dictionary.""" + ids = [int(x) for x in (_get(sample, "tokens") or [])] + mask = [float(x) for x in (_get(sample, "loss_mask") or [])] + issues: list[TrajectoryIssue] = [] + if len(ids) != len(mask): + issues.append( + TrajectoryIssue( + code="length_mismatch", + severity="error", + message=f"tokens/loss_mask length mismatch: {len(ids)} != {len(mask)}", + ) + ) + + tokens: list[TrajectoryToken] = [] + for idx, tok_id in enumerate(ids): + train_weight = mask[idx] if idx < len(mask) else 0.0 + tokens.append( + TrajectoryToken( + idx=idx, + token_id=tok_id, + decoded=_decode_one(tokenizer, tok_id), + role="assistant" if train_weight > 0.0 else "context", + turn_idx=None, + source="loss_mask", + train_weight=train_weight, + provenance="native_generated" if train_weight > 0.0 else "prompt_hard_append", + ) + ) + + return RolloutTrajectory( + source=source, + tokens=tokens, + issues=issues, + turns=[], + metadata=dict(metadata or {}), + ) + + +def _decode_one(tokenizer: Any | None, tok_id: int) -> str: + if tokenizer is None: + return "" + try: + return tokenizer.decode([tok_id], skip_special_tokens=False) + except Exception: # noqa: BLE001 + return "" + + +def _first_divergence(a: Sequence[int], b: Sequence[int]) -> int: + n = min(len(a), len(b)) + for idx in range(n): + if a[idx] != b[idx]: + return idx + return n + + +def _get(obj: Any, key: str) -> Any: + if isinstance(obj, Mapping): + return obj.get(key) + return getattr(obj, key, None) diff --git a/training/utils/rl/rollout/types.py b/training/utils/rl/rollout/types.py new file mode 100644 index 00000000..53ca4b03 --- /dev/null +++ b/training/utils/rl/rollout/types.py @@ -0,0 +1,232 @@ +"""Flat, mask-native rollout sample contract. + +The user's ``rollout_fn(sample_prompt)`` hands back one +:class:`RolloutSample` per call -- one trajectory. Each sample is three parallel lists +(``tokens``, ``logprobs``, ``loss_mask``) plus a scalar reward. Multi-turn +rollouts flatten into the same shape: turn boundaries are implicit in +``loss_mask`` transitions (0 on prompts / env feedback / tool responses, +1 on assistant-generated tokens). + +This matches the user-facing contract used by AReaL and slime. Group +assembly (collecting N samples per row and computing GRPO-style +advantages) happens framework-side via :class:`GroupAssembler` -- see +:mod:`training.utils.rl.rollout.group_assembler`. + +:class:`Rollout` is the internal group representation that the +:func:`rollout_to_prompt_group` adapter consumes after the assembler +joins N samples by row. The adapter emits ``tinker.Datum`` objects +with a per-token ``loss_mask`` in ``loss_fn_inputs``; the existing +loss kernels already honour this (see ``_get_loss_mask`` in +``training/utils/rl/common.py``). +""" + +from __future__ import annotations + +import logging +import math +from dataclasses import dataclass +from typing import Callable, List + +import tinker + +from training.utils.data import compute_advantages +from training.utils.rl.losses import PromptGroup +from training.utils.rl.router_replay import build_r3_routing_matrices + +logger = logging.getLogger(__name__) + +__all__ = [ + "RolloutSample", + "Rollout", + "rollout_to_prompt_group", +] + + +@dataclass +class RolloutSample: + """One completion's flat, trainer-ready data. + + The three parallel lists MUST have identical length. ``loss_mask`` + is ``1`` on assistant-generated positions (trained on) and ``0`` + everywhere else (prompt, user messages, tool responses, env feedback + injected between turns). ``logprobs`` is the per-token inference + logprob aligned with ``tokens``; use ``0.0`` on non-generated + positions since they carry no training signal. + """ + + tokens: List[int] + logprobs: List[float] + loss_mask: List[int] + reward: float + routing_matrices: List[str] | None = None + """Optional per-token MoE routing matrices captured from the inference + deployment (R3 / Router Replay). When set, the adapter aligns them to + ``tokens[:-1]`` via :func:`build_r3_routing_matrices` and threads them + through ``tinker.ModelInput.from_ints(routing_matrices=...)`` so the + training step replays the same expert routing decisions made at + inference time. Length should match ``tokens`` (echo mode) or the + completion suffix (legacy mode); the alignment helper handles both.""" + finish_reason: str = "stop" + text: str = "" + """Decoded assistant output. For logging only; not consumed by the + adapter or the trainer.""" + + +@dataclass +class Rollout: + """One row's worth of completions (one GRPO group).""" + + samples: List[RolloutSample] + row_meta: dict | None = None + + +def _validate(rollout: Rollout) -> None: + if not rollout.samples: + raise ValueError("Rollout.samples is empty") + for i, s in enumerate(rollout.samples): + n = len(s.tokens) + if len(s.logprobs) != n or len(s.loss_mask) != n: + raise ValueError( + f"Sample {i}: tokens/logprobs/loss_mask length mismatch " + f"({n} / {len(s.logprobs)} / {len(s.loss_mask)}). All three " + "lists must be the same length.", + ) + if n < 2: + raise ValueError(f"Sample {i}: tokens must have length >= 2.") + if not any(m > 0 for m in s.loss_mask): + raise ValueError( + f"Sample {i}: loss_mask is all zeros -- no tokens would be " + "trained on. Set loss_mask=1 on assistant-generated " + "positions.", + ) + + +def rollout_to_prompt_group( + rollout: Rollout, + *, + advantage_fn: Callable[[List[float]], List[float]] = compute_advantages, + with_reference: bool = False, + router_replay_completion_only: bool = False, +) -> PromptGroup | None: + """Pack one :class:`Rollout` into a :class:`PromptGroup`. + + Per-token ``loss_mask`` flows through ``tinker.Datum.loss_fn_inputs`` + so the existing kernels (``_get_loss_mask`` in + ``training/utils/rl/common.py``) mask prompt / env / tool tokens + correctly without any trainer-side changes. + + When a sample carries ``routing_matrices``, they are aligned to the + model_input shape via :func:`build_r3_routing_matrices` and passed + through ``tinker.ModelInput.from_ints(routing_matrices=...)`` so MoE + expert routing is replayed at training time (R3). + ``router_replay_completion_only=True`` zeroes out prompt-position + routing matrices so only completion-token routing is replayed (the + server picks its own routing for prompt tokens). Reference-side + datums never carry routing matrices. + + Returns ``None`` when the group has no samples; raises on structural + issues (mismatched lengths, empty samples, all-zero loss mask). + """ + _validate(rollout) + + rewards = [s.reward for s in rollout.samples] + advantages = list(advantage_fn(list(rewards))) + + # Drop on non-finite advantages (e.g., GRPO z-score on a length-1 group). + # Don't precheck on sample count -- REINFORCE (cpp=1, lambda r: r) is valid. + if any(not math.isfinite(a) for a in advantages): + logger.warning( + "rollout_to_prompt_group: dropping rollout (N=%d) -- " + "advantage_fn produced non-finite advantages %r. This " + "typically happens when the default GRPO-style " + "``compute_advantages`` z-score normalizer runs on a " + "single-sample group (std of a length-1 tensor is " + "undefined). For REINFORCE-style runs with " + "completions_per_prompt=1, pass a single-sample-safe " + "advantage_fn (e.g. ``lambda r: r``).", + len(rollout.samples), advantages, + ) + return None + + policy_data: List[tinker.Datum] = [] + reference_data: List[tinker.Datum] = [] + inf_logprobs_aligned: List[List[float]] = [] + completion_lens: List[int] = [] + truncated: List[bool] = [] + per_sample_prompt_lens: List[int] = [] + + # Datum predicts tokens[1:] from tokens[:-1]; shift both loss_mask and + # logprobs to match target positions. + for s in rollout.samples: + n = len(s.tokens) + target_len = n - 1 + + target_tokens = s.tokens[1:] + target_mask = s.loss_mask[1:] + target_logprobs = s.logprobs[1:] + + # Per-sample prompt boundary: index of the first assistant + # (loss_mask=1) token. Heterogeneous rollouts (multi-turn, + # tool branches) can have different prefix lengths per sample, + # so the single ``PromptGroup.prompt_len`` is wrong for them. + sample_prompt_len = next( + (i for i, m in enumerate(s.loss_mask) if m > 0), + 0, + ) + per_sample_prompt_lens.append(sample_prompt_len) + + rm = None + if s.routing_matrices is not None: + rm = build_r3_routing_matrices( + s.routing_matrices, + prompt_len=sample_prompt_len, + model_input_len=target_len, + completion_only=router_replay_completion_only, + ) + + policy_data.append(tinker.Datum( + model_input=tinker.ModelInput.from_ints(s.tokens[:-1], routing_matrices=rm), + loss_fn_inputs={ + "target_tokens": tinker.TensorData( + data=target_tokens, dtype="int64", shape=[target_len], + ), + "weights": tinker.TensorData( + data=target_mask, dtype="int64", shape=[target_len], + ), + }, + )) + + if with_reference: + reference_data.append(tinker.Datum( + model_input=tinker.ModelInput.from_ints(s.tokens[:-1]), + loss_fn_inputs={ + "target_tokens": tinker.TensorData( + data=target_tokens, dtype="int64", shape=[target_len], + ), + "loss_mask": tinker.TensorData( + data=target_mask, dtype="int64", shape=[target_len], + ), + }, + )) + + inf_logprobs_aligned.append(target_logprobs) + completion_lens.append(sum(1 for m in s.loss_mask if m > 0)) + truncated.append(s.finish_reason == "length") + + # ``prompt_len`` is the legacy scalar (back-compat); ``prompt_lens`` + # is the authoritative per-sample list for heterogeneous rollouts. + return PromptGroup( + data=policy_data, + ref_data=reference_data, + advantages=list(advantages), + ref_logprobs=None, + prompt_len=per_sample_prompt_lens[0], + rewards=rewards, + inf_logprobs=inf_logprobs_aligned, + completion_lens=completion_lens, + truncated=truncated, + prompt=None, + completions=None, + row_meta=dict(rollout.row_meta) if rollout.row_meta else None, + prompt_lens=per_sample_prompt_lens, + ) diff --git a/training/utils/rl/train.py b/training/utils/rl/train.py index ce7b6d1b..3da37d23 100644 --- a/training/utils/rl/train.py +++ b/training/utils/rl/train.py @@ -113,7 +113,7 @@ def _make_stats( "total_sampled": total_sampled, "filter_drops": filter_drops, "sample_fails": sample_fails, - "sample_wait_time": 0.0, + "trainer_wait_for_sampler_time": 0.0, "step_wall_time": wall, "all_raw_rewards": list(all_raw_rewards), } diff --git a/training/utils/timer.py b/training/utils/timer.py index fdc3d570..e1750c66 100644 --- a/training/utils/timer.py +++ b/training/utils/timer.py @@ -29,6 +29,7 @@ def save_model(...): from typing import Any from functools import wraps from contextlib import contextmanager +from dataclasses import dataclass logger = logging.getLogger(__name__) @@ -51,12 +52,13 @@ def start(self, name: str) -> None: logger.warning("Timer '%s' already running -- restarting", name) self._start_time[name] = time.time() - def end(self, name: str) -> None: + def end(self, name: str) -> float | None: if name not in self._start_time: logger.warning("Timer '%s' was never started -- ignoring end", name) - return + return None elapsed = time.time() - self._start_time.pop(name) self.timers[name] = self.timers.get(name, 0.0) + elapsed + return elapsed def add(self, name: str, elapsed: float) -> None: self.timers[name] = self.timers.get(name, 0.0) + elapsed @@ -106,6 +108,22 @@ def wrapper(*args, **kwargs): return wrapper +@dataclass +class TimerSpan: + elapsed: float = 0.0 + + +@contextmanager +def elapsed_timer(name: str): + span = TimerSpan() + t = Timer() + t.start(name) + try: + yield span + finally: + span.elapsed = t.end(name) or 0.0 + + @contextmanager def inverse_timer(name: str): """Measure the gap *between* operations (e.g. wait time). diff --git a/training/utils/validation.py b/training/utils/validation.py index 8449bc30..392429a3 100644 --- a/training/utils/validation.py +++ b/training/utils/validation.py @@ -41,12 +41,19 @@ def validate_output_model_id(_value: str) -> list[str]: def validate_config( base_model: str, - dataset: str, + dataset: str | None, hotload: WeightSyncConfig | None = None, deploy: DeployConfig | None = None, output_model_id: str | None = None, + *, + require_dataset: bool = True, ) -> None: - """Pre-flight validation. Catches misconfiguration before provisioning GPUs.""" + """Pre-flight validation. Catches misconfiguration before provisioning GPUs. + + ``require_dataset=False`` skips the dataset check so async recipes that + accept ``rows=`` directly (no JSONL path/URL) can still run the + base_model and output_model_id checks. + """ errors: list[str] = [] if not base_model: @@ -67,7 +74,7 @@ def validate_config( ) ) - if not dataset: + if require_dataset and not dataset: errors.append( format_sdk_error( "Missing dataset",