docs(training): align cookbook docs and skill references

skills/dev/SKILL.md

@@ -16,7 +16,7 @@ The cookbook is the reference implementation of the Fireworks Training SDK. Fork
 | "How do I set up / install the cookbook?" | [`references/setup.md`](references/setup.md) |
 | "I want to run something out of the box" | [`references/examples.md`](references/examples.md) |
 | "I want to fork a recipe and edit the Config" | [`references/recipes.md`](references/recipes.md) |
-| "How do I set the training / deployment shape?" | [`references/shapes.md`](references/shapes.md) |
+| "How do I set or override the training / deployment shape?" | [`references/shapes.md`](references/shapes.md) |
 | `RuntimeError: Failed to resolve latest validated training shape` | [`references/shapes.md`](references/shapes.md#when-resolve_training_profile-raises-failed-to-resolve-latest-validated-training-shape) — don't pin a version; retry or reach out |
 | "Can I run two deployments off one trainer (sampler + eval)?" | [`references/rl/hotload.md`](references/rl/hotload.md#two-deployments-one-trainer) |
 | "How does RL dispatch server-side vs client-side loss? What's the cost?" | [`references/rl/loss-paths.md`](references/rl/loss-paths.md) |
@@ -62,7 +62,7 @@ The requirement lives in the cookbook's `training/pyproject.toml` — look for t
 
 ```bash
 grep 'fireworks-ai\[training\]' cookbook/training/pyproject.toml
-# e.g. "fireworks-ai[training]>=1.0.0a62,<2"
+# e.g. "fireworks-ai[training]>=1.1.0a64,<2"
 
 pip show fireworks-ai | grep -i version
 ```
@@ -71,7 +71,7 @@ If the installed version doesn't satisfy the pin, upgrade first and retry. Only
 
 ## Non-negotiables
 
-1. **Shape first.** `cfg.infra.training_shape_id` is required. The deployment shape comes from the profile. Manual infra fields are a mistake; the backend will reject or ignore them. See [`references/shapes.md`](references/shapes.md).
+1. **Shape first.** Leave `cfg.infra.training_shape_id` unset to auto-select a validated shape, or set it to a bare training-shape resource to override. The deployment shape comes from the profile. Manual infra fields are a mistake; the backend will reject or ignore them. See [`references/shapes.md`](references/shapes.md).
 2. **`WeightSyncScope.PER_TRAINER` is the default.** Set `DeployConfig(weight_sync_scope=WeightSyncScope.PER_TRAINER)` (the default). Do not combine it with `hot_load_deployment_id` — that field belongs to `PER_DEPLOYMENT`. Pick one bucket scope. See [`references/rl/hotload.md`](references/rl/hotload.md#weight-sync-scope-per_trainer-vs-per_deployment).
 3. **Fork, don't reinvent.** Training loop plumbing lives in `training/recipes/`. Fork the file that matches the task; do not rewire `FiretitanTrainingClient` / `DeploymentManager` / `WeightSyncer` from scratch.
 4. **Validate `output_model_id` before promote.** Server cap is 63 chars, charset `[a-z0-9-]`. A rejected promote orphans the sampler blob; the same `checkpoint_id` returns "not found in GCS" after GC. See [`references/checkpoints.md`](references/checkpoints.md#output_model_id-validation).

skills/dev/references/checkpoints.md

@@ -63,15 +63,15 @@ Priority inside `TrainingCheckpoints.resume` (highest first):
 - `warm_start_from_adapter` and `init_from_checkpoint` are mutually exclusive.
 - Requires `lora_rank > 0`. Full-parameter continue-training uses `base_model` instead.
 
-Supported in all recipe loops: `sft_loop`, `dpo_loop`, `orpo_loop`, `rl_loop`, `igpo_loop`.
+Supported in all recipe loops: `sft_loop`, `dpo_loop`, `orpo_loop`, `rl_loop`, `async_rl_loop`.
 
 ## Cross-run resume
 
 Auto-resume (priority 2) is **scoped to one trainer job**. If you re-run with the same `log_path` but provision a fresh trainer, the new trainer's `list_checkpoints` is empty and resume falls through to fresh start.
 
 To resume across separate `main()` invocations, either:
 
-- Pin both runs to the same trainer via `cfg.trainer_job_id` (SFT/DPO/ORPO) or `cfg.policy_job_id` + `cfg.reference_job_id` (RL/IGPO). The second run reattaches and the CP rows are visible.
+- Pin both runs to the same trainer via `cfg.trainer_job_id` (SFT/DPO/ORPO) or `cfg.policy_job_id` + `cfg.reference_job_id` (RL). The second run reattaches and the CP rows are visible.
 - Or use `init_from_checkpoint=f"{prior_job_id}:step-N"` for explicit cross-job DCP load. This resets the step counter to 0 — fine for warm-start scenarios, not for "continue training and skip to step N".
 
 ---

skills/dev/references/examples.md

@@ -1,31 +1,36 @@
 # Out-of-the-box examples
 
-When the user wants "something that just runs", point them at `training/examples/`. Each subdirectory imports from `training/recipes/` and ships a ready-to-run `Config`.
+When the user wants "something that just runs", point them at `training/examples/`. Most examples import from `training/recipes/` and ship a ready-to-run `Config`; the larger agentic examples may own a custom loop plus rollout code.
 
 | Task | Path |
 |------|------|
-| Minimal SFT (single-file demo) | `training/examples/sft_getting_started/` |
 | SFT on real datasets | `training/examples/sft/` |
 | DPO | `training/examples/dpo/` |
 | ORPO | `training/examples/orpo/` |
-| GRPO on deepmath | `training/examples/deepmath_rl/` |
-| Tool-use RL (frozen lake) | `training/examples/frozen_lake/` |
-| Multi-hop QA RL | `training/examples/multihop_qa/` |
-| Generic RL wiring | `training/examples/rl/` |
+| GRPO on DeepMath | `training/examples/rl/deepmath/` |
+| Tool-use RL (Frozen Lake) | `training/examples/rl/frozen_lake/` |
+| Async single-turn RL (renderer-backed) | `training/examples/rl/single_turn_async/` |
+| Async single-turn RL on GSM8K | `training/examples/rl/gsm8k_async/` |
+| Multi-turn RL (token-native, no renderer) | `training/examples/rl/multi_turn_minimal/` |
+| Multi-turn renderer RL | `training/examples/rl/multi_turn_minimal_renderer/` |
+| Multi-turn tool RL | `training/examples/rl/multi_turn_tool/` |
+| Remote rollout service RL | `training/examples/rl/remote_rollout/` |
+| Eval-protocol remote grader RL | `training/examples/rl/ep_remote_grader/` |
+| Multi-hop QA IGPO | `training/examples/multihop_qa/` |
 
 ## What to read
 
-In each example directory, read the top-level `train_*.py` (or `main.py`) — it builds the `Config` and calls `main(config)` from the corresponding recipe. That's the whole script.
+In each example directory, read the top-level `train_*.py`, `train.py`, or `rollout.py` — it builds the `Config` or rollout function and calls the corresponding recipe/helper.
 
 ## Running
 
 ```bash
 cd cookbook/training
-pip install -e .
-python examples/sft_getting_started/<script>.py
+pip install --pre -e .
+python examples/sft/train_sft.py --output-model-id <model-id>
 ```
 
-Every example pulls `FIREWORKS_API_KEY` from the env (or `.env`). If the example needs a dataset path, it is listed in its `README.md`.
+Every example pulls `FIREWORKS_API_KEY` from the env (or `.env`). If an example needs a dataset path or extra arguments, check that script's `argparse` block, `README.md`, or `run.sh`.
 
 ## When the example isn't quite right
 

skills/dev/references/recipes.md

@@ -7,29 +7,28 @@ Each recipe is a single Python file in `training/recipes/` that wires the Traini
 | SFT | `training/recipes/sft_loop.py` |
 | DPO | `training/recipes/dpo_loop.py` |
 | 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) |
+| Async RL loop (rollout/train overlap, PPO inner minibatches, custom per-token losses such as IGPO) | `training/recipes/async_rl_loop.py` — see [`rl/async-rl.md`](rl/async-rl.md) |
 
 ## "Reference loop" means these files
 
 They are the canonical wiring of `FiretitanTrainingClient` + `DeploymentManager` + `WeightSyncer`. Do not rewrite — fork.
 
 ## What to fill in on `Config`
 
-Always required on `Config` + `InfraConfig`:
+Always required on `Config`:
 
 - `base_model` — `accounts/fireworks/models/<name>`
 - `dataset` — path to JSONL
 - `tokenizer_model` — HF model name
 - `log_path` — directory for `dataloader.json` and logs
-- `infra.training_shape_id` — **required**; do not set manual `accelerator_type` / `node_count` (see [`shapes.md`](shapes.md))
+- `infra.training_shape_id` — optional override. Leave unset to auto-select a validated shape; do not set manual `accelerator_type` / `node_count` (see [`shapes.md`](shapes.md)).
 
 RL-specific (in `rl_loop.py`'s `Config`): reward function, rollout batch sizes, deployment config (shape is auto-filled from the profile).
 
 ## Resume
 
-Auto-resume is scoped to one trainer. Pin both runs to the same trainer via `cfg.trainer_job_id` (SFT/DPO/ORPO) or `cfg.policy_job_id` + `cfg.reference_job_id` (RL/IGPO), keep the same `log_path`, and rerun. `TrainingCheckpoints.resume()` lists the trainer's checkpoints on the control plane, picks the newest resumable row, and restores the rollout cursor from `dataloader.json`. See [`checkpoints.md`](checkpoints.md) for the full priority order and constraints.
+Auto-resume is scoped to one trainer. Pin both runs to the same trainer via `cfg.trainer_job_id` (SFT/DPO/ORPO) or `cfg.policy_job_id` + `cfg.reference_job_id` (RL), keep the same `log_path`, and rerun. `TrainingCheckpoints.resume()` lists the trainer's checkpoints on the control plane, picks the newest resumable row, and restores the rollout cursor from `dataloader.json`. See [`checkpoints.md`](checkpoints.md) for the full priority order and constraints.
 
 ## Init from another job
 

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

@@ -32,13 +32,56 @@ loop, checkpoint plumbing — is handled by `recipes/async_rl_loop.py::main`.
 |---|---|
 | `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 |
+| (advanced) `prepare_prompt_groups_fn` + `client_loss_fn_builder` for per-token custom losses such as IGPO | Infra, async gate, checkpoints, weight sync, PPO minibatches |
 | 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`.
 
+## Advanced train-step hooks
+
+Use hooks only when the algorithm needs batch-level work after rollouts land.
+IGPO is the motivating case: score each completed trajectory for information
+gain, write per-token advantages into `PromptGroup.row_meta`, then build a
+custom `forward_backward_custom` loss from those per-token advantages.
+
+```python
+def prepare_prompt_groups_fn(groups, ctx) -> dict[str, float]:
+    # Mutate groups in place, e.g. groups[i].row_meta["per_token_advantages"] = ...
+    return {"igpo/mean_ig": mean_ig}
+
+def client_loss_fn_builder(
+    groups, data, advantages, ref_logprobs, prompt_lens,
+    inf_logprobs, old_policy_logprobs, ctx,
+):
+    # For PPO minibatches, ctx.minibatch_start / ctx.minibatch_end identify
+    # the slice of the flattened rollout batch represented by `data`.
+    return make_igpo_async_loss_fn(
+        groups,
+        ref_logprobs=ref_logprobs,
+        prompt_lens=prompt_lens,
+        inf_logprobs=inf_logprobs,
+        old_policy_logprobs=old_policy_logprobs,
+        kl_beta=ctx.config.kl_beta,
+        eps_clip=ctx.config.eps_clip,
+        start=ctx.minibatch_start,
+        end=ctx.minibatch_end,
+    )
+
+main(
+    cfg,
+    rollout_fn_factory=make_rollout_fn,
+    prepare_prompt_groups_fn=prepare_prompt_groups_fn,
+    client_loss_fn_builder=client_loss_fn_builder,
+)
+```
+
+Do not add a separate recipe loop for these algorithms. Keep rollout-specific
+logic in examples/utilities and reuse `async_rl_loop.py` for trainer lifecycle,
+checkpointing, and gate semantics.
+
 | | `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 |

skills/dev/references/setup.md

@@ -15,7 +15,7 @@ uv venv --python 3.12 && source .venv/bin/activate
 uv pip install --pre -e .
 ```
 
-`--pre` is required — the SDK (`fireworks-ai[training]`) is a prerelease. All dependencies (including `tinker-cookbook`) are pulled in automatically via `pyproject.toml`.
+`--pre` is required because the Training SDK dependency is pinned to a prerelease-compatible `fireworks-ai[training]` range. All dependencies are pulled in automatically via `pyproject.toml`.
 
 ## Credentials
 

skills/dev/references/shapes.md

@@ -1,10 +1,12 @@
 # Training and deployment shapes — always use a profile
 
-Shapes are the required entry point for both trainer and deployment. Never hand-set `accelerator_type`, `accelerator_count`, `node_count`, or `custom_image_tag` when a shape is in use — the backend will reject or silently ignore them.
+Shapes are the required entry point for both trainer and deployment. The cookbook can auto-select a validated training shape for the base model, LoRA mode, role, and context length. Set `cfg.infra.training_shape_id` only when you need to override that choice. Never hand-set `accelerator_type`, `accelerator_count`, `node_count`, or `custom_image_tag` when a shape is in use — the backend will reject or silently ignore them.
 
 ## Training shape
 
-Set `cfg.infra.training_shape_id`:
+Default: leave `cfg.infra.training_shape_id` unset and let `setup_infra` call `auto_select_training_shape(...)`.
+
+Override: set the bare training-shape resource:
 
 ```python
 cfg.infra.training_shape_id = "accounts/fireworks/trainingShapes/ts-qwen3-8b-policy"
@@ -20,7 +22,7 @@ profile = rlor_mgr.resolve_training_profile(cfg.infra.training_shape_id)
 # profile.accelerator_type, profile.node_count, ...  (read, do not copy to cfg)
 ```
 
-See `training/recipes/sft_loop.py` (search `resolve_training_profile`) and `training/recipes/rl_loop.py` (same — called once per policy, once per reference).
+See `training/utils/infra.py` (`_resolve_policy_shape`) and the recipe call sites in `training/recipes/sft_loop.py`, `training/recipes/dpo_loop.py`, `training/recipes/orpo_loop.py`, and `training/recipes/rl_loop.py`.
 
 ## Deployment shape
 
@@ -35,13 +37,13 @@ That is a **versioned** path (`accounts/fw/deploymentShapes/ds-x/versions/abc123
 
 ## Reference-model shape (RL / DPO)
 
-For **full-parameter** training with a frozen reference, set `cfg.infra.ref_training_shape_id` explicitly — there is no implicit fallback. It can share the same shape as the policy; the control plane appends `--forward-only` automatically.
+For **full-parameter** training with a frozen reference, leave `cfg.infra.ref_training_shape_id` unset to auto-select a validated forward-only reference shape, or set it explicitly to override. It can share the same shape as the policy; the control plane appends `--forward-only` automatically.
 
 For **LoRA** (`lora_rank > 0`), two valid options:
 - **Shared session (recommended, saves GPUs)**: leave `ref_training_shape_id` unset. `setup_infra` uses `policy.create_base_reference()` on the policy trainer for reference logprobs — no separate trainer, no extra GPUs.
 - **Separate LoRA-capable ref trainer**: set `ref_training_shape_id` to a `LORA_TRAINER` shape (typically the same as the policy shape). `setup_infra` provisions a forward-only LoRA ref trainer (its own GPUs) and forwards `lora_rank` to both the trainer request and the ref client so the gateway infers `trainer_mode=LORA_TRAINER` and matches the shape. CP's V2 DPO auto-resolver picks this path by default for LoRA DPO.
 
-The CI pattern for the saves-GPUs variant is `ref_shape = "" if lora_rank > 0 else <explicit shape>`.
+The CI pattern for the saves-GPUs LoRA variant is `ref_shape = ""` when `lora_rank > 0`.
 
 ## When to skip validation
 

skills/dev/references/tools.md

@@ -82,7 +82,7 @@ python training/examples/tools/list_checkpoints.py --job-id <job-id> --promotabl
 python training/examples/tools/list_checkpoints.py --job-id <job-id> --json              # machine-readable
 ```
 
-Requires `fireworks-ai[training] >= 1.0.0a62`. On older SDKs the method doesn't exist and the script will fail on import.
+Requires an installed `fireworks-ai[training]` version that satisfies `training/pyproject.toml` (currently `>=1.1.0a64,<2`). On older SDKs the method doesn't exist and the script will fail on import.
 
 ## `verify_logprobs.py`
 

training/README.md

@@ -10,6 +10,7 @@ Each recipe is a single Python file you can fork and customize.
 | Recipe | File | Description |
 | --- | --- | --- |
 | GRPO / IS / DAPO / DRO / GSPO / CISPO | `recipes/rl_loop.py` | On-policy RL with streaming rollouts. Set `policy_loss="grpo"`, `"importance_sampling"`, `"dapo"`, `"dro"`, `"gspo"`, or `"cispo"`. |
+| Async RL (any of the above losses, plus custom per-token losses such as IGPO) | `recipes/async_rl_loop.py` | Gate-native async RL: rollout/train overlap with bounded off-policy staleness. Strict superset of `rl_loop.py`; recommended for new RL work. |
 | DPO | `recipes/dpo_loop.py` | Direct preference optimization with cached reference logprobs. |
 | ORPO | `recipes/orpo_loop.py` | Odds-ratio preference optimization -- no reference model needed. |
 | SFT | `recipes/sft_loop.py` | Supervised fine-tuning with response-only cross-entropy loss. |
@@ -145,15 +146,23 @@ For detailed guides, configuration reference, and examples, see the official doc
 ## Directory layout
 
 ```
-recipes/                 Training loop scripts (fork these)
-utils/                   Shared config, data loading, loss functions, metrics
-examples/rl/deepmath/    Worked example: math reasoning with GRPO
-examples/rl/frozen_lake/ Worked example: Frozen Lake with tool-use RL
-examples/orpo/ifeval/    Worked example: IFEval with ORPO
-examples/sft/            Worked example: SFT getting started
-examples/dpo/            Worked example: DPO
-examples/snippets/       Standalone utility scripts
-tests/                   Unit and end-to-end tests
+recipes/                              Training loop scripts (fork these)
+utils/                                Shared config, data loading, loss functions, metrics
+examples/sft/                         Worked example: SFT getting started
+examples/dpo/                         Worked example: DPO
+examples/orpo/ifeval/                 Worked example: IFEval with ORPO
+examples/rl/deepmath/                 GRPO on DeepMath (rl_loop)
+examples/rl/frozen_lake/              Frozen Lake tool-use RL (custom loop)
+examples/rl/single_turn_async/        Async RL single-turn, renderer-backed
+examples/rl/gsm8k_async/              Async RL on GSM8K, renderer-backed
+examples/rl/multi_turn_minimal/       Async RL multi-turn, token-native (no renderer)
+examples/rl/multi_turn_minimal_renderer/  Async RL multi-turn, renderer-backed
+examples/rl/multi_turn_tool/          Async RL multi-turn with tool calls
+examples/rl/remote_rollout/           Async RL with a RolloutService
+examples/rl/ep_remote_grader/         Async RL with eval_protocol grading
+examples/multihop_qa/                 Multi-hop QA with IGPO utilities
+examples/tools/                       Standalone utility scripts
+tests/                                Unit and end-to-end tests
 ```
 
 ## Tests
@@ -174,7 +183,5 @@ pytest -q tests/unit tests/test_smoke_imports.py examples/rl/frozen_lake/test_ma
 python tests/coverage_summary.py coverage.json
 ```
 
-See [issues/training-script-coverage-baseline.md](./issues/training-script-coverage-baseline.md)
-for the current baseline and
-[issues/training-script-coverage-roadmap.md](./issues/training-script-coverage-roadmap.md)
-for the expansion plan.
+Use `training/tests/coverage_summary.py` to inspect the current recipe and
+example coverage baseline from `coverage.json`.