Add files using upload-large-folder tool

code/align/assets/tokenizer/goedel-prover-v2-8b/PROVENANCE.md

@@ -0,0 +1,7 @@
+Goedel-LM/Goedel-Prover-V2-8B tokenizer files (Qwen3 byte-level BPE; Apache-2.0), vendored so
+P2 tokenisation does not need the vLLM venv or the 16 GB model snapshot. Copied byte-for-byte
+from the HF snapshot dfd02e6271a58375dfbf3ece0175277cf6b6a89a. sha256:
+  tokenizer.json             aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
+  tokenizer_config.json      5d27a191a7ba6f93bcc2a8e78bc633764bbaaee263e9f3d9225e2c8621a5dd72
+  special_tokens_map.json    76862e765266b85aa9459767e33cbaf13970f327a0e88d1c65846c2ddd3a1ecd
+Tokenizer only: no config.json, no weights, no vocab.json/merges.txt (redundant with tokenizer.json).

code/align/assets/tokenizer/goedel-prover-v2-8b/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

code/align/assets/tokenizer/goedel-prover-v2-8b/tokenizer_config.json

@@ -0,0 +1,241 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if loop.index0 > ns.last_query_index %}\n            {%- if loop.last or (not loop.last and reasoning_content) %}\n                {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n            {%- else %}\n                {{- '<|im_start|>' + message.role + '\\n' + content }}\n            {%- endif %}\n        {%- else %}\n            {{- '<|im_start|>' + message.role + '\\n' + content }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if enable_thinking is defined and enable_thinking is false %}\n        {{- '<think>\\n\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

code/align/model/artifacts/logit_noise_floor.json

@@ -0,0 +1,20 @@
+{
+ "2": {
+  "max_abs_dlogit": 0.75586,
+  "mean_abs_dlogit": 0.015027,
+  "argmax_flips": 1,
+  "positions": 794
+ },
+ "4": {
+  "max_abs_dlogit": 0.78516,
+  "mean_abs_dlogit": 0.01436,
+  "argmax_flips": 1,
+  "positions": 794
+ },
+ "8": {
+  "max_abs_dlogit": 0.78516,
+  "mean_abs_dlogit": 0.014431,
+  "argmax_flips": 1,
+  "positions": 794
+ }
+}

code/align/model/artifacts/sdpa_baseline.json

@@ -0,0 +1,17 @@
+{
+ "512": {
+  "FLASH_ATTENTION": 0.18939885194413364,
+  "EFFICIENT_ATTENTION": NaN,
+  "MATH": 4.349638050189242
+ },
+ "1024": {
+  "FLASH_ATTENTION": 0.6236543995328248,
+  "EFFICIENT_ATTENTION": NaN,
+  "MATH": 16.252079099649563
+ },
+ "2048": {
+  "FLASH_ATTENTION": 2.281787298852578,
+  "EFFICIENT_ATTENTION": NaN,
+  "MATH": 61.56193719943985
+ }
+}

code/align/model/artifacts/triton_fa2_equivalence.json

@@ -0,0 +1,25 @@
+{
+  "environment": {
+    "compute_capability": "8.6",
+    "device_name": "NVIDIA A10G",
+    "torch": "2.11.0+cu128",
+    "triton": "3.6.0"
+  },
+  "kernel_sha256": "5f2cbe2088a7e583895268c59895366e39a4eb79d4dc8b719cbca2e57c2eed4b",
+  "schema_version": 1,
+  "shapes_covered": [
+    "gqa (n_q_heads, n_kv_heads) in [(8,2), (32,8)] plus MHA-degenerate (4,4)",
+    "head_dim in [64, 128]",
+    "seq in [192, 256, 512, 1000] (block-divisible AND non-divisible)",
+    "causal and non-causal",
+    "key_padding_mask, right-padded, non-block-aligned",
+    "every autotune config individually pinned (fwd + both bwd kernels)",
+    "dtype bfloat16 only (fp16/fp32 rejected)"
+  ],
+  "tolerances": {
+    "absolute_floor": 0.02,
+    "policy": "max |err| vs fp32 reference <= max(3.0 x torch-SDPA-bf16 error, 2e-2)",
+    "slack_over_sdpa": 3.0
+  },
+  "verified_at": "2026-08-18T13:18:04+00:00"
+}

code/align/rl/envs/__init__.py

@@ -0,0 +1,35 @@
+"""RL environments (opt-in; nothing in `align.rl` imports this package).
+
+    lean_multiturn.py   multi-turn Lean prover episodes with Goedel correction feedback,
+                        multi-turn loss masks and per-turn advantage assignment.
+"""
+
+from .lean_multiturn import (
+    EpisodeRecord,
+    EpisodeSpec,
+    GenRequest,
+    Generation,
+    LeanMultiTurnConfig,
+    LeanMultiTurnEnv,
+    PrefixBreak,
+    Problem,
+    SandboxNotWired,
+    TrainingRow,
+    TurnRecord,
+    assign_advantages,
+)
+
+__all__ = [
+    "EpisodeRecord",
+    "EpisodeSpec",
+    "GenRequest",
+    "Generation",
+    "LeanMultiTurnConfig",
+    "LeanMultiTurnEnv",
+    "PrefixBreak",
+    "Problem",
+    "SandboxNotWired",
+    "TrainingRow",
+    "TurnRecord",
+    "assign_advantages",
+]

code/align/rl/envs/lean_multiturn.py

@@ -0,0 +1,538 @@
+"""Multi-turn Lean prover environment (P4-B): proof attempt -> compiler feedback -> retry.
+
+An EPISODE is up to K turns on one formal statement. Turn 0 is the Goedel round-0 prompt.
+If the verifier rejects the proof with compile errors, turn k+1's prompt is the Goedel
+CORRECTION prompt: the history, the failed attempt as an assistant message, and a user
+message carrying the `<error></error>`-annotated compiler output. The episode ends at the
+first verified turn, at K turns, or at the first rejection that feedback cannot repair.
+
+Everything below exists to prevent one named failure; each is pinned by a test.
+
+* PROMPT DRIFT.  The prompts are NOT re-expressed here. `extras/vllm_rollout/lean_prompts.py`
+  (byte-faithful to Goedel-Prover-V2) is loaded BY FILE PATH, read-only, the way
+  tests/test_lean_gen_contract.py and scripts/p2/p2_common.py do. A second copy of the
+  correction template would drift by one newline and silently change the policy sampled.
+
+* LOSS ON FEEDBACK TOKENS.  The compiler's error text is environment output, not policy
+  output. If it carried loss the model would be trained to *write* error messages, and the
+  gradient would reward predicting them. Rows carry `response_mask` that is 1 ONLY on tokens
+  the generator returned; prompt, chat-template glue and feedback are 0. `turn_ids` labels
+  each token with its turn (-1 unlabeled) so tests and the loss can audit it.
+
+* RE-TOKENIZATION DRIFT.  Token ids are never re-derived from text for generated tokens.
+  layout="concatenated" builds ONE row per episode: context_k = context_{k-1} + generated_{k-1}
+  + encode(template delta). That is only valid when the chat template renders turn k's
+  prompt as a pure extension of (turn k-1's prompt + completion text); templates that rewrite
+  earlier assistant content (Qwen3 strips old <think> blocks) break it. We CHECK the prefix
+  property and raise `PrefixBreak` rather than train on tokens the policy never saw.
+  layout="per_turn" (the default) emits one row per turn from that turn's exact prompt ids;
+  it is always exact, at the price of re-running the shared prefix.
+
+* STATEMENT REWRITING IS NOT A SOLVE.  `problem_check` re-inserts the ORIGINAL statement over
+  whatever theorem the model displayed, so a model that states an easier theorem and proves
+  that is *accepted* if its tactics happen to close the original. Goedel's evaluation credits
+  it; a TRAINING reward must not (it teaches the model to misstate the problem). Before the
+  verifier is consulted, `scripts/p2/p2_common.displayed_statement_matches` (loaded by path)
+  must hold, else the turn is terminal with reason `statement_rewrite` and reward 0.
+
+* TRUNCATION IS NOT A SUBMISSION.  A generation that hit the length limit is rejected
+  (`truncated`) even if its plan contains an earlier closed code fence; the final answer was
+  cut off, and verifying a stale fence rewards an accident.
+
+* SPARSE REWARD, NO LENGTH IN THE LOSS.  Episode reward = turn_discount ** solved_turn
+  (0 if never solved) plus an optional DAPO soft-overlong term on TOTAL generated tokens,
+  weighted by `overlong_weight`. The penalty lives in reward composition (reusing
+  `align.rl.rewards.soft_overlong_punishment`), never in the loss.
+
+* CREDIT ASSIGNMENT.  advantage_mode="episode": every generated token gets the Dr.GRPO
+  group-centered episode reward (the repo's `compute_group_advantages`, reused).
+  advantage_mode="per_turn": turn t gets G_t - mean_group(G_t) with
+  G_t = sum_{s>=t} credit_gamma^(s-t) * u_s, u_s the turn's own outcome reward. With
+  credit_gamma=0 a failed first attempt gets no credit for the repair that followed (the
+  failure mode of episode-level sharing: rewarding the bad first draft); with gamma>0 it
+  shares in it, because its error context is what made the repair possible. The baseline for
+  turn t averages only episodes that HAVE a turn t.
+
+* DEAD VERIFIER != ZERO REWARD.  `LeanInfraError` from the verifier propagates unchanged. A
+  verifier that is dying SOFTLY (every proof comes back `crash`/`protocol`/`memory`/
+  `output_too_large` instead of raising) would otherwise be turned by this env into
+  `rejected:crash`, reward 0, a batch of zeros that reads as "the model got worse".
+  `LeanReward.accuracy_batch` refuses that; the env bypasses `LeanReward`, so it carries the
+  same guard: once at least `min_infra_batch` verifications have run in a `run_group`, a
+  cumulative infra-failure share above `max_infra_fraction` raises `LeanInfraError`.
+
+* ENGINE FINISH REASONS.  The P3 engine reports `eos|stop_token|stop_string|max_new_tokens|
+  max_seq_len`, the lean_gen contract reports `stop|length`. Only `length` used to count as
+  truncated, so an engine-capped generation (`max_new_tokens`/`max_seq_len`) was VERIFIED as
+  if it were a finished answer. Both vocabularies are accepted; an unknown reason raises
+  (a new engine reason must be classified by a human, not guessed).
+
+* CONTEXT BUDGET.  Compiler feedback is unbounded and each repair turn re-sends the whole
+  history. With `max_prompt_tokens` set, a repair prompt over the budget ends the episode
+  (`context_overflow`, reward 0 for the unsolved episode) instead of being handed to an
+  engine that would refuse it (or silently truncate the prompt from the left).
+
+* EPISODE KEY COLLISIONS.  `generate` adapters key seeds / KV state by `episode_key`; two
+  specs with the same (problem_id, sample_idx, statement) would share one stream. Refused.
+
+Sandbox: `LeanMultiTurnConfig.sandbox=True` is a pass-through request for the
+`align/verify/lean/sandbox.py` REPL. The pool is not wired to it yet, so the env raises
+`SandboxNotWired` unless the pool itself advertises a truthy `sandbox` attribute.
+
+Error positions: the Goedel format needs REPL message dicts (`pos`, `endPos`, `data`)
+indexing `compile_view(code)`. `StrictResult.diagnostics["repl_errors"]` is the contract for
+those; `StrictVerifier` does not populate it yet (it keeps only 3 truncated message strings).
+Absent positions the env degrades per `unpositioned_feedback` ("message_only" renders the
+messages with no code excerpt; "raise" refuses) instead of inventing an `<error>` location.
+
+stdlib-only (+ align.rl.rewards / align.rl.advantages); no torch tensors are built here.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.util
+import json
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Optional, Sequence
+
+from align.rl.advantages import compute_group_advantages
+from align.rl.rewards import soft_overlong_punishment
+
+_REPO = Path(__file__).resolve().parents[3]
+
+_MODS: dict[str, Any] = {}
+
+
+def _load_by_path(key: str, rel: str):
+    if key not in _MODS:
+        spec = importlib.util.spec_from_file_location(key, _REPO / rel)
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+        _MODS[key] = mod
+    return _MODS[key]
+
+
+def _lean_prompts():
+    return _load_by_path("align_envs_lean_prompts", "extras/vllm_rollout/lean_prompts.py")
+
+
+def _p2_common():
+    return _load_by_path("align_envs_p2_common", "scripts/p2/p2_common.py")
+
+
+#: finish reasons that mean "the generator ended on its own" / "the generator was cut off"
+STOP_REASONS = frozenset({"stop", "eos", "stop_token", "stop_string"})
+TRUNCATED_REASONS = frozenset({"length", "max_new_tokens", "max_seq_len"})
+
+
+class SandboxNotWired(RuntimeError):
+    """sandbox=True was requested but the verifier pool has no sandboxed REPL."""
+
+
+class PrefixBreak(RuntimeError):
+    """The chat template did not render turn k's prompt as an extension of turn k-1's
+    prompt + completion, so one concatenated token row would not be what the policy saw."""
+
+
+@dataclass(frozen=True)
+class Problem:
+    problem_id: str
+    formal_statement: str
+
+
+@dataclass(frozen=True)
+class EpisodeSpec:
+    problem: Problem
+    sample_idx: int = 0
+
+
+@dataclass(frozen=True)
+class GenRequest:
+    episode_key: str
+    turn: int
+    prompt_ids: tuple
+    prompt_text: str
+
+
+@dataclass(frozen=True)
+class Generation:
+    """`token_ids`: exactly what the engine sampled (incl. a trailing EOS if sampled).
+    `text`: the decoded completion WITHOUT the EOS string. `logprobs`: per generated token."""
+
+    token_ids: tuple
+    text: str
+    logprobs: Optional[tuple] = None
+    finish_reason: str = "stop"
+
+
+@dataclass(frozen=True)
+class LeanMultiTurnConfig:
+    max_turns: int = 2
+    layout: str = "per_turn"  # "per_turn" | "concatenated"
+    turn_discount: float = 1.0
+    overlong_max_len: Optional[int] = None
+    overlong_cache_len: Optional[int] = None
+    overlong_weight: float = 1.0
+    advantage_mode: str = "episode"  # "episode" | "per_turn"
+    credit_gamma: float = 0.0
+    retry_reasons: tuple = ("compile_error",)
+    reject_truncated: bool = True
+    #: ABLATION ONLY: accept a proof whose displayed theorem differs from the statement.
+    allow_statement_rewrite: bool = False
+    error_thres: bool = True
+    unpositioned_feedback: str = "message_only"  # "message_only" | "raise"
+    sandbox: bool = False
+    #: ceiling on the repair prompt length in tokens (None = unbounded, the old behaviour)
+    max_prompt_tokens: Optional[int] = None
+    #: soft-infra-failure guard (see module docstring); same defaults as LeanReward
+    max_infra_fraction: float = 0.5
+    min_infra_batch: int = 8
+    eos_token_id: Optional[int] = None
+    eos_text: str = "<|im_end|>"
+
+    def __post_init__(self) -> None:
+        if self.max_turns < 1:
+            raise ValueError(f"max_turns must be >= 1, got {self.max_turns}")
+        if self.layout not in ("per_turn", "concatenated"):
+            raise ValueError(f"layout must be 'per_turn' or 'concatenated', got {self.layout!r}")
+        if not 0.0 < self.turn_discount <= 1.0:
+            raise ValueError(f"turn_discount must be in (0, 1], got {self.turn_discount}")
+        if self.advantage_mode not in ("episode", "per_turn"):
+            raise ValueError(f"advantage_mode must be 'episode' or 'per_turn', got {self.advantage_mode!r}")
+        if not 0.0 <= self.credit_gamma <= 1.0:
+            raise ValueError(f"credit_gamma must be in [0, 1], got {self.credit_gamma}")
+        if self.max_prompt_tokens is not None and self.max_prompt_tokens < 1:
+            raise ValueError(f"max_prompt_tokens must be None or >= 1, got {self.max_prompt_tokens}")
+        if not 0.0 <= self.max_infra_fraction <= 1.0:
+            raise ValueError(f"max_infra_fraction must be in [0, 1], got {self.max_infra_fraction}")
+        if self.min_infra_batch < 1:
+            raise ValueError(f"min_infra_batch must be >= 1, got {self.min_infra_batch}")
+        if self.unpositioned_feedback not in ("message_only", "raise"):
+            raise ValueError(f"unpositioned_feedback must be 'message_only' or 'raise'")
+        # same construction-time refusals as CompositeReward: a half-armed penalty is silently off
+        if self.overlong_max_len is None and self.overlong_cache_len is not None:
+            raise ValueError("overlong_cache_len given without overlong_max_len: penalty would be silently OFF")
+        if self.overlong_max_len is not None:
+            if self.overlong_cache_len is None or not 0 < self.overlong_cache_len < self.overlong_max_len:
+                raise ValueError("need 0 < overlong_cache_len < overlong_max_len")
+
+
+@dataclass
+class TurnRecord:
+    turn: int
+    prompt_text: str
+    prompt_ids: list
+    completion_ids: list
+    completion_text: str
+    finish_reason: str
+    lean_code: Optional[str] = None
+    code_error: Optional[str] = None
+    displayed_statement_ok: Optional[bool] = None
+    reasons: tuple = ()
+    verified: bool = False
+    credited: bool = False
+    feedback_text: Optional[str] = None
+    feedback_positioned: Optional[bool] = None
+
+
+@dataclass
+class TrainingRow:
+    input_ids: list
+    #: 1 ONLY on tokens the generator returned (input space, like RolloutBatch.response_mask)
+    response_mask: list
+    #: turn index of each generated token, -1 elsewhere
+    turn_ids: list
+    #: per-token engine logprobs in input space (0.0 where unlabeled); None if not supplied
+    logprobs: Optional[list] = None
+    #: per-token advantage (0.0 where unlabeled); filled by `assign_advantages`
+    advantages: Optional[list] = None
+
+
+@dataclass
+class EpisodeRecord:
+    episode_key: str
+    problem_id: str
+    sample_idx: int
+    turns: list
+    rows: list
+    terminal: str = ""
+    solved_turn: Optional[int] = None
+    reward: dict = field(default_factory=dict)
+    turn_rewards: list = field(default_factory=list)  # u_t, the per-turn outcome rewards
+    n_generated: int = 0
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), sort_keys=True, ensure_ascii=False)
+
+
+# ------------------------------------------------------------------------- reward composition
+
+
+def compose_episode_reward(cfg: LeanMultiTurnConfig, solved_turn: Optional[int], n_turns: int, n_generated: int):
+    """(terms, u): itemised episode reward terms and the per-turn outcome rewards u_t.
+    The overlong term is attached to the LAST turn's u (the turn that ended the episode)."""
+    u = [0.0] * n_turns
+    acc = 0.0
+    if solved_turn is not None:
+        acc = cfg.turn_discount ** solved_turn
+        u[solved_turn] = acc
+    ov = 0.0
+    if cfg.overlong_max_len is not None:
+        ov = cfg.overlong_weight * soft_overlong_punishment(n_generated, cfg.overlong_max_len, cfg.overlong_cache_len)
+        u[-1] += ov
+    terms = {
+        "reward/accuracy": 1.0 if solved_turn is not None else 0.0,
+        "reward/turn_discount": acc,
+        "reward/overlong": ov,
+        "reward/total": acc + ov,
+    }
+    return terms, u
+
+
+def assign_advantages(episodes: Sequence[EpisodeRecord], cfg: LeanMultiTurnConfig) -> None:
+    """Fill `row.advantages` in place. Episodes of one problem_id form a group (input order
+    preserved). Groups must be equal-sized in "episode" mode (the shared helper's layout)."""
+    groups: dict[str, list] = {}
+    for ep in episodes:
+        groups.setdefault(ep.problem_id, []).append(ep)
+    for eps in groups.values():
+        n_turn_max = max(len(e.turns) for e in eps)
+        if cfg.advantage_mode == "episode":
+            import numpy as np
+
+            adv = compute_group_advantages(np.array([e.reward["reward/total"] for e in eps]), len(eps), "centered")
+            per_turn = [[float(a)] * n_turn_max for a in adv]
+        else:
+            G = []
+            for e in eps:
+                g, acc = [0.0] * len(e.turns), 0.0
+                for t in range(len(e.turns) - 1, -1, -1):
+                    acc = e.turn_rewards[t] + cfg.credit_gamma * acc
+                    g[t] = acc
+                G.append(g)
+            per_turn = [[0.0] * n_turn_max for _ in eps]
+            for t in range(n_turn_max):
+                have = [i for i, e in enumerate(eps) if t < len(e.turns)]
+                base = sum(G[i][t] for i in have) / len(have)
+                for i in have:
+                    per_turn[i][t] = G[i][t] - base
+        for e, adv_t in zip(eps, per_turn):
+            for row in e.rows:
+                row.advantages = [adv_t[tid] if tid >= 0 else 0.0 for tid in row.turn_ids]
+
+
+# ------------------------------------------------------------------------- the environment
+
+
+class _Ep:
+    def __init__(self, spec: EpisodeSpec, key: str):
+        self.spec, self.key = spec, key
+        self.turns: list[TurnRecord] = []
+        self.messages: list[dict] = []
+        self.prompt_text = ""
+        self.prompt_ids: list = []
+        self.ctx_ids: list = []  # concatenated layout: running context
+        self.gens: list[Generation] = []
+        self.done = False
+        self.terminal = ""
+        self.solved_turn: Optional[int] = None
+        self.pending: Optional[tuple] = None
+
+
+class LeanMultiTurnEnv:
+    """`verifier`: anything with `verify(statement, code) -> StrictResult` (and optionally
+    `verify_batch(items)`), e.g. `StrictVerifier`. `tokenizer`: `apply_chat_template(messages,
+    tokenize=False, add_generation_prompt=True)`. `encode`: text -> ids."""
+
+    def __init__(self, verifier, tokenizer, encode: Callable[[str], list], cfg: LeanMultiTurnConfig = LeanMultiTurnConfig()):
+        if cfg.sandbox and not getattr(getattr(verifier, "pool", verifier), "sandbox", None):
+            raise SandboxNotWired(
+                "sandbox=True requested but the verifier pool exposes no sandboxed REPL "
+                "(align/verify/lean/sandbox.py is not wired into repl_pool yet)"
+            )
+        self.verifier, self.tok, self.encode, self.cfg = verifier, tokenizer, encode, cfg
+
+    # ---- public
+    def run_episode(self, spec: EpisodeSpec, generate: Callable[[list], list]) -> EpisodeRecord:
+        return self.run_group([spec], generate)[0]
+
+    def run_group(self, specs: Sequence[EpisodeSpec], generate: Callable[[list], list]) -> list:
+        """Lockstep over episodes: each round, ONE `generate(list[GenRequest]) -> list[Generation]`
+        call for all active episodes, then one batched verification. Output order = input order."""
+        eps = [_Ep(s, self._key(s)) for s in specs]
+        if len({e.key for e in eps}) != len(eps):
+            dup = sorted({e.key for e in eps if sum(x.key == e.key for x in eps) > 1})
+            raise ValueError(f"duplicate episode keys {dup}: specs must differ in problem_id/sample_idx/statement")
+        n_ver = n_infra = 0
+        for ep in eps:
+            self._start(ep)
+        active = list(eps)
+        while active:
+            reqs = [GenRequest(e.key, len(e.turns), tuple(e.prompt_ids), e.prompt_text) for e in active]
+            gens = generate(reqs)
+            if len(gens) != len(active):
+                raise ValueError(f"generate returned {len(gens)} generations for {len(active)} requests")
+            for e, g in zip(active, gens):
+                self._absorb(e, g)
+            todo = [e for e in active if not e.done]
+            if todo:
+                items = [e.pending for e in todo]
+                vb = getattr(self.verifier, "verify_batch", None)
+                results = vb(items) if vb is not None else [self.verifier.verify(s, c) for s, c in items]
+                if len(results) != len(todo):
+                    raise ValueError(f"verifier returned {len(results)} results for {len(todo)} items")
+                n_ver += len(results)
+                n_infra += sum(1 for r in results if getattr(r, "infra_failure", False))
+                self._check_infra(n_infra, n_ver)
+                for e, r in zip(todo, results):
+                    self._judge(e, r)
+            active = [e for e in active if not e.done]
+        return [self._finish(e) for e in eps]
+
+    # ---- internals
+    def _check_infra(self, n_infra: int, n_ver: int) -> None:
+        cfg = self.cfg
+        if n_ver >= cfg.min_infra_batch and n_infra / n_ver > cfg.max_infra_fraction:
+            from align.verify.lean.strict import LeanInfraError
+
+            raise LeanInfraError(
+                f"{n_infra}/{n_ver} verifications died for infrastructure reasons: either the pool is "
+                f"dead or the policy is emitting REPL-killing proofs; refusing to emit episodes of "
+                f"reward 0 that read as 'model got worse'"
+            )
+
+    def _key(self, s: EpisodeSpec) -> str:
+        h = hashlib.sha256(f"{s.problem.problem_id}\0{s.sample_idx}\0{s.problem.formal_statement}".encode()).hexdigest()
+        return h[:16]
+
+    def _render(self, messages: list) -> str:
+        return self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    def _start(self, ep: _Ep) -> None:
+        text, msgs = _lean_prompts().initial_prompt(ep.spec.problem.formal_statement, self.tok)
+        ep.messages, ep.prompt_text = list(msgs), text
+        ep.prompt_ids = list(self.encode(text))
+        ep.ctx_ids = list(ep.prompt_ids)
+
+    def _absorb(self, ep: _Ep, g: Generation) -> None:
+        cfg, lp = self.cfg, _lean_prompts()
+        if g.logprobs is not None and len(g.logprobs) != len(g.token_ids):
+            raise ValueError(f"{len(g.logprobs)} logprobs for {len(g.token_ids)} generated tokens")
+        t = TurnRecord(len(ep.turns), ep.prompt_text, list(ep.prompt_ids), list(g.token_ids), g.text, g.finish_reason)
+        ep.turns.append(t)
+        ep.gens.append(g)
+        stmt = ep.spec.problem.formal_statement
+        if g.finish_reason not in STOP_REASONS and g.finish_reason not in TRUNCATED_REASONS:
+            raise ValueError(
+                f"unknown finish_reason {g.finish_reason!r}; classify it in STOP_REASONS or TRUNCATED_REASONS"
+            )
+        if cfg.reject_truncated and g.finish_reason in TRUNCATED_REASONS:
+            return self._end(ep, "truncated")
+        code, err = lp.lean_code_for_record(stmt, g.text)
+        t.lean_code, t.code_error = code, err
+        if code is None:
+            return self._end(ep, f"no_code:{err}")
+        if not cfg.allow_statement_rewrite:
+            t.displayed_statement_ok = bool(_p2_common().displayed_statement_matches(stmt, g.text))
+            if not t.displayed_statement_ok:
+                return self._end(ep, "statement_rewrite")
+        ep.pending = (stmt, code)
+
+    def _judge(self, ep: _Ep, r) -> None:
+        cfg, t = self.cfg, ep.turns[-1]
+        t.reasons, t.verified = tuple(r.reasons), bool(r.accepted)
+        if r.accepted:
+            t.credited, ep.solved_turn = True, t.turn
+            return self._end(ep, "verified")
+        if t.turn + 1 >= cfg.max_turns:
+            return self._end(ep, "max_turns")
+        if not all(x in cfg.retry_reasons for x in t.reasons):
+            return self._end(ep, f"rejected:{t.reasons[0] if t.reasons else 'unknown'}")
+        fb = self._feedback(ep, r)
+        if fb is None:
+            return self._end(ep, "no_feedback")
+        err_str, positioned = fb
+        t.feedback_text, t.feedback_positioned = err_str, positioned
+        lp = _lean_prompts()
+        text, msgs = lp.DeepSeekCoTHandler().generate_correction_prompt(
+            lp.goedel_statement(ep.spec.problem.formal_statement), ep.messages, t.completion_text,
+            err_str, self.tok, t.turn + 1,
+        )
+        if cfg.max_prompt_tokens is not None and len(self.encode(text)) > cfg.max_prompt_tokens:
+            return self._end(ep, "context_overflow")
+        ep.messages = msgs
+        prev_end = ep.prompt_text + t.completion_text
+        if cfg.layout == "concatenated":
+            if not text.startswith(prev_end):
+                raise PrefixBreak(
+                    f"turn {t.turn + 1} prompt does not extend turn {t.turn} prompt + completion "
+                    f"(template rewrites earlier assistant content, e.g. Qwen3 <think> stripping); "
+                    f"use layout='per_turn'"
+                )
+            delta = text[len(prev_end):]
+            gen = ep.gens[-1]
+            if (cfg.eos_token_id is not None and gen.token_ids and gen.token_ids[-1] == cfg.eos_token_id
+                    and delta.startswith(cfg.eos_text)):
+                delta = delta[len(cfg.eos_text):]  # the sampled EOS is already in the context
+            ep.ctx_ids = ep.ctx_ids + list(gen.token_ids) + list(self.encode(delta))
+            ep.prompt_ids = list(ep.ctx_ids)
+        else:
+            ep.prompt_ids = list(self.encode(text))
+        ep.prompt_text = text
+
+    def _feedback(self, ep: _Ep, r) -> Optional[tuple]:
+        lp = _lean_prompts()
+        diag = r.diagnostics or {}
+        t = ep.turns[-1]
+        errs = diag.get("repl_errors")
+        if errs and all(isinstance(e, dict) and e.get("pos") and "data" in e for e in errs):
+            view = diag.get("error_code_view") or lp.compile_view(t.lean_code)
+            return lp.get_error_str(view, errs, self.cfg.error_thres), True
+        msgs = diag.get("errors")
+        if not msgs and not errs:
+            return None
+        if self.cfg.unpositioned_feedback == "raise":
+            raise ValueError("verifier result has no positioned REPL errors (diagnostics['repl_errors'])")
+        raw = [e["data"] if isinstance(e, dict) else str(e) for e in (errs or msgs)]
+        return "".join(f"\nError {i + 1}:\n\nError Message: {m}\n" for i, m in enumerate(raw)), False
+
+    def _end(self, ep: _Ep, why: str) -> None:
+        ep.done, ep.terminal = True, why
+
+    def _finish(self, ep: _Ep) -> EpisodeRecord:
+        n_gen = sum(len(g.token_ids) for g in ep.gens)
+        terms, u = compose_episode_reward(self.cfg, ep.solved_turn, len(ep.turns), n_gen)
+        rows = self._rows(ep)
+        s = ep.spec
+        return EpisodeRecord(ep.key, s.problem.problem_id, s.sample_idx, ep.turns, rows, ep.terminal,
+                             ep.solved_turn, terms, u, n_gen)
+
+    def _rows(self, ep: _Ep) -> list:
+        if self.cfg.layout == "per_turn":
+            out = []
+            for t, g in zip(ep.turns, ep.gens):
+                n = len(g.token_ids)
+                out.append(TrainingRow(
+                    list(t.prompt_ids) + list(g.token_ids),
+                    [0] * len(t.prompt_ids) + [1] * n,
+                    [-1] * len(t.prompt_ids) + [t.turn] * n,
+                    ([0.0] * len(t.prompt_ids) + list(g.logprobs)) if g.logprobs is not None else None,
+                ))
+            return out
+        ids, mask, tids, lp, have_lp = [], [], [], [], all(g.logprobs is not None for g in ep.gens)
+        for t, g in zip(ep.turns, ep.gens):
+            n_ctx = len(t.prompt_ids) - len(ids)  # new context tokens (prompt / feedback glue)
+            ids += list(t.prompt_ids[len(ids):])
+            mask += [0] * n_ctx
+            tids += [-1] * n_ctx
+            lp += [0.0] * n_ctx
+            ids += list(g.token_ids)
+            mask += [1] * len(g.token_ids)
+            tids += [t.turn] * len(g.token_ids)
+            if have_lp:
+                lp += list(g.logprobs)
+            else:
+                lp += [0.0] * len(g.token_ids)
+        return [TrainingRow(ids, mask, tids, lp if have_lp else None)]

code/align/rl/envs/subgoal_repair.py

@@ -0,0 +1,694 @@
+"""APOLLO-style subgoal extraction and reassembly for failed Lean 4 proofs (K5 agentic arm, docs/K5_LEVERS.md).
+
+Pure Python (stdlib only) plus read-only Lean calls through the existing `align.verify.lean` API. Nothing here
+changes the verifier: the reassembled proof is only a CANDIDATE; the strict verifier (`StrictVerifier`) is the sole
+judge of acceptance, against the ORIGINAL statement. So a lemma that drifted from the goal it was cut from (the
+pretty-printer is not a round trip) can cost a wasted attempt, never a false accept.
+
+Pipeline (APOLLO, arXiv 2505.05758, section 3; DeepSeek-Prover-V2, arXiv 2504.21801, section 2.1: "decompose, prove the
+subgoals separately, recombine"):
+
+  1. `locate_targets`   map every Lean error of a failed proof to the smallest repairable unit that contains it:
+                        a failing `have name : T := <rhs>` (kind "have", innermost wins), a failing main-line tactic
+                        (kind "tactic"), or "the proof ended with goals left" (kind "tail"). Errors on a `have`
+                        HEADER (the statement itself does not elaborate) are not repairable and are reported.
+  2. `build_skeleton`   replace each target's text by `sorry` ("sorrifier"). Targets that contain another target are
+                        dropped (the inner failure is the finer cut; the outer one is picked up by the next round).
+  3. `plan_subgoals`    run the skeleton in Lean once: its `sorries` carry the exact goal state (hypotheses and
+                        goal) at each gap. Each state becomes a standalone `theorem k5_sg_<hash> (hyps) : goal`,
+                        which is checked to elaborate. These are ordinary problems for the prover.
+  4. `reassemble`       splice a prover-supplied tactic body back into the skeleton in place of the `sorry`
+                        (inline: the lemma's hypotheses carry the context's names; inaccessible ones are
+                        re-introduced with `rename_i`). Unproven gaps get their original text back.
+
+Coordinates. Lean errors index `compile_view(code)` (header lines removed, `p2_common.compile_view`), not the full
+code. `view_line_map` translates; all text surgery happens on the FULL code so the result can be verified as is.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+# --------------------------------------------------------------------------- lexical masking
+
+_OPEN = "([{⟨"
+_CLOSE = ")]}⟩"
+
+
+def mask(text: str) -> str:
+    """`text` with comments and string literals blanked (newlines kept), so indentation and bracket scans cannot be
+    fooled by prose. Same length as the input: offsets and (line, column) positions are preserved."""
+    out = list(text)
+    i, n = 0, len(text)
+    while i < n:
+        c = text[i]
+        if c == "-" and text.startswith("--", i):
+            j = text.find("\n", i)
+            j = n if j < 0 else j
+            for k in range(i, j):
+                out[k] = " "
+            i = j
+        elif c == "/" and text.startswith("/-", i):
+            depth, j = 1, i + 2
+            while j < n and depth:
+                if text.startswith("/-", j):
+                    depth += 1
+                    j += 2
+                elif text.startswith("-/", j):
+                    depth -= 1
+                    j += 2
+                else:
+                    j += 1
+            for k in range(i, min(j, n)):
+                if out[k] != "\n":
+                    out[k] = " "
+            i = j
+        elif c == '"':
+            j = i + 1
+            while j < n and text[j] != '"':
+                j += 2 if text[j] == "\\" else 1
+            for k in range(i + 1, min(j, n)):
+                if out[k] != "\n":
+                    out[k] = " "
+            i = j + 1
+        else:
+            i += 1
+    return "".join(out)
+
+
+def strip_comments(text: str) -> str:
+    """`text` with block and line comments deleted (used on headers only, which hold no string literals)."""
+    text = re.sub(r"/-.*?-/", "", text, flags=re.S)
+    return re.sub(r"--[^\n]*", "", text)
+
+
+def view_line_map(code: str) -> list[int]:
+    """`map[i]` = index in `code.split("\\n")` of line `i` of `compile_view(code)` (the REPL's coordinate system).
+    `compile_view` (the reference's `handle`) drops every line whose stripped text starts with import / set_option /
+    open, nothing else."""
+    out = []
+    for i, line in enumerate(code.split("\n")):
+        s = line.strip()
+        if s.startswith("import") or s.startswith("set_option") or s.startswith("open"):
+            continue
+        out.append(i)
+    return out
+
+
+# --------------------------------------------------------------------------- proof structure
+
+
+@dataclass
+class _Lines:
+    """Masked, bullet-normalised working copy of the code."""
+
+    raw: list[str]
+    m: list[str]  # masked lines (comments blanked); bullets `· ` replaced by two spaces
+    indent: list[int]
+    blank: list[bool]
+    depth_before: list[int]  # open-bracket depth at the start of each line
+    cont: list[bool]  # the line continues the previous statement at the same indent
+    ends: list[int] = field(default_factory=list)  # extent end line of a statement starting here (-1: not a start)
+
+
+_BULLET = re.compile(r"^(\s*)[·.](?=\s)")
+
+
+def _analyse(code: str) -> _Lines:
+    raw = code.split("\n")
+    mk = mask(code).split("\n")
+    m = []
+    for ln in mk:
+        mm = _BULLET.match(ln)
+        if mm:
+            ln = mm.group(1) + " " + ln[len(mm.group(1)) + 1:]
+        m.append(ln)
+    indent = [len(x) - len(x.lstrip(" ")) for x in m]
+    blank = [not x.strip() for x in m]
+    depth, db = 0, []
+    for x in m:
+        db.append(depth)
+        for ch in x:
+            if ch in _OPEN:
+                depth += 1
+            elif ch in _CLOSE:
+                depth = max(0, depth - 1)
+    cont = [False] * len(m)
+    prev = None
+    for i, x in enumerate(m):
+        if blank[i]:
+            continue
+        s = x.strip()
+        if prev is not None and (db[i] > 0 or s.startswith("<;>") or m[prev].rstrip().endswith("<;>")):
+            cont[i] = True
+        prev = i
+    L = _Lines(raw, m, indent, blank, db, cont)
+    L.ends = [-1] * len(m)
+    for i in range(len(m)):
+        if blank[i] or cont[i]:
+            continue
+        e = i
+        for j in range(i + 1, len(m)):
+            if blank[j]:
+                continue
+            if indent[j] > indent[i] or (indent[j] == indent[i] and cont[j]):
+                e = j
+            else:
+                break
+        L.ends[i] = e
+    return L
+
+
+_THEOREM = re.compile(r"^\s*(?:theorem|lemma)\b")
+
+
+@dataclass(frozen=True)
+class Have:
+    start: int  # first line
+    end: int  # last line of the unit (inclusive)
+    name: str
+    assign_line: int  # line holding the depth-0 `:=`
+    assign_col: int  # column of that `:=`
+
+
+def _first_assign(L: _Lines, s: int, e: int) -> Optional[tuple[int, int]]:
+    depth = 0
+    for i in range(s, e + 1):
+        x = L.m[i]
+        for c, ch in enumerate(x):
+            if ch in _OPEN:
+                depth += 1
+            elif ch in _CLOSE:
+                depth = max(0, depth - 1)
+            elif ch == ":" and depth == 0 and x.startswith(":=", c):
+                return i, c
+    return None
+
+
+def find_haves(code: str, L: Optional[_Lines] = None) -> list[Have]:
+    """Every `have <name> : T := <rhs>` unit (any nesting depth) that has a `:=`. Side-goal haves (no `:=`, closed by
+    following bullets) and `have ... from` are not units this module can cut."""
+    L = L or _analyse(code)
+    out = []
+    for i, x in enumerate(L.m):
+        if L.ends[i] < 0:
+            continue
+        s = x.strip()
+        if not (s.startswith("have ") or s.startswith("have:") or s.startswith("haveI ")):
+            continue
+        a = _first_assign(L, i, L.ends[i])
+        if a is None:
+            continue
+        toks = s.split()
+        name = toks[1] if len(toks) > 1 and toks[1] not in (":", ":=") and not toks[1].startswith(("⟨", "(", ":")) else "_"
+        out.append(Have(i, L.ends[i], name, a[0], a[1]))
+    return out
+
+
+def theorem_by_line(L: _Lines) -> Optional[tuple[int, int]]:
+    """(line, column just after `by`) of the main theorem's `:= by`; None when absent."""
+    for i, x in enumerate(L.m):
+        if _THEOREM.match(x):
+            a = _first_assign(L, i, min(len(L.m) - 1, i + 12))
+            if a is None:
+                return None
+            j = L.m[a[0]].find("by", a[1])
+            return (a[0], j + 2) if j >= 0 else None
+    return None
+
+
+# --------------------------------------------------------------------------- targets
+
+
+@dataclass
+class Target:
+    kind: str  # have | tactic | tail
+    start: tuple[int, int]  # (line, col) of the text to replace, 0-based, full-code coordinates
+    end: tuple[int, int]  # exclusive end (line, col); for "tail" start == end (insertion point)
+    original: str  # the replaced text, restorable verbatim
+    indent: int  # indentation of the statement the gap lives in
+    name: str = ""
+    errors: list = field(default_factory=list)
+
+
+def _cut(raw: list[str], a: tuple[int, int], b: tuple[int, int]) -> str:
+    if a[0] == b[0]:
+        return raw[a[0]][a[1]:b[1]]
+    parts = [raw[a[0]][a[1]:]] + raw[a[0] + 1:b[0]] + [raw[b[0]][:b[1]]]
+    return "\n".join(parts)
+
+
+def tail_target(L: _Lines, tb: tuple[int, int]) -> Target:
+    """Insertion point for a `sorry` after the proof's last line: the proof ended with goals left."""
+    body_lines = [i for i in range(tb[0] + 1, len(L.m)) if not L.blank[i]]
+    ind = min((L.indent[i] for i in body_lines), default=2)
+    last = max(body_lines, default=tb[0])
+    pos = (last, len(L.raw[last]))
+    return Target("tail", pos, pos, "", ind)
+
+
+def locate_targets(code: str, errors: list[dict], repl_code: Optional[str] = None) -> tuple[list[Target], list[dict]]:
+    """Map Lean errors (positions in `repl_code`, default `compile_view(code)`) to repairable targets.
+
+    Returns (targets sorted by position, unrepairable errors with a `why`). Only severity-error messages are used.
+    Rules, in order, for an error at line l, column c:
+      * inside the rhs of a `have`            -> kind "have" (innermost such have); target = text after its `:=`
+      * on the theorem's `:= by` line, or the proof ends with goals left -> kind "tail"
+      * on a `have` header (before its `:=`)  -> unrepairable "have_header" (the statement itself is wrong)
+      * otherwise                             -> kind "tactic": the statement containing the error, as one unit
+    Targets that contain another target are dropped; the finer cut wins."""
+    L = _analyse(code)
+    vmap = view_line_map(code)
+    if repl_code is not None and len(repl_code.split("\n")) != len(vmap):
+        raise ValueError("repl_code does not match compile_view(code): the error positions would index other lines")
+    haves = find_haves(code, L)
+    tb = theorem_by_line(L)
+    found: dict[tuple, Target] = {}
+    bad: list[dict] = []
+    for e in errors:
+        if e.get("severity", "error") != "error" or not e.get("pos"):
+            continue
+        vl, c = int(e["pos"]["line"]) - 1, int(e["pos"]["column"])
+        if not 0 <= vl < len(vmap):
+            bad.append({**e, "why": "position_outside_code"})
+            continue
+        l = vmap[vl]
+        inside = [h for h in haves if h.start <= l <= h.end and (l, c) > (h.assign_line, h.assign_col + 1)]
+        if inside:
+            h = max(inside, key=lambda x: x.start)
+            rhs_start = (h.assign_line, h.assign_col + 2)
+            end = (h.end, len(L.raw[h.end]))
+            key = ("have", rhs_start)
+            t = found.setdefault(key, Target("have", rhs_start, end, _cut(L.raw, rhs_start, end), L.indent[h.start], h.name))
+            t.errors.append(e)
+            continue
+        if tb is not None and l == tb[0]:
+            t = found.setdefault(("tail",), tail_target(L, tb))
+            t.errors.append(e)
+            continue
+        if any(h.start <= l <= h.assign_line and (l, c) <= (h.assign_line, h.assign_col + 1) for h in haves):
+            bad.append({**e, "why": "have_header"})
+            continue
+        starts = [i for i in range(l + 1) if L.ends[i] >= l and not L.blank[i]]
+        if not starts or (tb is not None and l < tb[0]):
+            bad.append({**e, "why": "no_statement"})
+            continue
+        s = max(starts)
+        end = (L.ends[s], len(L.raw[L.ends[s]]))
+        st = (s, L.indent[s])
+        t = found.setdefault(("tactic", st), Target("tactic", st, end, _cut(L.raw, st, end), L.indent[s]))
+        t.errors.append(e)
+    ts = sorted(found.values(), key=lambda t: (t.start, t.end))
+    keep = []
+    for t in ts:
+        if t.kind != "tail" and any(o is not t and o.kind != "tail" and t.start <= o.start and o.end <= t.end and (t.start, t.end) != (o.start, o.end)
+                                    for o in ts):
+            continue  # contains a finer target
+        keep.append(t)
+    return keep, bad
+
+
+# --------------------------------------------------------------------------- skeleton
+
+
+GAP = "⟪K5GAP{}⟫"
+
+
+def build_skeleton(code: str, targets: list[Target]) -> tuple[str, list[tuple[int, int]]]:
+    """(skeleton, [(line, col) of the i-th target's `sorry` in the skeleton]). Order of positions == order of
+    `targets` as given (they are applied bottom-up on a copy)."""
+    raw = code.split("\n")
+    order = sorted(range(len(targets)), key=lambda i: (targets[i].start, targets[i].end), reverse=True)
+    for i in order:
+        t = targets[i]
+        if t.kind == "have":
+            ins = " by " + GAP.format(i)
+        elif t.kind == "tactic":
+            ins = GAP.format(i)
+        else:  # tail: a new line after the last proof line
+            raw.insert(t.start[0] + 1, " " * t.indent + GAP.format(i))
+            continue
+        a, b = t.start, t.end
+        raw[a[0]:b[0] + 1] = [raw[a[0]][:a[1]] + ins + raw[b[0]][b[1]:]]
+    text = "\n".join(raw)
+    pos: dict[int, tuple[int, int]] = {}
+    out_lines = []
+    for ln_no, ln in enumerate(text.split("\n")):
+        while True:
+            m = re.search(r"⟪K5GAP(\d+)⟫", ln)
+            if not m:
+                break
+            pos[int(m.group(1))] = (ln_no, m.start())
+            ln = ln[:m.start()] + "sorry" + ln[m.end():]
+        out_lines.append(ln)
+    return "\n".join(out_lines), [pos[i] for i in range(len(targets))]
+
+
+# --------------------------------------------------------------------------- goal state -> lemma
+
+
+_INACC = re.compile(r"✝[⁰¹²³⁴⁵⁶⁷⁸⁹]*$")
+
+
+@dataclass
+class GoalState:
+    case: Optional[str]
+    hyps: list[tuple[list[str], str]]  # (names, type)
+    goal: str
+    has_let: bool = False
+
+
+def parse_goal_state(text: str) -> GoalState:
+    """Parse a Lean goal display (`case h` / `x y : T` / ... / `⊢ goal`). Wrapped lines (indented) are joined."""
+    lines = text.split("\n")
+    case = None
+    blocks: list[str] = []
+    goal_lines: list[str] = []
+    in_goal = False
+    for ln in lines:
+        if in_goal:
+            goal_lines.append(ln.strip())
+        elif ln.startswith("⊢"):
+            in_goal = True
+            goal_lines.append(ln[1:].strip())
+        elif ln.startswith("case ") and not blocks:
+            case = ln[5:].strip()
+        elif ln.startswith((" ", "\t")) and blocks:
+            blocks[-1] += " " + ln.strip()
+        elif ln.strip():
+            blocks.append(ln.strip())
+    hyps, has_let = [], False
+    for b in blocks:
+        names, sep, typ = b.partition(" : ")
+        if not sep:
+            continue
+        if re.search(r"\s:=\s", typ):
+            has_let = True
+        hyps.append((names.split(), typ.strip()))
+    return GoalState(case, hyps, " ".join(goal_lines), has_let)
+
+
+@dataclass
+class Lemma:
+    name: str
+    statement: str  # `theorem name (binders) :\n    goal := by sorry` (no header)
+    rename: list[str]  # `rename_i` argument list ("_" for instances); [] when no hypothesis is inaccessible
+
+
+def lemma_from_goal(gs: GoalState, name: str) -> Optional[Lemma]:
+    """Standalone lemma for a goal state, or None when unsupported (let-bound hypotheses).
+
+    Inaccessible hypotheses (`n✝`, shadowed or auto-named) get fresh names `k5h<i>` and every later mention of the old name
+    in a type or the goal is rewritten, so the lemma is closed; `rename` lists them (and `_` for instances) in context order
+    for the `rename_i` that gives the inlined proof the same names."""
+    if gs.has_let or not gs.goal:
+        return None
+    mp: dict[str, str] = {}
+    k = 0
+    for names, _ in gs.hyps:
+        inst = all(_INACC.search(n) and n.startswith("inst") for n in names)
+        if inst:
+            continue
+        for n in names:
+            if _INACC.search(n):
+                k += 1
+                mp[n] = f"k5h{k}"
+    pats = sorted(mp, key=len, reverse=True)
+
+    def fix(text: str) -> str:
+        for old in pats:
+            text = re.sub(r"(?<![\w.])" + re.escape(old) + r"(?![\w⁰¹²³⁴⁵⁶⁷⁸⁹])", mp[old], text)
+        return text
+
+    binders, rename = [], []
+    for names, typ in gs.hyps:
+        inst = all(_INACC.search(n) and n.startswith("inst") for n in names)
+        typ = fix(typ)
+        for n in names:
+            if _INACC.search(n):
+                rename.append("_" if inst else mp[n])
+        if inst:
+            binders += [f"[{typ}]" for _ in names]
+        else:
+            binders.append(f"({' '.join(mp.get(n, n) for n in names)} : {typ})")
+    head = " ".join(["theorem", name] + binders)
+    return Lemma(name, f"{head} :\n    {fix(gs.goal)} := by sorry", rename)
+
+
+# --------------------------------------------------------------------------- the plan
+
+
+@dataclass
+class Gap:
+    target: Target
+    sorry_pos: tuple[int, int]  # in the skeleton, full-code coordinates
+    goal_state: Optional[str] = None
+    lemma: Optional[Lemma] = None
+    why_not: Optional[str] = None
+
+
+@dataclass
+class SubgoalPlan:
+    code: str
+    skeleton: str
+    gaps: list[Gap]
+    unrepairable: list[dict]
+    residual_errors: list[dict]  # errors the skeleton still has (outer failures: next round)
+    header: str  # the header every lemma problem needs (comment-free)
+
+    @property
+    def lemmas(self) -> list[Gap]:
+        return [g for g in self.gaps if g.lemma is not None]
+
+    @property
+    def complete(self) -> bool:
+        """Every gap has a lemma and the skeleton has no error of its own."""
+        return bool(self.gaps) and all(g.lemma for g in self.gaps) and not self.residual_errors and not self.unrepairable
+
+
+Probe = Callable[[str], dict]
+"""`probe(code) -> {"errors": [...], "sorries": [{"pos": {"line","column"}, "goal": str}, ...], "view_lines": int}`: runs
+`compile_view(code)` in the statement's header environment (one REPL command) and returns the error messages and the
+`sorries` the REPL reports. `repl_probe` builds one from a pool worker."""
+
+MAX_GAPS = 6
+
+
+def header_of(statement: str) -> str:
+    """Everything before the first `theorem`/`lemma`, comments removed: the header every derived lemma must carry."""
+    m = re.search(r"\b(?:theorem|lemma)\b", statement)
+    head = statement[:m.start()] if m else statement
+    return strip_comments(head).rstrip() + "\n\n"
+
+
+def _sorry_map(resp: dict, skeleton: str) -> dict:
+    vmap = view_line_map(skeleton)
+    out = {}
+    for s in resp.get("sorries", []):
+        p = s.get("pos") or {}
+        if "line" in p and 0 <= int(p["line"]) - 1 < len(vmap):
+            out[(vmap[int(p["line"]) - 1], int(p["column"]))] = s.get("goal", "")
+    return out
+
+
+def plan_subgoals(statement: str, code: str, errors: list[dict], probe: Probe, parent_id: str = "p",
+                  repl_code: Optional[str] = None, max_gaps: int = MAX_GAPS) -> SubgoalPlan:
+    """Cut a failed proof into gaps (see the module docstring). Two passes: the have / tactic cuts first; the "proof ended
+    with goals left" tail cut only when the skeleton still shows that error (after a parse error Lean reports a cascade
+    of `unsolved goals` at the theorem's `by` that vanishes once the broken step is cut)."""
+    targets, bad = locate_targets(code, errors, repl_code)
+    header = header_of(statement)
+    if not targets or len(targets) > max_gaps:
+        return SubgoalPlan(code, code, [], bad + ([{"why": f"too_many_gaps:{len(targets)}"}] if targets else []), [], header)
+    main = [t for t in targets if t.kind != "tail"]
+    tails = [t for t in targets if t.kind == "tail"]
+    first = main or tails
+    skeleton, pos = build_skeleton(code, first)
+    resp = probe(skeleton)
+    tail_gap = None
+    if main and tails:
+        L2 = _analyse(skeleton)
+        tb2 = theorem_by_line(L2)
+        vmap2 = view_line_map(skeleton)
+        at_by = [e for e in resp.get("errors", []) if e.get("pos") and tb2 is not None
+                 and 0 <= int(e["pos"]["line"]) - 1 < len(vmap2) and vmap2[int(e["pos"]["line"]) - 1] == tb2[0]]
+        if at_by:
+            tt = tail_target(L2, tb2)
+            tt.errors = at_by
+            skeleton2, pos2 = build_skeleton(skeleton, [tt])
+            tail_gap = (tt, pos2[0])
+            skeleton = skeleton2
+            resp = probe(skeleton)
+    sorries = _sorry_map(resp, skeleton)
+    gaps = []
+    for i, t in enumerate(first):
+        gaps.append(_make_gap(t, pos[i], sorries, parent_id))
+    if tail_gap is not None:
+        gaps.append(_make_gap(tail_gap[0], tail_gap[1], sorries, parent_id))
+    resid = [e for e in resp.get("errors", []) if isinstance(e, dict)]
+    return SubgoalPlan(code, skeleton, gaps, bad, resid, header)
+
+
+def _make_gap(t: Target, pos: tuple[int, int], sorries: dict, parent_id: str) -> Gap:
+    g = Gap(t, pos)
+    goal = sorries.get(pos)
+    if goal is None:
+        g.why_not = "no_sorry_reported"
+        return g
+    g.goal_state = goal
+    h = hashlib.sha256((parent_id + "|" + goal).encode()).hexdigest()[:10]
+    lem = lemma_from_goal(parse_goal_state(goal), f"k5_sg_{h}")
+    if lem is None:
+        g.why_not = "unsupported_goal_state"
+    else:
+        g.lemma = lem
+    return g
+
+
+def lemma_problem_statement(plan: SubgoalPlan, gap: Gap) -> str:
+    """The reference statement for the prover: header + the lemma ending `:= by sorry` (what `lean_gen` problems carry)."""
+    assert gap.lemma is not None
+    return plan.header + gap.lemma.statement
+
+
+# --------------------------------------------------------------------------- reassembly
+
+
+_LEMMA_HEAD = re.compile(r"theorem\s+(\S+)[\s\S]*?:=\s*by", re.M)
+
+
+def lemma_body(lemma_code: str, lemma_name: str) -> Optional[str]:
+    """The tactic text after the lemma's own `:= by` in a prover-produced code (`lean_code` of a record), or None when
+    the code is not a proof of exactly that lemma (another declaration, no body). Comments are masked for the search."""
+    mk = mask(lemma_code)
+    m = None
+    for cand in re.finditer(r"theorem\s+" + re.escape(lemma_name) + r"\b[\s\S]*?:=\s*by", mk):
+        m = cand
+        break
+    if m is None:
+        return None
+    body = lemma_code[m.end():]
+    if re.search(r"^\s*(?:theorem|lemma)\s", mask(body), re.M):
+        return None  # a second declaration: not an inline-able body
+    return body
+
+
+def _indent_block(body: str, indent: int) -> list[str]:
+    ls = body.split("\n")
+    while ls and not ls[0].strip():
+        ls.pop(0)
+    while ls and not ls[-1].strip():
+        ls.pop()
+    if not ls:
+        return []
+    first = ls[0]
+    # `:= by simp` leaves the first tactic on the lemma's line: its text starts after one space
+    if body[:1] not in ("\n", "\r") and body.strip():
+        ls[0] = first.strip()
+        rest = [x for x in ls[1:]]
+        mi = min((len(x) - len(x.lstrip(" ")) for x in rest if x.strip()), default=0)
+        return [" " * indent + ls[0]] + [" " * indent + x[mi:] if x.strip() else "" for x in rest]
+    mi = min(len(x) - len(x.lstrip(" ")) for x in ls if x.strip())
+    return [" " * indent + x[mi:] if x.strip() else "" for x in ls]
+
+
+def reassemble(plan: SubgoalPlan, bodies: dict[int, str], restore_unproven: bool = True) -> Optional[str]:
+    """Full code with each proven gap's tactic body spliced in.
+
+    `bodies[i]` = tactic text for `plan.gaps[i]` (from `lemma_body`). An unproven gap gets its original text back when
+    `restore_unproven` (the result is then still a failing proof, but with every proven cut fixed: a better parent for
+    the next round), else the function returns None as soon as one is missing."""
+    raw = plan.skeleton.split("\n")
+    for i in sorted(range(len(plan.gaps)), key=lambda k: plan.gaps[k].sorry_pos, reverse=True):
+        g = plan.gaps[i]
+        ln, col = g.sorry_pos
+        line = raw[ln]
+        assert line[col:col + 5] == "sorry", "skeleton drifted from the plan"
+        body = bodies.get(i)
+        if body is None or not body.strip():
+            if not restore_unproven:
+                return None
+            if g.target.kind == "have":  # the skeleton holds `... := by sorry`: put the original rhs back
+                raw[ln] = line[:col - len(" by ")] + g.target.original + line[col + 5:]
+            elif g.target.kind == "tactic":
+                raw[ln] = line[:col] + g.target.original + line[col + 5:]
+            else:
+                del raw[ln]
+            continue
+        pre = ["rename_i " + " ".join(g.lemma.rename)] if g.lemma is not None and g.lemma.rename else []
+        if g.target.kind == "have":
+            ind = g.target.indent + 2
+            raw[ln:ln + 1] = [line[:col].rstrip()] + [" " * ind + p for p in pre] + _indent_block(body, ind)
+        elif g.target.kind == "tactic":
+            blk = [" " * g.target.indent + p for p in pre] + _indent_block(body, g.target.indent)
+            raw[ln:ln + 1] = [line[:col] + blk[0].lstrip()] + blk[1:]
+        else:  # tail: the inserted line holds only the gap
+            raw[ln:ln + 1] = [" " * g.target.indent + p for p in pre] + _indent_block(body, g.target.indent)
+    return "\n".join(raw)
+
+
+# --------------------------------------------------------------------------- Lean (read-only use of align.verify.lean)
+
+
+def repl_probe(worker: Any, statement: str, timeout_s: float = 120.0) -> Probe:
+    """A `Probe` over one pool worker (a `Worker` from `ReplPool.acquire`/`map`): the statement's header environment is
+    built (cached by the worker) and `compile_view(code)` is run in it, exactly as `verify_generations.collect_feedback`
+    does. Read-only: no change to the verifier, the pool or its limits."""
+    from align.verify.lean.strict import parse_statement
+
+    ref = parse_statement(statement)
+    if not ref.ok:
+        raise ValueError(f"statement not supported by the strict parser: {ref.problem}")
+
+    def probe(code: str) -> dict:
+        view = "\n".join(code.split("\n")[i] for i in view_line_map(code))
+        env = worker.header_env(ref.header_text, "")
+        resp = worker.command(view, env, timeout_s)
+        errs = [{"severity": "error", "pos": m.get("pos"), "endPos": m.get("endPos"), "data": m.get("data", "")}
+                for m in resp.get("messages", []) if isinstance(m, dict) and m.get("severity") == "error"]
+        return {"errors": errs, "sorries": resp.get("sorries", []), "view": view}
+
+    return probe
+
+
+def check_lemma_elaborates(worker: Any, plan: SubgoalPlan, gap: Gap, timeout_s: float = 120.0) -> bool:
+    """True iff `lemma_problem_statement` elaborates (no error; the `sorry` warning is expected). Gaps whose lemma does
+    not elaborate (pretty-printer round-trip failures) are dropped by the caller, not sent to the prover."""
+    from align.verify.lean.strict import parse_statement
+
+    stmt = lemma_problem_statement(plan, gap)
+    ref = parse_statement(stmt)
+    if not ref.ok:
+        return False
+    env = worker.header_env(ref.header_text, "")
+    view = stmt[len(plan.header):]
+    resp = worker.command(view, env, timeout_s)
+    return not any(isinstance(m, dict) and m.get("severity") == "error" for m in resp.get("messages", []))
+
+
+# --------------------------------------------------------------------------- (de)serialisation (route -> assemble hand-off)
+
+
+def plan_to_dict(plan: SubgoalPlan) -> dict:
+    return {
+        "code": plan.code, "skeleton": plan.skeleton, "header": plan.header,
+        "unrepairable": plan.unrepairable, "residual_errors": plan.residual_errors,
+        "gaps": [{
+            "kind": g.target.kind, "start": list(g.target.start), "end": list(g.target.end), "original": g.target.original,
+            "indent": g.target.indent, "name": g.target.name, "sorry_pos": list(g.sorry_pos), "goal_state": g.goal_state,
+            "why_not": g.why_not,
+            "lemma": None if g.lemma is None else {"name": g.lemma.name, "statement": g.lemma.statement, "rename": g.lemma.rename},
+        } for g in plan.gaps],
+    }
+
+
+def plan_from_dict(d: dict) -> SubgoalPlan:
+    gaps = []
+    for g in d["gaps"]:
+        t = Target(g["kind"], tuple(g["start"]), tuple(g["end"]), g["original"], g["indent"], g.get("name", ""))
+        lem = None if g["lemma"] is None else Lemma(g["lemma"]["name"], g["lemma"]["statement"], list(g["lemma"]["rename"]))
+        gaps.append(Gap(t, tuple(g["sorry_pos"]), g.get("goal_state"), lem, g.get("why_not")))
+    return SubgoalPlan(d["code"], d["skeleton"], gaps, d.get("unrepairable", []), d.get("residual_errors", []), d["header"])

code/align/rl/lean_rl.py

@@ -0,0 +1,659 @@
+"""Single-turn on-policy RL for Lean provers (lane R): config, rollout->batch, losses, gradient step.
+
+The run loop, checkpoints, ledger and PEFT export live in `align.rl.lean_rl_run`; the CLI is
+`scripts/p2/rl_lean.py`. Nothing here imports vLLM or runs Lean: generation and verification are
+file contracts (the P2 Generation / Verification JSONL records), so this module is testable on CPU
+with a fake generator and a fake verifier.
+
+THREE LOGPROB BASES, kept apart and stamped (Codex P13-32, audit U04-3):
+  behaviour  log pi_infer(token)  recorded by the sampler at sampling time (`sampling_logprobs`).
+             Requires top_p == 1.0 and top_k disabled: under a truncated sampler the sampled
+             distribution is NOT softmax(logits/T), so no trainer-side ratio against it is meaningful.
+  old        trainer logprobs at the rollout-time weights. Single optimizer step per iteration: this is
+             `current.detach()` (no extra forward); it is used ONLY for the engine-mismatch telemetry
+             |old - behaviour|, never as the importance denominator.
+  current    the differentiable forward of the step.
+The importance ratio is r = pi_current / pi_BEHAVIOUR, so the CISPO clamp (and the PPO clip of
+`dr_grpo`) is a trust region against the policy that actually produced the tokens, including checkpoint
+lag and the sampler/trainer kernel mismatch (U04-3). No separate rho factor exists to double count.
+
+LOSS MODES (all token-level, all normalised by an EXPLICIT global token count, see `norm_count`):
+  cispo    -A * sg(min(r, 1+eps_high)[, floor 1-eps_low]) * log pi      (MiniMax-M1 Eq. 4-5)
+  dr_grpo  PPO clip max(-A r, -A clip(r, 1-lo, 1+hi)), centered advantages, no std, no length norm
+  ei       -log pi on ACCEPTED samples only, weight 1: the same pipeline with negatives dropped
+           (K1: positive-only SFT vs signed updates on the SAME cached rollouts, Codex P13-12)
+Optional `neg_seq_ratio_floor` (cispo / dr_grpo): a sequence with A < 0 whose geometric-mean ratio is below
+the floor is masked whole (pushing down what the policy already made unlikely is the unstable direction).
+
+LOSS NORMALISATION (Codex P13-31). loss = sum_tokens(per-token loss) / N, N = the number of RESPONSE tokens
+of every sample the batch builder kept, over ALL micro-batches of ALL ranks (each rank all-reduces its local
+count and the result is checked against the count of the full list). Loss-time masks (negative-sequence
+mask, ratio clamp) do not shrink N: a masked token still counts, exactly as a clipped token does in DAPO.
+Consequence tested: the gradient is invariant to micro-batch partitioning and to the number of ranks.
+
+REWARDS. accepted -> 1.0, rejected -> 0.0. A verdict whose reason is in `no_signal_reasons` (timeout, and
+infra deaths) is MASKED: dropped from the batch AND from its group's baseline, never scored 0.
+Advantages are group-centered over the valid samples only (Dr.GRPO: no /std).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import hashlib
+import json
+import math
+import random
+from dataclasses import dataclass, field
+from typing import Any, Optional, Sequence
+
+import numpy as np
+import torch
+
+from align.rl.advantages import DEAD_ATOL, GroupFilterStats
+from align.rl.cispo import cispo_policy_loss, fp32_head_forward
+from align.rl.ops import log_softmax_and_gather
+
+LOSS_MODES = ("cispo", "dr_grpo", "ei")
+NORMS = ("global_token_mean", "constant")
+#: Verdict reasons that carry NO reward signal. `memory` / `output_too_large` are deliberately absent: a
+#: proof that kills its REPL is the policy's doing and scores 0 (scripts/p2/verify_generations.py).
+DEFAULT_NO_SIGNAL = ("timeout", "crash", "protocol")
+
+
+class RLConfigError(ValueError):
+    """A config that cannot launch."""
+
+
+# ---------------------------------------------------------------------------------------- config
+
+
+@dataclass
+class LossCfg:
+    mode: str = "cispo"
+    eps_high: float = 4.0  # CISPO IS-weight clamp (pre-registered sweep knob; paper does not state it)
+    eps_low: Optional[float] = None  # None = no lower clamp (the paper's setting)
+    clip_lower: float = 0.2  # dr_grpo PPO clip
+    clip_higher: float = 0.28
+    neg_seq_ratio_floor: Optional[float] = None
+    norm: str = "global_token_mean"
+    constant: float = 0.0  # denominator when norm == "constant"
+    fp32_head: bool = False
+    fp32_head_chunk_tokens: int = 2048
+    temperature: float = 1.0
+    max_grad_norm: float = 1.0
+
+    def validate(self) -> None:
+        if self.mode not in LOSS_MODES:
+            raise RLConfigError(f"loss.mode={self.mode!r}: one of {LOSS_MODES}")
+        if self.norm not in NORMS:
+            raise RLConfigError(f"loss.norm={self.norm!r}: one of {NORMS}")
+        if self.norm == "constant" and not self.constant > 0:
+            raise RLConfigError("loss.norm=constant needs loss.constant > 0")
+        if not (self.eps_high > 0 and math.isfinite(self.eps_high)):
+            raise RLConfigError("loss.eps_high must be finite and > 0")
+        if self.eps_low is not None and not (self.eps_low > 0 and math.isfinite(self.eps_low)):
+            raise RLConfigError("loss.eps_low must be None or finite > 0")
+        if self.neg_seq_ratio_floor is not None and not (0 < self.neg_seq_ratio_floor < 1):
+            raise RLConfigError("loss.neg_seq_ratio_floor must be in (0, 1)")
+        if self.neg_seq_ratio_floor is not None and self.mode == "ei":
+            raise RLConfigError("neg_seq_ratio_floor is meaningless in ei mode (no negatives)")
+        if not self.temperature > 0:
+            raise RLConfigError("loss.temperature must be > 0")
+
+
+@dataclass
+class BatchCfg:
+    group_size: int = 8  # G rollouts per prompt
+    min_valid_per_group: int = 2  # after masking timeouts; fewer -> group dropped
+    no_signal_reasons: tuple = DEFAULT_NO_SIGNAL
+    #: ei only: keep the positives of all-solved groups (rejection-sampling FT keeps them); signed modes always
+    #: drop dead groups (zero advantage = zero gradient, wasted tokens).
+    ei_keep_all_solved: bool = True
+    micro_tokens: int = 16384  # padded tokens per micro-batch (rows x max len)
+    max_total_len: int = 0  # >0: drop samples whose prompt+completion exceeds this (counted)
+
+
+@dataclass
+class RolloutCfg:
+    top_p: float = 1.0  # REQUIRED 1.0 (Codex P13-32)
+    top_k: int = 0
+    temperature: float = 1.0
+    cap_tokens: int = 16384
+    max_model_len: int = 20480
+    tp: int = 4
+    adapter_mode: str = "unmerged"  # "unmerged" (vLLM LoRARequest) | "merged_export" (fallback)
+    rollout_python: str = ".venv-rollout/bin/python"
+    gen_extra_args: tuple = ()
+
+    def validate(self) -> None:
+        if self.top_p != 1.0:
+            raise RLConfigError(f"rollout.top_p={self.top_p}: RL rollouts need top_p == 1.0 (Codex P13-32)")
+        if self.top_k not in (0, -1):
+            raise RLConfigError(f"rollout.top_k={self.top_k}: must be 0/-1 (disabled)")
+        if self.adapter_mode not in ("unmerged", "merged_export"):
+            raise RLConfigError(f"rollout.adapter_mode={self.adapter_mode!r}")
+
+
+@dataclass
+class OptimCfg:
+    lr: float = 1e-5
+    betas: tuple = (0.9, 0.95)
+    weight_decay: float = 0.0
+    eps: float = 1e-8
+    warmup_iters: int = 0
+    steps_per_iteration: int = 1  # optimizer steps taken on one rollout batch
+
+
+@dataclass
+class LoraCfg:
+    r: int = 64
+    alpha: float = 128.0
+    rslora: bool = False
+    targets: tuple = ("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj")
+    adapter: str = "policy"
+    param_dtype: str = "float32"
+
+
+@dataclass
+class RLConfig:
+    run_name: str = "p2-rl-lean"
+    run_dir: str = ""
+    model_snapshot: str = ""  # HF snapshot dir of the base (Goedel-8B / Kimina-Distill-8B)
+    problems: str = ""  # P2 problem jsonl (training pool)
+    iterations: int = 1
+    prompts_per_iter: int = 8
+    seed: int = 0
+    keep_checkpoints: int = 2
+    export_merged: bool = False
+    replay_from: str = ""
+    verify_workers: int = 8
+    verify_budget_s: float = 180.0
+    verify_python: str = ".venv/bin/python"
+    loss: LossCfg = field(default_factory=LossCfg)
+    batch: BatchCfg = field(default_factory=BatchCfg)
+    rollout: RolloutCfg = field(default_factory=RolloutCfg)
+    optim: OptimCfg = field(default_factory=OptimCfg)
+    lora: LoraCfg = field(default_factory=LoraCfg)
+
+    def validate(self) -> None:
+        self.loss.validate()
+        self.rollout.validate()
+        if self.loss.temperature != self.rollout.temperature:
+            raise RLConfigError(
+                f"loss.temperature={self.loss.temperature} != rollout.temperature={self.rollout.temperature}: the "
+                "trainer must score tokens under the distribution that sampled them"
+            )
+        if self.batch.group_size < 2:
+            raise RLConfigError("batch.group_size must be >= 2 (a group of 1 has no baseline)")
+        if self.batch.min_valid_per_group < 2:
+            raise RLConfigError("batch.min_valid_per_group must be >= 2")
+        if self.prompts_per_iter < 1 or self.iterations < 1:
+            raise RLConfigError("prompts_per_iter and iterations must be >= 1")
+        if self.optim.steps_per_iteration < 1:
+            raise RLConfigError("optim.steps_per_iteration must be >= 1")
+        if self.loss.mode == "cispo" and self.loss.fp32_head and self.loss.fp32_head_chunk_tokens < 1:
+            raise RLConfigError("fp32_head_chunk_tokens must be >= 1")
+
+    def fingerprint(self) -> str:
+        """Hash of everything that changes the experiment (not paths of scratch). A resume under a different
+        fingerprint is a different run and is refused."""
+        d = dataclasses.asdict(self)
+        for k in ("run_dir", "run_name", "iterations", "replay_from", "verify_workers", "keep_checkpoints", "export_merged"):
+            d.pop(k, None)
+        return hashlib.sha256(json.dumps(d, sort_keys=True, default=str).encode()).hexdigest()
+
+
+_SECTIONS = {"loss": LossCfg, "batch": BatchCfg, "rollout": RolloutCfg, "optim": OptimCfg, "lora": LoraCfg}
+
+
+def config_from_dict(d: dict) -> RLConfig:
+    """Strict: an unknown key anywhere is an error (a typo'd knob silently keeping its default is how a
+    pre-registered setting gets lost)."""
+    d = dict(d)
+    kw: dict[str, Any] = {}
+    for name, cls in _SECTIONS.items():
+        sec = dict(d.pop(name, {}))
+        names = {f.name for f in dataclasses.fields(cls)}
+        bad = sorted(set(sec) - names)
+        if bad:
+            raise RLConfigError(f"[{name}] unknown keys {bad}")
+        for k, v in list(sec.items()):
+            if isinstance(v, list):
+                sec[k] = tuple(v)
+        kw[name] = cls(**sec)
+    names = {f.name for f in dataclasses.fields(RLConfig)} - set(_SECTIONS)
+    bad = sorted(set(d) - names)
+    if bad:
+        raise RLConfigError(f"unknown top-level keys {bad}")
+    cfg = RLConfig(**d, **kw)
+    cfg.validate()
+    return cfg
+
+
+def load_config(path: str, overrides: Sequence[str] = ()) -> RLConfig:
+    import tomllib
+
+    with open(path, "rb") as f:
+        d = tomllib.load(f)
+    for ov in overrides:
+        key, _, raw = ov.partition("=")
+        try:
+            val = tomllib.loads(f"v = {raw}")["v"]
+        except tomllib.TOMLDecodeError:
+            val = raw
+        cur = d
+        parts = key.split(".")
+        for p in parts[:-1]:
+            cur = cur.setdefault(p, {})
+        cur[parts[-1]] = val
+    return config_from_dict(d)
+
+
+# ------------------------------------------------------------------------- prompt stream (U04-2)
+
+
+def prompt_block(ids: Sequence[str], seed: int, per_iter: int, k: int) -> list[str]:
+    """The prompt ids of iteration `k`: a PURE function of (ids, seed, per_iter, k).
+
+    Replays the stream from the start (cheap: id lists). Each pass over the set is an independent shuffle
+    seeded by (seed, pass index) -- epoch-indexed, so passes differ -- and the tail of one pass is carried
+    into the next block (nothing is dropped, every prompt is drawn once per pass). A block never repeats a
+    prompt unless it is larger than the pool.
+    """
+    n = len(ids)
+    if n == 0:
+        raise ValueError("empty prompt pool")
+    order: list[int] = []
+    cursor = 0
+    passes = 0
+    out: list[str] = []
+    for it in range(k + 1):
+        block: list[int] = []
+        while len(block) < per_iter:
+            if cursor >= len(order):
+                order = list(range(n))
+                random.Random(seed * 1000003 + passes * 7919 + 11).shuffle(order)
+                passes += 1
+                cursor = 0
+            i = order[cursor]
+            if per_iter <= n and i in block:
+                for j in range(cursor + 1, len(order)):
+                    if order[j] not in block:
+                        order[cursor], order[j] = order[j], order[cursor]
+                        i = order[cursor]
+                        break
+            cursor += 1
+            block.append(i)
+        out = [ids[i] for i in block]
+    return out
+
+
+def stream_position(n_pool: int, per_iter: int, k: int) -> dict:
+    """Position stored in the checkpoint (draws made, pass index); recomputable, kept as a cross-check."""
+    draws = k * per_iter
+    return {"draws": draws, "pass": draws // n_pool, "offset": draws % n_pool}
+
+
+# -------------------------------------------------------------------------------- rollout -> batch
+
+
+@dataclass
+class Sample:
+    key: tuple  # (problem_id, sample_idx, round, cap_tokens)
+    group: str  # problem_id
+    prompt_ids: list
+    completion_ids: list
+    behaviour_lp: list  # one per completion token
+    reward: float
+    adv: float = 0.0
+
+    @property
+    def n_resp(self) -> int:
+        return len(self.completion_ids)
+
+    @property
+    def total_len(self) -> int:
+        return len(self.prompt_ids) + len(self.completion_ids)
+
+
+def validate_rl_generation(rec: dict) -> None:
+    """The rollout-time contract an RL record must meet (raises with the key named)."""
+    key = (rec.get("problem_id"), rec.get("sample_idx"))
+    s = rec.get("sampling") or {}
+    if s.get("top_p") != 1.0:
+        raise RLConfigError(f"{key}: sampling.top_p={s.get('top_p')!r}; RL rollouts require top_p == 1.0")
+    if s.get("top_k") not in (0, -1, None):
+        raise RLConfigError(f"{key}: sampling.top_k={s.get('top_k')!r}; RL rollouts require top_k disabled")
+    lp = rec.get("sampling_logprobs")
+    comp = rec.get("completion_ids")
+    if not isinstance(lp, list) or not isinstance(comp, list) or len(lp) != len(comp) or not comp:
+        raise RLConfigError(f"{key}: sampling_logprobs must be a list aligned with completion_ids")
+    arr = np.asarray(lp, dtype=np.float64)
+    if not np.all(np.isfinite(arr)) or np.any(arr > 1e-4):
+        raise RLConfigError(f"{key}: sampling_logprobs must be finite and <= 0")
+
+
+def reward_of(ver: dict, no_signal: Sequence[str]) -> tuple[Optional[float], Optional[str]]:
+    """(reward, None) or (None, why-masked). Timeouts / infra deaths carry NO signal."""
+    reasons = list(ver.get("reasons") or [])
+    if ver.get("accepted"):
+        return 1.0, None
+    hit = [r for r in reasons if r in no_signal]
+    if hit:
+        return None, hit[0]
+    return 0.0, None
+
+
+def build_samples(gens: Sequence[dict], vers: dict, cfg: RLConfig) -> tuple[list[Sample], dict]:
+    """Join generations with verdicts -> trainable samples + telemetry.
+
+    Dead groups (all-rewards-equal over the VALID samples) are dropped in signed modes; `ei` keeps the
+    positives of all-solved groups when `batch.ei_keep_all_solved`. Masked (no-signal) samples are removed
+    before the baseline is computed.
+    """
+    b = cfg.batch
+    stats = {"samples": len(gens), "masked": collections_counter(), "groups": 0, "groups_too_few_valid": 0,
+             "dropped_overlength": 0, "accepted": 0, "valid": 0}
+    by_group: dict[str, list[Sample]] = {}
+    for g in gens:
+        validate_rl_generation(g)
+        key = (g["problem_id"], g["sample_idx"], g["round"], g["cap_tokens"])
+        v = vers.get(key)
+        if v is None:
+            raise RLConfigError(f"{key}: generation has no verification record (verify before training)")
+        r, why = reward_of(v, b.no_signal_reasons)
+        if r is None:
+            stats["masked"][why] += 1
+            continue
+        s = Sample(key, g["problem_id"], list(g["prompt_ids"]), list(g["completion_ids"]),
+                   [float(x) for x in g["sampling_logprobs"]], r)
+        if b.max_total_len and s.total_len > b.max_total_len:
+            stats["dropped_overlength"] += 1
+            continue
+        stats["valid"] += 1
+        stats["accepted"] += int(r == 1.0)
+        by_group.setdefault(g["problem_id"], []).append(s)
+    fstats = GroupFilterStats(max_possible_score=1.0)
+    out: list[Sample] = []
+    for pid in sorted(by_group):
+        grp = sorted(by_group[pid], key=lambda s: s.key)
+        stats["groups"] += 1
+        if len(grp) < b.min_valid_per_group:
+            stats["groups_too_few_valid"] += 1
+            continue
+        rewards = np.array([s.reward for s in grp])
+        dead = fstats.observe(rewards)
+        if cfg.loss.mode == "ei":
+            keep = [s for s in grp if s.reward == 1.0]
+            all_solved = bool(np.all(rewards == 1.0))
+            if not keep or (all_solved and not b.ei_keep_all_solved):
+                continue
+            for s in keep:
+                s.adv = 1.0
+            out.extend(keep)
+            continue
+        if dead:
+            continue
+        mean = float(rewards.mean())
+        for s in grp:
+            s.adv = s.reward - mean
+        out.extend(grp)
+    stats["masked"] = dict(stats["masked"])
+    stats["group_filter"] = fstats.as_log_dict()
+    stats["kept_samples"] = len(out)
+    stats["kept_resp_tokens"] = sum(s.n_resp for s in out)
+    return out, stats
+
+
+def collections_counter():
+    import collections
+
+    return collections.Counter()
+
+
+# ------------------------------------------------------------------------------- sharding / packing
+
+
+def shard_samples(samples: Sequence[Sample], world: int) -> list[list[Sample]]:
+    """Deterministic length-balanced split (longest-first onto the lightest rank)."""
+    order = sorted(samples, key=lambda s: (-s.total_len, s.key))
+    loads = [0] * world
+    shards: list[list[Sample]] = [[] for _ in range(world)]
+    for s in order:
+        r = min(range(world), key=lambda i: (loads[i], i))
+        shards[r].append(s)
+        loads[r] += s.total_len
+    return shards
+
+
+@dataclass
+class Micro:
+    input_ids: torch.Tensor  # [B, T]
+    attention_mask: torch.Tensor  # [B, T]
+    resp: torch.Tensor  # [B, T-1] bool, label space
+    behaviour: torch.Tensor  # [B, T-1]
+    adv: torch.Tensor  # [B]
+    dummy: bool = False
+
+    @property
+    def n_resp(self) -> int:
+        return 0 if self.dummy else int(self.resp.sum())
+
+
+def pack_micros(samples: Sequence[Sample], micro_tokens: int, pad_id: int = 0, dtype=torch.float32) -> list[Micro]:
+    """Token-budget micro-batches (rows x padded max length <= micro_tokens; a longer lone row gets its own)."""
+    rows = sorted(samples, key=lambda s: (s.total_len, s.key))
+    micros: list[Micro] = []
+    cur: list[Sample] = []
+    cur_max = 0
+    for s in rows:
+        m = max(cur_max, s.total_len)
+        if cur and m * (len(cur) + 1) > micro_tokens:
+            micros.append(_collate(cur, pad_id, dtype))
+            cur, m = [], s.total_len
+        cur.append(s)
+        cur_max = m
+    if cur:
+        micros.append(_collate(cur, pad_id, dtype))
+    return micros
+
+
+def _collate(rows: Sequence[Sample], pad_id: int, dtype) -> Micro:
+    T = max(s.total_len for s in rows)
+    B = len(rows)
+    ids = torch.full((B, T), pad_id, dtype=torch.int64)
+    att = torch.zeros((B, T), dtype=torch.int64)
+    resp = torch.zeros((B, T - 1), dtype=torch.bool)
+    beh = torch.zeros((B, T - 1), dtype=dtype)
+    adv = torch.zeros(B, dtype=dtype)
+    for i, s in enumerate(rows):
+        n = s.total_len
+        ids[i, :n] = torch.tensor(s.prompt_ids + s.completion_ids)
+        att[i, :n] = 1
+        p = len(s.prompt_ids)
+        resp[i, p - 1 : n - 1] = True  # label t+1 in the completion
+        beh[i, p - 1 : n - 1] = torch.tensor(s.behaviour_lp, dtype=dtype)
+        adv[i] = s.adv
+    return Micro(ids, att, resp, beh, adv)
+
+
+def dummy_micro(like: Optional[Micro], pad_id: int = 0, dtype=torch.float32) -> Micro:
+    """A micro that trains on nothing but runs the same forward/backward (so every rank issues the same
+    collectives): one row, empty response mask, zero advantage."""
+    if like is not None:
+        n = int(like.attention_mask[0].sum())
+        ids = like.input_ids[:1, :n].clone()
+    else:
+        ids = torch.full((1, 2), pad_id, dtype=torch.int64)
+    T = ids.shape[1]
+    return Micro(ids, torch.ones_like(ids), torch.zeros((1, T - 1), dtype=torch.bool),
+                 torch.zeros((1, T - 1), dtype=dtype), torch.zeros(1, dtype=dtype), dummy=True)
+
+
+# ---------------------------------------------------------------------------------------- the loss
+
+
+def token_loss(new_lp: torch.Tensor, behaviour: torch.Tensor, adv: torch.Tensor, resp: torch.Tensor,
+               cfg: LossCfg) -> tuple[torch.Tensor, dict]:
+    """Unreduced per-token loss [B, T-1] (zero off the response mask) + diagnostic tensors.
+
+    ratio = exp(new - behaviour): importance weight against the BEHAVIOUR policy (U04-3).
+    """
+    m = resp.to(new_lp.dtype)
+    A = adv.to(new_lp.dtype).unsqueeze(-1).expand_as(new_lp)
+    diff = (new_lp - behaviour).clamp(-20.0, 20.0)
+    ratio = torch.exp(diff)
+    clipfrac = torch.zeros_like(new_lp)
+    row_w = torch.ones(new_lp.shape[0], dtype=new_lp.dtype, device=new_lp.device)
+    if cfg.mode == "ei":
+        loss = -new_lp
+    else:
+        if cfg.mode == "cispo":
+            loss, clipfrac = cispo_policy_loss(new_lp, ratio, A, cfg.eps_high, cfg.eps_low)
+        else:  # dr_grpo
+            l1 = -A * ratio
+            l2 = -A * torch.clamp(ratio, 1.0 - cfg.clip_lower, 1.0 + cfg.clip_higher)
+            loss = torch.max(l1, l2)
+            clipfrac = (l2 > l1).to(new_lp.dtype)
+        if cfg.neg_seq_ratio_floor is not None:
+            n = m.sum(-1).clamp_min(1.0)
+            geo = torch.exp((diff.detach() * m).sum(-1) / n)
+            drop = (adv < 0) & (geo < cfg.neg_seq_ratio_floor)
+            row_w = (~drop).to(new_lp.dtype)
+    loss = loss * m * row_w.unsqueeze(-1)
+    return loss, {"ratio": ratio.detach(), "clipfrac": clipfrac.detach() * m, "row_w": row_w.detach()}
+
+
+def policy_logprobs(model, mb: Micro, cfg: LossCfg) -> torch.Tensor:
+    """Label-aligned [B, T-1] logprobs of the batch's own tokens under the CURRENT policy."""
+    if cfg.fp32_head:
+        return fp32_head_forward(model, mb.input_ids, mb.attention_mask, cfg.temperature, cfg.fp32_head_chunk_tokens)
+    out = model(input_ids=mb.input_ids, attention_mask=mb.attention_mask)
+    logits = getattr(out, "logits", out)
+    if cfg.temperature != 1.0:
+        logits = logits / cfg.temperature
+    return log_softmax_and_gather(logits[:, :-1], mb.input_ids[:, 1:])
+
+
+# ------------------------------------------------------------------------------- distributed bits
+
+
+def _dist_on() -> bool:
+    import torch.distributed as dist
+
+    return dist.is_available() and dist.is_initialized()
+
+
+def world_size() -> int:
+    import torch.distributed as dist
+
+    return dist.get_world_size() if _dist_on() else 1
+
+
+def all_reduce_ints(vals: list[int], op: str = "sum") -> list[int]:
+    if not _dist_on():
+        return [int(v) for v in vals]
+    import torch.distributed as dist
+
+    from align.utils.dist import gloo_group
+
+    g = gloo_group() or dist.group.WORLD
+    t = torch.tensor([int(v) for v in vals], dtype=torch.int64)
+    dist.all_reduce(t, op=dist.ReduceOp.SUM if op == "sum" else dist.ReduceOp.MAX, group=g)
+    return [int(v) for v in t.tolist()]
+
+
+def all_reduce_floats(vals: list[float]) -> list[float]:
+    if not _dist_on():
+        return [float(v) for v in vals]
+    import torch.distributed as dist
+
+    from align.utils.dist import gloo_group
+
+    g = gloo_group() or dist.group.WORLD
+    t = torch.tensor([float(v) for v in vals], dtype=torch.float64)
+    dist.all_reduce(t, op=dist.ReduceOp.SUM, group=g)
+    return [float(v) for v in t.tolist()]
+
+
+def norm_count(local_resp_tokens: int, cfg: LossCfg, expect_total: Optional[int] = None) -> float:
+    """The explicit loss denominator N (module docstring). Global = SUM over ranks of the local response-token
+    counts, verified against the full-list count when the caller can supply it."""
+    if cfg.norm == "constant":
+        return float(cfg.constant)
+    total = all_reduce_ints([local_resp_tokens])[0]
+    if expect_total is not None and total != expect_total:
+        raise RuntimeError(f"global token count {total} != full-batch count {expect_total}: ranks hold different data")
+    if total <= 0:
+        raise RuntimeError("global response-token count is 0")
+    return float(total)
+
+
+def _is_sharded(params: Sequence[torch.nn.Parameter]) -> bool:
+    return any(hasattr(p, "to_local") for p in params)
+
+
+def accumulate_gradients(model, micros: Sequence[Micro], cfg: LossCfg, denom: float,
+                         params: Sequence[torch.nn.Parameter], pad_id: int = 0) -> dict:
+    """Forward+backward over this rank's micro-batches with the GLOBAL denominator; returns summed stats.
+
+    Every rank runs the same number of micro-steps (agreed by MAX; short ranks run zero-weight dummies).
+    Gradient scale: loss_i / N. FSDP2 reduce-scatter AVERAGES over ranks, so a sharded model multiplies by the
+    world size; a replicated model (CPU tests, single GPU) sums gradients explicitly with an all-reduce SUM.
+    Either way grad = d/dtheta [ sum over ALL ranks/micros/tokens of per-token loss / N ].
+    """
+    world = world_size()
+    n_agreed = all_reduce_ints([len(micros)], "max")[0]
+    sharded = _is_sharded(params)
+    scale = float(world) if (sharded and world > 1) else 1.0
+    like = micros[0] if micros else None
+    dtype = like.behaviour.dtype if like is not None else torch.float32
+    acc = {"loss_sum": 0.0, "resp_tokens": 0.0, "clip_tokens": 0.0, "ratio_sum": 0.0, "rows": 0.0,
+           "rows_masked_negseq": 0.0, "mismatch_abs_sum": 0.0, "adv_pos_rows": 0.0, "entropy_proxy": 0.0}
+    for i in range(n_agreed):
+        mb = micros[i] if i < len(micros) else dummy_micro(like, pad_id, dtype)
+        new_lp = policy_logprobs(model, mb, cfg)
+        tok, d = token_loss(new_lp, mb.behaviour.to(new_lp.dtype), mb.adv, mb.resp, cfg)
+        loss = tok.sum() / denom * scale
+        loss.backward()
+        with torch.no_grad():
+            m = mb.resp.to(torch.float64)
+            acc["loss_sum"] += float(tok.detach().double().sum())
+            acc["resp_tokens"] += float(m.sum())
+            acc["clip_tokens"] += float((d["clipfrac"].double() * m).sum())
+            acc["ratio_sum"] += float((d["ratio"].double() * m).sum())
+            acc["mismatch_abs_sum"] += float(((new_lp.detach() - mb.behaviour.to(new_lp.dtype)).abs().double() * m).sum())
+            if not mb.dummy:
+                acc["rows"] += float(mb.input_ids.shape[0])
+                acc["rows_masked_negseq"] += float((d["row_w"] == 0).sum())
+                acc["adv_pos_rows"] += float((mb.adv > 0).sum())
+    if world > 1 and not sharded:
+        import torch.distributed as dist
+
+        from align.utils.dist import gloo_group
+
+        g = gloo_group() or dist.group.WORLD
+        for p in params:
+            if p.grad is None:
+                p.grad = torch.zeros_like(p)
+            dist.all_reduce(p.grad, op=dist.ReduceOp.SUM, group=g)
+    keys = sorted(acc)
+    red = all_reduce_floats([acc[k] for k in keys])
+    return dict(zip(keys, red))
+
+
+def summarize(stats: dict, denom: float, grad_norm: float, cfg: LossCfg) -> dict:
+    t = max(stats["resp_tokens"], 1.0)
+    return {
+        "loss": stats["loss_sum"] / denom,
+        "denominator": denom,
+        "resp_tokens": stats["resp_tokens"],
+        "rows": stats["rows"],
+        "clipfrac": stats["clip_tokens"] / t,
+        "ratio_mean": stats["ratio_sum"] / t,
+        # engine mismatch |log pi_current(=old at step 0) - log pi_behaviour|, mean over response tokens
+        "old_vs_behaviour_abs_logratio": stats["mismatch_abs_sum"] / t,
+        "negseq_masked_rows": stats["rows_masked_negseq"],
+        "grad_norm": grad_norm,
+        "loss_mode": cfg.mode,
+    }

code/align/rl/lean_rl_run.py

@@ -0,0 +1,716 @@
+"""Lean RL run loop: iteration stages, atomic checkpoints, PEFT export, budget ledger, backends.
+
+    iteration k (policy = ckpt/iter_k's adapter):
+        gen      generator.generate -> iters/iter_k/gens.jsonl      (CURRENT policy; behaviour logprobs recorded)
+        verify   verifier.verify    -> iters/iter_k/ver.jsonl
+        train    trainer.train(k)   -> ckpt/iter_{k+1}/             (adapter, optimizer, RNG, stream, ledger)
+
+Every stage is idempotent and commits with a marker (`gen.done`, `ver.done`) or a directory rename
+(checkpoints), so a kill at any point resumes into the same state: a killed `train` leaves only a `.tmp_*`
+directory, which is deleted, and the stage re-runs from the same gens/ver files and the same ckpt k.
+
+GPU HANDOFF. The default real flow keeps the loop process CPU-only: generation (vLLM, tp=4) and training
+(torchrun, FSDP2) are separate child processes that run one after the other, so neither holds the GPUs while the
+other runs; the loop itself never imports a model. `InProcessTrainer` (tests, single-device runs) calls
+`train_iteration` directly. This module imports torch/numpy/safetensors/stdlib/align only.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+from typing import Callable, Optional
+
+import torch
+
+from align.rl import lean_rl as L
+from align.utils.dist import capture_rng_state, restore_rng_state
+
+ADAPTER_NAME_PEFT = "adapter_model.safetensors"
+
+
+# ------------------------------------------------------------------------------------ fs helpers
+
+
+def fsync_dir(path: str) -> None:
+    fd = os.open(path, os.O_RDONLY)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def atomic_write_text(path: str, text: str) -> None:
+    tmp = path + ".partial"
+    with open(tmp, "w") as f:
+        f.write(text)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp, path)
+
+
+def read_jsonl(path: str) -> list[dict]:
+    out = []
+    with open(path, encoding="utf-8") as f:
+        for n, line in enumerate(f, 1):
+            if line.strip():
+                try:
+                    out.append(json.loads(line))
+                except json.JSONDecodeError as e:
+                    raise ValueError(f"{path}:{n}: {e}") from e
+    return out
+
+
+def write_jsonl_atomic(path: str, rows) -> None:
+    atomic_write_text(path, "".join(json.dumps(r, sort_keys=True) + "\n" for r in rows))
+
+
+def sha256_file(path: str) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for c in iter(lambda: f.read(1 << 20), b""):
+            h.update(c)
+    return h.hexdigest()
+
+
+# --------------------------------------------------------------------------------- PEFT export
+
+
+def peft_key_map(qcfg) -> dict[str, str]:
+    """align module path -> HF module path, for every projection LoRA can target."""
+    from align.model.qwen3_checkpoint import _table
+
+    out = {}
+    for hf_w, al_w in _table(qcfg).items():
+        if hf_w.endswith(".weight") and ".layers." in hf_w:
+            out[al_w[: -len(".weight")]] = hf_w[: -len(".weight")]
+    return out
+
+
+def lora_to_peft_tensors(lora_state: dict, qcfg) -> dict[str, torch.Tensor]:
+    m = peft_key_map(qcfg)
+    out = {}
+    for k, v in lora_state.items():
+        path, ab = k.rsplit(".", 1)
+        if path not in m:
+            raise KeyError(f"adapter path {path!r} has no HF counterpart")
+        out[f"base_model.model.{m[path]}.lora_{ab}.weight"] = v.detach().to("cpu").contiguous()
+    return out
+
+
+def peft_to_lora_state(tensors: dict, qcfg) -> dict[str, torch.Tensor]:
+    inv = {v: k for k, v in peft_key_map(qcfg).items()}
+    out = {}
+    for k, v in tensors.items():
+        assert k.startswith("base_model.model.")
+        body = k[len("base_model.model."):]
+        hf_path, lora_ab = body.rsplit(".lora_", 1)
+        out[f"{inv[hf_path]}.{lora_ab.split('.')[0]}"] = v
+    return out
+
+
+def write_peft_dir(out_dir: str, lora_state: dict, qcfg, lcfg: L.LoraCfg, base_name: str, dtype=torch.float32) -> None:
+    """A PEFT adapter directory (adapter_config.json + adapter_model.safetensors), written to a sibling temp dir
+    and renamed in, so a reader (vLLM) never sees half of it."""
+    import safetensors.torch as st
+
+    tmp = out_dir.rstrip("/") + ".partial"
+    shutil.rmtree(tmp, ignore_errors=True)
+    os.makedirs(tmp)
+    tens = {k: v.to(dtype) for k, v in lora_to_peft_tensors(lora_state, qcfg).items()}
+    st.save_file(tens, os.path.join(tmp, ADAPTER_NAME_PEFT), metadata={"format": "pt"})
+    cfg = {
+        "peft_type": "LORA", "task_type": "CAUSAL_LM", "r": lcfg.r, "lora_alpha": lcfg.alpha, "lora_dropout": 0.0,
+        "bias": "none", "use_rslora": bool(lcfg.rslora), "target_modules": list(lcfg.targets),
+        "base_model_name_or_path": base_name, "inference_mode": True, "fan_in_fan_out": False,
+        "init_lora_weights": True, "modules_to_save": None,
+    }
+    with open(os.path.join(tmp, "adapter_config.json"), "w") as f:
+        json.dump(cfg, f, indent=1, sort_keys=True)
+        f.flush()
+        os.fsync(f.fileno())
+    if os.path.isdir(out_dir):
+        shutil.rmtree(out_dir)
+    os.replace(tmp, out_dir)
+
+
+def read_peft_dir(path: str, qcfg) -> dict[str, torch.Tensor]:
+    import safetensors.torch as st
+
+    return peft_to_lora_state(st.load_file(os.path.join(path, ADAPTER_NAME_PEFT)), qcfg)
+
+
+def export_merged_checkpoint(model, adapter: str, qcfg, out_dir: str, tokenizer_dir: Optional[str] = None,
+                             dtype=torch.bfloat16) -> dict:
+    """FALLBACK for serving without un-merged LoRA: base+adapter folded into an HF checkpoint. Lossy by
+    construction (bf16 cast drops small delta elements; the fidelity report is returned). Collective under FSDP:
+    every rank must call; rank 0 writes."""
+    from align.model.qwen3_checkpoint import save_hf_qwen3
+    from align.ppo.lora import export_merged_state_dict
+    from align.utils.dist import is_rank0
+
+    if L._is_sharded(list(model.parameters())):
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+
+        sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+    else:
+        sd = {k: v.detach().to("cpu") for k, v in model.state_dict().items()}
+    report: dict = {}
+    if sd:
+        merged = export_merged_state_dict(model, sd, adapter, report=report)
+        if is_rank0():
+            tmp = out_dir.rstrip("/") + ".partial"
+            shutil.rmtree(tmp, ignore_errors=True)
+            save_hf_qwen3(merged, tmp, cfg=qcfg, dtype=dtype, tokenizer_dir=tokenizer_dir)
+            if os.path.isdir(out_dir):
+                shutil.rmtree(out_dir)
+            os.replace(tmp, out_dir)
+    return report
+
+
+# --------------------------------------------------------------------------------- checkpoints
+
+
+def ckpt_path(run_dir: str, k: int) -> str:
+    return os.path.join(run_dir, "ckpt", f"iter_{k:06d}")
+
+
+def latest_complete(run_dir: str) -> Optional[int]:
+    root = os.path.join(run_dir, "ckpt")
+    if not os.path.isdir(root):
+        return None
+    best = None
+    for name in os.listdir(root):
+        if name.startswith("iter_") and os.path.isfile(os.path.join(root, name, "COMPLETE")):
+            k = int(name[5:])
+            best = k if best is None else max(best, k)
+    return best
+
+
+def clean_partial(run_dir: str) -> list[str]:
+    """Delete half-written checkpoints (`.tmp_*`) and incomplete `iter_*` dirs: they look resumable and are not."""
+    root = os.path.join(run_dir, "ckpt")
+    gone = []
+    if not os.path.isdir(root):
+        return gone
+    for name in sorted(os.listdir(root)):
+        p = os.path.join(root, name)
+        if name.startswith(".tmp_") or (name.startswith("iter_") and not os.path.isfile(os.path.join(p, "COMPLETE"))):
+            shutil.rmtree(p, ignore_errors=True)
+            gone.append(name)
+    return gone
+
+
+def _rank() -> int:
+    import torch.distributed as dist
+
+    return dist.get_rank() if L._dist_on() else 0
+
+
+def _barrier() -> None:
+    import torch.distributed as dist
+
+    if L._dist_on():
+        dist.barrier()
+
+
+def save_checkpoint(run_dir: str, k: int, *, model, optimizer, lora_state: dict, qcfg, cfg: L.RLConfig, ledger: list,
+                    bases: dict, base_name: str, base_fingerprint: Optional[str]) -> str:
+    """Commit checkpoint k atomically: write everything into `.tmp_iter_k`, then rename. `COMPLETE` is the
+    last file written inside; the rename is the commit. `lora_state` must already be gathered (collective done
+    by the caller on every rank)."""
+    final = ckpt_path(run_dir, k)
+    tmp = os.path.join(run_dir, "ckpt", f".tmp_iter_{k:06d}")
+    r = _rank()
+    if r == 0:
+        os.makedirs(os.path.join(run_dir, "ckpt"), exist_ok=True)
+        shutil.rmtree(tmp, ignore_errors=True)
+        os.makedirs(tmp)
+    _barrier()
+    optim_sd = optimizer_state_for_save(model, optimizer)  # collective on every rank
+    if r == 0:
+        torch.save(_cpu_tree(optim_sd), os.path.join(tmp, "optim.pt"))
+        state = {
+            "iteration": k, "cfg_fingerprint": cfg.fingerprint(), "world_size": L.world_size(),
+            "adapter": {a: t.detach().cpu() for a, t in lora_state.items()}, "rng": capture_rng_state(),
+            "stream": L.stream_position(_pool_size(cfg), cfg.prompts_per_iter, k), "ledger": ledger,
+            "bases": bases, "base_fingerprint": base_fingerprint,
+        }
+        torch.save(state, os.path.join(tmp, "state.pt"))
+        write_peft_dir(os.path.join(tmp, "adapter"), lora_state, qcfg, cfg.lora, base_name)
+        write_jsonl_atomic(os.path.join(tmp, "ledger.jsonl"), ledger)
+        atomic_write_text(os.path.join(tmp, "COMPLETE"), json.dumps({"iteration": k, "t": time.time()}))
+    _barrier()
+    if r == 0:
+        if os.path.isdir(final):
+            shutil.rmtree(final)
+        fsync_dir(tmp)
+        os.replace(tmp, final)
+        fsync_dir(os.path.dirname(final))
+        write_jsonl_atomic(os.path.join(run_dir, "ledger.jsonl"), ledger)
+        _prune(run_dir, cfg.keep_checkpoints)
+    _barrier()
+    return final
+
+
+def _prune(run_dir: str, keep: int) -> None:
+    root = os.path.join(run_dir, "ckpt")
+    done = sorted(n for n in os.listdir(root) if n.startswith("iter_") and os.path.isfile(os.path.join(root, n, "COMPLETE")))
+    for n in done[:-max(keep, 1)]:
+        shutil.rmtree(os.path.join(root, n), ignore_errors=True)
+
+
+def load_checkpoint_state(run_dir: str, k: int, cfg: L.RLConfig) -> dict:
+    d = ckpt_path(run_dir, k)
+    if not os.path.isfile(os.path.join(d, "COMPLETE")):
+        raise FileNotFoundError(f"no complete checkpoint {d}")
+    st = torch.load(os.path.join(d, "state.pt"), weights_only=True)
+    if st["cfg_fingerprint"] != cfg.fingerprint():
+        raise RuntimeError("checkpoint was written under a different config fingerprint; refusing to resume")
+    return st
+
+
+def optimizer_state_for_save(model, optimizer) -> dict:
+    """Optimizer state in a world-size-independent form. Replicated model: the local state (identical on every
+    rank). FSDP2-sharded: the FULL state via DCP (a collective; rank 0 receives it, the others {})."""
+    if L._is_sharded(list(model.parameters())):
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, get_optimizer_state_dict
+
+        return get_optimizer_state_dict(model, optimizer, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+    return optimizer.state_dict()
+
+
+def load_optimizer(run_dir: str, k: int, model, optimizer) -> None:
+    p = os.path.join(ckpt_path(run_dir, k), "optim.pt")
+    sd = torch.load(p, weights_only=True)
+    if L._is_sharded(list(model.parameters())):
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_optimizer_state_dict
+
+        set_optimizer_state_dict(model, optimizer, optim_state_dict=sd,
+                                 options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=False))
+    else:
+        optimizer.load_state_dict(sd)
+
+
+def _cpu_tree(x):
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu()
+    if isinstance(x, dict):
+        return {k: _cpu_tree(v) for k, v in x.items()}
+    if isinstance(x, (list, tuple)):
+        return type(x)(_cpu_tree(v) for v in x)
+    return x
+
+
+def _pool_size(cfg: L.RLConfig) -> int:
+    return len(read_jsonl(cfg.problems))
+
+
+def load_problems(path: str) -> list[dict]:
+    rows = read_jsonl(path)
+    ids = [r["problem_id"] for r in rows]
+    if len(set(ids)) != len(ids):
+        raise ValueError(f"{path}: duplicate problem_id")
+    return rows
+
+
+# ------------------------------------------------------------------------------- iteration files
+
+
+def iter_dir(run_dir: str, k: int) -> str:
+    return os.path.join(run_dir, "iters", f"iter_{k:04d}")
+
+
+def stage_marker(d: str, name: str) -> str:
+    return os.path.join(d, f"{name}.done")
+
+
+def add_stage_time(d: str, stage: str, wall_s: float) -> None:
+    p = os.path.join(d, "stage_times.json")
+    cur = json.load(open(p)) if os.path.exists(p) else {}
+    cur[f"{stage}_wall_s"] = cur.get(f"{stage}_wall_s", 0.0) + wall_s
+    cur[f"{stage}_attempts"] = cur.get(f"{stage}_attempts", 0) + 1
+    atomic_write_text(p, json.dumps(cur, sort_keys=True))
+
+
+def ledger_entry(cfg: L.RLConfig, k: int, d: str, bstats: dict, metrics: dict, train_wall_s: float,
+                 padded_tokens: float, prev: list) -> dict:
+    gens = read_jsonl(os.path.join(d, "gens.jsonl"))
+    vers = read_jsonl(os.path.join(d, "ver.jsonl"))
+    st = json.load(open(os.path.join(d, "stage_times.json"))) if os.path.exists(os.path.join(d, "stage_times.json")) else {}
+    gen_tokens = sum(len(g["completion_ids"]) for g in gens)
+    ver_s = sum(float(v.get("verify_s", 0.0)) for v in vers)
+    e = {
+        "iteration": k, "loss_mode": cfg.loss.mode, "prompts": cfg.prompts_per_iter, "samples": len(gens),
+        "generated_tokens": gen_tokens, "prompt_tokens": sum(len(g["prompt_ids"]) for g in gens),
+        "verifier_seconds": ver_s, "verifier_wall_s": st.get("verify_wall_s", 0.0),
+        "generate_wall_s": st.get("gen_wall_s", 0.0), "train_wall_s": train_wall_s,
+        "train_tokens": metrics.get("resp_tokens_total", 0.0), "train_padded_tokens": padded_tokens,
+        "stage_attempts": {s: st.get(f"{s}_attempts", 0) for s in ("gen", "verify")},
+        "updated": bool(metrics.get("updated", False)), "batch": bstats, "metrics": metrics,
+        "bases": metrics.get("bases", {}),
+    }
+    e["wall_s"] = e["generate_wall_s"] + e["verifier_wall_s"] + train_wall_s
+    keys = ("generated_tokens", "verifier_seconds", "train_tokens", "train_padded_tokens", "wall_s")
+    e["cumulative"] = {c: (prev[-1]["cumulative"][c] if prev else 0.0) + e[c] for c in keys}
+    return e
+
+
+# ---------------------------------------------------------------------------- the train iteration
+
+
+def bases_stamp(cfg: L.RLConfig, gens: list) -> dict:
+    g0 = gens[0] if gens else {}
+    s = g0.get("sampling", {})
+    return {
+        "behaviour": {"source": "sampler", "field": "sampling_logprobs", "temperature": s.get("temperature"),
+                      "top_p": s.get("top_p"), "top_k": s.get("top_k"), "model_rev": g0.get("model_rev"),
+                      "serving": (g0.get("adapter") or {}).get("serving", "merged_or_base")},
+        "old": {"source": "trainer", "weights": "rollout-time (= current at the first optimizer step)",
+                "role": "mismatch telemetry only; NOT the importance denominator", "fp32_head": cfg.loss.fp32_head},
+        "current": {"source": "trainer", "differentiable": True, "fp32_head": cfg.loss.fp32_head,
+                    "temperature": cfg.loss.temperature},
+        "importance_ratio": "pi_current / pi_behaviour",
+    }
+
+
+def lr_at(cfg: L.RLConfig, k: int) -> float:
+    w = cfg.optim.warmup_iters
+    return cfg.optim.lr * (min(1.0, (k + 1) / w) if w > 0 else 1.0)
+
+
+def train_iteration(cfg: L.RLConfig, k: int, *, model, params, optimizer, qcfg, pad_id: int, base_name: str,
+                    base_fingerprint: Optional[str] = None, tokenizer_dir: Optional[str] = None) -> dict:
+    """One training iteration on iteration k's cached gens/ver, committing checkpoint k+1.
+
+    `model` carries the injected LoRA adapter `cfg.lora.adapter`; `params` are its trainable parameters;
+    `optimizer` is built over them. State is loaded from checkpoint k (adapter, optimizer, RNG)."""
+    from align.ppo.lora import load_lora_state_dict, lora_state_dict
+
+    run = cfg.run_dir
+    d = iter_dir(run, k)
+    st = load_checkpoint_state(run, k, cfg)
+    load_lora_state_dict(model, cfg.lora.adapter, st["adapter"], strict=True)
+    load_optimizer(run, k, model, optimizer)
+    restore_rng_state(st["rng"])
+    ledger = list(st["ledger"])
+    if len(ledger) != k:
+        raise RuntimeError(f"checkpoint {k} carries {len(ledger)} ledger rows")
+    pos = L.stream_position(_pool_size(cfg), cfg.prompts_per_iter, k)
+    if st["stream"] != pos:
+        raise RuntimeError(f"prompt stream position {st['stream']} != recomputed {pos}")
+
+    gens = read_jsonl(os.path.join(d, "gens.jsonl"))
+    vers = {(v["problem_id"], v["sample_idx"], v["round"], v["cap_tokens"]): v for v in read_jsonl(os.path.join(d, "ver.jsonl"))}
+    samples, bstats = L.build_samples(gens, vers, cfg)
+    world = L.world_size()
+    r = _rank()
+    t0 = time.perf_counter()
+    metrics: dict = {"updated": False, "bases": bases_stamp(cfg, gens), "resp_tokens_total": 0.0}
+    padded = 0.0
+    if samples:
+        dtype = next(iter(params)).dtype
+        steps = cfg.optim.steps_per_iteration
+        groups = sorted({s.group for s in samples})
+        chunks = [[s for s in samples if groups.index(s.group) % steps == i] for i in range(steps)]
+        step_metrics = []
+        for i, chunk in enumerate(chunks):
+            if not chunk:
+                continue
+            mine = L.shard_samples(chunk, world)[r]
+            micros = L.pack_micros(mine, cfg.batch.micro_tokens, pad_id, dtype)
+            expect = sum(s.n_resp for s in chunk)
+            denom = L.norm_count(sum(m.n_resp for m in micros), cfg.loss, expect)
+            for p in params:
+                p.grad = None
+            t_step = time.perf_counter()
+            acc = L.accumulate_gradients(model, micros, cfg.loss, denom, params, pad_id)
+            total = torch.nn.utils.clip_grad_norm_(list(params), cfg.loss.max_grad_norm)
+            if hasattr(total, "full_tensor"):
+                total = total.full_tensor()
+            for g in optimizer.param_groups:
+                g["lr"] = lr_at(cfg, k)
+            optimizer.step()
+            if torch.cuda.is_available() and torch.cuda.is_initialized():
+                torch.cuda.synchronize()
+            for p in params:
+                p.grad = None
+            padded += L.all_reduce_floats([float(sum(m.input_ids.numel() for m in micros))])[0]
+            sm = L.summarize(acc, denom, float(total), cfg.loss)
+            sm["step_wall_s"] = time.perf_counter() - t_step
+            sm["n_micro_this_rank"] = len(micros)
+            if torch.cuda.is_available() and torch.cuda.is_initialized():
+                sm["cuda_peak_alloc_gb"] = torch.cuda.max_memory_allocated() / 2**30
+                sm["cuda_peak_reserved_gb"] = torch.cuda.max_memory_reserved() / 2**30
+            step_metrics.append(sm)
+        metrics.update(step_metrics[-1])
+        metrics["steps"] = step_metrics
+        metrics["updated"] = True
+        metrics["resp_tokens_total"] = float(sum(m["resp_tokens"] for m in step_metrics))
+    train_wall = time.perf_counter() - t0
+    lora = lora_state_dict(model, cfg.lora.adapter)  # collective under FSDP
+    entry = ledger_entry(cfg, k, d, bstats, metrics, train_wall, padded, ledger)
+    ledger.append(entry)
+    if cfg.export_merged:
+        entry["merged_export"] = export_merged_checkpoint(model, cfg.lora.adapter, qcfg,
+                                                          os.path.join(ckpt_path(run, k + 1) + ".merged_tmp"), tokenizer_dir)
+    save_checkpoint(run, k + 1, model=model, optimizer=optimizer, lora_state=lora, qcfg=qcfg, cfg=cfg, ledger=ledger,
+                    bases=metrics["bases"], base_name=base_name, base_fingerprint=base_fingerprint)
+    if cfg.export_merged and _rank() == 0:
+        src = ckpt_path(run, k + 1) + ".merged_tmp"
+        if os.path.isdir(src):
+            os.replace(src, os.path.join(ckpt_path(run, k + 1), "merged"))
+    _barrier()
+    return entry
+
+
+def init_checkpoint(cfg: L.RLConfig, *, model, optimizer, qcfg, base_name: str, base_fingerprint: Optional[str] = None) -> None:
+    """Checkpoint 0 = the freshly initialised adapter (B = 0, so it IS the base policy) and a fresh optimizer."""
+    from align.ppo.lora import lora_state_dict
+
+    if latest_complete(cfg.run_dir) is not None:
+        return
+    save_checkpoint(cfg.run_dir, 0, model=model, optimizer=optimizer, lora_state=lora_state_dict(model, cfg.lora.adapter), qcfg=qcfg,
+                    cfg=cfg, ledger=[], bases={}, base_fingerprint=base_fingerprint, base_name=base_name)
+
+
+def make_optimizer(cfg: L.RLConfig, params):
+    return torch.optim.AdamW(list(params), lr=cfg.optim.lr, betas=tuple(cfg.optim.betas), eps=cfg.optim.eps,
+                             weight_decay=cfg.optim.weight_decay)
+
+
+# ----------------------------------------------------------------------------- real policy build
+
+
+def build_real_policy(cfg: L.RLConfig) -> tuple:
+    """Rank-serial load of the frozen bf16 base -> GPU -> LoRA (identical init on every rank, B = 0) -> adapter-only
+    trainable -> activation checkpointing -> FSDP2 (world > 1). Same recipe as `align.train.qwen3_stage`, without
+    its RunConfig. Returns the tuple `InProcessTrainer` expects."""
+    import gc
+
+    from align.model import qwen3 as q3
+    from align.model.checkpoint import load_hf_safetensors, resolve_snapshot_dir
+    from align.model.qwen3_checkpoint import hf_to_align, read_qwen3_config
+    from align.ppo.lora import LoRAConfig, inject_lora, mark_only_adapter_trainable
+    from align.train import fsdp
+    from align.train.qwen3_stage import base_fingerprint
+    from align.utils.dist import barrier, local_rank, rank, world_size
+
+    snapshot = resolve_snapshot_dir(cfg.model_snapshot)
+    qcfg = read_qwen3_config(snapshot).replace(attn_implementation="sdpa")
+    device = torch.device("cuda", local_rank())
+    torch.cuda.set_device(device)
+    model = None
+    for r in range(world_size()):
+        if rank() == r:
+            sd = hf_to_align(load_hf_safetensors(snapshot, dtype=torch.bfloat16), qcfg)
+            model = q3.Qwen3ForCausalLM.from_state_dict(qcfg, sd, device="cuda").to(device)
+            del sd
+            gc.collect()
+        barrier()
+    fp = base_fingerprint(model.state_dict())  # whole tensors, before inject/FSDP
+    lc = cfg.lora
+    gen = torch.Generator().manual_seed(cfg.seed)  # identical A on every rank: FSDP2 shards, never broadcasts
+    hit = inject_lora(model, LoRAConfig(r=lc.r, alpha=lc.alpha, target_modules=tuple(lc.targets), rslora=lc.rslora),
+                      adapter=lc.adapter, generator=gen)
+    expected = len(lc.targets) * qcfg.n_layers
+    if len(hit) != expected:
+        raise RuntimeError(f"LoRA hit {len(hit)} layers, expected {expected}")
+    mark_only_adapter_trainable(model, lc.adapter)
+    fsdp.apply_activation_checkpointing(model)
+    if world_size() > 1:
+        fsdp.apply_fsdp(model)
+    params = [p for p in model.parameters() if p.requires_grad]
+    pad = qcfg.pad_token_id if qcfg.pad_token_id is not None else (qcfg.eos_token_id or 0)
+    return model, params, qcfg, pad, snapshot, fp, snapshot
+
+
+# --------------------------------------------------------------------------------------- backends
+
+
+class InProcessTrainer:
+    """Trainer backend that runs `train_iteration` in this process. `builder()` -> (model, params, qcfg, pad_id,
+    base_name, base_fp, tokenizer_dir)."""
+
+    def __init__(self, cfg: L.RLConfig, builder: Callable[[], tuple]):
+        self.cfg, self.builder = cfg, builder
+
+    def _build(self):
+        model, params, qcfg, pad_id, base_name, base_fp, tok = self.builder()
+        return model, params, make_optimizer(self.cfg, params), qcfg, pad_id, base_name, base_fp, tok
+
+    def init(self) -> None:
+        model, params, opt, qcfg, pad_id, base_name, base_fp, _ = self._build()
+        init_checkpoint(self.cfg, model=model, optimizer=opt, qcfg=qcfg, base_name=base_name, base_fingerprint=base_fp)
+
+    def train(self, k: int) -> dict:
+        model, params, opt, qcfg, pad_id, base_name, base_fp, tok = self._build()
+        return train_iteration(self.cfg, k, model=model, params=params, optimizer=opt, qcfg=qcfg, pad_id=pad_id,
+                               base_name=base_name, base_fingerprint=base_fp, tokenizer_dir=tok)
+
+
+def repo_root() -> str:
+    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+class LeanGenCliGenerator:
+    """Real generator: `extras/vllm_rollout/lean_gen.py` as a child process (needs the genharness flags
+    `--lora-adapter`, `--top-p`, `--temperature` and `--record-logprobs`)."""
+
+    def __init__(self, cfg: L.RLConfig, config_path: str = ""):
+        self.cfg = cfg
+
+    def command(self, k: int, d: str, problems_path: str, adapter_dir: str, merged_dir: Optional[str], seed: int) -> list[str]:
+        c, ro = self.cfg, self.cfg.rollout
+        root = repo_root()
+        py = ro.rollout_python if os.path.isabs(ro.rollout_python) else os.path.join(root, ro.rollout_python)
+        model = c.model_snapshot
+        cmd = [py, os.path.join(root, "extras", "vllm_rollout", "lean_gen.py"), "--problems", problems_path,
+               "--out", os.path.join(d, "gens.jsonl"), "--n", str(c.batch.group_size), "--cap-tokens", str(ro.cap_tokens),
+               "--round", "0", "--tp", str(ro.tp), "--wait-lock", "--seed", str(seed), "--max-model-len", str(ro.max_model_len),
+               "--temperature", str(ro.temperature), "--top-p", "1.0", "--record-logprobs",
+               "--label", f"rl-lean-gen-{k}"]
+        if ro.adapter_mode == "unmerged":
+            cmd += ["--model", model, "--lora-adapter", adapter_dir]
+        else:
+            if not merged_dir:
+                raise L.RLConfigError("adapter_mode=merged_export needs export_merged=true")
+            cmd += ["--model", merged_dir]
+        return cmd + list(ro.gen_extra_args)
+
+    def generate(self, k: int, d: str, problems_path: str, adapter_dir: str, merged_dir: Optional[str], seed: int) -> None:
+        subprocess.run(self.command(k, d, problems_path, adapter_dir, merged_dir, seed), check=True)
+
+
+class VerifyCliVerifier:
+    def __init__(self, cfg: L.RLConfig):
+        self.cfg = cfg
+
+    def command(self, d: str, problems_path: str) -> list[str]:
+        c, root = self.cfg, repo_root()
+        py = c.verify_python if os.path.isabs(c.verify_python) else os.path.join(root, c.verify_python)
+        return [py, os.path.join(root, "scripts", "p2", "verify_generations.py"), "--generations", os.path.join(d, "gens.jsonl"),
+                "--problems", problems_path, "--out", os.path.join(d, "ver.jsonl"), "--workers", str(c.verify_workers),
+                "--budget-s", str(c.verify_budget_s)]
+
+    def verify(self, d: str, problems_path: str) -> None:
+        subprocess.run(self.command(d, problems_path), check=True)
+
+
+class ReplayBackend:
+    """Replay cached rollouts + verdicts of another run (K1: positive-only vs signed on the SAME rollouts).
+    Rollouts beyond iteration 0 were sampled by the SOURCE arm's policy, so they are off-policy for this arm;
+    the importance ratio against the recorded behaviour logprobs keeps the update well defined."""
+
+    def __init__(self, src_run_dir: str):
+        self.src = src_run_dir
+
+    def generate(self, k, d, problems_path, adapter_dir, merged_dir, seed) -> None:
+        shutil.copyfile(os.path.join(iter_dir(self.src, k), "gens.jsonl"), os.path.join(d, "gens.jsonl"))
+
+    def verify(self, d, problems_path) -> None:
+        # k is encoded in d's name; copy the sibling verdicts
+        k = int(os.path.basename(d).split("_")[1])
+        shutil.copyfile(os.path.join(iter_dir(self.src, k), "ver.jsonl"), os.path.join(d, "ver.jsonl"))
+
+
+def torchrun_command(cfg_path: str, k: int, nproc: int, overrides=()) -> list[str]:
+    return ["torchrun", "--standalone", f"--nproc_per_node={nproc}", os.path.join(repo_root(), "scripts", "p2", "rl_lean.py"),
+            "train-iter", cfg_path, "--iteration", str(k), *[f"--set={o}" for o in overrides]]
+
+
+class TorchrunTrainer:
+    def __init__(self, cfg: L.RLConfig, cfg_path: str, nproc: int, overrides=()):
+        self.cfg, self.cfg_path, self.nproc, self.overrides = cfg, cfg_path, nproc, tuple(overrides)
+
+    def init(self) -> None:
+        subprocess.run(torchrun_command(self.cfg_path, -1, self.nproc, self.overrides), check=True)
+
+    def train(self, k: int) -> dict:
+        subprocess.run(torchrun_command(self.cfg_path, k, self.nproc, self.overrides), check=True)
+        return {}
+
+
+# ------------------------------------------------------------------------------------- the loop
+
+
+def derive_seed(seed: int, k: int) -> int:
+    return int.from_bytes(hashlib.sha256(f"rl|{seed}|{k}".encode()).digest()[:8], "big") % (2**31)
+
+
+def run_loop(cfg: L.RLConfig, generator, verifier, trainer, hook: Optional[Callable[[str, int], None]] = None,
+             until: Optional[int] = None) -> list[dict]:
+    """Run (or resume) iterations up to `until` (default cfg.iterations). `hook(stage, k)` fires after each
+    committed stage (tests inject crashes there). Returns the committed ledger."""
+    run = cfg.run_dir
+    os.makedirs(run, exist_ok=True)
+    cfg_json = os.path.join(run, "run_config.json")
+    import dataclasses
+
+    fp = cfg.fingerprint()
+    if os.path.exists(cfg_json):
+        old = json.load(open(cfg_json))["fingerprint"]
+        if old != fp:
+            raise RuntimeError(f"{run} was started under config {old[:12]}, this is {fp[:12]}: refusing to mix runs")
+    else:
+        atomic_write_text(cfg_json, json.dumps({"fingerprint": fp, "config": dataclasses.asdict(cfg)}, sort_keys=True, default=str, indent=1))
+    clean_partial(run)
+    problems = load_problems(cfg.problems)
+    by_id = {p["problem_id"]: p for p in problems}
+    ids = [p["problem_id"] for p in problems]
+    if latest_complete(run) is None:
+        trainer.init()
+    k0 = latest_complete(run)
+    if k0 is None:
+        raise RuntimeError("no checkpoint after init")
+    hook = hook or (lambda stage, k: None)
+    end = cfg.iterations if until is None else min(until, cfg.iterations)
+    for k in range(k0, end):
+        d = iter_dir(run, k)
+        os.makedirs(d, exist_ok=True)
+        probs_path = os.path.join(d, "problems.jsonl")
+        if not os.path.exists(probs_path):
+            block = L.prompt_block(ids, cfg.seed, cfg.prompts_per_iter, k)
+            write_jsonl_atomic(probs_path, [by_id[i] for i in block])
+        if not os.path.exists(stage_marker(d, "gen")):
+            gp = os.path.join(d, "gens.jsonl")
+            adapter_dir = os.path.join(ckpt_path(run, k), "adapter")
+            merged = os.path.join(ckpt_path(run, k), "merged") if cfg.export_merged else None
+            t0 = time.perf_counter()
+            gen = ReplayBackend(cfg.replay_from) if cfg.replay_from else generator
+            gen.generate(k, d, probs_path, adapter_dir, merged, derive_seed(cfg.seed, k))
+            add_stage_time(d, "gen", time.perf_counter() - t0)
+            _check_generations(gp, cfg, [by_id_id["problem_id"] for by_id_id in read_jsonl(probs_path)])
+            atomic_write_text(stage_marker(d, "gen"), json.dumps({"sha256": sha256_file(gp)}))
+            hook("gen", k)
+        if not os.path.exists(stage_marker(d, "verify")):
+            vp = os.path.join(d, "ver.jsonl")
+            t0 = time.perf_counter()
+            (ReplayBackend(cfg.replay_from) if cfg.replay_from else verifier).verify(d, probs_path)
+            add_stage_time(d, "verify", time.perf_counter() - t0)
+            n_g, n_v = len(read_jsonl(os.path.join(d, "gens.jsonl"))), len(read_jsonl(vp))
+            if n_g != n_v:
+                raise RuntimeError(f"{n_v} verdicts for {n_g} generations")
+            atomic_write_text(stage_marker(d, "verify"), json.dumps({"sha256": sha256_file(vp)}))
+            hook("verify", k)
+        t0 = time.perf_counter()
+        trainer.train(k)
+        add_stage_time(d, "train_proc", time.perf_counter() - t0)  # whole child: model load + step + commit
+        hook("train", k)
+    last = latest_complete(run)
+    return read_jsonl(os.path.join(run, "ledger.jsonl")) if last else []
+
+
+def _check_generations(path: str, cfg: L.RLConfig, expect_ids: list[str]) -> None:
+    gens = read_jsonl(path)
+    want = len(expect_ids) * cfg.batch.group_size
+    if len(gens) != want:
+        raise RuntimeError(f"{path}: {len(gens)} generations, expected {want} ({len(expect_ids)} prompts x G={cfg.batch.group_size})")
+    for g in gens:
+        L.validate_rl_generation(g)

code/align/rl/lean_rl_testing.py

@@ -0,0 +1,174 @@
+"""TEST SUPPORT for the Lean RL loop (CPU only): a tiny Qwen3-shaped model, a FAKE generator that really
+samples from it (so the loop is genuinely on-policy), and a FAKE verifier. No vLLM, no Lean.
+
+Not used by any real run; it lives in the package so the tests and `scripts/p2/rl_lean.py --selftest` share it.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import random
+
+import torch
+
+from align.model.qwen3 import Qwen3Config, Qwen3ForCausalLM
+from align.ppo.lora import LoRAConfig, inject_lora, lora_modules, mark_only_adapter_trainable
+from align.rl import lean_rl as L
+from align.rl import lean_rl_run as R
+
+TINY = Qwen3Config(name="tiny", vocab_size=48, hidden_size=32, intermediate_size=64, n_layers=2, n_heads=4,
+                   n_kv_heads=2, head_dim=8, max_position_embeddings=256, bos_token_id=1, eos_token_id=2,
+                   pad_token_id=0)
+EOS = 2
+
+
+def tiny_base(seed: int = 0, dtype=torch.float32) -> Qwen3ForCausalLM:
+    torch.manual_seed(seed)
+    m = Qwen3ForCausalLM(TINY, attn_implementation="eager", device="cpu")
+    return m.to(dtype)
+
+
+def tiny_policy(lcfg: L.LoraCfg, seed: int = 0, dtype=torch.float32, adapter_seed: int = 1234):
+    """(model, params) with LoRA injected (B = 0 unless randomised) and only the adapter trainable."""
+    m = tiny_base(seed, torch.float32)
+    gen = torch.Generator().manual_seed(adapter_seed)
+    inject_lora(m, LoRAConfig(r=lcfg.r, alpha=lcfg.alpha, target_modules=tuple(lcfg.targets), rslora=lcfg.rslora),
+                adapter=lcfg.adapter, generator=gen)
+    mark_only_adapter_trainable(m, lcfg.adapter)
+    m = m.to(dtype)
+    params = [p for p in m.parameters() if p.requires_grad]
+    return m, params
+
+
+def randomize_adapter(model, adapter: str, seed: int, std: float = 0.05) -> None:
+    g = torch.Generator().manual_seed(seed)
+    with torch.no_grad():
+        for _, mod in lora_modules(model):
+            a = mod.adapters[adapter]
+            a.A.copy_(torch.randn(a.A.shape, generator=g) * std)
+            a.B.copy_(torch.randn(a.B.shape, generator=g) * std)
+
+
+def tiny_builder(cfg: L.RLConfig, dtype=torch.float32):
+    def build():
+        model, params = tiny_policy(cfg.lora, dtype=dtype)
+        return model, params, TINY, 0, "tiny-base", None, None
+
+    return build
+
+
+def synthetic_problems(n: int) -> list[dict]:
+    rows = []
+    for i in range(n):
+        fs = f"theorem t{i} : {i} = {i} := by"
+        rows.append({"problem_id": f"p{i:03d}", "source": "synthetic", "split": "train", "formal_statement": fs,
+                     "informal_prefix": "", "statement_sha": hashlib.sha256(fs.encode()).hexdigest()})
+    return rows
+
+
+def prompt_ids_for(pid: str) -> list[int]:
+    h = int.from_bytes(hashlib.sha256(pid.encode()).digest()[:4], "big")
+    rng = random.Random(h)
+    return [1] + [rng.randrange(3, TINY.vocab_size) for _ in range(rng.randrange(2, 6))]
+
+
+def _key_hash(*parts) -> int:
+    return int.from_bytes(hashlib.sha256("|".join(map(str, parts)).encode()).digest()[:8], "big")
+
+
+class FakeGenerator:
+    """Samples from the tiny model + the adapter in `adapter_dir` (the CURRENT policy), temperature 1, top_p 1,
+    recording the sampled-token logprobs. A real on-policy rollout, CPU-sized. `crash_in_gen` (set by a test)
+    kills the process mid-stage after writing nothing, like a vLLM crash before the output is complete."""
+
+    def __init__(self, cfg: L.RLConfig, max_new: int = 8, base_seed: int = 0, dtype=torch.float32):
+        self.cfg, self.max_new, self.base_seed, self.dtype = cfg, max_new, base_seed, dtype
+        self.calls = 0
+
+    def generate(self, k, d, problems_path, adapter_dir, merged_dir, seed) -> None:
+        from align.ppo.lora import load_lora_state_dict
+
+        self.calls += 1
+        model, _ = tiny_policy(self.cfg.lora, seed=self.base_seed, dtype=self.dtype)
+        load_lora_state_dict(model, self.cfg.lora.adapter, R.read_peft_dir(adapter_dir, TINY), strict=True)
+        model.eval()
+        recs = []
+        for prob in R.read_jsonl(problems_path):
+            pid = prob["problem_id"]
+            prompt = prompt_ids_for(pid)
+            for si in range(self.cfg.batch.group_size):
+                g = torch.Generator().manual_seed(_key_hash(seed, pid, si))
+                ids = list(prompt)
+                lps = []
+                finish = "length"
+                for _ in range(self.max_new):
+                    with torch.no_grad():
+                        logits = model(input_ids=torch.tensor([ids]))[0, -1].double()
+                    lp = torch.log_softmax(logits, -1)
+                    tok = int(torch.multinomial(lp.exp(), 1, generator=g))
+                    ids.append(tok)
+                    lps.append(float(lp[tok]))
+                    if tok == EOS:
+                        finish = "stop"
+                        break
+                comp = ids[len(prompt):]
+                recs.append({
+                    "problem_id": pid, "sample_idx": si, "round": 0, "cap_tokens": self.cfg.rollout.cap_tokens,
+                    "prompt_ids": prompt, "completion_ids": comp, "prompt_text": "", "completion_text": " ".join(map(str, comp)),
+                    "finish_reason": finish, "lean_code": "theorem x := by\n  trivial", "messages": [],
+                    "model_path": "tiny", "model_rev": "tiny+adapter",
+                    "sampling": {"temperature": 1.0, "top_p": 1.0, "top_k": 0, "n": 1, "max_tokens": self.max_new, "seed": seed},
+                    "context_clipped": False, "parent": None, "sampling_logprobs": lps,
+                    "adapter": {"serving": "unmerged_vllm_lora"},
+                })
+        R.write_jsonl_atomic(os.path.join(d, "gens.jsonl"), recs)
+
+
+class FakeVerifier:
+    """accepted iff the completion contains an even token > 10 in its first 3 tokens (learnable, ~50% under a
+    random policy); every 9th sample by key hash is a TIMEOUT (no signal)."""
+
+    def __init__(self, timeout_mod: int = 9):
+        self.timeout_mod = timeout_mod
+        self.calls = 0
+
+    def verify(self, d, problems_path) -> None:
+        self.calls += 1
+        out = []
+        for g in R.read_jsonl(os.path.join(d, "gens.jsonl")):
+            key = (g["problem_id"], g["sample_idx"], g["round"], g["cap_tokens"])
+            timeout = self.timeout_mod and _key_hash(*key) % self.timeout_mod == 0
+            ok = any(t % 2 == 0 and t > 10 for t in g["completion_ids"][:3])
+            out.append({"problem_id": key[0], "sample_idx": key[1], "round": key[2], "cap_tokens": key[3],
+                        "accepted": bool(ok and not timeout), "reasons": ["timeout"] if timeout else ([] if ok else ["compile_error"]),
+                        "verify_s": 0.25, "budget_s": 1.0, "verifier": {"fake": True},
+                        "code_sha256": hashlib.sha256(g["lean_code"].encode()).hexdigest(), "stage": "fake", "diagnostics": {}})
+        R.write_jsonl_atomic(os.path.join(d, "ver.jsonl"), out)
+
+
+def synthetic_rollouts(n_prompts: int, G: int, seed: int, timeouts: int = 0, dead_prompts: int = 0,
+                       cap: int = 64) -> tuple[list[dict], dict]:
+    """Random gens + verdicts with UNEVEN lengths, behaviour logprobs unrelated to any model (so the importance
+    ratio spans well below 1 and above the clamp), `dead_prompts` all-equal groups and `timeouts` masked samples."""
+    rng = random.Random(seed)
+    gens, vers = [], {}
+    for pi in range(n_prompts):
+        pid = f"p{pi:03d}"
+        prompt = prompt_ids_for(pid)
+        dead = pi < dead_prompts
+        for si in range(G):
+            n = rng.randrange(1, 10)
+            comp = [rng.randrange(3, TINY.vocab_size) for _ in range(n)]
+            lps = [-rng.uniform(0.05, 4.0) for _ in range(n)]
+            gens.append({"problem_id": pid, "sample_idx": si, "round": 0, "cap_tokens": cap, "prompt_ids": prompt,
+                         "completion_ids": comp, "finish_reason": "stop", "lean_code": "x",
+                         "sampling": {"temperature": 1.0, "top_p": 1.0, "top_k": 0}, "sampling_logprobs": lps})
+            acc = False if dead else (rng.random() < 0.5)
+            reasons = [] if acc else ["compile_error"]
+            vers[(pid, si, 0, cap)] = {"accepted": acc, "reasons": reasons, "verify_s": 0.1}
+    keys = list(vers)
+    for key in rng.sample([k for k in keys if int(k[0][1:]) >= dead_prompts], timeouts):
+        vers[key] = {"accepted": False, "reasons": ["timeout"], "verify_s": 5.0}
+    return gens, vers

code/align/rl/multiturn_source.py

@@ -0,0 +1,454 @@
+"""Multi-turn Lean rollouts -> GRPO rows with PER-TOKEN advantages (opt-in).
+
+`GRPOTrainModule` used to see one scalar advantage per row, so the credit the
+multi-turn env computes per turn could not reach the loss. This module is the
+producer side of the fix (the consumer side is `expand_advantages` in
+grpo_loss.py): it runs `align/rl/envs/lean_multiturn.py` episodes, assigns
+advantages, and emits rows whose ``advantages`` is [B, T] in INPUT space,
+nonzero only on tokens the policy generated.
+
+Every design choice below names the failure it prevents.
+
+ROW LAYOUT: one row per TURN (the env's ``layout="per_turn"``, the default).
+  * A packed row per episode is valid only when the chat template renders
+    turn k's prompt as an exact extension of turn k-1's prompt + completion.
+    Qwen3-family templates (Goedel-Prover-V2 is one) strip earlier <think>
+    blocks, so the env raises `PrefixBreak` on them; per-turn rows are exact on
+    every template because each row carries the exact prompt ids the policy
+    saw at that turn.
+  * Memory: a per-turn row is bounded by one turn's prompt + completion, so
+    micro-batch padding is bounded by the single longest TURN. A packed row is
+    the SUM of all turns' context and pads every micro-batch row to that sum.
+    Per-turn rows pay for it by re-running the shared prefix (~K/2 x extra
+    forward tokens at K turns) -- accepted: it is bounded and the alternative
+    cannot always be built.
+  * Rows are independent in the loss. The Dr.GRPO denominator is a constant
+    and advantages are already per-token, so no statistic crosses rows: a
+    micro-slice or rank boundary that splits one episode's rows changes
+    nothing (a test pins that the summed loss is invariant to row order and
+    split). The packed layout is still accepted (the env's ``layout=
+    "concatenated"``) when the template permits it; the same masks and
+    advantages come out.
+  * DENOMINATOR: the constant D is per OPTIMIZER STEP. A multi-turn episode
+    can generate max_turns x max_response_len tokens, so D should be sized
+    prompts_per_step x G x max_turns x max_response_len
+    (`recommended_loss_denominator`); sizing it for one turn silently scales
+    every gradient by up to max_turns. The A3 gate still only requires a
+    positive finite float; the value is the caller's pre-registered choice.
+
+CREDIT (`MultiTurnCreditConfig`):
+  mode="episode" (DEFAULT): one number per episode, A_ep = R_ep - mean_group(R),
+    the Dr.GRPO centered advantage over the episode's TOTAL reward, broadcast
+    to EVERY generated token of EVERY turn. Prevents: a baseline computed on a
+    thinned sample (see below) and the variance of fitting a per-turn value
+    function from G <= 16 episodes. It is the exact single-turn algebra, so
+    group centering, dead-group semantics and A3 all carry over unchanged.
+    Its cost is credit blur: a failed first draft is rewarded for the repair
+    that followed. That is wrong only if the draft did not make the repair
+    possible; with turn_discount < 1 the solved-late episode already scores
+    lower than a solved-first one, which is the ordering we want taught.
+  mode="per_turn_loo": turn t gets G_t - b_t, G_t = sum_{s>=t} gamma^(s-t) u_s
+    (the env's return; gamma=0 -> each turn judged on its own outcome).
+      turn 0   : b_0 = mean over the WHOLE group. Every episode has a turn 0,
+                 so this is the ordinary group baseline and the turn-0
+                 advantages are exactly centered.
+      turn t>=1: b_t = mean over the OTHER episodes of the group that reached
+                 turn t (leave-one-out). The old rule averaged over ALL
+                 episodes that reached turn t INCLUDING self; with one repair
+                 in the group the baseline was the sample's own return, so the
+                 lone repair that SOLVED the problem got advantage exactly 0
+                 (and one that failed also 0). LOO removes self-inclusion, and
+                 the LOO advantages of a cell sum to exactly 0 (x_i - mean_{-i}
+                 = n/(n-1) (x_i - mean)).
+      lone cell: if no other episode of the group reached turn t there is no
+                 within-group peer. The baseline is then the mean of the other
+                 episodes at turn t POOLED across the other groups of the same
+                 rollout step (>= ``min_pool_peers`` of them; still leave-one-
+                 out and independent of the sample's own action, so the
+                 estimator stays unbiased). With fewer, the cell's advantage
+                 is 0 and counted (`credit/cells_no_baseline`) -- a baseline
+                 invented from nothing would be a made-up gradient.
+    Pooled cells trade within-group centering for a baseline that exists; the
+    telemetry counts how many cells took each path.
+
+DEAD GROUPS. "episode": dead iff the group's episode rewards have zero
+variance, classified all_zero / all_solved / nonzero_std0 by the shared
+`GroupFilterStats` (max_possible_score=1.0: the best episode reward). "per_turn
+_loo": dead iff EVERY token advantage of the group is exactly 0 -- the group
+contributes no gradient, so keeping it only spends compute; classified from
+the rewards the same way. Dead groups are dropped whole, as in single-turn
+assembly; an all-dead step returns None (the driver resamples).
+
+No CUDA, no Lean here: the env, generator and verifier are injected.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from collections import OrderedDict
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional, Sequence
+
+import numpy as np
+import torch
+
+from align.rl.advantages import GroupFilterStats, compute_group_advantages
+from align.rl.batch_source import GrpoBatchSource
+from align.rl.envs.lean_multiturn import (
+    EpisodeRecord,
+    EpisodeSpec,
+    LeanMultiTurnEnv,
+    Problem,
+)
+
+CREDIT_MODES = ("episode", "per_turn_loo")
+
+
+@dataclass(frozen=True)
+class MultiTurnCreditConfig:
+    mode: str = "episode"
+    #: per-turn return discount (per_turn_loo only): 0 = a turn is judged on its own outcome,
+    #: >0 = an earlier turn shares in the later turns' outcomes
+    gamma: float = 0.0
+    #: pooled-baseline floor for a turn no group-mate reached (per_turn_loo only)
+    min_pool_peers: int = 2
+
+    def __post_init__(self) -> None:
+        if self.mode not in CREDIT_MODES:
+            raise ValueError(f"mode must be one of {CREDIT_MODES}, got {self.mode!r}")
+        if not 0.0 <= self.gamma <= 1.0:
+            raise ValueError(f"gamma must be in [0, 1], got {self.gamma}")
+        if self.min_pool_peers < 1:
+            raise ValueError(f"min_pool_peers must be >= 1, got {self.min_pool_peers}")
+        if self.mode == "episode" and self.gamma != 0.0:
+            raise ValueError("gamma is read only by mode='per_turn_loo'; it would be silently ignored")
+
+
+def recommended_loss_denominator(
+    prompts_per_step: int, n_samples: int, max_turns: int, max_response_len: int
+) -> float:
+    """Dr.GRPO constant sized for the worst-case GENERATED tokens of one optimizer step."""
+    return float(prompts_per_step * n_samples * max_turns * max_response_len)
+
+
+# ---------------------------------------------------------------------------- credit
+
+
+def turn_returns(u: Sequence[float], gamma: float) -> list[float]:
+    """G_t = sum_{s>=t} gamma^(s-t) u_s (the env's definition)."""
+    out, acc = [0.0] * len(u), 0.0
+    for t in range(len(u) - 1, -1, -1):
+        acc = u[t] + gamma * acc
+        out[t] = acc
+    return out
+
+
+@dataclass
+class CreditResult:
+    #: per episode (input order), per turn advantage
+    turn_adv: list
+    #: path taken per (episode, turn) cell: "episode" | "group_mean" | "group_loo" | "pooled" | "none"
+    cell_path: list
+    counts: dict = field(default_factory=dict)
+
+
+def compute_turn_credit(episodes: Sequence[EpisodeRecord], credit: MultiTurnCreditConfig) -> CreditResult:
+    """Advantage per (episode, turn). Episodes of one problem_id form a group."""
+    n = len(episodes)
+    groups: "OrderedDict[str, list[int]]" = OrderedDict()
+    for i, e in enumerate(episodes):
+        groups.setdefault(e.problem_id, []).append(i)
+
+    ep_adv = [0.0] * n
+    for idx in groups.values():
+        a = compute_group_advantages(
+            np.array([episodes[i].reward["reward/total"] for i in idx]), len(idx), "centered"
+        )
+        for i, v in zip(idx, a):
+            ep_adv[i] = float(v)
+
+    n_turns = [len(e.turns) for e in episodes]
+    if credit.mode == "episode":
+        return CreditResult(
+            [[ep_adv[i]] * n_turns[i] for i in range(n)],
+            [["episode"] * n_turns[i] for i in range(n)],
+            {"credit/cells_episode": sum(n_turns)},
+        )
+
+    G = [turn_returns(e.turn_rewards, credit.gamma) for e in episodes]
+    turn_adv = [[0.0] * n_turns[i] for i in range(n)]
+    path = [[""] * n_turns[i] for i in range(n)]
+    counts = {"credit/cells_group_mean": 0, "credit/cells_group_loo": 0,
+              "credit/cells_pooled": 0, "credit/cells_no_baseline": 0}
+    for idx in groups.values():
+        for i in idx:  # turn 0: the ordinary group baseline (every episode has a turn 0)
+            base = sum(G[j][0] for j in idx) / len(idx)
+            turn_adv[i][0] = G[i][0] - base
+            path[i][0] = "group_mean"
+            counts["credit/cells_group_mean"] += 1
+    for gidx in groups.values():
+        for i in gidx:
+            for t in range(1, n_turns[i]):
+                peers = [G[j][t] for j in gidx if j != i and n_turns[j] > t]
+                if peers:
+                    turn_adv[i][t] = G[i][t] - sum(peers) / len(peers)
+                    path[i][t] = "group_loo"
+                    counts["credit/cells_group_loo"] += 1
+                    continue
+                pool = [G[j][t] for j in range(n) if j != i and n_turns[j] > t and j not in gidx]
+                if len(pool) >= credit.min_pool_peers:
+                    turn_adv[i][t] = G[i][t] - sum(pool) / len(pool)
+                    path[i][t] = "pooled"
+                    counts["credit/cells_pooled"] += 1
+                else:
+                    path[i][t] = "none"
+                    counts["credit/cells_no_baseline"] += 1
+    return CreditResult(turn_adv, path, counts)
+
+
+# ---------------------------------------------------------------------------- assembly
+
+
+@dataclass
+class AssembledMultiTurn:
+    """Duck-types `AssembledRollout` (the fields `GrpoBatchSource._produce_step` reads),
+    except ``advantages`` is [B, T] in input space."""
+
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    response_mask: torch.Tensor
+    advantages: torch.Tensor
+    vllm_logprobs: torch.Tensor
+    group_ids: torch.Tensor
+    rewards: dict
+    turn_ids: torch.Tensor
+    stats: dict = field(default_factory=dict)
+    kept_groups: int = 0
+    dropped_groups: int = 0
+
+
+#: A group is dead below this spread. `np.std(...) == 0` (the single-turn rule) is an EXACT test, and
+#: with discounted multi-turn rewards it misfires: seven episodes that all solved on turn 1 at
+#: discount 0.9 have std 1.1e-16, not 0, because ``mean(x, x, ..., x) != x`` in float64. Such a group
+#: was kept with advantages of +-1e-16 (a zero-gradient group paid for in full, and counted as
+#: ``groups/kept``). Rewards live in [-1, 1] and real differences are >= ~1e-3, so 1e-9 separates
+#: rounding noise from signal by six orders of magnitude on each side.
+DEAD_ATOL = 1e-9
+
+
+def _observe_group(stats: GroupFilterStats, rewards: np.ndarray, live: bool) -> None:
+    """Record one group in the three-way telemetry with an externally decided liveness."""
+    if live:
+        stats.kept += 1
+        return
+    if np.all(np.abs(rewards) <= DEAD_ATOL):
+        stats.all_zero += 1
+    elif np.all(np.abs(rewards - stats.max_possible_score) <= DEAD_ATOL):
+        stats.all_solved += 1
+    else:
+        stats.nonzero_std0 += 1
+
+
+def assemble_multiturn(
+    episodes: Sequence[EpisodeRecord],
+    credit: MultiTurnCreditConfig,
+    pad_token_id: int,
+    filter_stats: Optional[GroupFilterStats] = None,
+) -> Optional[AssembledMultiTurn]:
+    """Episodes -> one training batch with per-token advantages; None if every group is dead."""
+    if not episodes:
+        raise ValueError("no episodes")
+    stats = filter_stats if filter_stats is not None else GroupFilterStats(max_possible_score=1.0)
+    res = compute_turn_credit(episodes, credit)
+
+    groups: "OrderedDict[str, list[int]]" = OrderedDict()
+    for i, e in enumerate(episodes):
+        groups.setdefault(e.problem_id, []).append(i)
+
+    keep: list[list[int]] = []
+    for idx in groups.values():
+        rewards = np.array([episodes[i].reward["reward/total"] for i in idx], dtype=np.float64)
+        if credit.mode == "episode":
+            # the single-turn rule (zero reward variance) with a rounding tolerance, see DEAD_ATOL
+            dead = float(np.ptp(rewards)) <= DEAD_ATOL
+        else:
+            dead = not any(abs(a) > DEAD_ATOL for i in idx for a in res.turn_adv[i])
+        _observe_group(stats, rewards, live=not dead)
+        if not dead:
+            keep.append(idx)
+    if not keep:
+        return None
+
+    rows = []  # (group_index, episode, row, per-token adv)
+    for gi, idx in enumerate(keep):
+        for i in idx:
+            ep = episodes[i]
+            for row in ep.rows:
+                if row.logprobs is None:
+                    raise ValueError(
+                        f"episode {ep.episode_key}: generation carries no logprobs; the rho "
+                        "correction cannot align without the engine's per-token logprobs"
+                    )
+                adv = [res.turn_adv[i][tid] if tid >= 0 else 0.0 for tid in row.turn_ids]
+                rows.append((gi, ep, row, adv))
+    B = len(rows)
+    T = max(len(r.input_ids) for _, _, r, _ in rows)
+    input_ids = torch.full((B, T), pad_token_id, dtype=torch.int64)
+    attention_mask = torch.zeros((B, T), dtype=torch.int64)
+    response_mask = torch.zeros((B, T), dtype=torch.int64)
+    turn_ids = torch.full((B, T), -1, dtype=torch.int64)
+    advantages = torch.zeros((B, T), dtype=torch.float32)
+    vllm_lp = torch.zeros((B, T - 1), dtype=torch.float32)
+    group_ids = torch.zeros((B,), dtype=torch.int64)
+    term_names = list(rows[0][1].reward)
+    reward_tensors = {k: torch.zeros(B, dtype=torch.float32) for k in term_names}
+    for b, (gi, ep, row, adv) in enumerate(rows):
+        L = len(row.input_ids)
+        if not (len(row.response_mask) == len(row.turn_ids) == len(row.logprobs) == L):
+            raise ValueError(f"row {b} of episode {ep.episode_key} is ragged across ids/mask/turn_ids/logprobs")
+        # Credit is placed by ``turn_ids`` and trained on by ``response_mask``. If the two disagree
+        # by even one position, the loss-side stray check (`expand_advantages`) catches it only
+        # when the misplaced credit lands on an excluded position that exists: a row whose last
+        # generated token is its last column (every micro-batch's longest per-turn row) has no
+        # position past it, so a right-shifted span slips through. Refuse here, where it is exact.
+        if [int(t >= 0) for t in row.turn_ids] != [int(m) for m in row.response_mask]:
+            raise ValueError(
+                f"row {b} of episode {ep.episode_key}: turn_ids >= 0 does not coincide with "
+                f"response_mask; per-token credit would be misaligned with the trained tokens"
+            )
+        input_ids[b, :L] = torch.tensor(row.input_ids, dtype=torch.int64)
+        attention_mask[b, :L] = 1
+        response_mask[b, :L] = torch.tensor(row.response_mask, dtype=torch.int64)
+        turn_ids[b, :L] = torch.tensor(row.turn_ids, dtype=torch.int64)
+        advantages[b, :L] = torch.tensor(adv, dtype=torch.float32)
+        # label position j <-> input token j+1
+        vllm_lp[b, : L - 1] = torch.tensor(row.logprobs[1:], dtype=torch.float32)
+        group_ids[b] = gi
+        for k in term_names:
+            reward_tensors[k][b] = float(ep.reward[k])
+
+    kept_eps = [episodes[i] for idx in keep for i in idx]
+    gsz = [len(idx) for idx in keep]
+    first = [e.solved_turn == 0 for e in kept_eps]
+    failed_first = [e for e in kept_eps if e.solved_turn != 0]
+    stat = {
+        "multiturn/episodes": float(len(kept_eps)),
+        "multiturn/rows": float(B),
+        "multiturn/turns_mean": float(np.mean([len(e.turns) for e in kept_eps])),
+        "multiturn/solved_turn0": float(np.mean(first)),
+        "multiturn/solved": float(np.mean([e.solved_turn is not None for e in kept_eps])),
+        "multiturn/repair_rate": (
+            float(np.mean([e.solved_turn is not None for e in failed_first])) if failed_first else float("nan")
+        ),
+        "multiturn/generated_tokens_mean": float(np.mean([e.n_generated for e in kept_eps])),
+        "multiturn/group_size_min": float(min(gsz)),
+        "multiturn/group_size_max": float(max(gsz)),
+        **{k: float(v) for k, v in res.counts.items()},
+        **{k: float(v) for k, v in stats.as_log_dict().items()},
+    }
+    return AssembledMultiTurn(
+        input_ids, attention_mask, response_mask, advantages, vllm_lp, group_ids, reward_tensors,
+        turn_ids, stat, kept_groups=len(keep), dropped_groups=len(groups) - len(keep),
+    )
+
+
+# ---------------------------------------------------------------------------- the batch source
+
+
+class MultiTurnBatchSource(GrpoBatchSource):
+    """`GrpoBatchSource` whose rollout is a batch of multi-turn Lean episodes.
+
+    Everything downstream of rollout production (rank-0 produce + broadcast,
+    row slicing, no-grad old_logprobs, the validated batch, the collapse
+    tripwire, the draw-position resume state) is inherited unchanged; only
+    `_rollout_once` differs. Rows are sliced along dim 0 and the [B, T]
+    advantages slice with them.
+
+    ``prompt_rows``: dicts with ``formal_statement`` (or ``problem``) and a
+    UNIQUE ``problem_id`` (the group key; duplicates are refused -- two
+    problems sharing an id would be centered against each other).
+    ``generate_factory(seed=..., weights_dir=...)`` returns the env's
+    ``generate(list[GenRequest]) -> list[Generation]`` for this step.
+    """
+
+    def __init__(
+        self,
+        prompt_rows: Sequence[dict[str, Any]],
+        *,
+        env: LeanMultiTurnEnv,
+        credit: MultiTurnCreditConfig = MultiTurnCreditConfig(),
+        prompts_per_step: int,
+        n_samples_per_prompt: int,
+        max_response_len: int,
+        temperature: float,
+        optimizer_steps: int,
+        grad_accum: int,
+        seed: int,
+        dead_group_retries: int = 5,
+    ):
+        rows, seen = [], set()
+        for i, r in enumerate(prompt_rows):
+            stmt = r.get("formal_statement", r.get("problem"))
+            if not isinstance(stmt, str) or not stmt:
+                raise ValueError(f"prompt row {i} lacks formal_statement/problem: {sorted(r)}")
+            pid = str(r.get("problem_id") or hashlib.sha256(stmt.encode()).hexdigest()[:12])
+            if pid in seen:
+                raise ValueError(f"duplicate problem_id {pid!r}: it is the advantage group key")
+            seen.add(pid)
+            rows.append({"problem": stmt, "answer": "", "problem_id": pid})
+        super().__init__(
+            rows,
+            prompts_per_step=prompts_per_step,
+            n_samples_per_prompt=n_samples_per_prompt,
+            max_response_len=max_response_len,
+            temperature=temperature,
+            optimizer_steps=optimizer_steps,
+            grad_accum=grad_accum,
+            seed=seed,
+            prompt_builder=lambda p: p,  # unused: the env builds the Goedel prompts
+            dead_group_retries=dead_group_retries,
+        )
+        self.env = env
+        self.credit = credit
+        # Multi-turn episode reward is <= 1.0 (turn_discount ** 0); the single-turn
+        # source's 1.1 (accuracy + format bonus) would misfile all-solved groups.
+        self.filter_stats.max_possible_score = 1.0
+        self.last_rollout_stats: dict[str, float] = {}
+
+    def bind(  # noqa: D401 - deliberately narrower than the parent's (no client/reward/verifier)
+        self,
+        *,
+        generate_factory: Callable[..., Callable[[list], list]],
+        old_logprob_fn,
+        weights_dir_fn,
+        rank: int,
+        world_size: int,
+        device: str = "cpu",
+        canary=None,
+        canary_generate_fn=None,
+        collapse_tripwire=None,
+        trainer=None,
+        canary_log_path=None,
+    ) -> None:
+        self.generate_factory = generate_factory
+        super().bind(
+            client=None, reward=None, verifier=None, old_logprob_fn=old_logprob_fn,
+            weights_dir_fn=weights_dir_fn, rank=rank, world_size=world_size, device=device,
+            canary=canary, canary_generate_fn=canary_generate_fn, collapse_tripwire=collapse_tripwire,
+            trainer=trainer, canary_log_path=canary_log_path,
+        )
+
+    def _rollout_once(self, indices: list[int], optimizer_step: int = 0) -> Optional[AssembledMultiTurn]:
+        specs = [
+            EpisodeSpec(Problem(self.prompt_rows[i]["problem_id"], self.prompt_rows[i]["problem"]), s)
+            for i in indices
+            for s in range(self.n_samples_per_prompt)
+        ]
+        generate = self.generate_factory(
+            seed=self.rollout_seed(optimizer_step), weights_dir=self.weights_dir_fn()
+        )
+        episodes = self.env.run_group(specs, generate)
+        assembled = assemble_multiturn(episodes, self.credit, self.pad_token_id, self.filter_stats)
+        if assembled is not None:
+            self.last_rollout_stats = assembled.stats
+        return assembled

code/align/rl/ops.py

@@ -0,0 +1,99 @@
+"""Pure tensor ops for the GRPO layer.
+
+Vendored from allenai/open-instruct @ 8759c91:
+  masked_mean            <- open_instruct/rl_utils.py:426-448
+  entropy_from_logits    <- open_instruct/model_utils.py:326-334
+  log_softmax_and_gather <- open_instruct/model_utils.py:588-601
+  estimate_kl            <- open_instruct/model_utils.py:793-819
+
+Bodies are verbatim; only the module docstrings/imports differ. Any edit to
+these functions breaks parity with the upstream line citations in
+docs/research/nmo-research/docs/04_ALLENAI_CODEBASE_MAP.md and must be
+recorded as a deliberate divergence, not slipped in.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+
+
+def masked_mean(
+    values: torch.Tensor,
+    mask: torch.Tensor,
+    axis: Optional[int] = None,
+    denominator: Optional[float] = None,
+) -> torch.Tensor:
+    """Compute mean of tensor with masked values.
+
+    Returns 0 if mask is empty (no valid elements) to avoid division by zero.
+    This can happen with sequence parallel when a chunk contains only
+    query/padding tokens.
+
+    A3 note: passing a constant float ``denominator`` is what realises the
+    Dr.GRPO constant-global-denominator form; leaving it None realises the
+    DAPO token-mean. align's conformance gate requires the float.
+    """
+    extra_dims = values.ndim - mask.ndim
+    if axis is None:
+        sum_dims: int | tuple[int, ...] = tuple(range(extra_dims, values.ndim))
+    elif axis >= 0:
+        sum_dims = axis + extra_dims
+    else:
+        sum_dims = axis
+    numerator = (values * mask).sum(dim=sum_dims)
+    denom = mask.sum(dim=axis) if denominator is None else denominator
+    if isinstance(denom, torch.Tensor):
+        result = torch.where(denom > 0, numerator / denom, torch.zeros_like(numerator))
+    else:
+        result = numerator / denom if denom > 0 else torch.zeros_like(numerator)
+    return result.flatten(extra_dims).mean(-1) if result.ndim > extra_dims else result
+
+
+def entropy_from_logits(logits: torch.Tensor) -> torch.Tensor:
+    """Compute the entropy of the logits.
+
+    Upstream credit chain: open-instruct borrowed this from verl
+    (verl/utils/torch_functional.py). logsumexp form avoids materialising
+    log-probabilities next to probabilities.
+    """
+    pd = torch.nn.functional.softmax(logits, dim=-1)
+    entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)
+    return entropy
+
+
+def log_softmax_and_gather(logits: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
+    """`log_softmax -> gather` without allocating a [B, T, V] logprob tensor.
+
+    Upstream compiles this (open-instruct PR #584); align leaves compilation to
+    the caller because torch.compile on this box requires the vendored-header
+    toolchain fix (align/kernels/_toolchain.py) and is opted into per-site.
+    """
+    logprobs = logits.log_softmax(dim=-1)
+    return torch.gather(logprobs, dim=-1, index=index.unsqueeze(-1)).squeeze(-1)
+
+
+def estimate_kl(ref_logprobs_diff: torch.Tensor, ratio: torch.Tensor) -> torch.Tensor:
+    """Compute 4 KL divergence estimators between current and reference policies.
+
+    Args:
+        ref_logprobs_diff: (new_logprobs - ref_logprobs), pre-clamped to
+            [-40, 40] by the caller for numerical stability.
+        ratio: exp(new_logprobs - old_logprobs) vs the step-start policy.
+
+    Returns:
+        Tensor of shape [4, *ref_logprobs_diff.shape]:
+            [0] linear approximation
+            [1] quadratic approximation (diff^2 / 2)
+            [2] numerically stable form (expm1(-diff) + diff)  -- default
+            [3] importance-weighted (ratio * diff)
+    """
+    return torch.stack(
+        [
+            ref_logprobs_diff,
+            (ref_logprobs_diff) ** 2 / 2,
+            torch.expm1(-ref_logprobs_diff) + ref_logprobs_diff,
+            ratio * ref_logprobs_diff,
+        ]
+    )

code/align/rl/rewards.py

@@ -0,0 +1,98 @@
+"""Reward composition -- the ONE genuinely net-new A3 piece.
+
+The corpus's conformance table (04_ALLENAI_CODEBASE_MAP.md item 9) found soft
+overlong punishment absent from open-instruct entirely (its
+``mask_truncated_completions`` is DAPO's HARD binary filter, a different
+mechanism). Per the same table's architectural note: this lives in the
+REWARD-COMPOSITION layer, never in the loss -- keeping compute_grpo_loss
+provably unmodified from upstream is an auditability property of the
+pre-registration, not a style preference.
+
+soft_overlong_punishment implements DAPO eq. (arXiv:2503.14476, "Soft
+Overlong Punishment"): zero inside the budget, linear ramp to -1 across the
+final ``cache`` tokens, -1 beyond max length.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Sequence
+
+
+def soft_overlong_punishment(response_len: int, max_len: int, cache_len: int) -> float:
+    """DAPO soft length penalty in [-1, 0].
+
+    0                                for len <= max_len - cache_len
+    ((max_len - cache_len) - len)/cache_len   inside the buffer (linear to -1)
+    -1                               for len > max_len
+    """
+    if cache_len <= 0 or cache_len >= max_len:
+        raise ValueError(f"need 0 < cache_len < max_len, got cache={cache_len} max={max_len}")
+    soft_start = max_len - cache_len
+    if response_len <= soft_start:
+        return 0.0
+    if response_len > max_len:
+        return -1.0
+    return (soft_start - response_len) / cache_len
+
+
+@dataclass
+class CompositeReward:
+    """A3 reward: verifier accuracy + soft format term + soft overlong penalty.
+
+    ``accuracy`` and ``format_score`` arrive from the verifier layer (rule
+    verifier only -- no PRM/neural RM, A3's hardest lock and the field's
+    strongest cross-lab consensus per doc 08). This class only composes and
+    itemises them so history.jsonl can log each term separately -- the
+    format/structural-consistency instrument (doc 09 amendment) reads the
+    per-term stream, not the sum.
+    """
+
+    format_weight: float = 0.1
+    overlong_max_len: Optional[int] = None
+    overlong_cache_len: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        # Fail at CONSTRUCTION, not at the first rollout. The dangerous half is
+        # cache_len set with max_len unset: __call__ takes the "overlong off"
+        # branch and the soft length penalty silently never applies while the
+        # config reads as if it were armed (a disarmed-instrument shape).
+        if self.overlong_max_len is None and self.overlong_cache_len is not None:
+            raise ValueError(
+                f"overlong_cache_len={self.overlong_cache_len} given without overlong_max_len: "
+                "the soft overlong penalty would be silently OFF"
+            )
+        if self.overlong_max_len is not None:
+            if self.overlong_cache_len is None:
+                raise ValueError("overlong_cache_len required when overlong_max_len is set")
+            if not 0 < self.overlong_cache_len < self.overlong_max_len:
+                raise ValueError(
+                    f"need 0 < overlong_cache_len < overlong_max_len, got "
+                    f"cache={self.overlong_cache_len} max={self.overlong_max_len}"
+                )
+
+    def __call__(self, accuracy: float, format_score: float, response_len: int) -> dict[str, float]:
+        terms: dict[str, float] = {
+            "reward/accuracy": float(accuracy),
+            "reward/format": self.format_weight * float(format_score),
+        }
+        if self.overlong_max_len is not None:
+            if self.overlong_cache_len is None:
+                raise ValueError("overlong_cache_len required when overlong_max_len is set")
+            terms["reward/overlong"] = soft_overlong_punishment(
+                response_len, self.overlong_max_len, self.overlong_cache_len
+            )
+        else:
+            terms["reward/overlong"] = 0.0
+        terms["reward/total"] = (
+            terms["reward/accuracy"] + terms["reward/format"] + terms["reward/overlong"]
+        )
+        return terms
+
+    def batch(
+        self, accuracies: Sequence[float], format_scores: Sequence[float], lengths: Sequence[int]
+    ) -> list[float]:
+        return [
+            self(a, f, n)["reward/total"]
+            for a, f, n in zip(accuracies, format_scores, lengths, strict=True)
+        ]

code/align/rl/stage_config.py

@@ -0,0 +1,253 @@
+"""Typed stage-b config: configs/grpo_stageb.toml -> dataclasses + launch refusals.
+
+Deliberately NOT align.config.RunConfig: that is the trainer's surface, this is
+the stage's pre-registration. Collapsing them would let a trainer default
+silently become an experimental parameter (the same separation dpo_stage
+draws). The trainer still gets a RunConfig -- built FROM this, explicitly,
+in scripts/grpo.py.
+
+Every check in `validate()` is a launch refusal for a failure phase-1 paid
+for: the denominator arithmetic (obligation a's config half), the
+save-cadence/weight-sync alignment (a drifted engine trains against stale
+rho baselines silently), the ceiling receipt (the 0/100-perfect-policy
+trap), max_grad_norm (R3), a live canary cadence (R4).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+
+from align.rl.grpo_loss import GRPOLossConfig, GRPOLossType
+
+
+@dataclass
+class RolloutShape:
+    prompts_per_step: int
+    max_response_len: int
+    temperature: float = 1.0
+    overlong_cache_len: int = 128
+    #: Where completions come from. "spool" = the external vLLM worker (own venv,
+    #: extras/vllm_rollout; the backend the pre-registration describes) and is
+    #: the default so an existing config keeps its generation engine. "native" =
+    #: the in-repo torch sampler on a STANDALONE engine copy on a dedicated GPU
+    #: (NMO_ROLLOUT_GPU), synced by reloading the latest checkpoint; GPU-unverified
+    #: at 7B at the time of the port, so it is opt-in. A silent switch changes
+    #: throughput by an order of magnitude and the logprob basis (native logprobs
+    #: are of the temperature-scaled distribution), hence never inferred.
+    backend: str = "spool"
+    #: INFORMATIONAL ONLY: configs/grpo_stageb.toml carries `vllm_env` as a note of the
+    #: worker environment the pre-registration ran under. Nothing reads it -- the env is
+    #: pinned in code (align/rollout/client.py) -- so it is declared here, rather than
+    #: left to the unknown-key refusal, purely so that refusal can be strict for every
+    #: OTHER key. A value here does NOT change the worker environment.
+    vllm_env: str = ""
+
+
+@dataclass
+class RewardCfg:
+    format_weight: float = 0.1
+
+
+@dataclass
+class DataCfg:
+    band_file: str
+    banding_manifest: str
+
+
+@dataclass
+class TrainCfg:
+    lr: float
+    optimizer_steps: int
+    warmup_frac: float = 0.03
+    micro_batch_size: int = 1
+    save_every_micro_steps: int = 100
+    keep_checkpoints: int = 2
+    #: Per-rank optimizer shards (TrainConfig.per_rank_optimizer_state); forwarded
+    #: by scripts/grpo.py::_run_config. Default False = gathered layout.
+    per_rank_optimizer_state: bool = False
+    max_grad_norm: float | None = None
+    seed: int = 41
+    optimizer_name: str = "adamw8bit"
+    grad_accum: int = 8
+    attn_implementation: str = "sdpa"
+
+    # -- run-health bounds (G0), forwarded to TrainConfig by the stage script's
+    # _run_config: a field that is not copied trains at the default while the
+    # TOML says otherwise. Semantics: align/config.py::TrainConfig.
+    max_wall_seconds: float | None = None
+    max_consecutive_skipped_steps: int = 3
+    max_skipped_fraction: float = 0.05
+
+
+@dataclass
+class CanaryCfg:
+    every: int
+    n_problems: int
+    source: str = "gsm8k"
+
+
+@dataclass
+class GrpoStageConfig:
+    sft_checkpoint: str
+    run_dir: str
+    grpo: GRPOLossConfig
+    rollout: RolloutShape
+    reward: RewardCfg
+    data: DataCfg
+    train: TrainCfg
+    canary: CanaryCfg
+    armed_tripwires: tuple[str, ...] = ("grad_norm",)
+    ceiling_receipt: str = "outputs/eval/ceiling_test.json"
+    rollout_timeout_s: float = 1800.0
+    dead_group_retries: int = 5
+    extras: dict = field(default_factory=dict)
+
+    def validate(self) -> None:
+        g = self.grpo
+        if self.rollout.backend not in ("spool", "native"):
+            raise ValueError(
+                f"rollout.backend {self.rollout.backend!r}: choose 'spool' (external vLLM "
+                "worker) or 'native' (in-repo torch sampler)"
+            )
+        expected = float(
+            self.rollout.prompts_per_step
+            * g.num_samples_per_prompt_rollout
+            * self.rollout.max_response_len
+        )
+        if not isinstance(g.loss_denominator, float) or g.loss_denominator != expected:
+            raise ValueError(
+                f"loss_denominator {g.loss_denominator!r} != prompts_per_step x G x "
+                f"max_response_len = {self.rollout.prompts_per_step} x "
+                f"{g.num_samples_per_prompt_rollout} x {self.rollout.max_response_len} "
+                f"= {expected}. The Dr.GRPO constant is DERIVED, never hand-typed."
+            )
+        if not os.path.isfile(self.data.band_file):
+            raise FileNotFoundError(
+                f"band file missing: {self.data.band_file}. Run "
+                "scripts/build_stageb_data.py first; the driver never builds data."
+            )
+        if self.train.save_every_micro_steps % self.train.grad_accum != 0:
+            raise ValueError(
+                f"save_every_micro_steps={self.train.save_every_micro_steps} is not a "
+                f"multiple of grad_accum={self.train.grad_accum}: checkpoint saves ARE "
+                "the weight sync (the engine reloads the latest checkpoint), so a save "
+                "landing mid-accumulation would sync half-stepped weights."
+            )
+        if self.train.max_grad_norm is None or self.train.max_grad_norm <= 0:
+            raise ValueError(
+                "train.max_grad_norm must be a positive float: R3 (2026-08-23 audit) "
+                "-- the recorded phase-1 rerun blew up to grad_norm 355,830 with no "
+                "clipping anywhere. Unclipped RL is not a configuration this stage "
+                "accepts."
+            )
+        if self.canary.every <= 0 or self.canary.n_problems <= 0:
+            raise ValueError(
+                "canary cadence/n must be positive: R4 -- 727 optimizer steps of "
+                "phase-1 preference training ran with zero task-accuracy signal."
+            )
+        if "grad_norm" not in self.armed_tripwires:
+            raise ValueError("stage-b requires the grad_norm HALT tripwire armed (R3)")
+        if not g.record_entropy:
+            raise ValueError("record_entropy must be true: the collapse tripwire reads entropy")
+
+    def validate_ceiling_receipt(self) -> dict:
+        """Separate from validate(): tests exercise config logic without a receipt."""
+        if not os.path.isfile(self.ceiling_receipt):
+            raise FileNotFoundError(
+                f"ceiling receipt missing: {self.ceiling_receipt}. No GRPO rollout "
+                "launches until the reward path is certified (the 0/100-perfect-policy "
+                "incident, mistral-prm-ppo 2026-08-14)."
+            )
+        with open(self.ceiling_receipt) as f:
+            receipt = json.load(f)
+        if not receipt.get("overall_pass"):
+            raise RuntimeError(f"ceiling receipt {self.ceiling_receipt} is not a PASS")
+        return receipt
+
+    def validate_world(self, world_size: int) -> None:
+        """Checks that need the launch topology."""
+        rows_per_step = self.rollout.prompts_per_step * self.grpo.num_samples_per_prompt_rollout
+        per_micro = rows_per_step // (world_size * self.train.grad_accum)
+        if per_micro < 1:
+            raise ValueError(
+                f"rollout shape too small for the topology: {rows_per_step} completions "
+                f"per step cannot fill world_size({world_size}) x "
+                f"grad_accum({self.train.grad_accum}) micro-slices"
+            )
+
+
+def load_stage_config(path: str, smoke: bool = False) -> GrpoStageConfig:
+    from align.config import load_toml
+
+    raw = load_toml(path)
+    g = dict(raw["grpo"])
+    g["loss_fn"] = GRPOLossType(g.get("loss_fn", "dapo"))
+    # TOML floats survive; TOML ints for float fields do not -- coerce the one
+    # field whose TYPE is load-bearing (the conformance gate refuses non-float).
+    if "loss_denominator" in g:
+        g["loss_denominator"] = float(g["loss_denominator"])
+    t = dict(raw["train"])
+    t["optimizer_name"] = raw.get("optimizer", {}).get("name", t.pop("optimizer_name", "adamw8bit"))
+    t["grad_accum"] = raw.get("optimizer", {}).get("grad_accum", t.pop("grad_accum", 8))
+    # STRUCT MODE for [train]. This used to filter through
+    # TrainCfg.__dataclass_fields__, which made a typo'd or not-yet-plumbed key
+    # (e.g. per_rank_optimizer_state before TrainCfg had it) ACCEPTED AND
+    # SILENTLY IGNORED: the TOML looked armed and the run was not. Raise at
+    # launch, naming the key, like DpoStageConfig's struct-mode loader.
+    unknown = sorted(k for k in t if k not in TrainCfg.__dataclass_fields__)
+    if unknown:
+        raise ValueError(
+            f"unknown key(s) {unknown} in [train] are not fields of TrainCfg "
+            f"(struct mode: a silently ignored key is a disarmed config). "
+            f"Known fields: {sorted(TrainCfg.__dataclass_fields__)}"
+        )
+    # STRUCT MODE for [rollout] too. The same filter-through-__dataclass_fields__
+    # idiom that [train] shed above survived here: `[rollout] max_reponse_len = 512`
+    # (typo) or a not-yet-plumbed `top_p = 0.9` was DROPPED, the shape kept its
+    # default, and `validate()` then checked the loss denominator against the
+    # default -- the pre-registered rollout looked armed and was not.
+    r = dict(raw["rollout"])
+    unknown_rollout = sorted(k for k in r if k not in RolloutShape.__dataclass_fields__)
+    if unknown_rollout:
+        raise ValueError(
+            f"unknown key(s) {unknown_rollout} in [rollout] are not fields of RolloutShape "
+            f"(struct mode: a silently ignored key is a disarmed config). "
+            f"Known fields: {sorted(RolloutShape.__dataclass_fields__)}"
+        )
+    cfg = GrpoStageConfig(
+        sft_checkpoint=raw["sft_checkpoint"],
+        run_dir=raw["run_dir"],
+        grpo=GRPOLossConfig(**g),
+        rollout=RolloutShape(**r),
+        reward=RewardCfg(**raw.get("reward", {})),
+        data=DataCfg(**raw["data"]),
+        train=TrainCfg(**t),
+        canary=CanaryCfg(**raw["canary"]),
+        armed_tripwires=tuple(raw.get("tripwires", {}).get("armed", ("grad_norm",))),
+    )
+    if smoke:
+        cfg = _smokeify(cfg)
+    cfg.validate()
+    return cfg
+
+
+def _smokeify(cfg: GrpoStageConfig) -> GrpoStageConfig:
+    """Tiny everything; the derived denominator is recomputed, not hand-edited."""
+    cfg.rollout.prompts_per_step = 2
+    cfg.grpo.num_samples_per_prompt_rollout = 8  # A3 floor is G >= 8, smoke included
+    # 256 tokens, not 32: the sixth smoke run proved a cap real answers
+    # cannot fit under produces zero reward variance -> every group dead ->
+    # the (correct) band-mischosen refusal. The smoke needs a policy with
+    # SOME task signal (the real SFT substrate) and room to finish answers.
+    cfg.rollout.max_response_len = 256
+    cfg.rollout.overlong_cache_len = 64
+    cfg.grpo.loss_denominator = float(2 * 8 * 256)
+    cfg.train.optimizer_steps = 10
+    cfg.train.grad_accum = 2
+    cfg.train.save_every_micro_steps = 10  # sync twice in a 20-micro smoke
+    cfg.canary.every = 5
+    cfg.canary.n_problems = 4
+    cfg.run_dir = cfg.run_dir + "_smoke"
+    return cfg

code/align/rl/tripwire.py

@@ -0,0 +1,202 @@
+"""Entropy/length collapse tripwire -- the second net-new A3 piece.
+
+Upstream has the loggable INPUTS (record_entropy=True) but no tripwire logic
+(corpus conformance table item 7: "the knob exists, the tripwire mechanism
+does not"). This is the GRPO sibling of the DPO chosen-logprob tripwire that
+fired -- and autonomously launched its pre-registered response -- on Arm B
+(2026-08-22, runs/dpo_arm_b_response/).
+
+Design contract, identical to the DPO one:
+  * The tripwire DETECTS and RECORDS. The stage script owns the response,
+    and the only admissible response is the pre-registered one
+    (KL=0 collapse -> ONE rerun with beta=1e-4, recorded AS the result;
+    no third configuration). No improvisation lives here by construction.
+  * Fed from the same metrics stream that history.jsonl records, so a fire
+    is always reproducible from the run's own log after the fact.
+
+Signals (cross-lab consensus per doc 08: entropy collapse is the 2025
+field-wide GRPO failure mode; length collapse is its usual companion):
+  E: policy entropy trend -- median over the trailing window falls below
+     ``entropy_floor_frac`` x the median over the baseline window.
+  L: mean response length trend -- same construction, lengths shrinking
+     below ``length_floor_frac`` of baseline (degenerate short answers).
+Either alone fires; both are cheap medians over deques, no model access.
+
+Resume (2026-10-12 review). The instrument is STATEFUL -- a baseline measured
+over the first ``baseline_steps`` observations, a trailing window, a step
+counter and a latch -- so a resume that rebuilds it fresh silently changes the
+experiment: the baseline is re-measured on the POST-resume (possibly already
+collapsed) policy, which makes the collapse the new "normal", and
+``min_steps_before_fire`` restarts at 0, blinding the guard for ~100 steps.
+``state_dict()`` / ``load_state_dict()`` carry all of it through the Trainer's
+``extra_state`` slot in resume.pt (``GRPOTrainModule.extra_state``).
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from collections import deque
+from dataclasses import dataclass, field
+from statistics import median
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CollapseTripwire:
+    baseline_steps: int = 50
+    window_steps: int = 50
+    entropy_floor_frac: float = 0.2
+    length_floor_frac: float = 0.4
+    min_steps_before_fire: int = 100
+
+    _baseline_entropy: list[float] = field(default_factory=list)
+    _baseline_length: list[float] = field(default_factory=list)
+    _entropy: deque = field(default_factory=deque)
+    _length: deque = field(default_factory=deque)
+    _steps: int = 0
+    fired: bool = False
+    fire_reason: Optional[str] = None
+    fire_step: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        self._entropy = deque(maxlen=self.window_steps)
+        self._length = deque(maxlen=self.window_steps)
+
+    def observe(self, step: int, entropy: float, mean_response_len: float) -> bool:
+        """Feed one optimizer step's metrics; returns True the step it fires.
+
+        Once fired it stays fired (a latched instrument, like the DPO one) --
+        the stage script reads ``fired``/``fire_reason`` and executes the
+        pre-registered response exactly once.
+        """
+        if not (math.isfinite(entropy) and math.isfinite(mean_response_len)):
+            # NaN compares False against every threshold below (`nan > 0`,
+            # `nan < x`), so a NaN entropy stream would leave the entropy arm
+            # permanently inert while the object reads as armed -- the
+            # disarmed-tripwire shape (audit D1). `_on_step_boundary` feeds NaN
+            # when a metric key is missing. Refuse loudly instead.
+            raise ValueError(
+                f"CollapseTripwire.observe(step={step}): non-finite input "
+                f"(entropy={entropy}, mean_response_len={mean_response_len}); a guard "
+                "fed NaN can never fire"
+            )
+        self._steps += 1
+        if len(self._baseline_entropy) < self.baseline_steps:
+            self._baseline_entropy.append(entropy)
+            self._baseline_length.append(mean_response_len)
+            return False
+        self._entropy.append(entropy)
+        self._length.append(mean_response_len)
+
+        if (
+            self.fired
+            or self._steps < self.min_steps_before_fire
+            or len(self._entropy) < self.window_steps
+        ):
+            return False
+
+        base_e = median(self._baseline_entropy)
+        base_l = median(self._baseline_length)
+        cur_e = median(self._entropy)
+        cur_l = median(self._length)
+        reasons = []
+        if base_e > 0 and cur_e < self.entropy_floor_frac * base_e:
+            reasons.append(
+                f"entropy collapse: trailing median {cur_e:.4f} < {self.entropy_floor_frac} x baseline {base_e:.4f}"
+            )
+        if base_l > 0 and cur_l < self.length_floor_frac * base_l:
+            reasons.append(
+                f"length collapse: trailing median {cur_l:.1f} < {self.length_floor_frac} x baseline {base_l:.1f}"
+            )
+        if reasons:
+            self.fired = True
+            self.fire_step = step
+            self.fire_reason = "; ".join(reasons)
+            logger.error(
+                "TRIPWIRE (GRPO collapse) at step %d: %s -- pre-registered response: ONE beta=1e-4 "
+                "rerun recorded AS the result; no third configuration.",
+                step,
+                self.fire_reason,
+            )
+            return True
+        return False
+
+    #: Bump on any change to the persisted layout; a mismatch refuses to load.
+    STATE_VERSION = 1
+
+    def _config_dict(self) -> dict:
+        return {
+            "baseline_steps": self.baseline_steps,
+            "window_steps": self.window_steps,
+            "entropy_floor_frac": self.entropy_floor_frac,
+            "length_floor_frac": self.length_floor_frac,
+            "min_steps_before_fire": self.min_steps_before_fire,
+        }
+
+    def state_dict(self) -> dict:
+        """Everything `observe` reads, as primitives only.
+
+        Plain dict/list/float/int/bool/None/str so it survives the Trainer's
+        ``torch.load(weights_only=True)`` of resume.pt -- unpickling arbitrary
+        objects to restore a counter would be a code-execution path for nothing.
+        The thresholds are included so `load_state_dict` can refuse a state cut
+        for a different configuration.
+        """
+        return {
+            "version": self.STATE_VERSION,
+            "config": self._config_dict(),
+            "baseline_entropy": [float(x) for x in self._baseline_entropy],
+            "baseline_length": [float(x) for x in self._baseline_length],
+            "entropy": [float(x) for x in self._entropy],
+            "length": [float(x) for x in self._length],
+            "steps": int(self._steps),
+            "fired": bool(self.fired),
+            "fire_reason": self.fire_reason,
+            "fire_step": self.fire_step,
+        }
+
+    def load_state_dict(self, sd: dict) -> None:
+        """Restore `state_dict()` output; refuse anything that is not exactly it.
+
+        A silently half-restored tripwire is the disarmed-instrument shape, so
+        every inconsistency raises: wrong version, thresholds that differ from
+        this instance's (a different experiment -- rebuild it deliberately),
+        or windows longer than the configured sizes. Validation happens BEFORE
+        any field is assigned, so a refused load leaves the instance untouched.
+        """
+        if sd.get("version") != self.STATE_VERSION:
+            raise ValueError(
+                f"CollapseTripwire state version {sd.get('version')!r} != {self.STATE_VERSION}"
+            )
+        if sd.get("config") != self._config_dict():
+            raise ValueError(
+                f"CollapseTripwire state was saved under {sd.get('config')!r} but this "
+                f"instance is configured {self._config_dict()!r}; resuming under different "
+                "thresholds is a different experiment."
+            )
+        base_e, base_l = list(sd["baseline_entropy"]), list(sd["baseline_length"])
+        ent, length = list(sd["entropy"]), list(sd["length"])
+        if len(base_e) != len(base_l) or len(ent) != len(length):
+            raise ValueError("CollapseTripwire state has mismatched entropy/length series")
+        if len(base_e) > self.baseline_steps or len(ent) > self.window_steps:
+            raise ValueError("CollapseTripwire state windows exceed the configured sizes")
+        if not all(math.isfinite(x) for x in base_e + base_l + ent + length):
+            raise ValueError("CollapseTripwire state contains non-finite observations")
+        self._baseline_entropy = [float(x) for x in base_e]
+        self._baseline_length = [float(x) for x in base_l]
+        self._entropy = deque((float(x) for x in ent), maxlen=self.window_steps)
+        self._length = deque((float(x) for x in length), maxlen=self.window_steps)
+        self._steps = int(sd["steps"])
+        self.fired = bool(sd["fired"])
+        self.fire_reason = sd["fire_reason"]
+        self.fire_step = sd["fire_step"]
+
+    def as_log_dict(self) -> dict:
+        return {
+            "tripwire/fired": self.fired,
+            "tripwire/fire_step": self.fire_step,
+            "tripwire/fire_reason": self.fire_reason,
+        }

code/align/rl/types.py

@@ -0,0 +1,84 @@
+"""Rollout data types for stage-b GRPO.
+
+The rollout ORCHESTRATOR (M-D, separate build) produces these; the
+GRPOTrainModule consumes them. Alignment convention, stated once: per-token
+logprob tensors are LABEL-ALIGNED -- position t holds the logprob of
+input_ids[:, t+1], so they have length T-1 and pair with
+`shifted_response_mask`. Getting this half-open boundary wrong is the same
+off-by-one family as the decode-position incident (section-09) and the
+completion-mask defect (audit R1); the validating constructor exists so the
+mismatch fails at batch construction, not as a silently shifted objective.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+
+@dataclass
+class RolloutBatch:
+    """One training micro-batch of rollout data.
+
+    Shapes (B = completions in the batch, T = padded sequence length):
+      input_ids       [B, T] int64  prompt + completion, right-padded
+      attention_mask  [B, T] int64  1 on real tokens
+      response_mask   [B, T] int64  1 on COMPLETION tokens only (input space)
+      advantages      [B]  float    per-sequence Dr.GRPO centered advantage
+      old_logprobs    [B, T-1]      policy-at-rollout logprobs (label-aligned)
+      vllm_logprobs   [B, T-1]      inference-engine logprobs (label-aligned)
+      ref_logprobs    [B, T-1]?     reference-policy logprobs; REQUIRED iff
+                                    the loss config has beta > 0 (the KL
+                                    ablation), absent otherwise
+      rewards         {name: [B]}   per-completion reward terms (reward/total,
+                                    reward/accuracy, ... -- itemised so the
+                                    format-consistency instrument reads terms,
+                                    not sums)
+      group_ids       [B] int64     prompt-group id of each completion
+    """
+
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    response_mask: torch.Tensor
+    advantages: torch.Tensor
+    old_logprobs: torch.Tensor
+    vllm_logprobs: torch.Tensor
+    group_ids: torch.Tensor
+    rewards: dict[str, torch.Tensor] = field(default_factory=dict)
+    ref_logprobs: torch.Tensor | None = None
+
+    def __post_init__(self) -> None:
+        B, T = self.input_ids.shape
+        for name, t, shape in (
+            ("attention_mask", self.attention_mask, (B, T)),
+            ("response_mask", self.response_mask, (B, T)),
+            ("advantages", self.advantages, (B,)),
+            ("old_logprobs", self.old_logprobs, (B, T - 1)),
+            ("vllm_logprobs", self.vllm_logprobs, (B, T - 1)),
+            ("group_ids", self.group_ids, (B,)),
+        ):
+            if tuple(t.shape) != shape:
+                raise ValueError(
+                    f"RolloutBatch.{name} has shape {tuple(t.shape)}, expected {shape} "
+                    f"(B={B}, T={T}; logprob tensors are label-aligned, length T-1)"
+                )
+        if self.ref_logprobs is not None and tuple(self.ref_logprobs.shape) != (B, T - 1):
+            raise ValueError(
+                f"RolloutBatch.ref_logprobs has shape {tuple(self.ref_logprobs.shape)}, "
+                f"expected {(B, T - 1)}"
+            )
+        for name, t in self.rewards.items():
+            if tuple(t.shape) != (B,):
+                raise ValueError(
+                    f"RolloutBatch.rewards[{name!r}] has shape {tuple(t.shape)}, expected {(B,)}"
+                )
+
+    @property
+    def shifted_response_mask(self) -> torch.Tensor:
+        """Response mask in LABEL space [B, T-1], pairing with the logprobs."""
+        return self.response_mask[:, 1:]
+
+    @property
+    def n_response_tokens(self) -> int:
+        return int(self.shifted_response_mask.sum())

code/align/rollout/__init__.py

@@ -0,0 +1,74 @@
+"""Rollout: request/result shapes, the native sampler, and the optional spool client.
+
+Backends (selected by `align.rollout.client.make_rollout_client`):
+
+  'native' (default)  `native.NativeRolloutBackend` -- in-process torch KV-cache
+                      sampler; zero dependencies beyond torch.
+  'spool'             `client.RolloutClient` -- file protocol to the optional
+                      vLLM worker in extras/vllm_rollout/ (its own venv).
+
+Nothing under align/ imports vllm or ray, and nothing under align/ imports
+extras/ (tests/test_dependency_firewall.py enforces both). The extras worker
+imports `align.rollout.protocol`, which is stdlib-only BY CONTRACT; this
+package `__init__` therefore imports ONLY that module eagerly. The torch-using
+pieces (`NativeRolloutBackend`, `make_rollout_client`) resolve lazily so the
+worker's venv never pays for, or is constrained by, them.
+
+Per the 2026-08-22 sovereignty amendment (RESULTS.md dated entry): vLLM is
+permitted for GENERATION only, in an isolated venv (.venv-rollout), behind a
+process boundary. The v1 topology (DESIGN v3 section B.1b, decided 2026-08-24):
+3 FSDP train ranks + 1 dedicated engine GPU; weight handoff via checkpoint
+directory + engine reload every `sync_every` optimizer steps;
+use_rho_correction=True absorbs the resulting off-policyness -- that is the
+term's entire purpose. The native backend honours the same contract
+(`load_weights_fn` on a changed weights_dir).
+"""
+
+from align.rollout.protocol import (
+    RolloutBackend,
+    RolloutRequest,
+    RolloutResult,
+    read_request,
+    read_result,
+    write_request,
+    write_result,
+)
+
+__all__ = [
+    "PagedEngine",
+    "PagedEngineConfig",
+    "PagedResult",
+    "SamplingParams",
+    "NativeRolloutBackend",
+    "NativeRolloutConfig",
+    "RolloutBackend",
+    "RolloutRequest",
+    "RolloutResult",
+    "checkpoint_weight_loader",
+    "make_rollout_client",
+    "read_request",
+    "read_result",
+    "write_request",
+    "write_result",
+]
+
+_LAZY = {
+    # P3 opt-in paged continuous-batching sampler (docs/P3_ENGINE.md). Nothing
+    # in the default GRPO/rollout path imports these.
+    "PagedEngine": "align.rollout.engine",
+    "PagedEngineConfig": "align.rollout.engine",
+    "PagedResult": "align.rollout.engine",
+    "SamplingParams": "align.rollout.engine",
+    "NativeRolloutBackend": "align.rollout.native",
+    "NativeRolloutConfig": "align.rollout.native",
+    "checkpoint_weight_loader": "align.rollout.native",
+    "make_rollout_client": "align.rollout.client",
+}
+
+
+def __getattr__(name):
+    if name in _LAZY:
+        import importlib
+
+        return getattr(importlib.import_module(_LAZY[name]), name)
+    raise AttributeError(f"module 'align.rollout' has no attribute {name!r}")

code/align/rollout/client.py

@@ -0,0 +1,234 @@
+"""Trainer-side rollout clients and the backend selector.
+
+Two backends serve the same `RolloutRequest -> RolloutResult` contract
+(align/rollout/protocol.py), chosen by `make_rollout_client(backend=...)`:
+
+  'native' (DEFAULT) -- `native.NativeRolloutBackend`: in-process torch
+      KV-cache sampler on the training/passed model. Zero extra dependencies;
+      this is what makes GRPO runnable on torch + triton + numpy + safetensors.
+  'spool'            -- `RolloutClient` below: file protocol to the optional
+      vLLM worker (extras/vllm_rollout/rollout_worker.py) running in its OWN
+      venv. A throughput option, never a requirement.
+
+This module lives in the TRAINING venv and imports nothing from vLLM -- the
+process boundary is the sovereignty line. `spawn_worker` launches the extras
+worker under the ISOLATED rollout venv with the full required environment
+baked in, so a caller cannot forget the flashinfer workaround or point a cache
+at the 13GB root filesystem.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from typing import Optional
+
+from align.rollout.protocol import (
+    RolloutBackend,
+    RolloutRequest,
+    RolloutResult,
+    wait_for_result,
+    write_request,
+)
+
+#: The selectable backends. Data, so error messages and CLIs iterate it.
+BACKENDS: tuple[str, ...] = ("native", "spool")
+DEFAULT_BACKEND = "native"
+
+#: Repo root, derived from this file's location -- the client and the worker
+#: script it spawns always come from the same checkout.
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+ROLLOUT_PYTHON = os.path.join(REPO_ROOT, ".venv-rollout", "bin", "python")
+#: The vLLM worker is an EXTRA (own venv, see extras/vllm_rollout/README.md),
+#: not part of the align package: nothing under align/ may import it, but a
+#: path string to launch it as a subprocess is the intended coupling.
+WORKER_SCRIPT = os.path.join(REPO_ROOT, "extras", "vllm_rollout", "rollout_worker.py")
+
+
+class RolloutClient:
+    """The 'spool' backend: submit rollout requests and block on their results.
+
+    Timeouts are the caller's budget decision (stage-b sets them from the
+    measured pilot throughput); the stale-heartbeat threshold distinguishes
+    a slow request from a dead worker -- see protocol.wait_for_result.
+    """
+
+    def __init__(self, spool: str, timeout_s: float, stale_heartbeat_s: float = 600.0):
+        self.spool = spool
+        self.timeout_s = float(timeout_s)
+        self.stale_heartbeat_s = float(stale_heartbeat_s)
+
+    def submit(self, request: RolloutRequest) -> str:
+        write_request(self.spool, request)
+        return request.request_id
+
+    def collect(self, request_id: str) -> RolloutResult:
+        return wait_for_result(
+            self.spool,
+            request_id,
+            timeout_s=self.timeout_s,
+            stale_heartbeat_s=self.stale_heartbeat_s,
+        )
+
+
+#: torchrun/torchelastic variables that MUST NOT reach the worker: vLLM's
+#: engine initialises torch.distributed internally, and a child that inherits
+#: the trainer's rendezvous (MASTER_ADDR/RANK/WORLD_SIZE/...) tries to JOIN
+#: the trainer's process group -- observed on the 2026-08-24 smoke as the
+#: engine core dialing torchrun's TCP store while ranks 1-2 timed out waiting
+#: for rank 0. The worker is a separate world of size 1, always.
+_POISON_ENV_KEYS = (
+    "MASTER_ADDR",
+    "MASTER_PORT",
+    "RANK",
+    "WORLD_SIZE",
+    "LOCAL_RANK",
+    "LOCAL_WORLD_SIZE",
+    "GROUP_RANK",
+    "GROUP_WORLD_SIZE",
+    "ROLE_RANK",
+    "ROLE_WORLD_SIZE",
+    "ROLE_NAME",
+    "NODE_RANK",
+)
+_POISON_ENV_PREFIXES = ("TORCHELASTIC_", "PET_", "TORCH_NCCL_")
+
+
+def _worker_env() -> dict:
+    """A copy of os.environ with every trainer-distributed variable scrubbed."""
+    return {
+        k: v
+        for k, v in os.environ.items()
+        if k not in _POISON_ENV_KEYS and not k.startswith(_POISON_ENV_PREFIXES)
+    }
+
+
+def spawn_worker(
+    spool: str,
+    gpu: int,
+    log_path: str,
+    gpu_mem_util: float = 0.85,
+    tp: int = 1,
+    idle_exit_s: float = 0.0,
+    rollout_python: Optional[str] = None,
+) -> subprocess.Popen:
+    """Launch the vLLM worker on ONE dedicated GPU with the required env.
+
+    The env block is the launch checklist, encoded: the flashinfer sampler
+    is disabled (its JIT ninja build fails on this box's toolchain, measured
+    2026-08-23), every cache/tmp path lands on /apps, and the HF hub cache is
+    the rlhf one the stage-b snapshots live in.
+    """
+    python = rollout_python or ROLLOUT_PYTHON
+    if not os.path.isfile(python):
+        raise FileNotFoundError(
+            f"rollout venv python not found at {python}; build .venv-rollout first "
+            "(uv venv + torch==2.11.0 triton==3.6.0 vllm==0.26.0 ray pyarrow)."
+        )
+    env = _worker_env()
+    env.update(
+        {
+            "CUDA_VISIBLE_DEVICES": str(gpu),
+            "VLLM_USE_FLASHINFER_SAMPLER": "0",
+            "TMPDIR": os.path.join(REPO_ROOT, ".tmp"),
+            "XDG_CACHE_HOME": os.path.join(REPO_ROOT, ".cache"),
+            "HF_HUB_CACHE": "/apps/development/rlhf/hf_cache/hub",
+            # worker logs must flush per line or a killed worker leaves an
+            # empty log (first smoke run, 2026-08-24)
+            "PYTHONUNBUFFERED": "1",
+        }
+    )
+    os.makedirs(os.path.dirname(os.path.abspath(log_path)), exist_ok=True)
+    log = open(log_path, "a")
+    return subprocess.Popen(
+        [
+            python,
+            WORKER_SCRIPT,
+            "--spool",
+            spool,
+            "--gpu-mem-util",
+            str(gpu_mem_util),
+            "--tp",
+            str(tp),
+            "--idle-exit-s",
+            str(idle_exit_s),
+        ],
+        stdout=log,
+        stderr=subprocess.STDOUT,
+        env=env,
+        cwd=REPO_ROOT,
+        # own process GROUP so teardown can reap vLLM's EngineCore child --
+        # terminating only the parent orphaned a 19.6GB engine core on GPU 3
+        # (fifth smoke run, 2026-08-24)
+        start_new_session=True,
+    )
+
+
+def terminate_worker(proc: subprocess.Popen, timeout_s: float = 30.0) -> None:
+    """Kill the worker's WHOLE process group (worker + vLLM engine core).
+
+    proc.terminate() reaches only the parent; the engine core reparents to
+    init and keeps its VRAM. Group-SIGTERM, grace, then group-SIGKILL.
+    """
+    import signal
+
+    try:
+        pgid = os.getpgid(proc.pid)
+    except ProcessLookupError:
+        return
+    try:
+        os.killpg(pgid, signal.SIGTERM)
+        proc.wait(timeout=timeout_s)
+    except subprocess.TimeoutExpired:
+        os.killpg(pgid, signal.SIGKILL)
+        proc.wait(timeout=10)
+    except ProcessLookupError:
+        pass
+
+
+#: Explicit name for the spool backend; `RolloutClient` is kept as the
+#: original spelling because existing callers and tests import it.
+SpoolRolloutClient = RolloutClient
+
+
+def make_rollout_client(
+    backend: str = DEFAULT_BACKEND,
+    *,
+    # -- native -----------------------------------------------------------------
+    model=None,
+    tokenizer=None,
+    native_config=None,
+    load_weights_fn=None,
+    device=None,
+    # -- spool ------------------------------------------------------------------
+    spool: Optional[str] = None,
+    timeout_s: float = 3600.0,
+    stale_heartbeat_s: float = 600.0,
+) -> RolloutBackend:
+    """Build the selected rollout backend.
+
+    Each backend REFUSES the other's arguments' absence loudly instead of
+    falling back: 'native' without a model is not "spool by default", and
+    'spool' without a spool dir is not "native by default". A silent switch of
+    generation backend changes throughput by an order of magnitude and
+    off-policyness by a sync interval, and would show up first as a mystery in
+    the step time.
+    """
+    if backend == "native":
+        if model is None or tokenizer is None:
+            raise ValueError(
+                "backend='native' needs model= and tokenizer= (the in-process sampler runs "
+                "on the model it is given); pass backend='spool' to use the external worker"
+            )
+        # lazy: keeps `import align.rollout.client` free of torch/engine for
+        # spool-only callers
+        from align.rollout.native import NativeRolloutBackend
+
+        return NativeRolloutBackend(
+            model, tokenizer, native_config, load_weights_fn=load_weights_fn, device=device
+        )
+    if backend == "spool":
+        if not spool:
+            raise ValueError("backend='spool' needs spool= (the request/result spool directory)")
+        return RolloutClient(spool, timeout_s=timeout_s, stale_heartbeat_s=stale_heartbeat_s)
+    raise ValueError(f"unknown rollout backend {backend!r}; choose one of {BACKENDS}")

code/align/rollout/engine.py

@@ -0,0 +1,674 @@
+"""Continuous-batching sampler over a paged KV cache (P3, opt-in).
+
+This is the native-engine step ROADMAP Track L calls P3: many sequences of very
+different lengths in flight at once, admitted and retired independently,
+backed by `align.rollout.paged`'s block tables. It is model-agnostic over
+align's causal LMs (Mistral, Qwen3, OLMoE) WITHOUT touching them: the engine
+drives the model's own submodules (`embed_tokens`, per-block `attn_norm` /
+`attn.{q,k,v,o}_proj` / optional `q_norm`,`k_norm` / `mlp_norm` / `mlp`,
+`norm`, `lm_head`, `rope`) layer by layer and replaces only the attention
+kernel, which is the one thing a paged cache changes. Qwen3's per-head QK-norm
+and OLMoE's full-width QK-norm (+ `clip_qkv`) are told apart by the norm weight
+width, the same distinction the model files' docstrings warn about.
+
+THE SCHEDULER, in one paragraph. Every sequence is `tokens` (prompt + sampled)
+and `num_computed` (how many have KV in the cache). A step schedules, per
+sequence, the next `n_new = min(len(tokens) - num_computed, chunk_size,
+budget)` tokens. Prefill, chunked prefill, decode and post-preemption
+recompute are therefore ONE code path; a sequence samples exactly when its
+scheduled tokens reach the end of `tokens`. Mixed prefill chunks and decode
+tokens share a step (one flat ragged batch), so a long prompt cannot stall
+the decoders behind it (the head-of-line blocking a prefill-then-decode engine
+has).
+
+MEMORY PRESSURE. Blocks are allocated lazily as sequences grow. When a running
+sequence cannot get its next block the YOUNGEST unscheduled running sequence is
+preempted by RECOMPUTE: its blocks are freed, `num_computed` returns to 0, and
+it re-enters the front of the waiting queue with its tokens (and its sampler
+state) intact, so it re-prefills and carries on. Recompute rather than swap
+because there is no host swap space on this box by design. The generator is
+only advanced when a token is actually sampled, so a preemption changes WHEN a
+token is drawn, never WHICH.
+
+BATCH INVARIANCE (the property RL rollouts care about). A sequence's tokens
+and logprobs should not depend on who else is in the batch, when it was
+admitted, how its prompt was chunked, or whether it was preempted: otherwise
+`pi_infer_old` (the logprobs that feed the rho correction) is a function of the
+scheduler's mood, and two reruns of the same rollout seed disagree. Floating
+point breaks this through (a) GEMM kernels whose reduction order depends on the
+number of rows, and (b) attention whose reduction length depends on padding or
+chunk size. `numerics='invariant'` removes both by construction: every
+row-wise module call (projections, norms, MLP, lm_head) is issued on one row
+at a time, and attention is one SDPA call per query row over exactly its own
+keys. Every operand the kernel sees then depends only on the sequence itself.
+It is slow and is the reference/RL-safe mode. `numerics='fast'` batches rows
+and pads decode attention: faster, equal to 'invariant' to rounding only
+(fp64: ~1e-15 on logits and tokens identical; bf16: expect occasional
+argmax/sample flips on near-ties -- docs/P3_ENGINE.md has the measured rate).
+
+SAMPLING. One `torch.Generator` PER SEQUENCE, seeded from `SamplingParams.seed`
+(or a stable hash of the request id), advanced only by that sequence's own
+draws. A shared engine generator (what `align.inference.Engine` uses) makes a
+sampled token depend on how many other rows drew before it in the step, i.e. on
+batch composition. Temperature, top-k, then top-p (minimal nucleus) are applied
+to the fp32 logits; the RETURNED logprob is `log_softmax(logits/T)[token]`
+UNTRUNCATED, the same basis as `align.rollout.native.LOGPROB_BASIS`, so the
+rho correction contract is unchanged.
+
+PREFIX SHARING (opt-in). Requests whose prompts share leading FULL blocks with
+a running sequence adopt that sequence's blocks by reference instead of
+recomputing them (GRPO: G samples of one prompt prefill it once). A follower
+waits until the leader has computed the shared region, never shares a partial
+block, and always recomputes at least its last prompt token (it needs that
+row's logits). The allocator's refcounts keep shared blocks alive until the
+last owner leaves; `PagedKVCache.slots` refuses writes into a shared block.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import time
+from dataclasses import dataclass, field
+from typing import Callable, Optional, Sequence
+
+import torch
+import torch.nn as nn
+
+from align.inference.engine import _sampling_distribution, find_stop
+from align.model.mistral import apply_rotary, kv_head_index_map
+from align.rollout.paged import (
+    OutOfBlocks,
+    PagedAttention,
+    PagedKVCache,
+    Segment,
+    StepBatch,
+    TorchPagedAttention,
+)
+
+FINISH_REASONS = ("eos", "stop_token", "stop_string", "max_new_tokens", "max_seq_len")
+
+
+# ---------------------------------------------------------------------------
+# Public dataclasses
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class SamplingParams:
+    """Per-sequence sampling. `temperature` 0/None = greedy."""
+
+    max_tokens: int = 64
+    temperature: Optional[float] = 0.0
+    top_p: float = 1.0
+    #: 0 disables. Applied BEFORE top-p (vLLM's order).
+    top_k: int = 0
+    #: None -> stable hash of the request id (NOT admission order, which would
+    #: make "unseeded" depend on batch composition).
+    seed: Optional[int] = None
+    stop_token_ids: tuple[int, ...] = ()
+    #: text stops need a tokenizer; the completion keeps the token that
+    #: completed the stop string (token ids and logprobs stay aligned).
+    stop: tuple[str, ...] = ()
+
+    def validate(self) -> None:
+        if self.max_tokens < 1:
+            raise ValueError(f"max_tokens={self.max_tokens} must be >= 1")
+        t = self.temperature
+        if t is not None and (math.isnan(t) or t < 0):
+            raise ValueError(f"temperature={t} must be >= 0")
+        if math.isnan(self.top_p) or not (0.0 < self.top_p <= 1.0):
+            raise ValueError(f"top_p={self.top_p} must be in (0, 1]")
+        if self.top_k < 0:
+            raise ValueError(f"top_k={self.top_k} must be >= 0")
+
+    @property
+    def greedy(self) -> bool:
+        return self.temperature is None or self.temperature == 0.0
+
+
+@dataclass
+class PagedEngineConfig:
+    #: Tokens per KV block. 16 is vLLM's default; smaller wastes less tail,
+    #: larger gathers fewer indices.
+    block_size: int = 16
+    num_blocks: int = 256
+    #: Max sequences resident (admitted) at once.
+    max_num_seqs: int = 32
+    #: Per-step token budget across all segments (decode tokens + prefill chunks).
+    max_num_batched_tokens: int = 512
+    #: Max prompt tokens one sequence contributes to a step. None -> the budget.
+    chunk_size: Optional[int] = None
+    #: Hard cap on prompt + completion. None -> the model's max_position_embeddings.
+    max_seq_len: Optional[int] = None
+    #: 'invariant' (bit-reproducible across batch compositions; slow) or 'fast'.
+    numerics: str = "invariant"
+    prefix_sharing: bool = False
+    eos_token_id: Optional[int] = None
+    #: KV dtype. None -> the model's parameter dtype.
+    kv_dtype: Optional[torch.dtype] = None
+    #: Replace the attention backend (e.g. a future Triton kernel).
+    attention: Optional[PagedAttention] = None
+
+    def validate(self) -> None:
+        if self.numerics not in ("invariant", "fast"):
+            raise ValueError(f"numerics={self.numerics!r}; expected 'invariant' or 'fast'")
+        if min(self.block_size, self.num_blocks, self.max_num_seqs, self.max_num_batched_tokens) < 1:
+            raise ValueError(f"bad PagedEngineConfig: {self}")
+        if self.chunk_size is not None and self.chunk_size < 1:
+            raise ValueError(f"chunk_size={self.chunk_size} must be >= 1")
+
+
+@dataclass
+class PagedResult:
+    rid: str
+    prompt_ids: list[int]
+    token_ids: list[int]
+    logprobs: list[float]
+    finish_reason: str
+    preemptions: int = 0
+    shared_prefix_tokens: int = 0
+    text: Optional[str] = None
+    step_logits: Optional[list[torch.Tensor]] = None
+
+
+@dataclass
+class _Seq:
+    rid: str
+    prompt: list[int]
+    params: SamplingParams
+    order: int  # monotonically increasing arrival index (age for preemption)
+    generator: torch.Generator
+    tokens: list[int] = field(default_factory=list)
+    num_computed: int = 0
+    logprobs: list[float] = field(default_factory=list)
+    finish: Optional[str] = None
+    preemptions: int = 0
+    shared_tokens: int = 0
+    logits_log: Optional[list[torch.Tensor]] = None
+
+    @property
+    def remaining(self) -> int:
+        return len(self.tokens) - self.num_computed
+
+    @property
+    def n_generated(self) -> int:
+        return len(self.tokens) - len(self.prompt)
+
+
+# ---------------------------------------------------------------------------
+# Sampling
+# ---------------------------------------------------------------------------
+
+
+def top_k_filter(logits: torch.Tensor, k: int) -> torch.Tensor:
+    """Keep the k largest logits of a [V] / [B,V] tensor, others -> -inf.
+
+    Ties at the k-th value are resolved by `torch.topk`'s index order, which is
+    deterministic; the cut is by COUNT, never by value, so exactly k survive."""
+    if k <= 0 or k >= logits.shape[-1]:
+        return logits
+    idx = torch.topk(logits, k, dim=-1).indices
+    out = torch.full_like(logits, float("-inf"))
+    return out.scatter(-1, idx, logits.gather(-1, idx))
+
+
+def row_distribution(logits: torch.Tensor, p: SamplingParams) -> torch.Tensor:
+    """The distribution a non-greedy draw is made from: softmax(logits/T) over the
+    top-k survivors, truncated to the minimal top-p nucleus, renormalised. [V] fp32.
+    Split out so tests can compare it with an independent fp64 reference instead
+    of inferring it from draw frequencies."""
+    masked = top_k_filter(logits.float(), p.top_k)
+    return _sampling_distribution(masked.unsqueeze(0), float(p.temperature), p.top_p)[0]
+
+
+def sample_row(
+    logits: torch.Tensor, p: SamplingParams, gen: torch.Generator
+) -> tuple[int, float]:
+    """One row [V] -> (token, logprob). Greedy = argmax. The logprob is of the
+    UNTRUNCATED tempered distribution (T=1 for greedy): top-k/top-p decide which
+    token is drawn, they are not part of the distribution the training objective
+    models (see align.rollout.native for why folding them in breaks rho)."""
+    x = logits.float()
+    if p.greedy:
+        tok = int(x.argmax())
+        return tok, float(torch.log_softmax(x, dim=-1)[tok])
+    lp = torch.log_softmax(x / float(p.temperature), dim=-1)
+    probs = row_distribution(logits, p)
+    tok = int(torch.multinomial(probs.unsqueeze(0), 1, generator=gen)[0, 0])
+    return tok, float(lp[tok])
+
+
+def _stable_seed(rid: str) -> int:
+    return int.from_bytes(hashlib.sha256(rid.encode()).digest()[:8], "little") & ((1 << 62) - 1)
+
+
+# ---------------------------------------------------------------------------
+# Model adapter
+# ---------------------------------------------------------------------------
+
+
+class _ModelView:
+    """Reads align's causal-LM structure and runs it layer by layer on a flat
+    token batch. Raises at construction on anything it does not recognise, so a
+    new architecture fails at build, not as wrong logits."""
+
+    def __init__(self, model: nn.Module):
+        try:
+            inner = model.model
+            self.embed, self.layers, self.final_norm, self.rope = (
+                inner.embed_tokens, list(inner.layers), inner.norm, inner.rope,
+            )
+            self.lm_head = model.lm_head
+        except AttributeError as e:
+            raise TypeError(
+                f"{type(model).__name__} does not look like an align causal LM "
+                f"(needs .model.{{embed_tokens,layers,norm,rope}} and .lm_head): {e}"
+            ) from None
+        a0 = self.layers[0].attn
+        for attr in ("n_heads", "n_kv_heads", "head_dim", "scale", "q_proj", "k_proj", "v_proj", "o_proj"):
+            if not hasattr(a0, attr):
+                raise TypeError(f"attention module lacks `{attr}`; unsupported architecture")
+        self.n_layers = len(self.layers)
+        self.n_heads, self.n_kv_heads, self.head_dim, self.scale = (
+            a0.n_heads, a0.n_kv_heads, a0.head_dim, a0.scale,
+        )
+        # The paged backends use SDPA's enable_gqa, i.e. the canonical BLOCKED map
+        # q_head // group -> kv_head. Anything else would be silently wrong.
+        group = self.n_heads // self.n_kv_heads
+        idx = kv_head_index_map(self.n_heads, self.n_kv_heads, device=torch.device("cpu"))
+        if not torch.equal(idx, torch.arange(self.n_heads) // group):
+            raise NotImplementedError("non-canonical kv-head map is not supported by the paged engine")
+        self.qk_norm = self._qk_norm_kind(a0)
+        self.clip = getattr(a0, "clip_qkv", None)
+
+    def _qk_norm_kind(self, attn) -> Optional[str]:
+        if not hasattr(attn, "q_norm"):
+            return None
+        w = attn.q_norm.weight.numel()
+        if w == self.head_dim:
+            return "head"  # Qwen3
+        if w == self.n_heads * self.head_dim:
+            return "full"  # OLMoE
+        raise NotImplementedError(f"q_norm width {w} matches neither head_dim nor n_heads*head_dim")
+
+    def qkv(self, attn, h: torch.Tensor):
+        """h [1,t,D] -> q [1,t,H,Dh], k/v [1,t,Hkv,Dh] BEFORE RoPE."""
+        b, t, _ = h.shape
+        q, k, v = attn.q_proj(h), attn.k_proj(h), attn.v_proj(h)
+        if self.qk_norm == "full":
+            q, k = attn.q_norm(q), attn.k_norm(k)
+        if self.clip is not None:
+            c = self.clip
+            q, k, v = q.clamp(-c, c), k.clamp(-c, c), v.clamp(-c, c)
+        q = q.view(b, t, self.n_heads, self.head_dim)
+        k = k.view(b, t, self.n_kv_heads, self.head_dim)
+        v = v.view(b, t, self.n_kv_heads, self.head_dim)
+        if self.qk_norm == "head":
+            q, k = attn.q_norm(q), attn.k_norm(k)
+        return q, k, v
+
+
+def _assert_finite_logits(logits: torch.Tensor, rids: list[str]) -> None:
+    """Raise FloatingPointError if any sampled row holds NaN or +inf.
+
+    Sampled rows would fail loudly inside torch.multinomial, but GREEDY rows do not:
+    argmax of a row containing NaN returns the NaN's index and `log_softmax` returns
+    NaN, so a bf16 overflow silently becomes a plausible-looking token with a NaN
+    logprob that then poisons the rho correction. -inf is allowed (a masked vocab).
+    Called BEFORE `num_computed` advances, so the step is retryable and no sequence has
+    consumed a draw from its generator."""
+    ok = (logits == logits) & (logits != float("inf"))
+    if logits.numel() and not bool(ok.all()):
+        bad = [rids[i] for i in range(len(rids)) if not bool(ok[i].all())]
+        raise FloatingPointError(f"non-finite logits (NaN/+inf) for request(s) {bad}")
+
+
+def _rows(fn: Callable, x: torch.Tensor, rowwise: bool):
+    """Apply `fn` to x [1,T,D]; if `rowwise`, once per row [1,1,D] and re-join.
+
+    This is the whole invariance trick for non-attention ops: a GEMM with M=1
+    runs the same kernel on the same operands whatever the step contains."""
+    if not rowwise or x.shape[1] == 1:
+        return fn(x)
+    outs = [fn(x[:, i : i + 1]) for i in range(x.shape[1])]
+    if isinstance(outs[0], tuple):
+        return tuple(torch.cat([o[j] for o in outs], dim=1) for j in range(len(outs[0])))
+    return torch.cat(outs, dim=1)
+
+
+# ---------------------------------------------------------------------------
+# Engine
+# ---------------------------------------------------------------------------
+
+
+class PagedEngine:
+    """Continuous-batching sampler. Usage:
+
+        eng = PagedEngine(model, PagedEngineConfig(num_blocks=64))
+        eng.add_request("a", prompt_ids, SamplingParams(max_tokens=32, seed=1))
+        while eng.has_unfinished():
+            eng.step()               # may add_request between steps
+        res = eng.results["a"]
+    """
+
+    def __init__(self, model: nn.Module, cfg: Optional[PagedEngineConfig] = None, *,
+                 tokenizer=None, return_logits: bool = False):
+        self.cfg = cfg if cfg is not None else PagedEngineConfig()
+        self.cfg.validate()
+        self.model = model
+        self.view = _ModelView(model)
+        self.tokenizer = tokenizer
+        self.return_logits = return_logits
+        p0 = next(model.parameters())
+        self.device = p0.device
+        kv_dtype = self.cfg.kv_dtype or p0.dtype
+        mcfg = getattr(model, "cfg", None)
+        self.max_seq_len = self.cfg.max_seq_len or getattr(mcfg, "max_position_embeddings", 2048)
+        self.cache = PagedKVCache(
+            self.view.n_layers, self.cfg.num_blocks, self.cfg.block_size,
+            self.view.n_kv_heads, self.view.head_dim, kv_dtype, self.device,
+        )
+        self.rowwise = self.cfg.numerics == "invariant"
+        self.attention: PagedAttention = self.cfg.attention or TorchPagedAttention(
+            "row" if self.rowwise else "segment"
+        )
+        self.waiting: list[_Seq] = []
+        self.running: list[_Seq] = []
+        self.results: dict[str, PagedResult] = {}
+        self._known: set[str] = set()
+        self._arrivals = 0
+        self.stats = dict(steps=0, tokens_forward=0, preemptions=0, prefix_tokens_shared=0,
+                          peak_blocks_used=0, tokens_sampled=0, forward_seconds=0.0)
+
+    # -- requests ---------------------------------------------------------------
+
+    def add_request(self, rid: str, prompt_ids: Sequence[int], params: Optional[SamplingParams] = None) -> None:
+        params = params if params is not None else SamplingParams()
+        params.validate()
+        if rid in self._known:
+            raise ValueError(f"duplicate request id {rid!r}")
+        prompt = [int(t) for t in prompt_ids]
+        if not prompt:
+            raise ValueError(f"{rid!r}: empty prompt")
+        # An out-of-range id fails INSIDE the embedding lookup during `step()`, long after
+        # add_request returned: on CPU it is an IndexError that re-fires on every later step
+        # (the poisoned request stays in `running`, so one bad prompt wedges every other
+        # request), on CUDA a device-side assert that corrupts the whole context. Refuse it here.
+        vocab = self.view.embed.weight.shape[0]
+        bad = next((t for t in prompt if not 0 <= t < vocab), None)
+        if bad is not None:
+            raise ValueError(f"{rid!r}: token id {bad} outside the model vocabulary [0, {vocab})")
+        if len(prompt) >= self.max_seq_len:
+            raise ValueError(f"{rid!r}: prompt of {len(prompt)} tokens reaches max_seq_len={self.max_seq_len}")
+        if params.stop and self.tokenizer is None:
+            raise ValueError("text stop sequences need a tokenizer")
+        # The LAST sampled token is never fed back, so it never gets KV: a request
+        # that ends with `worst` tokens holds `worst - 1` in the cache. Sizing the guard
+        # on `worst` refuses a request that fits exactly (prompt 3 + 6 new, block 4 ->
+        # 8 cached tokens = 2 blocks, yet a 2-block pool was rejected as needing 3).
+        worst = min(len(prompt) + params.max_tokens, self.max_seq_len) - 1
+        if self.cache.blocks_needed(worst) > self.cache.num_blocks:
+            raise ValueError(
+                f"{rid!r} can need {self.cache.blocks_needed(worst)} blocks but the pool has "
+                f"{self.cache.num_blocks}: it could never finish even running alone"
+            )
+        seed = params.seed if params.seed is not None else _stable_seed(rid)
+        gen = torch.Generator(device=self.device)
+        gen.manual_seed(int(seed))
+        seq = _Seq(rid=rid, prompt=prompt, params=params, order=self._arrivals, generator=gen,
+                   tokens=list(prompt), logits_log=[] if self.return_logits else None)
+        self._arrivals += 1
+        self._known.add(rid)
+        self.waiting.append(seq)
+
+    def has_unfinished(self) -> bool:
+        return bool(self.waiting or self.running)
+
+    # -- scheduling -------------------------------------------------------------
+
+    def _lcp(self, a: list[int], b: list[int]) -> int:
+        n = min(len(a), len(b))
+        i = 0
+        while i < n and a[i] == b[i]:
+            i += 1
+        return i
+
+    def _find_share(self, seq: _Seq) -> Optional[tuple[list[int], bool]]:
+        """-> (blocks to adopt, defer?) for the best leader among running seqs."""
+        bs = self.cfg.block_size
+        cap_blocks = (len(seq.tokens) - 1) // bs  # always recompute the last token
+        best: list[int] = []
+        defer = False
+        for src in self.running:
+            lcp_blocks = min(self._lcp(seq.tokens, src.tokens), len(src.prompt)) // bs
+            lcp_blocks = min(lcp_blocks, cap_blocks)
+            if lcp_blocks == 0:
+                continue
+            ready = min(lcp_blocks, src.num_computed // bs)
+            if ready < lcp_blocks and src.num_computed < len(src.prompt):
+                defer = True  # the leader is still computing blocks we could adopt
+            if ready > len(best):
+                best = list(self.cache.tables[src.rid][:ready])
+        return best, defer
+
+    @staticmethod
+    def _pick_victim(victims: list[_Seq]) -> _Seq:
+        """The YOUNGEST (latest arrival): it has the least computed work to throw
+        away, and older sequences keep making progress, so the oldest request is
+        never starved by newer ones."""
+        return max(victims, key=lambda s: s.order)
+
+    def _preempt(self, seq: _Seq) -> None:
+        self.cache.free_seq(seq.rid)
+        self.running.remove(seq)
+        seq.num_computed = 0
+        seq.shared_tokens = 0
+        seq.preemptions += 1
+        self.stats["preemptions"] += 1
+        self.waiting.append(seq)
+        self.waiting.sort(key=lambda s: s.order)
+
+    def _schedule(self) -> list[tuple[_Seq, int]]:
+        cfg = self.cfg
+        chunk = cfg.chunk_size or cfg.max_num_batched_tokens
+        budget = cfg.max_num_batched_tokens
+        sched: list[tuple[_Seq, int]] = []
+        scheduled: set[str] = set()
+
+        # 1) resident sequences: decodes first, then prefill chunks, oldest first
+        for seq in sorted(self.running, key=lambda s: (s.remaining > 1, s.order)):
+            if seq not in self.running or budget <= 0:
+                continue
+            n = min(seq.remaining, chunk, budget)
+            while True:
+                try:
+                    self.cache.ensure_capacity(seq.rid, seq.num_computed + n)
+                    break
+                except OutOfBlocks:
+                    victims = [s for s in self.running if s is not seq and s.rid not in scheduled]
+                    if not victims:
+                        n = 0
+                        break
+                    self._preempt(self._pick_victim(victims))
+            if n == 0:
+                continue
+            sched.append((seq, n))
+            scheduled.add(seq.rid)
+            budget -= n
+
+        # 2) admit waiting sequences (FIFO, deferred followers skipped)
+        for seq in list(self.waiting):
+            if budget <= 0 or len(self.running) >= cfg.max_num_seqs:
+                break
+            shared: list[int] = []
+            if cfg.prefix_sharing:
+                got = self._find_share(seq)
+                shared, defer = got
+                if defer:
+                    continue
+            start = len(shared) * cfg.block_size
+            n = min(len(seq.tokens) - start, chunk, budget)
+            need = self.cache.blocks_needed(len(seq.tokens)) - len(shared)
+            if not self.cache.allocator.can_allocate(need):
+                break  # keep FIFO: do not let a small request leapfrog a starving big one
+            self.waiting.remove(seq)
+            self.cache.add_seq(seq.rid, shared)
+            seq.num_computed = start
+            seq.shared_tokens = start
+            self.stats["prefix_tokens_shared"] += start
+            self.cache.ensure_capacity(seq.rid, start + n)
+            self.running.append(seq)
+            sched.append((seq, n))
+            scheduled.add(seq.rid)
+            budget -= n
+        return sched
+
+    # -- one step ---------------------------------------------------------------
+
+    def step(self) -> list[tuple[str, int, Optional[str]]]:
+        """Run one scheduler step. Returns [(rid, token, finish_reason|None)] for
+        every token sampled this step."""
+        sched = self._schedule()
+        if not sched:
+            if self.running or self.waiting:
+                raise RuntimeError(
+                    "scheduler made no progress: no sequence can get the blocks it needs "
+                    f"(free={self.cache.allocator.num_free}, running={len(self.running)}, "
+                    f"waiting={len(self.waiting)})"
+                )
+            return []
+        batch = StepBatch()
+        toks: list[int] = []
+        pos: list[int] = []
+        off = 0
+        for seq, n in sched:
+            batch.segments.append(Segment(seq.rid, off, n, seq.num_computed))
+            toks.extend(seq.tokens[seq.num_computed : seq.num_computed + n])
+            pos.extend(range(seq.num_computed, seq.num_computed + n))
+            off += n
+        # rows whose logits we need: the last row of every segment that reaches
+        # the end of its sequence's tokens
+        sample_rows = [
+            (seq, seg.offset + seg.n_new - 1)
+            for (seq, n), seg in zip(sched, batch.segments)
+            if seq.num_computed + n == len(seq.tokens)
+        ]
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            logits = self._forward(
+                torch.tensor(toks, dtype=torch.long, device=self.device),
+                torch.tensor(pos, dtype=torch.long, device=self.device),
+                batch,
+                [r for _, r in sample_rows],
+            )
+        # The finiteness check reads a device scalar, i.e. it synchronises: placing it BEFORE
+        # the timer stops makes `forward_seconds` include the GPU work (CUDA launches are
+        # async, so without a sync here the stat measured kernel-launch time only).
+        _assert_finite_logits(logits, [seq.rid for seq, _ in sample_rows])
+        self.stats["forward_seconds"] += time.perf_counter() - t0
+        self.stats["steps"] += 1
+        self.stats["tokens_forward"] += off
+        self.stats["peak_blocks_used"] = max(self.stats["peak_blocks_used"], self.cache.allocator.num_used)
+
+        for seq, n in sched:
+            seq.num_computed += n
+            self.cache.lengths[seq.rid] = seq.num_computed
+
+        events = []
+        for i, (seq, _) in enumerate(sample_rows):
+            row = logits[i]
+            tok, lp = sample_row(row, seq.params, seq.generator)
+            if seq.logits_log is not None:
+                seq.logits_log.append(row.detach().float().cpu())
+            seq.tokens.append(tok)
+            seq.logprobs.append(lp)
+            self.stats["tokens_sampled"] += 1
+            reason = self._finish_reason(seq)
+            events.append((seq.rid, tok, reason))
+            if reason is not None:
+                self._retire(seq, reason)
+        return events
+
+    def _finish_reason(self, seq: _Seq) -> Optional[str]:
+        p, gen = seq.params, seq.tokens[len(seq.prompt):]
+        last = gen[-1]
+        if self.cfg.eos_token_id is not None and last == self.cfg.eos_token_id:
+            return "eos"
+        if last in p.stop_token_ids:
+            return "stop_token"
+        if p.stop:
+            text = self.tokenizer.decode(gen, skip_special_tokens=True)
+            if find_stop(text, p.stop) is not None:
+                return "stop_string"
+        if len(gen) >= p.max_tokens:
+            return "max_new_tokens"
+        if len(seq.tokens) >= self.max_seq_len:
+            return "max_seq_len"
+        return None
+
+    def _retire(self, seq: _Seq, reason: str) -> None:
+        seq.finish = reason
+        self.cache.free_seq(seq.rid)
+        self.running.remove(seq)
+        gen = seq.tokens[len(seq.prompt):]
+        text = self.tokenizer.decode(gen, skip_special_tokens=True) if self.tokenizer is not None else None
+        self.results[seq.rid] = PagedResult(
+            rid=seq.rid, prompt_ids=list(seq.prompt), token_ids=list(gen), logprobs=list(seq.logprobs),
+            finish_reason=reason, preemptions=seq.preemptions, shared_prefix_tokens=seq.shared_tokens,
+            text=text, step_logits=seq.logits_log,
+        )
+
+    # -- the model, layer by layer -------------------------------------------------
+
+    def _forward(self, tokens: torch.Tensor, positions: torch.Tensor, batch: StepBatch,
+                 sample_row_ids: list[int]) -> torch.Tensor:
+        v, rw = self.view, self.rowwise
+        x = v.embed(tokens).unsqueeze(0)  # [1, T, D]
+        cos, sin = v.rope(positions.unsqueeze(0))
+        for li, blk in enumerate(v.layers):
+            attn = blk.attn
+            h = _rows(blk.attn_norm, x, rw)
+            q, k, val = _rows(lambda z: v.qkv(attn, z), h, rw)
+            qh, kh = apply_rotary(q.transpose(1, 2), k.transpose(1, 2), cos, sin)  # [1,H,T,Dh]
+            out = self.attention(
+                li, qh[0].transpose(0, 1), kh[0].transpose(0, 1), val[0],
+                self.cache, batch, v.scale,
+            )  # [T, H, Dh]
+            x = x + _rows(attn.o_proj, out.reshape(1, out.shape[0], -1), rw)
+
+            def mlp(z, blk=blk):
+                o = blk.mlp(blk.mlp_norm(z))
+                return o[0] if isinstance(o, tuple) else o
+
+            x = x + _rows(mlp, x, rw)
+        if not sample_row_ids:
+            return x.new_zeros((0, 1))
+        sel = x[:, torch.tensor(sample_row_ids, device=x.device)]
+        sel = _rows(v.final_norm, sel, rw)
+        return _rows(v.lm_head, sel, rw)[0]
+
+    # -- convenience ----------------------------------------------------------------
+
+    def generate(
+        self,
+        requests: Sequence[tuple[str, Sequence[int], SamplingParams]],
+        arrival_step: Optional[dict[str, int]] = None,
+        max_steps: int = 1_000_000,
+    ) -> dict[str, PagedResult]:
+        """Run `requests` to completion. `arrival_step[rid] = s` admits that
+        request to the queue just before step s (0 = immediately), which is how
+        tests stage 'admitted at different steps'."""
+        arrival_step = arrival_step or {}
+        pending = sorted(requests, key=lambda r: arrival_step.get(r[0], 0))
+        i, s = 0, 0
+        while i < len(pending) or self.has_unfinished():
+            while i < len(pending) and arrival_step.get(pending[i][0], 0) <= s:
+                self.add_request(*pending[i])
+                i += 1
+            if self.has_unfinished():
+                self.step()
+            s += 1
+            if s > max_steps:
+                raise RuntimeError("generate exceeded max_steps")
+        return {r[0]: self.results[r[0]] for r in requests}

code/align/rollout/native.py

@@ -0,0 +1,372 @@
+"""Native rollout: GRPO generation on torch alone, in-process, zero extra deps.
+
+This is the CORE rollout backend. vLLM (extras/vllm_rollout/) is an optional
+throughput upgrade behind a process boundary; this module is what makes the
+GRPO stage runnable with exactly torch + triton + numpy + safetensors. It
+satisfies the same `RolloutBackend` surface (submit/collect) and returns the
+same `RolloutResult` the spool client returns, so `align.rl.batch_source`
+cannot tell which backend served it.
+
+What it does, in one paragraph: tokenize each prompt ONCE (those ids are
+returned in `extras["prompt_token_ids"]` and are what assembly concatenates --
+prompts are never re-tokenized), expand to P x G `GenerationRequest`s, run
+them through `align.inference.engine.Engine` (varlen batched prefill,
+contiguous per-slot KV, completion-based refill) with
+`return_logprobs=True`, and regroup into completions[i][j] / logprobs[i][j] /
+texts[i][j]. The per-token logprobs are captured AT SAMPLING TIME from the
+very logits the token was drawn from; they are never recomputed from text.
+
+THE RHO-CORRECTION CONTRACT (the reason logprobs travel with the rollout).
+The GRPO module's `old_logprobs` is a teacher-forced pass over
+(prompt + completion): `log_softmax(logits / T)` gathered at the label.
+`vllm_logprobs` (pi^infer_old) must be the same quantity as seen by the
+sampler. With the native engine they differ only by floating-point noise
+(KV-cache decode vs full-sequence forward), so rho = exp(old - infer) == 1 up
+to ~1e-6 in fp32 -- which `tests/test_rollout_native.py` pins by recomputing
+the teacher-forced gather on the same model. The basis is therefore
+log_softmax(logits / T), UNTRUNCATED: top-p shapes which token gets drawn but
+is not a part of the distribution the training objective models, and folding
+it into the logprob would make rho != 1 on-policy for every top_p < 1.
+
+TOPOLOGY CAVEAT. The backend drives a model it can call standalone on its own
+device. That is a dedicated engine copy (the v1 topology: one engine GPU
+beside the FSDP train ranks) or a single-process run (CPU tests, 1-GPU
+smoke). It is NOT valid to point it at an FSDP2-sharded module from one rank:
+the decode loop's forwards are collectives and the other ranks would never
+join them. `GrpoBatchSource` only generates on rank 0, which is exactly why the
+sharded case needs a separate engine-side copy plus `load_weights_fn`.
+
+WEIGHT SYNC. `request.weights_dir` is provenance, not an input, unless the
+caller supplies `load_weights_fn`: then a CHANGED weights_dir triggers
+`load_weights_fn(weights_dir)` (and `engine_reloaded=True`), mirroring the
+spool worker's "checkpoint saves ARE the weight sync" contract. With no hook
+the model is assumed live (shared parameters), so there is nothing to sync.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import time
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import torch
+
+from align.inference.engine import Engine, GenerationRequest, InferenceConfig
+from align.rollout.protocol import RolloutRequest, RolloutResult
+
+logger = logging.getLogger(__name__)
+
+#: How the returned logprobs are defined; stamped into `extras` so a result
+#: file is self-describing about the one thing the rho correction depends on.
+LOGPROB_BASIS = "log_softmax(logits/T)[token], fp32, no top-p truncation (T=1 for greedy)"
+
+
+def checkpoint_weight_loader(model, dtype: Optional[torch.dtype] = None) -> Callable[[str], None]:
+    """A `load_weights_fn` that refreshes a STANDALONE engine model from a
+    trainer checkpoint directory ("checkpoint saves ARE the weight sync").
+
+    Reads the HF-layout safetensors the Trainer wrote (single file or sharded
+    snapshot), converts to the align key layout and copies INTO the model's
+    existing parameters (`assign=False`): the Engine and any CUDA graphs hold
+    references to those tensors, so rebinding them would silently leave the
+    sampler on stale weights. `strict=True`: a checkpoint that does not cover
+    every parameter is a wrong sync, not a partial one.
+    """
+    from align.model import checkpoint as ckpt
+
+    def _load(weights_dir: str) -> None:
+        if dtype is None:
+            target = next(model.parameters()).dtype
+        else:
+            target = dtype
+        state = ckpt.hf_to_nmo(ckpt.load_hf_safetensors(weights_dir, dtype=target))
+        with torch.no_grad():
+            model.load_state_dict(state, strict=True)
+
+    return _load
+
+
+@dataclass
+class NativeRolloutConfig:
+    """Knobs of the native backend, one docstring per field."""
+
+    #: Concurrent sequences = KV slots. Each slot owns `max_seq_len` tokens of
+    #: KV for the engine's life (128 KiB/token at 7B bf16), so this multiplies
+    #: the cache budget directly. The engine is built with
+    #: min(n_slots, sequences-in-the-request) so a small request does not pay
+    #: for a big cache.
+    n_slots: int = 32
+    #: HARD cap on per-slot sequence length (prompt + completion). The engine
+    #: is sized to the request's actual need (longest prompt + max_tokens)
+    #: rounded up to `seq_len_granularity` and never above this; a sequence
+    #: that reaches the cap ends with finish_reason 'max_seq_len' (recorded in
+    #: extras), while a PROMPT that alone reaches it is refused loudly -- the
+    #: truncation of a prompt must be visible upstream, not done here.
+    max_seq_len_cap: int = 2048
+    #: Engine rebuilds are expensive (the KV is preallocated), so the length is
+    #: rounded up to a multiple of this to make growth rare.
+    seq_len_granularity: int = 64
+    #: Passed to the tokenizer when encoding prompts. True matches what vLLM's
+    #: `LLM.generate(prompts)` does (prepends BOS for Mistral), which keeps the
+    #: native and spool backends token-identical on the same prompt text.
+    add_special_tokens: bool = True
+    #: Used when a request carries `seed=None`: the n-th unseeded request gets
+    #: default_seed + n, so even "unseeded" rollouts are reproducible per run.
+    default_seed: int = 0
+    #: Ends a sequence when sampled. None -> take it from the tokenizer.
+    eos_token_id: Optional[int] = None
+    #: Filler for right-padded prefill rows; never read back.
+    pad_token_id: int = 0
+    #: 'length' admits shortest-first (measured default, see InferenceConfig).
+    admission_order: str = "length"
+    #: Run the engine's VRAM preflight when the KV cache is allocated on CUDA.
+    preflight: bool = True
+
+
+def _make_encoder(tokenizer, add_special_tokens: bool) -> Callable[[str], list[int]]:
+    """prompt text -> token ids, over the minimal tokenizer contract.
+
+    The contract is `encode(text) -> list[int]` and
+    `decode(ids, skip_special_tokens=True) -> str`. If `encode` takes an
+    `add_special_tokens` keyword it is passed; otherwise the tokenizer is
+    taken to apply its own convention and the flag is not forwarded. Decided
+    once from the signature, not by catching TypeError per call -- a TypeError
+    raised INSIDE a tokenizer must surface, not be mistaken for "no such kwarg".
+    """
+    enc = tokenizer.encode
+    try:
+        params = inspect.signature(enc).parameters
+    except (TypeError, ValueError):  # builtins / odd callables
+        params = {}
+    accepts = "add_special_tokens" in params or any(
+        p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()
+    )
+
+    def encode(text: str) -> list[int]:
+        out = enc(text, add_special_tokens=add_special_tokens) if accepts else enc(text)
+        ids = getattr(out, "ids", out)  # tolerate an Encoding-like return
+        return [int(t) for t in ids]
+
+    return encode
+
+
+class NativeRolloutBackend:
+    """In-process rollout on a torch model via the repo's own KV-cache engine.
+
+    `submit(request)` runs the whole request synchronously and parks the
+    result; `collect(request_id)` hands it over. That keeps the two-call
+    surface of the spool client (so `GrpoBatchSource` is backend-agnostic)
+    while being honest that there is no worker to wait on.
+    """
+
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        cfg: Optional[NativeRolloutConfig] = None,
+        *,
+        load_weights_fn: Optional[Callable[[str], None]] = None,
+        device: Optional[torch.device] = None,
+    ):
+        if tokenizer is None:
+            raise ValueError(
+                "native rollout needs a tokenizer: prompts must be encoded here (and their "
+                "ids returned) and completions decoded to the exact text the reward path scores"
+            )
+        self.model = model
+        self.tokenizer = tokenizer
+        self.cfg = cfg if cfg is not None else NativeRolloutConfig()
+        if self.cfg.n_slots < 1 or self.cfg.max_seq_len_cap < 2 or self.cfg.seq_len_granularity < 1:
+            raise ValueError(f"bad NativeRolloutConfig: {self.cfg}")
+        self.load_weights_fn = load_weights_fn
+        self.device = device
+        eos = self.cfg.eos_token_id
+        self.eos_token_id = eos if eos is not None else getattr(tokenizer, "eos_token_id", None)
+        self._encode = _make_encoder(tokenizer, self.cfg.add_special_tokens)
+
+        self._engine: Optional[Engine] = None
+        self._results: dict[str, RolloutResult] = {}
+        self._loaded_weights_dir: Optional[str] = None
+        self._unseeded = 0
+
+    # -- the RolloutBackend surface --------------------------------------------
+
+    def submit(self, request: RolloutRequest) -> str:
+        if request.request_id in self._results:
+            # same refusal as the spool protocol's write_request
+            raise FileExistsError(f"duplicate request_id {request.request_id!r}")
+        self._results[request.request_id] = self.generate(request)
+        return request.request_id
+
+    def collect(self, request_id: str) -> RolloutResult:
+        try:
+            return self._results.pop(request_id)
+        except KeyError:
+            raise KeyError(
+                f"no pending native result for {request_id!r} (never submitted, or already "
+                "collected: results are single-use, as the spool's consumed files are)"
+            ) from None
+
+    # -- core -------------------------------------------------------------------
+
+    def generate(self, request: RolloutRequest) -> RolloutResult:
+        """One request -> one RolloutResult. Exceptions propagate directly: there
+        is no worker boundary to convert them into an `error` result across."""
+        request.validate(check_weights_dir=False)
+        reloaded = self._maybe_sync_weights(request.weights_dir)
+
+        prompt_ids = [self._encode(p) for p in request.prompts]
+        n, g = len(request.prompts), request.n_samples_per_prompt
+        engine = self._engine_for(
+            need_len=max(len(p) for p in prompt_ids) + request.max_tokens, n_sequences=n * g
+        )
+
+        seed = request.seed
+        if seed is None:
+            seed = self.cfg.default_seed + self._unseeded
+            self._unseeded += 1
+        # Re-seed per request: the sample stream is a pure function of
+        # (weights, prompts, seed), not of how many requests came before.
+        engine.generator.manual_seed(int(seed))
+
+        gen_requests = [
+            GenerationRequest(
+                rid=f"{request.request_id}:{i}:{j}",
+                prompt_ids=prompt_ids[i],
+                max_new_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+            )
+            for i in range(n)
+            for j in range(g)
+        ]
+
+        t0 = time.perf_counter()
+        with _rollout_mode(self.model):
+            results = engine.generate(gen_requests, return_logprobs=True)
+        gen_s = time.perf_counter() - t0
+
+        completions, logprobs, texts, finish = [], [], [], []
+        for i in range(n):
+            row = results[i * g : (i + 1) * g]
+            for r in row:
+                if r.logprobs is None or len(r.logprobs) != len(r.token_ids):
+                    raise RuntimeError(
+                        f"{r.rid}: engine returned {None if r.logprobs is None else len(r.logprobs)}"
+                        f" logprobs for {len(r.token_ids)} tokens; the rho correction cannot run"
+                    )
+                if r.text is None:  # unreachable with a tokenizer; cheap to keep honest
+                    raise RuntimeError(f"{r.rid}: engine returned no decoded text")
+            completions.append([r.token_ids for r in row])
+            logprobs.append([r.logprobs for r in row])
+            texts.append([r.text for r in row])
+            finish.append([r.finish_reason for r in row])
+
+        return RolloutResult(
+            request_id=request.request_id,
+            completions=completions,
+            logprobs=logprobs,
+            texts=texts,
+            weights_dir=request.weights_dir,
+            engine_reloaded=reloaded,
+            generate_seconds=gen_s,
+            extras={
+                # THE tokenization of each prompt that the completions continue
+                # from; assembly forbids re-deriving it (audit R1 class).
+                "prompt_token_ids": [list(p) for p in prompt_ids],
+                "finish_reasons": finish,
+                "backend": "native",
+                "seed": int(seed),
+                "logprob_basis": LOGPROB_BASIS,
+                "engine_max_seq_len": engine.cfg.max_seq_len,
+                "engine_n_slots": engine.cfg.n_slots,
+            },
+        )
+
+    # -- internals ----------------------------------------------------------------
+
+    def _maybe_sync_weights(self, weights_dir: str) -> bool:
+        if self.load_weights_fn is None or not weights_dir:
+            return False
+        if weights_dir == self._loaded_weights_dir:
+            return False
+        t0 = time.perf_counter()
+        self.load_weights_fn(weights_dir)
+        self._loaded_weights_dir = weights_dir
+        logger.info("native rollout: weights synced from %s in %.1fs", weights_dir,
+                    time.perf_counter() - t0)
+        return True
+
+    def _engine_for(self, need_len: int, n_sequences: int) -> Engine:
+        """The cached engine if it is big enough, else a (re)built one.
+
+        Grows only: a shorter later request reuses the bigger cache, so a run
+        converges to one engine after a few steps instead of thrashing on
+        prompt-length variance. The old engine's KV is freed BEFORE the new
+        one is allocated -- both alive at once would double the cache.
+        """
+        gran = self.cfg.seq_len_granularity
+        want_len = min(self.cfg.max_seq_len_cap, -(-need_len // gran) * gran)
+        want_slots = min(self.cfg.n_slots, n_sequences)
+        eng = self._engine
+        if eng is not None and eng.cfg.max_seq_len >= want_len and eng.cfg.n_slots >= want_slots:
+            return eng
+        if eng is not None:
+            want_len = max(want_len, eng.cfg.max_seq_len)
+            want_slots = max(want_slots, eng.cfg.n_slots)
+            self._engine = None
+            del eng
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        icfg = InferenceConfig(
+            n_slots=want_slots,
+            max_seq_len=want_len,
+            max_new_tokens=1,  # always overridden per request
+            temperature=1.0,
+            top_p=1.0,
+            seed=self.cfg.default_seed,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.cfg.pad_token_id,
+            admission_order=self.cfg.admission_order,
+        )
+        self._engine = Engine(
+            self.model,
+            icfg,
+            tokenizer=self.tokenizer,
+            device=self.device,
+            preflight=self.cfg.preflight,
+        )
+        return self._engine
+
+
+class _rollout_mode:
+    """Borrow a TRAINING model for sampling, then give it back untouched.
+
+    The engine flips every attention module's backend per phase (prefill vs
+    decode) and runs the model as-is; on a model that is also being trained
+    that would leak: a decode-phase 'sdpa' left on a `triton_fa2` model would
+    silently change the next TRAIN step's kernel, and a model left in eval()
+    would stay there. Both are restored in `__exit__`, including on error.
+    """
+
+    def __init__(self, model):
+        self.model = model
+
+    def __enter__(self):
+        m = self.model
+        self.was_training = m.training
+        layers = getattr(getattr(m, "model", None), "layers", None)
+        self.backends = (
+            [(blk.attn, blk.attn.attn_implementation) for blk in layers] if layers else []
+        )
+        m.eval()
+        return m
+
+    def __exit__(self, *exc):
+        for attn, impl in self.backends:
+            attn.attn_implementation = impl
+        if self.was_training:
+            self.model.train()
+        return False

code/align/rollout/paged.py

@@ -0,0 +1,494 @@
+"""Paged KV cache for the native continuous-batching sampler (P3).
+
+WHY PAGING (the failure mode this replaces). `align.inference.engine.Engine`
+gives every slot a CONTIGUOUS `max_seq_len` strip of KV for the engine's whole
+life. For a prover-RL rollout the length distribution is brutal: most
+completions stop at a few thousand tokens, a few run to the 16k/40k cap. A
+strip sized for the cap wastes ~all of it on the short ones, so concurrency is
+`budget / cap` regardless of what is actually in flight. On the 23 GB A10G
+that is the difference between "a handful of sequences" and "as many as the
+live tokens allow" (numbers: `plan_capacity`, docs/P3_ENGINE.md). Paging cuts
+the cache into fixed `block_size`-token blocks handed out on demand from a
+free list; a sequence owns a *block table* (list of block ids) and the waste
+is bounded by one partially filled block per sequence.
+
+What lives here:
+
+  * `BlockAllocator`   free-list allocator with refcounts. The refcount is what
+                       makes prefix sharing safe: a block returns to the free
+                       list only when its LAST owner releases it.
+  * `PagedKVCache`     per-layer `[num_blocks, block_size, Hkv, D]` tensors +
+                       block tables. Append is COPY-FREE: the new token's K/V is
+                       written in place at (block, offset); nothing is moved or
+                       re-concatenated when a sequence grows (the dense-cache
+                       `torch.cat` path is O(L) per token per layer).
+  * `StepBatch`        the flat, ragged description of one engine step.
+  * `TorchPagedAttention`  the pure-torch attention backend (gather from the
+                       block table, SDPA on the gathered K/V). Correctness
+                       first. Three granularities, see its docstring.
+  * `TritonPagedAttention` INTERFACE ONLY (raises NotImplementedError): the
+                       signature a fused kernel must satisfy, so the engine
+                       does not change when the kernel lands.
+  * `kv_bytes_per_token`, `plan_capacity`  the memory accounting.
+
+Nothing here imports anything outside torch + stdlib (+ `align.model` for the
+GQA map check), so the dependency firewall is untouched.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Optional, Protocol, Sequence
+
+import torch
+import torch.nn.functional as F
+
+
+class OutOfBlocks(RuntimeError):
+    """The free list cannot satisfy an allocation. The engine catches exactly this
+    to decide on preemption; any other error is a bug and must propagate."""
+
+
+# ---------------------------------------------------------------------------
+# Allocator
+# ---------------------------------------------------------------------------
+
+
+class BlockAllocator:
+    """Free-list block allocator with reference counts.
+
+    LIFO free list: a just-freed block is the next one handed out, which keeps
+    the working set of block ids small and -- more usefully for tests --
+    deterministic, so a failing property run replays exactly.
+
+    Failure modes each check prevents:
+      * double free  -> the block would sit on the free list TWICE and be handed
+        to two sequences at once; their KV would silently overwrite each other.
+        `free` raises on refcount 0.
+      * free of an out-of-range / never-issued id -> same corruption, one hop
+        later.
+      * leaks -> `check_consistent` proves `free + referenced == total` and that
+        the free list has no duplicates; the property test calls it after every
+        event.
+    """
+
+    def __init__(self, num_blocks: int):
+        if num_blocks < 1:
+            raise ValueError(f"num_blocks={num_blocks} must be >= 1")
+        self.num_blocks = int(num_blocks)
+        # reversed so pop() yields block 0 first
+        self._free: list[int] = list(range(self.num_blocks - 1, -1, -1))
+        self._ref: list[int] = [0] * self.num_blocks
+
+    @property
+    def num_free(self) -> int:
+        return len(self._free)
+
+    @property
+    def num_used(self) -> int:
+        return self.num_blocks - len(self._free)
+
+    def refcount(self, block: int) -> int:
+        return self._ref[block]
+
+    def can_allocate(self, n: int) -> bool:
+        return n <= len(self._free)
+
+    def allocate(self) -> int:
+        if not self._free:
+            raise OutOfBlocks("no free KV blocks")
+        b = self._free.pop()
+        self._ref[b] = 1
+        return b
+
+    def incref(self, block: int) -> None:
+        if not (0 <= block < self.num_blocks) or self._ref[block] <= 0:
+            raise ValueError(f"incref of block {block} which is not allocated")
+        self._ref[block] += 1
+
+    def free(self, block: int) -> None:
+        if not (0 <= block < self.num_blocks):
+            raise ValueError(f"free of out-of-range block {block}")
+        if self._ref[block] <= 0:
+            raise ValueError(f"double free of block {block}")
+        self._ref[block] -= 1
+        if self._ref[block] == 0:
+            self._free.append(block)
+
+    def check_consistent(self) -> None:
+        """Raise AssertionError unless the books balance. O(num_blocks)."""
+        free = self._free
+        assert len(set(free)) == len(free), "duplicate block on the free list"
+        for b in free:
+            assert self._ref[b] == 0, f"block {b} is free but has refcount {self._ref[b]}"
+        referenced = sum(1 for r in self._ref if r > 0)
+        assert referenced + len(free) == self.num_blocks, (
+            f"leak: {referenced} referenced + {len(free)} free != {self.num_blocks}"
+        )
+        assert all(r >= 0 for r in self._ref), "negative refcount"
+
+
+# ---------------------------------------------------------------------------
+# KV cache
+# ---------------------------------------------------------------------------
+
+
+class PagedKVCache:
+    """Block-table KV cache. One K and one V tensor per layer, each
+    `[num_blocks, block_size, n_kv_heads, head_dim]`.
+
+    A sequence is `(block_table, length)`. Token `p` of a sequence lives at
+    block `table[p // block_size]`, offset `p % block_size`. Block 0 is NOT
+    reserved: nothing pads a table with a sentinel, because every reader slices
+    by `length` instead.
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        num_blocks: int,
+        block_size: int,
+        n_kv_heads: int,
+        head_dim: int,
+        dtype: torch.dtype = torch.float32,
+        device: Optional[torch.device | str] = None,
+    ):
+        if block_size < 1:
+            raise ValueError(f"block_size={block_size} must be >= 1")
+        self.n_layers, self.block_size = n_layers, int(block_size)
+        self.n_kv_heads, self.head_dim = n_kv_heads, head_dim
+        self.dtype, self.device = dtype, torch.device(device) if device is not None else torch.device("cpu")
+        self.allocator = BlockAllocator(num_blocks)
+        shape = (num_blocks, self.block_size, n_kv_heads, head_dim)
+        self.k = [torch.zeros(shape, dtype=dtype, device=self.device) for _ in range(n_layers)]
+        self.v = [torch.zeros(shape, dtype=dtype, device=self.device) for _ in range(n_layers)]
+        self.tables: dict[str, list[int]] = {}
+        self.lengths: dict[str, int] = {}
+
+    # -- accounting -----------------------------------------------------------
+
+    @property
+    def num_blocks(self) -> int:
+        return self.allocator.num_blocks
+
+    def blocks_needed(self, n_tokens: int) -> int:
+        return -(-n_tokens // self.block_size)
+
+    def bytes_total(self) -> int:
+        itemsize = torch.empty((), dtype=self.dtype).element_size()
+        return self.n_layers * 2 * self.num_blocks * self.block_size * self.n_kv_heads * self.head_dim * itemsize
+
+    # -- sequence lifecycle ----------------------------------------------------
+
+    def add_seq(self, seq_id: str, shared_blocks: Sequence[int] = ()) -> None:
+        """Register a sequence, optionally adopting FULL blocks of another one.
+
+        Only whole blocks can be shared: a partially filled tail block is still
+        being appended to by its owner, and a second writer (or even a reader
+        racing the owner's append) would see/clobber tokens that are not part
+        of the shared prefix. The shared blocks are incref'd, and the length
+        starts at the end of the last shared block.
+        """
+        if seq_id in self.tables:
+            raise ValueError(f"sequence {seq_id!r} already registered")
+        for b in shared_blocks:
+            self.allocator.incref(b)
+        self.tables[seq_id] = list(shared_blocks)
+        self.lengths[seq_id] = len(shared_blocks) * self.block_size
+
+    def free_seq(self, seq_id: str) -> int:
+        """Release every block of a sequence (shared blocks only drop a ref)."""
+        table = self.tables.pop(seq_id)
+        self.lengths.pop(seq_id)
+        for b in table:
+            self.allocator.free(b)
+        return len(table)
+
+    def ensure_capacity(self, seq_id: str, total_tokens: int) -> None:
+        """Grow the table to hold `total_tokens`. ALL-OR-NOTHING: if the free
+        list cannot cover the whole growth nothing is allocated, so a failed
+        reservation never strands blocks on a sequence that is about to be
+        preempted (which would make the preemption free them again -- fine --
+        but a *skipped* sequence would hold them forever)."""
+        table = self.tables[seq_id]
+        need = self.blocks_needed(total_tokens) - len(table)
+        if need <= 0:
+            return
+        if not self.allocator.can_allocate(need):
+            raise OutOfBlocks(f"need {need} blocks for {seq_id!r}, {self.allocator.num_free} free")
+        for _ in range(need):
+            table.append(self.allocator.allocate())
+
+    def slots(self, seq_id: str, start: int, n: int) -> torch.Tensor:
+        """Flat slot ids (block * block_size + offset) for positions [start, start+n).
+
+        Refuses to hand out a slot inside a block somebody else also references:
+        a write there would corrupt the sharer's KV with no error anywhere.
+        """
+        table = self.tables[seq_id]
+        bs = self.block_size
+        if start + n > len(table) * bs:
+            raise ValueError(f"{seq_id!r}: positions up to {start + n} exceed reserved {len(table) * bs}")
+        pos = torch.arange(start, start + n)
+        blocks = torch.tensor(table, dtype=torch.long)[pos // bs]
+        for b in set(blocks.tolist()):
+            if self.allocator.refcount(b) != 1:
+                raise RuntimeError(
+                    f"{seq_id!r}: write into shared block {b} (refcount {self.allocator.refcount(b)})"
+                )
+        return (blocks * bs + pos % bs).to(self.device)
+
+    # -- data movement ----------------------------------------------------------
+
+    def write(self, layer: int, slots: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None:
+        """In-place scatter of `k`,`v` [n, Hkv, D] at flat `slots`. No allocation,
+        no copy of existing KV (the 'copy-free append')."""
+        self.k[layer].view(-1, self.n_kv_heads, self.head_dim)[slots] = k.to(self.dtype)
+        self.v[layer].view(-1, self.n_kv_heads, self.head_dim)[slots] = v.to(self.dtype)
+
+    def gather(self, layer: int, seq_id: str, length: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """K, V of the first `length` tokens as contiguous `[1, Hkv, length, D]`.
+
+        Contiguity is deliberate: SDPA can pick different kernels / reduction
+        orders for strided inputs, and the batch-invariance guarantee needs the
+        SAME bytes in the SAME layout whatever else is in the batch."""
+        table = self.tables[seq_id]
+        idx = torch.tensor(table[: self.blocks_needed(length)], dtype=torch.long, device=self.device)
+        k = self.k[layer][idx].reshape(-1, self.n_kv_heads, self.head_dim)[:length]
+        v = self.v[layer][idx].reshape(-1, self.n_kv_heads, self.head_dim)[:length]
+        return (
+            k.permute(1, 0, 2).unsqueeze(0).contiguous(),
+            v.permute(1, 0, 2).unsqueeze(0).contiguous(),
+        )
+
+    def block_table_tensor(self, seq_ids: Sequence[str], pad: int = 0) -> torch.Tensor:
+        """`[B, max_blocks]` int32 table (padded with `pad`; readers slice by
+        context length). This is the argument a fused kernel consumes."""
+        width = max((len(self.tables[s]) for s in seq_ids), default=0)
+        out = torch.full((len(seq_ids), max(width, 1)), pad, dtype=torch.int32)
+        for i, s in enumerate(seq_ids):
+            t = self.tables[s]
+            if t:
+                out[i, : len(t)] = torch.tensor(t, dtype=torch.int32)
+        return out.to(self.device)
+
+
+# ---------------------------------------------------------------------------
+# One engine step, flat and ragged
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Segment:
+    """`n_new` consecutive tokens of one sequence in the flat batch.
+
+    `ctx` is the number of tokens of this sequence whose KV is ALREADY in the
+    cache before this step; the segment's token j sits at position `ctx + j`
+    and attends to positions `[0, ctx + j]`. A decode token is a segment with
+    n_new == 1; a prefill chunk is n_new > 1; they coexist in one step.
+    """
+
+    seq_id: str
+    offset: int  # row of the segment's first token in the flat batch
+    n_new: int
+    ctx: int
+
+
+@dataclass
+class StepBatch:
+    segments: list[Segment] = field(default_factory=list)
+
+    @property
+    def n_tokens(self) -> int:
+        return sum(s.n_new for s in self.segments)
+
+
+# ---------------------------------------------------------------------------
+# Attention backends
+# ---------------------------------------------------------------------------
+
+
+class PagedAttention(Protocol):
+    """The seam between the engine and an attention implementation.
+
+    `q` is `[T, H, D]` (flat rows, RoPE already applied), `k_new`/`v_new` are
+    `[T, Hkv, D]` for the same rows. The backend must (1) write the new K/V into
+    `cache` for `layer` at the segments' positions, then (2) return the
+    attention output `[T, H, D]` where row r attends to the cache contents
+    `[0, position(r)]` of its own sequence. Splitting write from read inside the
+    backend lets a fused kernel do both in one launch.
+    """
+
+    def __call__(
+        self, layer: int, q: torch.Tensor, k_new: torch.Tensor, v_new: torch.Tensor,
+        cache: PagedKVCache, batch: StepBatch, scale: float,
+    ) -> torch.Tensor: ...
+
+
+class TorchPagedAttention:
+    """Pure-torch paged attention: gather K/V from the block table, run SDPA.
+
+    `mode`:
+      'row'      one SDPA call per QUERY ROW, over exactly that row's keys
+                 `[0, pos]`, no mask. The operands of every call depend only on
+                 the sequence's own tokens, never on what else is in the step or
+                 how the prompt was chunked, so the output bits are identical
+                 across batch compositions and chunk schedules. Slow; it is the
+                 correctness/invariance reference and the mode RL rollouts that
+                 need bit-reproducible logprobs should use.
+      'segment'  one SDPA per segment (a prefill chunk is one causal call with
+                 an explicit bool mask; a decode token is mask-free) -- but all
+                 decode segments of a step are *padded and batched* into one
+                 call. Faster; the softmax/matmul reduction length now depends on
+                 the longest sequence in the batch (padding) or the chunk size,
+                 so results match 'row' to rounding only (see docs/P3_ENGINE.md
+                 for the measured bf16 expectation).
+    """
+
+    def __init__(self, mode: str = "row"):
+        if mode not in ("row", "segment"):
+            raise ValueError(f"mode={mode!r}; expected 'row' or 'segment'")
+        self.mode = mode
+
+    def __call__(self, layer, q, k_new, v_new, cache, batch, scale):
+        group = q.shape[1] // k_new.shape[1]
+        gqa = group > 1
+        out = torch.empty_like(q)
+        # 1) copy-free append of every segment's new K/V
+        for s in batch.segments:
+            sl = slice(s.offset, s.offset + s.n_new)
+            cache.write(layer, cache.slots(s.seq_id, s.ctx, s.n_new), k_new[sl], v_new[sl])
+        # 2) read
+        decode = [s for s in batch.segments if s.n_new == 1] if self.mode == "segment" else []
+        decode_ids = {id(s) for s in decode}
+        for s in batch.segments:
+            if id(s) in decode_ids:
+                continue
+            total = s.ctx + s.n_new
+            kk, vv = cache.gather(layer, s.seq_id, total)
+            if kk.dtype != q.dtype:  # a cache stored narrower than the model (kv_dtype=bf16/fp8): SDPA
+                kk, vv = kk.to(q.dtype), vv.to(q.dtype)  # insists on one dtype, so dequantise on read
+            if self.mode == "row":
+                for j in range(s.n_new):
+                    ln = s.ctx + j + 1
+                    qj = q[s.offset + j].view(1, q.shape[1], 1, q.shape[2])
+                    o = F.scaled_dot_product_attention(
+                        qj, kk[:, :, :ln].contiguous(), vv[:, :, :ln].contiguous(),
+                        enable_gqa=gqa, scale=scale,
+                    )
+                    out[s.offset + j] = o[0, :, 0]
+            else:
+                qs = q[s.offset : s.offset + s.n_new].permute(1, 0, 2).unsqueeze(0)
+                qpos = torch.arange(s.ctx, total, device=q.device)[:, None]
+                keep = qpos >= torch.arange(total, device=q.device)[None, :]
+                o = F.scaled_dot_product_attention(qs, kk, vv, attn_mask=keep, enable_gqa=gqa, scale=scale)
+                out[s.offset : s.offset + s.n_new] = o[0].permute(1, 0, 2)
+        if decode:
+            lens = [s.ctx + 1 for s in decode]
+            lmax = max(lens)
+            kp = q.new_zeros(len(decode), k_new.shape[1], lmax, q.shape[2])
+            vp = torch.zeros_like(kp)
+            keep = torch.zeros(len(decode), 1, 1, lmax, dtype=torch.bool, device=q.device)
+            for i, s in enumerate(decode):
+                kk, vv = cache.gather(layer, s.seq_id, lens[i])
+                kp[i, :, : lens[i]] = kk[0]
+                vp[i, :, : lens[i]] = vv[0]
+                keep[i, ..., : lens[i]] = True
+            rows = torch.tensor([s.offset for s in decode], device=q.device)
+            qd = q[rows].unsqueeze(2)  # [B, H, 1, D]
+            o = F.scaled_dot_product_attention(qd, kp, vp, attn_mask=keep, enable_gqa=gqa, scale=scale)
+            out[rows] = o[:, :, 0]
+        return out
+
+
+class TritonPagedAttention:
+    """INTERFACE ONLY. A later step replaces `TorchPagedAttention` with a fused
+    kernel; this class pins what that kernel receives so the engine is untouched.
+
+    Kernel contract (vLLM-style "paged attention v1/v2" over a ragged batch):
+
+        q            [T, H, D]            flat query rows (RoPE applied)
+        k_cache      [NB, BS, Hkv, D]     this layer's K blocks
+        v_cache      [NB, BS, Hkv, D]
+        block_tables [B, max_blocks] i32  from `PagedKVCache.block_table_tensor`
+        cu_seqlens_q [B+1] i32            row offsets of each segment in q
+        ctx_lens     [B] i32              tokens already cached BEFORE this step
+        scale        float
+        out          [T, H, D]
+
+    Row r of segment b (position ctx_lens[b] + j) attends to cached positions
+    [0, ctx_lens[b] + j]; the new K/V for the step must be written to the cache
+    BEFORE the launch (or fused in the same kernel via `slot_mapping`
+    [T] i64 = block * BS + offset). Bit-invariance requirement for RL use: the
+    reduction over keys for a given row must not depend on B, T or on other
+    segments (fixed KV-split count per context length, no atomics).
+    """
+
+    def __call__(self, layer, q, k_new, v_new, cache, batch, scale):  # pragma: no cover
+        raise NotImplementedError(
+            "TritonPagedAttention is an interface placeholder; use TorchPagedAttention. "
+            "See the class docstring for the kernel contract."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Memory accounting
+# ---------------------------------------------------------------------------
+
+
+def kv_bytes_per_token(n_layers: int, n_kv_heads: int, head_dim: int, dtype: torch.dtype = torch.bfloat16,
+                       tp_degree: int = 1) -> int:
+    """Bytes of KV one token costs on ONE device: layers x 2 (K,V) x kv-heads/tp x head_dim x itemsize.
+
+    `tp_degree` shards the KV heads (vLLM's tensor parallelism does exactly this),
+    which is why a TP=4 probe's "tokens in cache" is not comparable to a
+    single-device engine's without it. Qwen3-8B (36, 8, 128, bf16) -> 147456 B
+    = 144 KiB = the '147 KB/token' figure.
+    """
+    if n_kv_heads % tp_degree:
+        raise ValueError(f"n_kv_heads={n_kv_heads} not divisible by tp_degree={tp_degree}")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    return n_layers * 2 * (n_kv_heads // tp_degree) * head_dim * itemsize
+
+
+@dataclass(frozen=True)
+class CapacityPlan:
+    kv_bytes_per_token: int
+    kv_budget_bytes: int
+    num_blocks: int
+    block_size: int
+    kv_tokens: int
+    max_concurrent_at_cap: int  # sequences that can each hold `cap` tokens simultaneously
+    cap: int
+
+
+def plan_capacity(
+    device_total_bytes: int,
+    weights_bytes: int,
+    kv_bytes_per_tok: int,
+    *,
+    utilization: float = 0.9,
+    activation_bytes: int = 0,
+    other_bytes: int = 0,
+    block_size: int = 16,
+    cap: int = 16384,
+) -> CapacityPlan:
+    """How much KV fits, and how many `cap`-token sequences that sustains.
+
+    budget = utilization * device_total - weights - activation peak - other.
+    Paging rounds each sequence up to whole blocks, so `max_concurrent_at_cap`
+    = num_blocks // ceil(cap / block_size) (a dense per-slot strip would give
+    the same number for sequences that all reach the cap; paging's advantage is
+    for the ones that do not, which is why the engine also reports live usage).
+    """
+    budget = int(utilization * device_total_bytes) - weights_bytes - activation_bytes - other_bytes
+    budget = max(budget, 0)
+    per_block = kv_bytes_per_tok * block_size
+    nb = budget // per_block
+    per_seq_blocks = math.ceil(cap / block_size)
+    return CapacityPlan(
+        kv_bytes_per_token=kv_bytes_per_tok, kv_budget_bytes=budget, num_blocks=int(nb),
+        block_size=block_size, kv_tokens=int(nb * block_size),
+        max_concurrent_at_cap=int(nb // per_seq_blocks), cap=cap,
+    )

code/align/rollout/protocol.py

@@ -0,0 +1,227 @@
+"""Request/result shapes for rollout, plus the file protocol the spool worker speaks.
+
+DEPENDENCY-FREE ON PURPOSE: this module is imported by the training venv
+(torch 2.11/cu128), by the native in-process backend (`native.py`), AND by the
+optional vLLM worker in extras/vllm_rollout/ (its own venv, vllm/cu130); it
+may import nothing beyond the standard library, or the environments' pin sets
+start constraining each other -- the exact coupling the isolated-venv
+design exists to prevent.
+
+`RolloutRequest` / `RolloutResult` are the backend-neutral contract. The
+default backend (`native.NativeRolloutBackend`) fills them in-process with no
+file I/O; the 'spool' backend moves the SAME dataclasses through the files
+described below. The trainer cannot tell which one served it, by design.
+
+Mechanics: one JSON file per message, written atomically (.partial ->
+os.replace) into a spool directory. The worker polls `requests/`, writes
+into `results/`, and touches `heartbeat` every poll so the trainer can
+distinguish "generating" from "dead" (mtime-based liveness, per the house
+monitor-mtime-not-PIDs rule). Completions travel as token id lists plus
+the engine's per-token logprobs -- the trainer needs vllm_logprobs for the
+rho correction and must NOT re-derive them from text (re-tokenization is
+not identity).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from dataclasses import asdict, dataclass, field
+from typing import Optional, Protocol, runtime_checkable
+
+PROTOCOL_VERSION = 1
+
+
+@dataclass
+class RolloutRequest:
+    request_id: str
+    weights_dir: str  # HF-loadable checkpoint dir; worker reloads on change
+    prompts: list[str]
+    n_samples_per_prompt: int
+    max_tokens: int
+    temperature: float = 1.0
+    top_p: float = 1.0
+    seed: Optional[int] = None
+    protocol_version: int = PROTOCOL_VERSION
+
+    def validate(self, *, check_weights_dir: bool = True) -> None:
+        """Shape checks always; the weights_dir existence check only for
+        backends that LOAD from it. The native backend samples from a live
+        in-process model, where weights_dir is provenance, not an input -- a
+        step-0 request legitimately names a checkpoint that has not been
+        written yet."""
+        if not self.prompts:
+            raise ValueError("empty prompt list")
+        if self.n_samples_per_prompt < 1 or self.max_tokens < 1:
+            raise ValueError(
+                f"bad shape: n={self.n_samples_per_prompt}, max_tokens={self.max_tokens}"
+            )
+        if check_weights_dir and not os.path.isdir(self.weights_dir):
+            raise FileNotFoundError(f"weights_dir does not exist: {self.weights_dir}")
+
+
+@dataclass
+class RolloutResult:
+    """One request's samples, in the form `align.rl.assembly` consumes.
+
+    Whatever the backend, the contract the trainer relies on is:
+    logprobs[i][j][t] is the logprob of completions[i][j][t] under the
+    SAMPLING-TEMPERATURE-scaled, untruncated policy distribution, so
+    rho = exp(old - infer) is ~1 when the weights agree. And
+    `extras["prompt_token_ids"]` (one id list per prompt, request order) is
+    REQUIRED: completions continue from exactly those ids, and assembly
+    refuses to re-tokenize prompt text.
+    """
+
+    request_id: str
+    #: completions[i][j] = j-th sample for prompt i: token ids WITHOUT the prompt.
+    completions: list[list[list[int]]]
+    #: logprobs[i][j][t] = engine logprob of completions[i][j][t] (the rho
+    #: correction's pi^infer_old side; label-aligned with the completion ids).
+    logprobs: list[list[list[float]]]
+    #: decoded completion texts, same nesting -- the reward path consumes text
+    #: and must see EXACTLY what the engine's tokenizer decodes, not a re-decode.
+    texts: list[list[str]]
+    weights_dir: str
+    engine_reloaded: bool
+    generate_seconds: float
+    error: Optional[str] = None
+    protocol_version: int = PROTOCOL_VERSION
+    extras: dict = field(default_factory=dict)
+
+
+@runtime_checkable
+class RolloutBackend(Protocol):
+    """What `GrpoBatchSource` consumes: submit a request, collect its result.
+
+    Implemented by `native.NativeRolloutBackend` (in-process, synchronous: the
+    work happens in `submit`) and `client.RolloutClient` (file spool to an
+    external worker: the work happens between `submit` and `collect`).
+    """
+
+    def submit(self, request: RolloutRequest) -> str: ...
+
+    def collect(self, request_id: str) -> RolloutResult: ...
+
+
+def _atomic_write_json(payload: dict, path: str) -> None:
+    tmp = path + ".partial"
+    with open(tmp, "w") as f:
+        json.dump(payload, f)
+    os.replace(tmp, path)
+
+
+def spool_dirs(spool: str) -> tuple[str, str]:
+    req = os.path.join(spool, "requests")
+    res = os.path.join(spool, "results")
+    os.makedirs(req, exist_ok=True)
+    os.makedirs(res, exist_ok=True)
+    return req, res
+
+
+def write_request(spool: str, request: RolloutRequest) -> str:
+    request.validate()
+    req_dir, _ = spool_dirs(spool)
+    path = os.path.join(req_dir, f"{request.request_id}.json")
+    if os.path.exists(path):
+        raise FileExistsError(f"duplicate request_id {request.request_id!r}")
+    _atomic_write_json(asdict(request), path)
+    return path
+
+
+def read_request(path: str) -> RolloutRequest:
+    with open(path) as f:
+        payload = json.load(f)
+    if payload.get("protocol_version") != PROTOCOL_VERSION:
+        raise RuntimeError(
+            f"protocol version mismatch: request carries "
+            f"{payload.get('protocol_version')!r}, this side speaks {PROTOCOL_VERSION}"
+        )
+    return RolloutRequest(**payload)
+
+
+def write_result(spool: str, result: RolloutResult) -> str:
+    _, res_dir = spool_dirs(spool)
+    path = os.path.join(res_dir, f"{result.request_id}.json")
+    _atomic_write_json(asdict(result), path)
+    return path
+
+
+def read_result(spool: str, request_id: str) -> Optional[RolloutResult]:
+    """Non-blocking: None until the result file exists (atomically complete)."""
+    _, res_dir = spool_dirs(spool)
+    path = os.path.join(res_dir, f"{request_id}.json")
+    if not os.path.exists(path):
+        return None
+    with open(path) as f:
+        payload = json.load(f)
+    if payload.get("protocol_version") != PROTOCOL_VERSION:
+        raise RuntimeError(
+            f"protocol version mismatch on result {request_id!r}: "
+            f"{payload.get('protocol_version')!r} vs {PROTOCOL_VERSION}"
+        )
+    return RolloutResult(**payload)
+
+
+def touch_heartbeat(spool: str) -> None:
+    path = os.path.join(spool, "heartbeat")
+    with open(path, "a"):
+        os.utime(path, None)
+
+
+def heartbeat_age_seconds(spool: str) -> float:
+    """Inf when no heartbeat exists yet (worker never started != worker slow)."""
+    path = os.path.join(spool, "heartbeat")
+    if not os.path.exists(path):
+        return float("inf")
+    return time.time() - os.path.getmtime(path)
+
+
+def wait_for_result(
+    spool: str,
+    request_id: str,
+    timeout_s: float,
+    stale_heartbeat_s: float = 600.0,
+    poll_s: float = 2.0,
+    startup_grace_s: float = 300.0,
+) -> RolloutResult:
+    """Block until the result lands; raise on timeout, dead worker, or error.
+
+    THREE distinct failure signatures, kept distinct on purpose: a TIMEOUT
+    with a live heartbeat means the request is too big for the budget; a
+    STALE heartbeat means the worker died mid-service; and NO heartbeat at
+    all means the worker has not finished booting -- which is fatal only
+    after `startup_grace_s`, measured from this wait's start. The first
+    smoke run (2026-08-24) found the missing third case the hard way: a
+    step-0 canary raced the worker's ~30s vLLM import, age=inf compared
+    `> stale_heartbeat_s` as instantly dead, and the teardown killed a
+    perfectly healthy worker before its first log line flushed.
+    """
+    start = time.time()
+    while True:
+        result = read_result(spool, request_id)
+        if result is not None:
+            if result.error:
+                raise RuntimeError(f"rollout worker reported error: {result.error}")
+            return result
+        age = heartbeat_age_seconds(spool)
+        if age == float("inf"):
+            if time.time() - start > startup_grace_s:
+                raise RuntimeError(
+                    f"rollout worker produced NO heartbeat within {startup_grace_s:.0f}s "
+                    f"of waiting on request {request_id!r}: it never started. Check its "
+                    "log (spawn env, venv, import errors)."
+                )
+        elif age > stale_heartbeat_s:
+            raise RuntimeError(
+                f"rollout worker heartbeat is {age:.0f}s stale "
+                f"(> {stale_heartbeat_s:.0f}s) while request {request_id!r} is pending: "
+                "worker presumed dead. Check its log; do not resubmit blindly."
+            )
+        if time.time() - start > timeout_s:
+            raise TimeoutError(
+                f"rollout request {request_id!r} exceeded {timeout_s:.0f}s with a live "
+                "worker: the request shape is too big for the time budget."
+            )
+        time.sleep(poll_s)

code/align/rollout/tp_engine.py

@@ -0,0 +1,169 @@
+"""Tensor-parallel driver for `PagedEngine` (P3, opt-in): same scheduler, same sampler, a TP-sharded model.
+
+WHAT CHANGES vs the single-device engine, and what deliberately does not.
+
+  * The model is `align.model.tp`'s shard. `PagedEngine` drives it UNCHANGED through its own submodules, so the
+    paged KV cache is sharded by kv head for free: `_ModelView` reads the LOCAL `n_kv_heads` (Qwen3-8B at tp=4:
+    8 -> 2) and sizes `PagedKVCache` with it. Per-rank KV bytes are exactly 1/tp (tested), the block tables,
+    refcounts, prefix sharing and preemption logic are the original code, and attention is purely local (a rank's
+    q heads read only kv heads it owns, see `align.model.tp`).
+
+  * REPLICATED SCHEDULER, NOT A BROADCAST ONE. The task offered "scheduler on rank 0 broadcasting decisions".
+    Chosen instead: every rank runs the identical deterministic scheduler on identical inputs, and rank 0 is
+    authoritative for the INPUTS (the request stream, broadcast once by `generate`) and the VERIFIER of the
+    replication. Why:
+      - Broadcasting decisions means shipping a ragged object (segments, block tables, the sampled-row list) every
+        step: a pickle + host sync on the critical path, per step, for information every rank can recompute for
+        free from state it already holds.
+      - The scheduler is a pure function of (request stream, sampled tokens). Sampled tokens are identical on
+        every rank because the logits are: the lm_head ALL-GATHER hands every rank the same bytes, and sampling
+        (per-sequence seeded generators, `_stable_seed` = sha256 not `hash()`) is a deterministic function of the
+        logits row. So replicated sampling needs no extra collective; rank-0-only sampling would need a token
+        broadcast between every forward and the next (a serial hop) PLUS the full-logit row on rank 0 anyway.
+      - The one real risk of replication is SILENT DIVERGENCE (a single differing token or schedule makes ranks
+        issue different collectives: an NCCL hang for the timeout, or all-reduces of mismatched shapes). That is
+        closed by two tiny consensus all-gathers per step (8 bytes/rank each, vs 2*n_layers+1 all-reduces of
+        [rows, hidden]): `_schedule` checks a fingerprint of the step's schedule BEFORE any model collective, and
+        `has_unfinished` checks the previous step's events (tokens, finish reasons) and the loop
+        predicate. The check is SYMMETRIC (all-gather, every rank sees every fingerprint), so a divergence raises
+        `TPDesyncError` on ALL ranks instead of leaving the agreeing ones blocked in the next collective.
+
+  * Residual rounding. Under bf16 the row-parallel all-reduce sums per-rank partials each rounded to bf16, so TP
+    logits differ from the unsharded model's by bf16 rounding (measured in docs/P3_ENGINE.md); in fp64 they agree
+    to ~1e-15. `numerics='invariant'` keeps its meaning under TP: row-wise calls keep every all-reduce the same
+    shape whatever the batch holds, so a sequence's bits do not depend on its companions (tested bitwise).
+
+Every rank must call the same methods with the same arguments in the same order (`add_request`, `step`, `generate`);
+`generate` takes its requests from rank 0 and broadcasts them, so only rank 0 needs the real list.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import struct
+from typing import Optional, Sequence
+
+import torch
+
+from align.model.tp import TPGroup
+from align.rollout.engine import PagedEngine, PagedEngineConfig, PagedResult, SamplingParams
+
+
+class TPDesyncError(RuntimeError):
+    """Ranks disagree about the schedule, the sampled tokens or whether work remains. Raised on EVERY rank."""
+
+
+class _FullVocabEmbed:
+    """What `PagedEngine.add_request` sees as the embedding: a `.weight` whose shape[0] is the FULL vocabulary
+    (the real weight is this rank's shard, whose row count would reject every valid id >= V/tp), and the sharded
+    lookup as the call. The meta tensor allocates nothing."""
+
+    def __init__(self, emb):
+        self._emb = emb
+        self.weight = torch.empty((emb.full_vocab, emb.weight.shape[1]), device="meta")
+
+    def __call__(self, ids):
+        return self._emb(ids)
+
+
+def _digest(parts: Sequence[bytes]) -> int:
+    h = hashlib.sha256()
+    for p in parts:
+        h.update(struct.pack("<Q", len(p)))
+        h.update(p)
+    return int.from_bytes(h.digest()[:8], "little") & ((1 << 62) - 1)
+
+
+def schedule_fingerprint(sched) -> int:
+    """Hash of everything that decides the step's collectives and shapes: for each scheduled sequence its id, how
+    many tokens it already has computed, how many it adds, and the token ids it feeds. Module-level so a test can
+    substitute a constant (divergence then goes undetected until a collective hangs)."""
+    parts = []
+    for seq, n in sched:
+        parts.append(f"{seq.rid}|{seq.num_computed}|{n}".encode())
+        parts.append(struct.pack(f"<{n}q", *seq.tokens[seq.num_computed : seq.num_computed + n]))
+    return _digest(parts)
+
+
+def events_fingerprint(events) -> int:
+    """Hash of a step's sampled events: request id, token and finish reason. (Logprobs need no hash of their own: they are
+    a deterministic function of the all-gathered logits row and the token, both already agreed.)"""
+    return _digest([f"{rid}|{tok}|{reason}".encode() for rid, tok, reason in events])
+
+
+class TPPagedEngine(PagedEngine):
+    """`PagedEngine` over a `align.model.tp` shard. `verify=False` drops the two consensus all-gathers (a
+    measured-cost opt-out for a job that has proven itself; the default keeps the hang-proofing)."""
+
+    def __init__(self, model: torch.nn.Module, tp: Optional[TPGroup] = None, cfg: Optional[PagedEngineConfig] = None, *,
+                 tokenizer=None, return_logits: bool = False, verify: bool = True):
+        tp = tp if tp is not None else getattr(model, "tp", None)
+        if tp is None:
+            raise ValueError("TPPagedEngine needs a TPGroup (or a model built by align.model.tp, which carries one)")
+        super().__init__(model, cfg, tokenizer=tokenizer, return_logits=return_logits)
+        self.tp, self.verify = tp, verify and tp.world_size > 1
+        emb = self.view.embed
+        if not hasattr(emb, "full_vocab"):
+            raise TypeError("model is not a tensor-parallel shard (embed_tokens has no `full_vocab`)")
+        emb.check_range = False  # add_request already validates ids; skip the per-step device sync
+        self.view.embed = _FullVocabEmbed(emb)
+        self._last_events = 0
+        self.stats.update(consensus_checks=0)
+
+    # -- replication guards ----------------------------------------------------------
+
+    def _agree(self, label: str, values: list[int]) -> None:
+        got = self.tp.all_gather_i64(values, self.device)
+        self.stats["consensus_checks"] += 1
+        if any(g != got[0] for g in got):
+            raise TPDesyncError(
+                f"tensor-parallel ranks diverged at {label}: per-rank fingerprints {got} (this is rank {self.tp.rank}). "
+                "Every rank must see the same requests in the same order with the same config."
+            )
+
+    def has_unfinished(self) -> bool:
+        local = super().has_unfinished()
+        if self.verify:
+            self._agree("loop predicate / previous step's events", [int(local), self._last_events])
+        return local
+
+    def _schedule(self):
+        sched = super()._schedule()
+        if self.verify:
+            self._agree("the step schedule", [schedule_fingerprint(sched)])
+        return sched
+
+    def step(self):
+        events = super().step()
+        self._last_events = events_fingerprint(events)
+        return events
+
+    # -- driver ---------------------------------------------------------------------------
+
+    def generate(self, requests=None, arrival_step: Optional[dict[str, int]] = None, max_steps: int = 1_000_000
+                 ) -> dict[str, PagedResult]:
+        """As `PagedEngine.generate`, but rank 0's `requests`/`arrival_step` are broadcast to every rank (other
+        ranks may pass None): the request stream has exactly one author."""
+        payload = None
+        if self.tp.rank == 0:
+            # Built BEFORE the broadcast, so a malformed stream would raise on rank 0 alone and leave every other
+            # rank blocked in the broadcast until the process-group timeout (minutes to hours on NCCL). The error
+            # is therefore shipped as data and re-raised everywhere.
+            try:
+                payload = ("ok", ([(rid, list(ids), p) for rid, ids, p in requests], dict(arrival_step or {})))
+            except Exception as e:  # noqa: BLE001
+                payload = ("err", f"{type(e).__name__}: {e}")
+        tag, body = self.tp.broadcast_object(payload)
+        if tag == "err":
+            raise ValueError(f"rank 0 could not build the request stream: {body}")
+        reqs, arrivals = body
+        reqs = [(rid, ids, p if isinstance(p, SamplingParams) else SamplingParams(**p)) for rid, ids, p in reqs]
+        return super().generate(reqs, arrivals, max_steps)
+
+    # -- accounting --------------------------------------------------------------------------
+
+    def kv_bytes_per_rank(self) -> int:
+        return self.cache.bytes_total()
+
+
+__all__ = ["TPPagedEngine", "TPDesyncError", "schedule_fingerprint", "events_fingerprint"]

code/align/tokenizer/__init__.py

@@ -0,0 +1,85 @@
+"""align.tokenizer -- pure-Python (stdlib-only) Mistral-7B-v0.3 BPE tokenizer.
+
+    from align.tokenizer import Tokenizer, load_tokenizer
+    tok = load_tokenizer("/path/to/dir/with/tokenizer.json")
+    tok("a\\nb")["input_ids"]          # [1, 1032, <0x0A id>, 29494]   (<s> ▁a <0x0A> b)
+    tok.decode(ids, skip_special_tokens=True)
+
+It replaces `transformers.AutoTokenizer` for this codebase (dependency policy:
+torch/triton/numpy/safetensors only). Parity is certified against the legacy
+transformers tokenizer by the `oracle`-marked test in tests/test_tokenizer.py.
+
+Two places pure reimplementations usually get wrong -- both are load-bearing:
+
+1. Metaspace prepend_scheme="first".
+   The pre-tokenizer replaces ' ' with U+2581 and prepends U+2581 to split #0 of
+   the whole input -- and only if that split is NOT an added token. Added tokens
+   (<s>, [INST], [control_N], ...) are carved out of the raw text before the
+   pre-tokenizer runs; they occupy a split index but skip the pre-tokenizer, so
+   "first" is neither "first character of the string" nor "first segment after
+   each special token" nor "first non-added segment":
+       "hello<s>world"  -> ▁hello <s> world  (prefix on split 0 only)
+       "<s>hello"       -> <s> <s> hello     (bare 'hello': the leading <s> used
+                                              up split 0 -- measured, not guessed)
+       " hello"         -> ▁hello            (already starts with U+2581: no
+                                              second prefix)
+   With split=false the whole non-added segment is a single BPE "word", so
+   merges can run across spaces and newlines ('\\n' is not touched by Metaspace;
+   it has no vocab piece and falls back to the byte token <0x0A>, the id the vocab assigns it).
+
+2. Strip(" ", start=1) at the END of the decoder chain.
+   decode = Replace(U+2581 -> ' ') / ByteFallback / Fuse / Strip. Strip removes
+   at most ONE leading space from the FUSED string. Consequences: decode of the
+   ids for "hello" is "hello" (the prepended ▁ is stripped, round-trip exact),
+   but decoding a mid-sequence chunk [▁world] gives "world", not " world" -- so
+   decoding token deltas and concatenating drops spaces. Decode the whole
+   sequence and diff the TEXT (the inference engine does exactly that). Also:
+   ByteFallback decodes each maximal run of <0xNN> pieces as one UTF-8 unit, and
+   an invalid run yields one U+FFFD per byte (not Python's errors="replace").
+
+The loader refuses (NotImplementedError) any tokenizer.json whose pipeline
+differs from the one above, rather than silently producing wrong ids.
+
+Byte-level BPE (Qwen2/Qwen3, e.g. Goedel-Prover-V2-8B) lives in `.bytelevel`
+(`ByteLevelTokenizer`); see that module's docstring. `load_tokenizer` dispatches on
+tokenizer.json (model.type + pre_tokenizer): Metaspace -> the Mistral `Tokenizer`
+above, completely unchanged; Split+ByteLevel -> `ByteLevelTokenizer`; anything else
+-> NotImplementedError (never a guess).
+"""
+from __future__ import annotations
+
+import json
+import os
+from typing import Union
+
+from .bytelevel import ByteLevelTokenizer
+from .core import Tokenizer, _resolve_dir
+
+__all__ = ["Tokenizer", "ByteLevelTokenizer", "load_tokenizer"]
+
+
+def _pre_tokenizer_kinds(pt) -> set:
+    if not pt:
+        return set()
+    if pt.get("type") == "Sequence":
+        return {p.get("type") for p in pt.get("pretokenizers", [])}
+    return {pt.get("type")}
+
+
+def load_tokenizer(path: Union[str, os.PathLike]):
+    """Local-only loader. Mistral (Metaspace) -> Tokenizer; Qwen (ByteLevel) ->
+    ByteLevelTokenizer. The choice is made from tokenizer.json, never from the
+    directory name, and every other family fails loudly instead of being fed to a
+    pipeline that would encode it wrong."""
+    d = _resolve_dir(os.fspath(path))
+    with open(os.path.join(d, "tokenizer.json"), "rb") as f:
+        tj = json.loads(f.read())
+    kinds = _pre_tokenizer_kinds(tj.get("pre_tokenizer"))
+    if "ByteLevel" in kinds:
+        return ByteLevelTokenizer.from_pretrained(d)
+    if kinds == {"Metaspace"}:
+        return Tokenizer.from_pretrained(d)
+    raise NotImplementedError(
+        f"no tokenizer implementation for pre_tokenizer {sorted(k for k in kinds if k)!r} "
+        f"in {d!r}; supported: Metaspace (Mistral), Split+ByteLevel (Qwen2/3)"
+    )

code/align/tokenizer/_qwen2_unicode.py

@@ -0,0 +1,875 @@
+"""GENERATED by tests/data/make_qwen2_unicode_tables.py -- do not edit.
+
+Probed from transformers 5.19.0 / tokenizers 0.23.3 against Python unicodedata 15.0.0.
+See that script for why these are frozen.
+"""
+PYTHON_UNIDATA_VERSION = '15.0.0'  # the unicodedata these deltas were measured against
+
+#: \p{L} / \p{N} exactly as the Rust regex engine sees them (inclusive ranges).
+L_RANGES = (
+    (0x41, 0x5a),
+    (0x61, 0x7a),
+    (0xaa, 0xaa),
+    (0xb5, 0xb5),
+    (0xba, 0xba),
+    (0xc0, 0xd6),
+    (0xd8, 0xf6),
+    (0xf8, 0x2c1),
+    (0x2c6, 0x2d1),
+    (0x2e0, 0x2e4),
+    (0x2ec, 0x2ec),
+    (0x2ee, 0x2ee),
+    (0x370, 0x374),
+    (0x376, 0x377),
+    (0x37a, 0x37d),
+    (0x37f, 0x37f),
+    (0x386, 0x386),
+    (0x388, 0x38a),
+    (0x38c, 0x38c),
+    (0x38e, 0x3a1),
+    (0x3a3, 0x3f5),
+    (0x3f7, 0x481),
+    (0x48a, 0x52f),
+    (0x531, 0x556),
+    (0x559, 0x559),
+    (0x560, 0x588),
+    (0x5d0, 0x5ea),
+    (0x5ef, 0x5f2),
+    (0x620, 0x64a),
+    (0x66e, 0x66f),
+    (0x671, 0x6d3),
+    (0x6d5, 0x6d5),
+    (0x6e5, 0x6e6),
+    (0x6ee, 0x6ef),
+    (0x6fa, 0x6fc),
+    (0x6ff, 0x6ff),
+    (0x710, 0x710),
+    (0x712, 0x72f),
+    (0x74d, 0x7a5),
+    (0x7b1, 0x7b1),
+    (0x7ca, 0x7ea),
+    (0x7f4, 0x7f5),
+    (0x7fa, 0x7fa),
+    (0x800, 0x815),
+    (0x81a, 0x81a),
+    (0x824, 0x824),
+    (0x828, 0x828),
+    (0x840, 0x858),
+    (0x860, 0x86a),
+    (0x870, 0x887),
+    (0x889, 0x88e),
+    (0x8a0, 0x8c9),
+    (0x904, 0x939),
+    (0x93d, 0x93d),
+    (0x950, 0x950),
+    (0x958, 0x961),
+    (0x971, 0x980),
+    (0x985, 0x98c),
+    (0x98f, 0x990),
+    (0x993, 0x9a8),
+    (0x9aa, 0x9b0),
+    (0x9b2, 0x9b2),
+    (0x9b6, 0x9b9),
+    (0x9bd, 0x9bd),
+    (0x9ce, 0x9ce),
+    (0x9dc, 0x9dd),
+    (0x9df, 0x9e1),
+    (0x9f0, 0x9f1),
+    (0x9fc, 0x9fc),
+    (0xa05, 0xa0a),
+    (0xa0f, 0xa10),
+    (0xa13, 0xa28),
+    (0xa2a, 0xa30),
+    (0xa32, 0xa33),
+    (0xa35, 0xa36),
+    (0xa38, 0xa39),
+    (0xa59, 0xa5c),
+    (0xa5e, 0xa5e),
+    (0xa72, 0xa74),
+    (0xa85, 0xa8d),
+    (0xa8f, 0xa91),
+    (0xa93, 0xaa8),
+    (0xaaa, 0xab0),
+    (0xab2, 0xab3),
+    (0xab5, 0xab9),
+    (0xabd, 0xabd),
+    (0xad0, 0xad0),
+    (0xae0, 0xae1),
+    (0xaf9, 0xaf9),
+    (0xb05, 0xb0c),
+    (0xb0f, 0xb10),
+    (0xb13, 0xb28),
+    (0xb2a, 0xb30),
+    (0xb32, 0xb33),
+    (0xb35, 0xb39),
+    (0xb3d, 0xb3d),
+    (0xb5c, 0xb5d),
+    (0xb5f, 0xb61),
+    (0xb71, 0xb71),
+    (0xb83, 0xb83),
+    (0xb85, 0xb8a),
+    (0xb8e, 0xb90),
+    (0xb92, 0xb95),
+    (0xb99, 0xb9a),
+    (0xb9c, 0xb9c),
+    (0xb9e, 0xb9f),
+    (0xba3, 0xba4),
+    (0xba8, 0xbaa),
+    (0xbae, 0xbb9),
+    (0xbd0, 0xbd0),
+    (0xc05, 0xc0c),
+    (0xc0e, 0xc10),
+    (0xc12, 0xc28),
+    (0xc2a, 0xc39),
+    (0xc3d, 0xc3d),
+    (0xc58, 0xc5a),
+    (0xc5d, 0xc5d),
+    (0xc60, 0xc61),
+    (0xc80, 0xc80),
+    (0xc85, 0xc8c),
+    (0xc8e, 0xc90),
+    (0xc92, 0xca8),
+    (0xcaa, 0xcb3),
+    (0xcb5, 0xcb9),
+    (0xcbd, 0xcbd),
+    (0xcdd, 0xcde),
+    (0xce0, 0xce1),
+    (0xcf1, 0xcf2),
+    (0xd04, 0xd0c),
+    (0xd0e, 0xd10),
+    (0xd12, 0xd3a),
+    (0xd3d, 0xd3d),
+    (0xd4e, 0xd4e),
+    (0xd54, 0xd56),
+    (0xd5f, 0xd61),
+    (0xd7a, 0xd7f),
+    (0xd85, 0xd96),
+    (0xd9a, 0xdb1),
+    (0xdb3, 0xdbb),
+    (0xdbd, 0xdbd),
+    (0xdc0, 0xdc6),
+    (0xe01, 0xe30),
+    (0xe32, 0xe33),
+    (0xe40, 0xe46),
+    (0xe81, 0xe82),
+    (0xe84, 0xe84),
+    (0xe86, 0xe8a),
+    (0xe8c, 0xea3),
+    (0xea5, 0xea5),
+    (0xea7, 0xeb0),
+    (0xeb2, 0xeb3),
+    (0xebd, 0xebd),
+    (0xec0, 0xec4),
+    (0xec6, 0xec6),
+    (0xedc, 0xedf),
+    (0xf00, 0xf00),
+    (0xf40, 0xf47),
+    (0xf49, 0xf6c),
+    (0xf88, 0xf8c),
+    (0x1000, 0x102a),
+    (0x103f, 0x103f),
+    (0x1050, 0x1055),
+    (0x105a, 0x105d),
+    (0x1061, 0x1061),
+    (0x1065, 0x1066),
+    (0x106e, 0x1070),
+    (0x1075, 0x1081),
+    (0x108e, 0x108e),
+    (0x10a0, 0x10c5),
+    (0x10c7, 0x10c7),
+    (0x10cd, 0x10cd),
+    (0x10d0, 0x10fa),
+    (0x10fc, 0x1248),
+    (0x124a, 0x124d),
+    (0x1250, 0x1256),
+    (0x1258, 0x1258),
+    (0x125a, 0x125d),
+    (0x1260, 0x1288),
+    (0x128a, 0x128d),
+    (0x1290, 0x12b0),
+    (0x12b2, 0x12b5),
+    (0x12b8, 0x12be),
+    (0x12c0, 0x12c0),
+    (0x12c2, 0x12c5),
+    (0x12c8, 0x12d6),
+    (0x12d8, 0x1310),
+    (0x1312, 0x1315),
+    (0x1318, 0x135a),
+    (0x1380, 0x138f),
+    (0x13a0, 0x13f5),
+    (0x13f8, 0x13fd),
+    (0x1401, 0x166c),
+    (0x166f, 0x167f),
+    (0x1681, 0x169a),
+    (0x16a0, 0x16ea),
+    (0x16f1, 0x16f8),
+    (0x1700, 0x1711),
+    (0x171f, 0x1731),
+    (0x1740, 0x1751),
+    (0x1760, 0x176c),
+    (0x176e, 0x1770),
+    (0x1780, 0x17b3),
+    (0x17d7, 0x17d7),
+    (0x17dc, 0x17dc),
+    (0x1820, 0x1878),
+    (0x1880, 0x1884),
+    (0x1887, 0x18a8),
+    (0x18aa, 0x18aa),
+    (0x18b0, 0x18f5),
+    (0x1900, 0x191e),
+    (0x1950, 0x196d),
+    (0x1970, 0x1974),
+    (0x1980, 0x19ab),
+    (0x19b0, 0x19c9),
+    (0x1a00, 0x1a16),
+    (0x1a20, 0x1a54),
+    (0x1aa7, 0x1aa7),
+    (0x1b05, 0x1b33),
+    (0x1b45, 0x1b4c),
+    (0x1b83, 0x1ba0),
+    (0x1bae, 0x1baf),
+    (0x1bba, 0x1be5),
+    (0x1c00, 0x1c23),
+    (0x1c4d, 0x1c4f),
+    (0x1c5a, 0x1c7d),
+    (0x1c80, 0x1c8a),
+    (0x1c90, 0x1cba),
+    (0x1cbd, 0x1cbf),
+    (0x1ce9, 0x1cec),
+    (0x1cee, 0x1cf3),
+    (0x1cf5, 0x1cf6),
+    (0x1cfa, 0x1cfa),
+    (0x1d00, 0x1dbf),
+    (0x1e00, 0x1f15),
+    (0x1f18, 0x1f1d),
+    (0x1f20, 0x1f45),
+    (0x1f48, 0x1f4d),
+    (0x1f50, 0x1f57),
+    (0x1f59, 0x1f59),
+    (0x1f5b, 0x1f5b),
+    (0x1f5d, 0x1f5d),
+    (0x1f5f, 0x1f7d),
+    (0x1f80, 0x1fb4),
+    (0x1fb6, 0x1fbc),
+    (0x1fbe, 0x1fbe),
+    (0x1fc2, 0x1fc4),
+    (0x1fc6, 0x1fcc),
+    (0x1fd0, 0x1fd3),
+    (0x1fd6, 0x1fdb),
+    (0x1fe0, 0x1fec),
+    (0x1ff2, 0x1ff4),
+    (0x1ff6, 0x1ffc),
+    (0x2071, 0x2071),
+    (0x207f, 0x207f),
+    (0x2090, 0x209c),
+    (0x2102, 0x2102),
+    (0x2107, 0x2107),
+    (0x210a, 0x2113),
+    (0x2115, 0x2115),
+    (0x2119, 0x211d),
+    (0x2124, 0x2124),
+    (0x2126, 0x2126),
+    (0x2128, 0x2128),
+    (0x212a, 0x212d),
+    (0x212f, 0x2139),
+    (0x213c, 0x213f),
+    (0x2145, 0x2149),
+    (0x214e, 0x214e),
+    (0x2183, 0x2184),
+    (0x2c00, 0x2ce4),
+    (0x2ceb, 0x2cee),
+    (0x2cf2, 0x2cf3),
+    (0x2d00, 0x2d25),
+    (0x2d27, 0x2d27),
+    (0x2d2d, 0x2d2d),
+    (0x2d30, 0x2d67),
+    (0x2d6f, 0x2d6f),
+    (0x2d80, 0x2d96),
+    (0x2da0, 0x2da6),
+    (0x2da8, 0x2dae),
+    (0x2db0, 0x2db6),
+    (0x2db8, 0x2dbe),
+    (0x2dc0, 0x2dc6),
+    (0x2dc8, 0x2dce),
+    (0x2dd0, 0x2dd6),
+    (0x2dd8, 0x2dde),
+    (0x2e2f, 0x2e2f),
+    (0x3005, 0x3006),
+    (0x3031, 0x3035),
+    (0x303b, 0x303c),
+    (0x3041, 0x3096),
+    (0x309d, 0x309f),
+    (0x30a1, 0x30fa),
+    (0x30fc, 0x30ff),
+    (0x3105, 0x312f),
+    (0x3131, 0x318e),
+    (0x31a0, 0x31bf),
+    (0x31f0, 0x31ff),
+    (0x3400, 0x4dbf),
+    (0x4e00, 0xa48c),
+    (0xa4d0, 0xa4fd),
+    (0xa500, 0xa60c),
+    (0xa610, 0xa61f),
+    (0xa62a, 0xa62b),
+    (0xa640, 0xa66e),
+    (0xa67f, 0xa69d),
+    (0xa6a0, 0xa6e5),
+    (0xa717, 0xa71f),
+    (0xa722, 0xa788),
+    (0xa78b, 0xa7cd),
+    (0xa7d0, 0xa7d1),
+    (0xa7d3, 0xa7d3),
+    (0xa7d5, 0xa7dc),
+    (0xa7f2, 0xa801),
+    (0xa803, 0xa805),
+    (0xa807, 0xa80a),
+    (0xa80c, 0xa822),
+    (0xa840, 0xa873),
+    (0xa882, 0xa8b3),
+    (0xa8f2, 0xa8f7),
+    (0xa8fb, 0xa8fb),
+    (0xa8fd, 0xa8fe),
+    (0xa90a, 0xa925),
+    (0xa930, 0xa946),
+    (0xa960, 0xa97c),
+    (0xa984, 0xa9b2),
+    (0xa9cf, 0xa9cf),
+    (0xa9e0, 0xa9e4),
+    (0xa9e6, 0xa9ef),
+    (0xa9fa, 0xa9fe),
+    (0xaa00, 0xaa28),
+    (0xaa40, 0xaa42),
+    (0xaa44, 0xaa4b),
+    (0xaa60, 0xaa76),
+    (0xaa7a, 0xaa7a),
+    (0xaa7e, 0xaaaf),
+    (0xaab1, 0xaab1),
+    (0xaab5, 0xaab6),
+    (0xaab9, 0xaabd),
+    (0xaac0, 0xaac0),
+    (0xaac2, 0xaac2),
+    (0xaadb, 0xaadd),
+    (0xaae0, 0xaaea),
+    (0xaaf2, 0xaaf4),
+    (0xab01, 0xab06),
+    (0xab09, 0xab0e),
+    (0xab11, 0xab16),
+    (0xab20, 0xab26),
+    (0xab28, 0xab2e),
+    (0xab30, 0xab5a),
+    (0xab5c, 0xab69),
+    (0xab70, 0xabe2),
+    (0xac00, 0xd7a3),
+    (0xd7b0, 0xd7c6),
+    (0xd7cb, 0xd7fb),
+    (0xf900, 0xfa6d),
+    (0xfa70, 0xfad9),
+    (0xfb00, 0xfb06),
+    (0xfb13, 0xfb17),
+    (0xfb1d, 0xfb1d),
+    (0xfb1f, 0xfb28),
+    (0xfb2a, 0xfb36),
+    (0xfb38, 0xfb3c),
+    (0xfb3e, 0xfb3e),
+    (0xfb40, 0xfb41),
+    (0xfb43, 0xfb44),
+    (0xfb46, 0xfbb1),
+    (0xfbd3, 0xfd3d),
+    (0xfd50, 0xfd8f),
+    (0xfd92, 0xfdc7),
+    (0xfdf0, 0xfdfb),
+    (0xfe70, 0xfe74),
+    (0xfe76, 0xfefc),
+    (0xff21, 0xff3a),
+    (0xff41, 0xff5a),
+    (0xff66, 0xffbe),
+    (0xffc2, 0xffc7),
+    (0xffca, 0xffcf),
+    (0xffd2, 0xffd7),
+    (0xffda, 0xffdc),
+    (0x10000, 0x1000b),
+    (0x1000d, 0x10026),
+    (0x10028, 0x1003a),
+    (0x1003c, 0x1003d),
+    (0x1003f, 0x1004d),
+    (0x10050, 0x1005d),
+    (0x10080, 0x100fa),
+    (0x10280, 0x1029c),
+    (0x102a0, 0x102d0),
+    (0x10300, 0x1031f),
+    (0x1032d, 0x10340),
+    (0x10342, 0x10349),
+    (0x10350, 0x10375),
+    (0x10380, 0x1039d),
+    (0x103a0, 0x103c3),
+    (0x103c8, 0x103cf),
+    (0x10400, 0x1049d),
+    (0x104b0, 0x104d3),
+    (0x104d8, 0x104fb),
+    (0x10500, 0x10527),
+    (0x10530, 0x10563),
+    (0x10570, 0x1057a),
+    (0x1057c, 0x1058a),
+    (0x1058c, 0x10592),
+    (0x10594, 0x10595),
+    (0x10597, 0x105a1),
+    (0x105a3, 0x105b1),
+    (0x105b3, 0x105b9),
+    (0x105bb, 0x105bc),
+    (0x105c0, 0x105f3),
+    (0x10600, 0x10736),
+    (0x10740, 0x10755),
+    (0x10760, 0x10767),
+    (0x10780, 0x10785),
+    (0x10787, 0x107b0),
+    (0x107b2, 0x107ba),
+    (0x10800, 0x10805),
+    (0x10808, 0x10808),
+    (0x1080a, 0x10835),
+    (0x10837, 0x10838),
+    (0x1083c, 0x1083c),
+    (0x1083f, 0x10855),
+    (0x10860, 0x10876),
+    (0x10880, 0x1089e),
+    (0x108e0, 0x108f2),
+    (0x108f4, 0x108f5),
+    (0x10900, 0x10915),
+    (0x10920, 0x10939),
+    (0x10980, 0x109b7),
+    (0x109be, 0x109bf),
+    (0x10a00, 0x10a00),
+    (0x10a10, 0x10a13),
+    (0x10a15, 0x10a17),
+    (0x10a19, 0x10a35),
+    (0x10a60, 0x10a7c),
+    (0x10a80, 0x10a9c),
+    (0x10ac0, 0x10ac7),
+    (0x10ac9, 0x10ae4),
+    (0x10b00, 0x10b35),
+    (0x10b40, 0x10b55),
+    (0x10b60, 0x10b72),
+    (0x10b80, 0x10b91),
+    (0x10c00, 0x10c48),
+    (0x10c80, 0x10cb2),
+    (0x10cc0, 0x10cf2),
+    (0x10d00, 0x10d23),
+    (0x10d4a, 0x10d65),
+    (0x10d6f, 0x10d85),
+    (0x10e80, 0x10ea9),
+    (0x10eb0, 0x10eb1),
+    (0x10ec2, 0x10ec4),
+    (0x10f00, 0x10f1c),
+    (0x10f27, 0x10f27),
+    (0x10f30, 0x10f45),
+    (0x10f70, 0x10f81),
+    (0x10fb0, 0x10fc4),
+    (0x10fe0, 0x10ff6),
+    (0x11003, 0x11037),
+    (0x11071, 0x11072),
+    (0x11075, 0x11075),
+    (0x11083, 0x110af),
+    (0x110d0, 0x110e8),
+    (0x11103, 0x11126),
+    (0x11144, 0x11144),
+    (0x11147, 0x11147),
+    (0x11150, 0x11172),
+    (0x11176, 0x11176),
+    (0x11183, 0x111b2),
+    (0x111c1, 0x111c4),
+    (0x111da, 0x111da),
+    (0x111dc, 0x111dc),
+    (0x11200, 0x11211),
+    (0x11213, 0x1122b),
+    (0x1123f, 0x11240),
+    (0x11280, 0x11286),
+    (0x11288, 0x11288),
+    (0x1128a, 0x1128d),
+    (0x1128f, 0x1129d),
+    (0x1129f, 0x112a8),
+    (0x112b0, 0x112de),
+    (0x11305, 0x1130c),
+    (0x1130f, 0x11310),
+    (0x11313, 0x11328),
+    (0x1132a, 0x11330),
+    (0x11332, 0x11333),
+    (0x11335, 0x11339),
+    (0x1133d, 0x1133d),
+    (0x11350, 0x11350),
+    (0x1135d, 0x11361),
+    (0x11380, 0x11389),
+    (0x1138b, 0x1138b),
+    (0x1138e, 0x1138e),
+    (0x11390, 0x113b5),
+    (0x113b7, 0x113b7),
+    (0x113d1, 0x113d1),
+    (0x113d3, 0x113d3),
+    (0x11400, 0x11434),
+    (0x11447, 0x1144a),
+    (0x1145f, 0x11461),
+    (0x11480, 0x114af),
+    (0x114c4, 0x114c5),
+    (0x114c7, 0x114c7),
+    (0x11580, 0x115ae),
+    (0x115d8, 0x115db),
+    (0x11600, 0x1162f),
+    (0x11644, 0x11644),
+    (0x11680, 0x116aa),
+    (0x116b8, 0x116b8),
+    (0x11700, 0x1171a),
+    (0x11740, 0x11746),
+    (0x11800, 0x1182b),
+    (0x118a0, 0x118df),
+    (0x118ff, 0x11906),
+    (0x11909, 0x11909),
+    (0x1190c, 0x11913),
+    (0x11915, 0x11916),
+    (0x11918, 0x1192f),
+    (0x1193f, 0x1193f),
+    (0x11941, 0x11941),
+    (0x119a0, 0x119a7),
+    (0x119aa, 0x119d0),
+    (0x119e1, 0x119e1),
+    (0x119e3, 0x119e3),
+    (0x11a00, 0x11a00),
+    (0x11a0b, 0x11a32),
+    (0x11a3a, 0x11a3a),
+    (0x11a50, 0x11a50),
+    (0x11a5c, 0x11a89),
+    (0x11a9d, 0x11a9d),
+    (0x11ab0, 0x11af8),
+    (0x11bc0, 0x11be0),
+    (0x11c00, 0x11c08),
+    (0x11c0a, 0x11c2e),
+    (0x11c40, 0x11c40),
+    (0x11c72, 0x11c8f),
+    (0x11d00, 0x11d06),
+    (0x11d08, 0x11d09),
+    (0x11d0b, 0x11d30),
+    (0x11d46, 0x11d46),
+    (0x11d60, 0x11d65),
+    (0x11d67, 0x11d68),
+    (0x11d6a, 0x11d89),
+    (0x11d98, 0x11d98),
+    (0x11ee0, 0x11ef2),
+    (0x11f02, 0x11f02),
+    (0x11f04, 0x11f10),
+    (0x11f12, 0x11f33),
+    (0x11fb0, 0x11fb0),
+    (0x12000, 0x12399),
+    (0x12480, 0x12543),
+    (0x12f90, 0x12ff0),
+    (0x13000, 0x1342f),
+    (0x13441, 0x13446),
+    (0x13460, 0x143fa),
+    (0x14400, 0x14646),
+    (0x16100, 0x1611d),
+    (0x16800, 0x16a38),
+    (0x16a40, 0x16a5e),
+    (0x16a70, 0x16abe),
+    (0x16ad0, 0x16aed),
+    (0x16b00, 0x16b2f),
+    (0x16b40, 0x16b43),
+    (0x16b63, 0x16b77),
+    (0x16b7d, 0x16b8f),
+    (0x16d40, 0x16d6c),
+    (0x16e40, 0x16e7f),
+    (0x16f00, 0x16f4a),
+    (0x16f50, 0x16f50),
+    (0x16f93, 0x16f9f),
+    (0x16fe0, 0x16fe1),
+    (0x16fe3, 0x16fe3),
+    (0x17000, 0x187f7),
+    (0x18800, 0x18cd5),
+    (0x18cff, 0x18d08),
+    (0x1aff0, 0x1aff3),
+    (0x1aff5, 0x1affb),
+    (0x1affd, 0x1affe),
+    (0x1b000, 0x1b122),
+    (0x1b132, 0x1b132),
+    (0x1b150, 0x1b152),
+    (0x1b155, 0x1b155),
+    (0x1b164, 0x1b167),
+    (0x1b170, 0x1b2fb),
+    (0x1bc00, 0x1bc6a),
+    (0x1bc70, 0x1bc7c),
+    (0x1bc80, 0x1bc88),
+    (0x1bc90, 0x1bc99),
+    (0x1d400, 0x1d454),
+    (0x1d456, 0x1d49c),
+    (0x1d49e, 0x1d49f),
+    (0x1d4a2, 0x1d4a2),
+    (0x1d4a5, 0x1d4a6),
+    (0x1d4a9, 0x1d4ac),
+    (0x1d4ae, 0x1d4b9),
+    (0x1d4bb, 0x1d4bb),
+    (0x1d4bd, 0x1d4c3),
+    (0x1d4c5, 0x1d505),
+    (0x1d507, 0x1d50a),
+    (0x1d50d, 0x1d514),
+    (0x1d516, 0x1d51c),
+    (0x1d51e, 0x1d539),
+    (0x1d53b, 0x1d53e),
+    (0x1d540, 0x1d544),
+    (0x1d546, 0x1d546),
+    (0x1d54a, 0x1d550),
+    (0x1d552, 0x1d6a5),
+    (0x1d6a8, 0x1d6c0),
+    (0x1d6c2, 0x1d6da),
+    (0x1d6dc, 0x1d6fa),
+    (0x1d6fc, 0x1d714),
+    (0x1d716, 0x1d734),
+    (0x1d736, 0x1d74e),
+    (0x1d750, 0x1d76e),
+    (0x1d770, 0x1d788),
+    (0x1d78a, 0x1d7a8),
+    (0x1d7aa, 0x1d7c2),
+    (0x1d7c4, 0x1d7cb),
+    (0x1df00, 0x1df1e),
+    (0x1df25, 0x1df2a),
+    (0x1e030, 0x1e06d),
+    (0x1e100, 0x1e12c),
+    (0x1e137, 0x1e13d),
+    (0x1e14e, 0x1e14e),
+    (0x1e290, 0x1e2ad),
+    (0x1e2c0, 0x1e2eb),
+    (0x1e4d0, 0x1e4eb),
+    (0x1e5d0, 0x1e5ed),
+    (0x1e5f0, 0x1e5f0),
+    (0x1e7e0, 0x1e7e6),
+    (0x1e7e8, 0x1e7eb),
+    (0x1e7ed, 0x1e7ee),
+    (0x1e7f0, 0x1e7fe),
+    (0x1e800, 0x1e8c4),
+    (0x1e900, 0x1e943),
+    (0x1e94b, 0x1e94b),
+    (0x1ee00, 0x1ee03),
+    (0x1ee05, 0x1ee1f),
+    (0x1ee21, 0x1ee22),
+    (0x1ee24, 0x1ee24),
+    (0x1ee27, 0x1ee27),
+    (0x1ee29, 0x1ee32),
+    (0x1ee34, 0x1ee37),
+    (0x1ee39, 0x1ee39),
+    (0x1ee3b, 0x1ee3b),
+    (0x1ee42, 0x1ee42),
+    (0x1ee47, 0x1ee47),
+    (0x1ee49, 0x1ee49),
+    (0x1ee4b, 0x1ee4b),
+    (0x1ee4d, 0x1ee4f),
+    (0x1ee51, 0x1ee52),
+    (0x1ee54, 0x1ee54),
+    (0x1ee57, 0x1ee57),
+    (0x1ee59, 0x1ee59),
+    (0x1ee5b, 0x1ee5b),
+    (0x1ee5d, 0x1ee5d),
+    (0x1ee5f, 0x1ee5f),
+    (0x1ee61, 0x1ee62),
+    (0x1ee64, 0x1ee64),
+    (0x1ee67, 0x1ee6a),
+    (0x1ee6c, 0x1ee72),
+    (0x1ee74, 0x1ee77),
+    (0x1ee79, 0x1ee7c),
+    (0x1ee7e, 0x1ee7e),
+    (0x1ee80, 0x1ee89),
+    (0x1ee8b, 0x1ee9b),
+    (0x1eea1, 0x1eea3),
+    (0x1eea5, 0x1eea9),
+    (0x1eeab, 0x1eebb),
+    (0x20000, 0x2a6df),
+    (0x2a700, 0x2b739),
+    (0x2b740, 0x2b81d),
+    (0x2b820, 0x2cea1),
+    (0x2ceb0, 0x2ebe0),
+    (0x2ebf0, 0x2ee5d),
+    (0x2f800, 0x2fa1d),
+    (0x30000, 0x3134a),
+    (0x31350, 0x323af),
+)
+N_RANGES = (
+    (0x30, 0x39),
+    (0xb2, 0xb3),
+    (0xb9, 0xb9),
+    (0xbc, 0xbe),
+    (0x660, 0x669),
+    (0x6f0, 0x6f9),
+    (0x7c0, 0x7c9),
+    (0x966, 0x96f),
+    (0x9e6, 0x9ef),
+    (0x9f4, 0x9f9),
+    (0xa66, 0xa6f),
+    (0xae6, 0xaef),
+    (0xb66, 0xb6f),
+    (0xb72, 0xb77),
+    (0xbe6, 0xbf2),
+    (0xc66, 0xc6f),
+    (0xc78, 0xc7e),
+    (0xce6, 0xcef),
+    (0xd58, 0xd5e),
+    (0xd66, 0xd78),
+    (0xde6, 0xdef),
+    (0xe50, 0xe59),
+    (0xed0, 0xed9),
+    (0xf20, 0xf33),
+    (0x1040, 0x1049),
+    (0x1090, 0x1099),
+    (0x1369, 0x137c),
+    (0x16ee, 0x16f0),
+    (0x17e0, 0x17e9),
+    (0x17f0, 0x17f9),
+    (0x1810, 0x1819),
+    (0x1946, 0x194f),
+    (0x19d0, 0x19da),
+    (0x1a80, 0x1a89),
+    (0x1a90, 0x1a99),
+    (0x1b50, 0x1b59),
+    (0x1bb0, 0x1bb9),
+    (0x1c40, 0x1c49),
+    (0x1c50, 0x1c59),
+    (0x2070, 0x2070),
+    (0x2074, 0x2079),
+    (0x2080, 0x2089),
+    (0x2150, 0x2182),
+    (0x2185, 0x2189),
+    (0x2460, 0x249b),
+    (0x24ea, 0x24ff),
+    (0x2776, 0x2793),
+    (0x2cfd, 0x2cfd),
+    (0x3007, 0x3007),
+    (0x3021, 0x3029),
+    (0x3038, 0x303a),
+    (0x3192, 0x3195),
+    (0x3220, 0x3229),
+    (0x3248, 0x324f),
+    (0x3251, 0x325f),
+    (0x3280, 0x3289),
+    (0x32b1, 0x32bf),
+    (0xa620, 0xa629),
+    (0xa6e6, 0xa6ef),
+    (0xa830, 0xa835),
+    (0xa8d0, 0xa8d9),
+    (0xa900, 0xa909),
+    (0xa9d0, 0xa9d9),
+    (0xa9f0, 0xa9f9),
+    (0xaa50, 0xaa59),
+    (0xabf0, 0xabf9),
+    (0xff10, 0xff19),
+    (0x10107, 0x10133),
+    (0x10140, 0x10178),
+    (0x1018a, 0x1018b),
+    (0x102e1, 0x102fb),
+    (0x10320, 0x10323),
+    (0x10341, 0x10341),
+    (0x1034a, 0x1034a),
+    (0x103d1, 0x103d5),
+    (0x104a0, 0x104a9),
+    (0x10858, 0x1085f),
+    (0x10879, 0x1087f),
+    (0x108a7, 0x108af),
+    (0x108fb, 0x108ff),
+    (0x10916, 0x1091b),
+    (0x109bc, 0x109bd),
+    (0x109c0, 0x109cf),
+    (0x109d2, 0x109ff),
+    (0x10a40, 0x10a48),
+    (0x10a7d, 0x10a7e),
+    (0x10a9d, 0x10a9f),
+    (0x10aeb, 0x10aef),
+    (0x10b58, 0x10b5f),
+    (0x10b78, 0x10b7f),
+    (0x10ba9, 0x10baf),
+    (0x10cfa, 0x10cff),
+    (0x10d30, 0x10d39),
+    (0x10d40, 0x10d49),
+    (0x10e60, 0x10e7e),
+    (0x10f1d, 0x10f26),
+    (0x10f51, 0x10f54),
+    (0x10fc5, 0x10fcb),
+    (0x11052, 0x1106f),
+    (0x110f0, 0x110f9),
+    (0x11136, 0x1113f),
+    (0x111d0, 0x111d9),
+    (0x111e1, 0x111f4),
+    (0x112f0, 0x112f9),
+    (0x11450, 0x11459),
+    (0x114d0, 0x114d9),
+    (0x11650, 0x11659),
+    (0x116c0, 0x116c9),
+    (0x116d0, 0x116e3),
+    (0x11730, 0x1173b),
+    (0x118e0, 0x118f2),
+    (0x11950, 0x11959),
+    (0x11bf0, 0x11bf9),
+    (0x11c50, 0x11c6c),
+    (0x11d50, 0x11d59),
+    (0x11da0, 0x11da9),
+    (0x11f50, 0x11f59),
+    (0x11fc0, 0x11fd4),
+    (0x12400, 0x1246e),
+    (0x16130, 0x16139),
+    (0x16a60, 0x16a69),
+    (0x16ac0, 0x16ac9),
+    (0x16b50, 0x16b59),
+    (0x16b5b, 0x16b61),
+    (0x16d70, 0x16d79),
+    (0x16e80, 0x16e96),
+    (0x1ccf0, 0x1ccf9),
+    (0x1d2c0, 0x1d2d3),
+    (0x1d2e0, 0x1d2f3),
+    (0x1d360, 0x1d378),
+    (0x1d7ce, 0x1d7ff),
+    (0x1e140, 0x1e149),
+    (0x1e2f0, 0x1e2f9),
+    (0x1e4f0, 0x1e4f9),
+    (0x1e5f1, 0x1e5fa),
+    (0x1e8c7, 0x1e8cf),
+    (0x1e950, 0x1e959),
+    (0x1ec71, 0x1ecab),
+    (0x1ecad, 0x1ecaf),
+    (0x1ecb1, 0x1ecb4),
+    (0x1ed01, 0x1ed2d),
+    (0x1ed2f, 0x1ed3d),
+    (0x1f100, 0x1f10c),
+    (0x1fbf0, 0x1fbf9),
+)
+
+#: Code points the Rust NFC treats as combining-class 0 although unicodedata says otherwise,
+#: plus second members of compositions the Rust side lacks. NFC is applied around them.
+NFC_BARRIERS = (
+    (0x7fd, 0x7fd),
+    (0x898, 0x89f),
+    (0x8ca, 0x8d3),
+    (0x9fe, 0x9fe),
+    (0xc3c, 0xc3c),
+    (0xd3b, 0xd3c),
+    (0xeba, 0xeba),
+    (0x1715, 0x1715),
+    (0x1abf, 0x1ace),
+    (0x1df6, 0x1dfa),
+    (0xa82c, 0xa82c),
+    (0x10d24, 0x10d27),
+    (0x10eab, 0x10eac),
+    (0x10efd, 0x10eff),
+    (0x10f46, 0x10f50),
+    (0x10f82, 0x10f85),
+    (0x11070, 0x11070),
+    (0x1133b, 0x1133b),
+    (0x1145e, 0x1145e),
+    (0x11839, 0x1183a),
+    (0x11930, 0x11930),
+    (0x1193d, 0x1193e),
+    (0x11943, 0x11943),
+    (0x119e0, 0x119e0),
+    (0x11a34, 0x11a34),
+    (0x11a47, 0x11a47),
+    (0x11a99, 0x11a99),
+    (0x11d42, 0x11d42),
+    (0x11d44, 0x11d45),
+    (0x11d97, 0x11d97),
+    (0x11f41, 0x11f42),
+    (0x16ff0, 0x16ff1),
+    (0x1e08f, 0x1e08f),
+    (0x1e130, 0x1e136),
+    (0x1e2ae, 0x1e2ae),
+    (0x1e2ec, 0x1e2ef),
+    (0x1e4ec, 0x1e4ef),
+)

code/align/tokenizer/bytelevel.py

@@ -0,0 +1,693 @@
+"""Pure-Python (stdlib only) byte-level BPE tokenizer for Qwen2/Qwen3 (Goedel-Prover-V2-8B).
+
+Why this exists. P2 (LoRA SFT on Goedel-Prover-V2-8B's own verified samples) needs
+token ids for prompts and completions. Before this module the only way to get
+them for a Qwen3 vocabulary was `transformers.AutoTokenizer`, i.e. the isolated
+vLLM venv: the training process would have had to import (or shell out to) a
+second Python environment just to turn text into ints, and a version skew
+between that venv and the one that trained would silently change the ids
+(measured elsewhere in this repo: transformers 5.x REBUILDS pipelines from the
+config rather than from tokenizer.json). This module reads tokenizer.json
+directly and is certified byte-identical (ids AND decode) against the HF tokenizer
+on a frozen oracle corpus by tests/test_tokenizer_bytelevel.py.
+
+Pipeline (encode), exactly what tokenizer.json declares and nothing else:
+
+    text
+      -> split on ADDED tokens (<|im_start|>, <think>, ...), leftmost-longest,
+         on the RAW text (all are normalized=false)
+      -> per remaining segment: NFC normalizer
+      -> Split(Isolated) with the Qwen2 regex
+      -> ByteLevel: UTF-8 bytes of each piece mapped through GPT-2 bytes_to_unicode
+      -> BPE (rank-ordered merges) per piece
+      -> post-processor ByteLevel: adds nothing (no BOS/EOS)
+
+Decode: ids -> pieces (optionally drop special=true added tokens) -> join ->
+byte-decode -> UTF-8 with U+FFFD replacement.
+
+The three places a reimplementation gets this wrong -- each has a mutant in the
+test file that must make parity FAIL:
+
+1. NFC per SEGMENT, not per string. Added tokens are carved out first, so a
+   combining mark right after "<|im_end|>" has nothing to compose with. NFC on the
+   whole string would compose across the special token. Dropping NFC entirely
+   changes the ids of every decomposed-form string (e.g. "e" + U+0301).
+2. The regex has `\\p{N}` (ONE digit per piece, so "2024" is four pieces) and
+   `\\p{L}+` / `[^\\r\\n\\p{L}\\p{N}]?` (a single non-letter prefix glues onto a word).
+   Python `re` has no \\p{..}: the classes are explicit ranges (_class_body)
+   built from the frozen Rust-side tables in _qwen2_unicode.py. A "\\d+" or "\\w+" shortcut is
+   not equivalent (superscripts are \\p{N} but not \\d; "_" is \\w but not \\p{L}).
+3. BPE tie-break: lowest merge RANK wins, ties go to the LEFTMOST pair. Using
+   vocab id (or longest-match) instead of rank changes tokens inside long
+   identifiers and whitespace runs, silently.
+
+Unicode tables (measured, pinned in tests): the Rust tokenizer's Unicode tables are not
+those of any one Python. \\p{L}/\\p{N} are NEWER than CPython 3.12's unicodedata (5004
+Unicode-16 code points are letters/digits there and unassigned here) and the NFC
+normalizer is OLDER (108 combining marks of Unicode 10-15 keep combining class 0, and
+U+11935 U+11930 does not compose). A unicodedata-derived translation therefore split
+"X's" differently for those code points and canonically reordered marks HF leaves alone.
+Both are now FROZEN from the real tokenizer (_qwen2_unicode.py, generated by
+tests/data/make_qwen2_unicode_tables.py), so the result does not depend on the Python
+version for the regex classes; NFC still uses unicodedata for everything else and is pinned
+to the certified unicodedata version by a test.
+
+The loader VALIDATES the whole pipeline configuration and raises
+NotImplementedError on anything not implemented (a different regex, a missing
+NFC, ignore_merges, normalized added tokens, clean_up_tokenization_spaces,
+an unknown chat template ...). A different tokenizer.json must fail loudly, not
+silently encode wrong ids.
+"""
+from __future__ import annotations
+
+import hashlib
+import heapq
+import json
+import os
+import re
+import unicodedata
+from typing import Iterable, Optional, Sequence, Union
+
+from .core import _TOKENIZER_FILES, _as_list, _content, _expect, _resolve_dir
+from ._qwen2_unicode import L_RANGES, N_RANGES, NFC_BARRIERS
+
+__all__ = ["ByteLevelTokenizer", "QWEN2_PRETOKENIZE_REGEX", "bytes_to_unicode", "render_qwen3_chat"]
+
+#: The pre_tokenizer Split pattern the Qwen2/Qwen2.5/Qwen3 tokenizer.json files ship.
+#: The loader requires tokenizer.json to carry EXACTLY this string: the Python
+#: translation below (_build_pattern) is only valid for it.
+QWEN2_PRETOKENIZE_REGEX = (
+    r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*"
+    r"|\s*[\r\n]+|\s+(?!\S)|\s+"
+)
+
+#: sha256 of the Qwen3 chat_template string that render_qwen3_chat reimplements
+#: (the subset without tools). A config carrying any other template makes
+#: apply_chat_template raise instead of rendering with the wrong one.
+QWEN3_CHAT_TEMPLATE_SHA256 = "a55ee1b1660128b7098723e0abcd92caa0788061051c62d51cbe87d9cf1974d8"
+
+_CACHE_MAX_ENTRIES = 200_000
+
+# Unicode White_Space property (what the Rust regex `\s` means). NOT Python's
+# str.isspace(): that also accepts U+001C..U+001F, which are not White_Space.
+_WHITE_SPACE = (
+    (0x09, 0x0D), (0x20, 0x20), (0x85, 0x85), (0xA0, 0xA0), (0x1680, 0x1680),
+    (0x2000, 0x200A), (0x2028, 0x2029), (0x202F, 0x202F), (0x205F, 0x205F), (0x3000, 0x3000),
+)
+
+
+def bytes_to_unicode() -> dict[int, str]:
+    """GPT-2 byte -> printable-char map (the ByteLevel alphabet): printable Latin-1
+    bytes map to themselves, the other 68 bytes to U+0100 and up, in byte order."""
+    keep = list(range(0x21, 0x7F)) + list(range(0xA1, 0xAD)) + list(range(0xAE, 0x100))
+    out: dict[int, str] = {}
+    n = 0
+    for b in range(256):
+        if b in keep:
+            out[b] = chr(b)
+        else:
+            out[b] = chr(256 + n)
+            n += 1
+    return out
+
+
+def _ranges(pred) -> list[tuple[int, int]]:
+    out: list[tuple[int, int]] = []
+    start = None
+    for cp in range(0x110000):
+        ok = pred(cp)
+        if ok and start is None:
+            start = cp
+        elif not ok and start is not None:
+            out.append((start, cp - 1))
+            start = None
+    if start is not None:
+        out.append((start, 0x10FFFF))
+    return out
+
+
+def _class_body(ranges: Iterable[tuple[int, int]]) -> str:
+    """Body of a regex character class (without the brackets) for explicit ranges."""
+    parts = []
+    for a, b in ranges:
+        parts.append(f"\\U{a:08x}" if a == b else f"\\U{a:08x}-\\U{b:08x}")
+    return "".join(parts)
+
+
+_PATTERN: Optional[re.Pattern] = None
+_BARRIER_RE: Optional[re.Pattern] = None
+
+
+def _build_pattern() -> re.Pattern:
+    """Python-`re` translation of QWEN2_PRETOKENIZE_REGEX. \\p{L} / \\p{N} are the FROZEN
+    Rust-side tables (L_RANGES / N_RANGES), \\s = White_Space. The contraction group is
+    spelled out instead of (?i:...): Python's IGNORECASE is Unicode-aware and would also
+    accept U+017F (long s) for 's -- the oracle test pins what the reference does for it."""
+    global _PATTERN
+    if _PATTERN is not None:
+        return _PATTERN
+    L = _class_body(L_RANGES)
+    N = _class_body(N_RANGES)
+    S = _class_body(_WHITE_SPACE)
+    pat = (
+        r"'(?:[sS]|[tT]|[rR][eE]|[vV][eE]|[mM]|[lL][lL]|[dD])"
+        rf"|[^\r\n{L}{N}]?[{L}]+"
+        rf"|[{N}]"
+        rf"| ?[^{S}{L}{N}]+[\r\n]*"
+        rf"|[{S}]*[\r\n]+"
+        rf"|[{S}]+(?![^{S}])"
+        rf"|[{S}]+"
+    )
+    _PATTERN = re.compile(pat)
+    return _PATTERN
+
+
+def _barrier_re() -> re.Pattern:
+    global _BARRIER_RE
+    if _BARRIER_RE is None:
+        _BARRIER_RE = re.compile("([" + _class_body(NFC_BARRIERS) + "])")
+    return _BARRIER_RE
+
+
+class ByteLevelTokenizer:
+    """Qwen2/Qwen3 byte-level BPE, API-compatible with align.tokenizer.Tokenizer
+    for the calls this repo makes (encode / __call__ / decode / ids / save).
+
+    Build with `ByteLevelTokenizer.from_pretrained(dir)` or `load_tokenizer(dir)`.
+    """
+
+    def __init__(
+        self,
+        tokenizer_json: dict,
+        tokenizer_config: Optional[dict] = None,
+        special_tokens_map: Optional[dict] = None,
+        *,
+        files: Optional[dict] = None,
+    ):
+        self._files: dict[str, bytes] = dict(files or {})
+        self._config = dict(tokenizer_config or {})
+        self._special_map = dict(special_tokens_map or {})
+        self._tokenizer_json = tokenizer_json
+        self._validate_pipeline(tokenizer_json, self._config)
+
+        model = tokenizer_json["model"]
+        base_vocab: dict[str, int] = dict(model["vocab"])
+        self._base_vocab = base_vocab
+        vocab = dict(base_vocab)
+
+        self._added: dict[str, int] = {}
+        self._special_contents: set[str] = set()
+        for t in tokenizer_json.get("added_tokens", []):
+            _expect(
+                not (t.get("lstrip") or t.get("rstrip") or t.get("single_word")),
+                f"added token {t['content']!r} with lstrip/rstrip/single_word",
+            )
+            _expect(
+                not t.get("normalized"),
+                f"added token {t['content']!r} with normalized=true (NFC-before-match not implemented)",
+            )
+            vocab[t["content"]] = int(t["id"])
+            self._added[t["content"]] = int(t["id"])
+            if t.get("special"):
+                self._special_contents.add(t["content"])
+        self._vocab = vocab
+        size = max(vocab.values()) + 1
+        self._id_to_token: list[Optional[str]] = [None] * size
+        for tok, i in vocab.items():
+            self._id_to_token[i] = tok
+
+        self._b2u = bytes_to_unicode()
+        self._u2b = {c: b for b, c in self._b2u.items()}
+        missing = [b for b, c in self._b2u.items() if c not in base_vocab]
+        _expect(not missing, f"vocab lacks ByteLevel alphabet symbols for bytes {missing[:5]}")
+
+        self._ranks: dict[tuple[str, str], int] = {}
+        for rank, m in enumerate(model["merges"]):
+            a, b = (m[0], m[1]) if isinstance(m, (list, tuple)) else m.split(" ")
+            if (a + b) not in base_vocab or a not in base_vocab or b not in base_vocab:
+                raise ValueError(f"merge #{rank} {a!r}+{b!r} references a token outside the vocab")
+            self._ranks[(a, b)] = rank  # last duplicate wins (Rust HashMap semantics)
+
+        if self._added:
+            alts = sorted(self._added, key=len, reverse=True)
+            self._added_re: Optional[re.Pattern] = re.compile("|".join(re.escape(a) for a in alts))
+        else:
+            self._added_re = None
+        self._cache: dict[str, tuple[int, ...]] = {}
+        self._id_bytes: list[Optional[bytes]] = [None] * size
+
+        def pick(name: str, default: Optional[str] = None) -> Optional[str]:
+            v = _content(self._config.get(f"{name}_token"))
+            if v is None:
+                v = _content(self._special_map.get(f"{name}_token"))
+            if v is None:
+                v = default
+            return v if (v is None or v in self._vocab) else None
+
+        self.bos_token = pick("bos")
+        self.eos_token = pick("eos")
+        self.unk_token = pick("unk")
+        self.pad_token = pick("pad")
+
+    # ------------------------------------------------------------------ config
+
+    @staticmethod
+    def _validate_pipeline(tj: dict, cfg: dict) -> None:
+        model = tj["model"]
+        _expect(model.get("type") == "BPE", f"model.type={model.get('type')!r}")
+        _expect(not model.get("dropout"), "BPE dropout")
+        _expect(not model.get("continuing_subword_prefix"), "continuing_subword_prefix")
+        _expect(not model.get("end_of_word_suffix"), "end_of_word_suffix")
+        _expect(not model.get("byte_fallback"), "byte_fallback")
+        _expect(not model.get("ignore_merges"), "ignore_merges")
+        _expect(not model.get("fuse_unk"), "fuse_unk")
+        _expect(model.get("unk_token") is None, "an unk_token")
+        _expect(tj.get("normalizer") == {"type": "NFC"}, f"normalizer {tj.get('normalizer')!r} (only NFC)")
+        pt = tj.get("pre_tokenizer") or {}
+        _expect(pt.get("type") == "Sequence", f"pre_tokenizer {pt.get('type')!r}")
+        seq = pt.get("pretokenizers") or []
+        _expect(len(seq) == 2, "pre_tokenizer Sequence other than [Split, ByteLevel]")
+        sp, bl = seq
+        _expect(sp.get("type") == "Split", f"pre_tokenizer[0] {sp.get('type')!r}")
+        _expect(
+            (sp.get("pattern") or {}).get("Regex") == QWEN2_PRETOKENIZE_REGEX,
+            "a Split regex other than the Qwen2 one",
+        )
+        _expect(sp.get("behavior") == "Isolated", f"Split behavior {sp.get('behavior')!r}")
+        _expect(sp.get("invert") is False, "Split invert=true")
+        _expect(bl.get("type") == "ByteLevel", f"pre_tokenizer[1] {bl.get('type')!r}")
+        _expect(not bl.get("add_prefix_space"), "ByteLevel add_prefix_space=true")
+        _expect(bl.get("use_regex") is False, "ByteLevel use_regex=true (GPT-2 regex on top)")
+        dec = tj.get("decoder") or {}
+        _expect(dec.get("type") == "ByteLevel", f"decoder {dec.get('type')!r}")
+        pp = tj.get("post_processor")
+        if pp is not None:
+            _expect(
+                pp.get("type") == "ByteLevel" and not pp.get("add_prefix_space") and not pp.get("trim_offsets"),
+                f"post_processor {pp.get('type')!r} other than a no-op ByteLevel",
+            )
+        # tokenizer.json can carry truncation/padding; the Rust encoder applies them
+        # (silently clipping every encode). Neither is implemented.
+        _expect(tj.get("truncation") is None, "a tokenizer.json truncation setting")
+        _expect(tj.get("padding") is None, "a tokenizer.json padding setting")
+        ByteLevelTokenizer._check_added_tokens_decoder(tj, cfg)
+        # transformers applies these from the CONFIG on top of tokenizer.json.
+        _expect(not cfg.get("add_prefix_space"), "tokenizer_config.json add_prefix_space=true")
+        _expect(not cfg.get("add_bos_token"), "tokenizer_config.json add_bos_token=true")
+        _expect(
+            not cfg.get("clean_up_tokenization_spaces"),
+            "tokenizer_config.json clean_up_tokenization_spaces=true",
+        )
+        _expect(not cfg.get("split_special_tokens"), "tokenizer_config.json split_special_tokens=true")
+        _expect(cfg.get("errors", "replace") == "replace", f"decode errors={cfg.get('errors')!r}")
+
+    @staticmethod
+    def _check_added_tokens_decoder(tj: dict, cfg: dict) -> None:
+        """tokenizer_config.json `added_tokens_decoder` must AGREE with tokenizer.json's
+        added_tokens (same ids, contents and lstrip/rstrip/single_word/normalized/special).
+
+        Incident this prevents: transformers builds the added-token matcher from the
+        CONFIG as well as tokenizer.json. Measured (transformers 5.19.0, Goedel snapshot
+        with only `<think>.lstrip=true` flipped in tokenizer_config.json): HF encodes
+        "a <think> b" as [64, 151667, 293] (the space before the token is eaten) while a
+        reader of tokenizer.json alone gives [64, 220, 151667, 293]. A config/json pair
+        that disagrees must not load silently."""
+        adt = cfg.get("added_tokens_decoder")
+        if adt is None:
+            return
+        declared = {int(k): v for k, v in adt.items()}
+        in_json = {int(t["id"]): t for t in tj.get("added_tokens", [])}
+        _expect(
+            set(declared) == set(in_json),
+            "tokenizer_config.json added_tokens_decoder ids differ from tokenizer.json added_tokens",
+        )
+        for i, t in in_json.items():
+            c = declared[i]
+            _expect(c.get("content") == t["content"], f"added token {i} content differs between config and json")
+            for flag in ("special", "normalized", "lstrip", "rstrip", "single_word"):
+                _expect(
+                    bool(c.get(flag)) == bool(t.get(flag)),
+                    f"added token {t['content']!r}: {flag} differs between tokenizer_config.json and tokenizer.json",
+                )
+
+    # ------------------------------------------------------------- vocabulary
+
+    @property
+    def bos_token_id(self) -> Optional[int]:
+        return None if self.bos_token is None else self._vocab[self.bos_token]
+
+    @property
+    def eos_token_id(self) -> Optional[int]:
+        return None if self.eos_token is None else self._vocab[self.eos_token]
+
+    @property
+    def pad_token_id(self) -> Optional[int]:
+        return None if self.pad_token is None else self._vocab[self.pad_token]
+
+    @property
+    def unk_token_id(self) -> Optional[int]:
+        return None
+
+    @property
+    def vocab_size(self) -> int:
+        """BPE base vocabulary only (HF semantics): 151643 for Qwen3; len(tok) adds
+        the added tokens. The model's embedding rows (151936) are padding on top."""
+        return len(self._base_vocab)
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return sorted(self._special_contents, key=lambda t: self._vocab[t])
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return [self._vocab[t] for t in self.all_special_tokens]
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self._vocab)
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return dict(self._added)
+
+    def __len__(self) -> int:
+        return len(self._vocab)
+
+    def __repr__(self) -> str:
+        return f"ByteLevelTokenizer(n={len(self)}, merges={len(self._ranks)}, added={len(self._added)})"
+
+    def _piece(self, i: int) -> Optional[str]:
+        if i < 0:
+            raise ValueError(f"negative token id {i}: a label-mask value leaked into a token-id path")
+        return self._id_to_token[i] if i < len(self._id_to_token) else None
+
+    def convert_ids_to_tokens(self, ids: Union[int, Iterable[int]]):
+        if isinstance(ids, int):
+            return self._piece(ids)
+        return [self._piece(int(i)) for i in _as_list(ids)]
+
+    def convert_tokens_to_ids(self, tokens: Union[str, Iterable[str]]):
+        if isinstance(tokens, str):
+            return self._vocab.get(tokens)
+        return [self._vocab.get(t) for t in tokens]
+
+    # ------------------------------------------------------------------ encode
+
+    def __call__(self, text: Union[str, Sequence[str]], add_special_tokens: bool = True) -> dict:
+        if isinstance(text, str):
+            return {"input_ids": self.encode(text, add_special_tokens=add_special_tokens)}
+        return {"input_ids": [self.encode(t, add_special_tokens=add_special_tokens) for t in text]}
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
+        """`add_special_tokens` is accepted for API parity; Qwen3 adds none."""
+        if not isinstance(text, str):
+            raise TypeError(f"encode() takes str, got {type(text).__name__}")
+        try:
+            text.encode("utf-8")
+        except UnicodeEncodeError as e:
+            raise ValueError(f"text is not valid Unicode (lone surrogate at index {e.start})") from e
+        ids: list[int] = []
+        for is_added, seg in self._split_added(text):
+            if is_added:
+                ids.append(self._added[seg])
+            else:
+                ids.extend(self._encode_segment(self._normalize(seg)))
+        return ids
+
+    @staticmethod
+    def _normalize(seg: str) -> str:
+        """NFC as the Rust normalizer computes it. ASCII is NFC-invariant. A code point in
+        NFC_BARRIERS is a starter there (never reordered, never composed) but not in
+        Python's newer table, so the text is normalized AROUND it, never across it.
+
+        Incident this prevents: "\\u0301\\U00011d45" (acute, then Masaram Gondi virama) is
+        kept in that order by HF, but unicodedata.normalize moves the virama (ccc 9) in
+        front of the acute (ccc 230): same text, different ids."""
+        if seg.isascii():
+            return seg
+        rx = _barrier_re()
+        if rx.search(seg) is None:
+            return unicodedata.normalize("NFC", seg)
+        parts = rx.split(seg)           # [text, barrier, text, barrier, ..., text]
+        return "".join(p if k & 1 else unicodedata.normalize("NFC", p) for k, p in enumerate(parts))
+
+    def _split_added(self, text: str):
+        if self._added_re is None or not text:
+            if text:
+                yield False, text
+            return
+        pos = 0
+        for m in self._added_re.finditer(text):
+            if m.start() > pos:
+                yield False, text[pos : m.start()]
+            yield True, m.group(0)
+            pos = m.end()
+        if pos < len(text):
+            yield False, text[pos:]
+
+    def _pretokenize(self, seg: str):
+        """Split(Isolated): matches AND the gaps between them become pieces (the
+        regex's last alternatives cover every character, so gaps cannot occur for
+        the Qwen2 pattern; kept so a pattern change cannot silently drop text)."""
+        pos = 0
+        for m in _build_pattern().finditer(seg):
+            if m.start() > pos:
+                yield seg[pos : m.start()]
+            yield m.group(0)
+            pos = m.end()
+        if pos < len(seg):
+            yield seg[pos:]
+
+    def _encode_segment(self, seg: str) -> list[int]:
+        out: list[int] = []
+        cache = self._cache
+        for piece in self._pretokenize(seg):
+            r = cache.get(piece)
+            if r is None:
+                r = self._bpe(piece)
+                if len(cache) >= _CACHE_MAX_ENTRIES:
+                    cache.clear()
+                cache[piece] = r
+            out.extend(r)
+        return out
+
+    def _bpe(self, piece: str) -> tuple[int, ...]:
+        b2u = self._b2u
+        syms = [b2u[b] for b in piece.encode("utf-8")]
+        return self._merge(syms)
+
+    def _merge(self, syms: list[str]) -> tuple[int, ...]:
+        """Rank-ordered pair merging (linked list + lazy-deletion min-heap, O(n log n)
+        so a multi-thousand-char whitespace run is not quadratic). Tie-break:
+        lowest merge rank, then leftmost position."""
+        n = len(syms)
+        vocab = self._base_vocab
+        if n < 2:
+            return tuple(vocab[s] for s in syms)
+        ranks = self._ranks
+        nxt = list(range(1, n + 1))
+        nxt[-1] = -1
+        prv = list(range(-1, n - 1))
+        heap: list[tuple[int, int, str, str]] = []
+        for i in range(n - 1):
+            r = ranks.get((syms[i], syms[i + 1]))
+            if r is not None:
+                heap.append((r, i, syms[i], syms[i + 1]))
+        heapq.heapify(heap)
+        pop, push = heapq.heappop, heapq.heappush
+        while heap:
+            _, i, left, right = pop(heap)
+            if syms[i] != left:
+                continue
+            j = nxt[i]
+            if j == -1 or syms[j] != right:
+                continue
+            new = left + right
+            syms[i] = new
+            syms[j] = None  # type: ignore[call-overload]
+            nj = nxt[j]
+            nxt[i] = nj
+            if nj != -1:
+                prv[nj] = i
+            p = prv[i]
+            if p != -1:
+                r = ranks.get((syms[p], new))
+                if r is not None:
+                    push(heap, (r, p, syms[p], new))
+            if nj != -1:
+                r = ranks.get((new, syms[nj]))
+                if r is not None:
+                    push(heap, (r, i, new, syms[nj]))
+        return tuple(vocab[s] for s in syms if s is not None)
+
+    # ------------------------------------------------------------------ decode
+
+    def _bytes_of(self, i: int) -> Optional[bytes]:
+        b = self._id_bytes[i]
+        if b is None:
+            t = self._id_to_token[i]
+            u2b = self._u2b
+            # Same rule as the Rust ByteLevel decoder: a token whose every char is in
+            # the alphabet is byte-decoded; anything else passes through as literal text.
+            if all(c in u2b for c in t):  # type: ignore[union-attr]
+                b = bytes(u2b[c] for c in t)  # type: ignore[union-attr]
+            else:
+                b = t.encode("utf-8")  # type: ignore[union-attr]
+            self._id_bytes[i] = b
+        return b
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:
+        """Join all tokens' bytes THEN decode UTF-8 once: a multi-byte character split
+        across tokens decodes correctly, a truncated one yields U+FFFD (never an
+        exception). Ids outside the vocab are skipped (HF behaviour). With
+        skip_special_tokens only added tokens flagged special=true are dropped
+        (<think>/</think>/<tool_call> are special=false and stay)."""
+        special = self._special_contents
+        buf = bytearray()
+        for i in _as_list(ids):
+            i = int(i)
+            t = self._piece(i)
+            if t is None:
+                continue
+            if skip_special_tokens and t in special:
+                continue
+            buf += self._bytes_of(i)  # type: ignore[arg-type]
+        return bytes(buf).decode("utf-8", errors="replace")
+
+    # ------------------------------------------------------------ chat template
+
+    def apply_chat_template(
+        self,
+        messages: Sequence[dict],
+        *,
+        tokenize: bool = False,
+        add_generation_prompt: bool = False,
+        enable_thinking: Optional[bool] = None,
+        tools=None,
+    ):
+        """Qwen3 chat template, the system/user/assistant subset. See render_qwen3_chat.
+        Raises NotImplementedError if the shipped chat_template is not the one this
+        implements (checked by sha256) or if tools are requested.
+
+        `tokenize=True` returns a plain list[int]. transformers 5.x returns a BatchEncoding
+        ({"input_ids", "attention_mask"}) there (measured, 5.19.0), so code written against
+        HF's tokenize=True result must index ["input_ids"] on HF but not here; callers in
+        this repo use tokenize=False + encode()."""
+        tpl = self._config.get("chat_template")
+        _expect(
+            tpl is not None and hashlib.sha256(tpl.encode("utf-8")).hexdigest() == QWEN3_CHAT_TEMPLATE_SHA256,
+            "a chat_template other than the Qwen3 one (sha256 mismatch)",
+        )
+        text = render_qwen3_chat(
+            messages, add_generation_prompt=add_generation_prompt, enable_thinking=enable_thinking, tools=tools
+        )
+        return self.encode(text, add_special_tokens=False) if tokenize else text
+
+    # -------------------------------------------------------------------- save
+
+    def save_pretrained(self, save_directory: str) -> tuple[str, ...]:
+        """Copy-through of the source files (atomic tmp + os.replace per file)."""
+        os.makedirs(save_directory, exist_ok=True)
+        files = dict(self._files)
+        if "tokenizer.json" not in files:
+            files["tokenizer.json"] = json.dumps(self._tokenizer_json, ensure_ascii=False).encode("utf-8")
+        if "tokenizer_config.json" not in files and self._config:
+            files["tokenizer_config.json"] = json.dumps(self._config, indent=2).encode("utf-8")
+        if "special_tokens_map.json" not in files and self._special_map:
+            files["special_tokens_map.json"] = json.dumps(self._special_map, indent=2).encode("utf-8")
+        written = []
+        for name, blob in files.items():
+            dst = os.path.join(save_directory, name)
+            tmp = dst + f".tmp.{os.getpid()}"
+            with open(tmp, "wb") as f:
+                f.write(blob)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp, dst)
+            written.append(dst)
+        return tuple(written)
+
+    @classmethod
+    def from_pretrained(cls, path: Union[str, os.PathLike]) -> "ByteLevelTokenizer":
+        """LOCAL ONLY: a dir holding tokenizer.json, a tokenizer.json path, or an HF
+        cache model dir. Never touches the network."""
+        d = _resolve_dir(os.fspath(path))
+        files: dict[str, bytes] = {}
+        for name in _TOKENIZER_FILES:
+            p = os.path.join(d, name)
+            if os.path.isfile(p):
+                with open(p, "rb") as f:
+                    files[name] = f.read()
+        tj = json.loads(files["tokenizer.json"])
+        cfg = json.loads(files["tokenizer_config.json"]) if "tokenizer_config.json" in files else None
+        smap = json.loads(files["special_tokens_map.json"]) if "special_tokens_map.json" in files else None
+        return cls(tj, cfg, smap, files=files)
+
+
+def render_qwen3_chat(
+    messages: Sequence[dict],
+    *,
+    add_generation_prompt: bool = False,
+    enable_thinking: Optional[bool] = None,
+    tools=None,
+) -> str:
+    """Render the Qwen3 chat template for system/user/assistant messages.
+
+    Incident this prevents: P2 trains on (prompt, completion) pairs whose prompt
+    bytes must equal what Goedel-Prover sampled under; correction rounds replay
+    earlier assistant turns, and the template REWRITES them (earlier turns lose
+    their <think> block; only turns after the last user query keep
+    "<think>\\n..\\n</think>\\n\\n"). Hand-rolled "<|im_start|>" concatenation
+    gets the multi-turn case wrong without any error. Verified against the HF
+    template output in tests/test_tokenizer_bytelevel.py.
+
+    Deliberate deviations from the Jinja (each raises instead of silently differing):
+    tools / role "tool" / tool_calls -> NotImplementedError; a role outside
+    system/user/assistant -> ValueError (Jinja drops the message); non-str content
+    -> TypeError (Jinja renders ''); empty messages -> ValueError.
+    """
+    if tools:
+        raise NotImplementedError("Qwen3 chat template with tools is not implemented")
+    if not messages:
+        raise ValueError("messages is empty")
+    for k, m in enumerate(messages):
+        role = m.get("role")
+        if role == "tool" or m.get("tool_calls"):
+            raise NotImplementedError("tool messages / tool_calls are not implemented")
+        if role not in ("system", "user", "assistant"):
+            raise ValueError(f"message {k}: unsupported role {role!r}")
+        if not isinstance(m.get("content"), str):
+            raise TypeError(f"message {k}: content must be str, got {type(m.get('content')).__name__}")
+    out: list[str] = []
+    if messages[0]["role"] == "system":
+        out.append("<|im_start|>system\n" + messages[0]["content"] + "<|im_end|>\n")
+    last_query_index = len(messages) - 1
+    for idx in range(len(messages) - 1, -1, -1):
+        c = messages[idx]["content"]
+        if messages[idx]["role"] == "user" and not (c.startswith("<tool_response>") and c.endswith("</tool_response>")):
+            last_query_index = idx
+            break
+    n = len(messages)
+    for idx, m in enumerate(messages):
+        role, content = m["role"], m["content"]
+        if role == "user" or (role == "system" and idx != 0):
+            out.append("<|im_start|>" + role + "\n" + content + "<|im_end|>\n")
+        elif role == "assistant":
+            reasoning = ""
+            rc = m.get("reasoning_content")
+            if isinstance(rc, str):
+                reasoning = rc
+            elif "</think>" in content:
+                reasoning = content.split("</think>")[0].rstrip("\n").split("<think>")[-1].lstrip("\n")
+                content = content.split("</think>")[-1].lstrip("\n")
+            is_last = idx == n - 1
+            if idx > last_query_index and (is_last or reasoning):
+                out.append(
+                    "<|im_start|>assistant\n<think>\n" + reasoning.strip("\n") + "\n</think>\n\n" + content.lstrip("\n")
+                )
+            else:
+                out.append("<|im_start|>assistant\n" + content)
+            out.append("<|im_end|>\n")
+    if add_generation_prompt:
+        out.append("<|im_start|>assistant\n")
+        if enable_thinking is False:
+            out.append("<think>\n\n</think>\n\n")
+    return "".join(out)

code/align/tokenizer/core.py

@@ -0,0 +1,615 @@
+"""Pure-Python (stdlib only) BPE tokenizer for Mistral-7B-v0.3.
+
+This is a from-scratch reimplementation of exactly the pipeline that
+`tokenizers`/`transformers` run for the v0.3 `tokenizer.json`, and nothing
+else. It exists so the training stack has zero dependency on `transformers`
+(see docs/DEPENDENCY_POLICY.md); parity with the legacy stack is certified by
+the `oracle`-marked test in tests/test_tokenizer.py (hundreds of strings,
+ids AND decode byte-identical).
+
+Pipeline (encode):
+
+    text
+      -> split on ADDED tokens (771: <unk> <s> </s> [INST] ... [control_N]),
+         leftmost-longest, raw text, BEFORE any BPE        (_split_added)
+      -> per non-added segment: Metaspace pre-tokenizer     (_metaspace)
+           ' ' -> U+2581, prepend U+2581 per prepend_scheme
+      -> BPE over the WHOLE segment (split=false => ONE "word")  (_bpe)
+           chars not in vocab -> <0xNN> byte-fallback pieces
+      -> post-processor: prepend <s> (id 1) iff add_special_tokens
+
+Decode: ids -> pieces, optionally drop special pieces, then the decoder chain
+Replace(U+2581 -> ' ') / ByteFallback / Fuse / Strip(' ', start=1).
+
+The loader VALIDATES the pipeline configuration and raises
+NotImplementedError on anything this module does not implement. A different
+tokenizer.json must fail loudly, not silently encode wrong ids.
+"""
+from __future__ import annotations
+
+import heapq
+import json
+import os
+import re
+from typing import Iterable, Optional, Sequence, Union
+
+__all__ = ["Tokenizer", "load_tokenizer"]
+
+#: U+2581 LOWER ONE EIGHTH BLOCK, the SentencePiece "space" glyph.
+_SP = "▁"
+_HEX = frozenset("0123456789abcdefABCDEF")
+#: Cap on the per-segment encode cache. With split=false a "word" is a whole
+#: text segment, so the cache only pays off for short, repeated strings
+#: (stop strings, role tags, "\n"); long segments are never cached.
+_CACHE_MAX_ENTRIES = 50_000
+_CACHE_MAX_SEGMENT_CHARS = 128
+
+_TOKENIZER_FILES = ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json")
+#: Optional passthrough: AutoTokenizer on a LlamaTokenizer-class config wants it.
+_OPTIONAL_FILES = ("tokenizer.model",)
+
+
+def _content(tok) -> Optional[str]:
+    """HF configs store a special token as either a bare string or a dict
+    {"content": ...}; None means 'unset'."""
+    if tok is None:
+        return None
+    if isinstance(tok, dict):
+        return tok.get("content")
+    return str(tok)
+
+
+def _expect(cond: bool, what: str) -> None:
+    if not cond:
+        raise NotImplementedError(
+            f"tokenizer.json uses a configuration this pure-Python tokenizer does not "
+            f"implement: {what}. Refusing to load rather than silently encode wrong ids."
+        )
+
+
+class Tokenizer:
+    """Drop-in for the subset of the HF tokenizer API this repo calls.
+
+    Build with `Tokenizer.from_pretrained(dir)` / `load_tokenizer(dir)`.
+    """
+
+    def __init__(
+        self,
+        tokenizer_json: dict,
+        tokenizer_config: Optional[dict] = None,
+        special_tokens_map: Optional[dict] = None,
+        *,
+        files: Optional[dict] = None,
+    ):
+        #: raw bytes of the source files, kept so save_pretrained is a true
+        #: copy-through that survives the source dir being pruned/rotated.
+        self._files: dict[str, bytes] = dict(files or {})
+        self._config = dict(tokenizer_config or {})
+        self._special_map = dict(special_tokens_map or {})
+        self._tokenizer_json = tokenizer_json
+        self._validate_pipeline(tokenizer_json)
+        # transformers REBUILDS the Metaspace pre-tokenizer from the config's
+        # add_prefix_space (measured on 5.15: False changes the ids, while
+        # add_bos_token/add_eos_token are ignored by the tokenizer.json path). The
+        # pure-Python pipeline follows tokenizer.json only, so refuse loudly rather
+        # than silently disagree with the legacy stack. Every v0.3 config has True.
+        _expect(
+            self._config.get("add_prefix_space") is not False,
+            "tokenizer_config.json add_prefix_space=false",
+        )
+
+        model = tokenizer_json["model"]
+        vocab: dict[str, int] = dict(model["vocab"])
+
+        # -- added tokens: override/extend the model vocab, and are matched first.
+        added = list(tokenizer_json.get("added_tokens", []))
+        self._added: dict[str, int] = {}
+        self._special_contents: set[str] = set()
+        for t in added:
+            _expect(
+                not (t.get("lstrip") or t.get("rstrip") or t.get("single_word")),
+                f"added token {t['content']!r} with lstrip/rstrip/single_word",
+            )
+            vocab[t["content"]] = int(t["id"])
+            self._added[t["content"]] = int(t["id"])
+            if t.get("special"):
+                self._special_contents.add(t["content"])
+        self._vocab = vocab
+
+        size = max(vocab.values()) + 1
+        self._id_to_token: list[Optional[str]] = [None] * size
+        for tok, i in vocab.items():
+            self._id_to_token[i] = tok
+        self._base_vocab = model["vocab"]  # the BPE model's own vocab (ids for symbols)
+
+        # -- BPE tables. Merges may be [a, b] pairs (older saves) or "a b" strings.
+        self._ranks: dict[tuple[str, str], int] = {}
+        for rank, m in enumerate(model["merges"]):
+            a, b = (m[0], m[1]) if isinstance(m, (list, tuple)) else m.split(" ")
+            # Rust's BPE builder errors if the merge result is not in the vocab.
+            if (a + b) not in self._base_vocab or a not in self._base_vocab or b not in self._base_vocab:
+                raise ValueError(f"merge #{rank} {a!r}+{b!r} references a token outside the vocab")
+            # LAST duplicate wins: Rust collects merges into a HashMap, so a repeated
+            # pair keeps its final rank (measured against tokenizers 0.22). v0.3 has none.
+            self._ranks[(a, b)] = rank
+        self._byte_fallback = bool(model.get("byte_fallback"))
+        self._byte_pieces: list[Optional[str]] = [
+            (f"<0x{b:02X}>" if f"<0x{b:02X}>" in self._base_vocab else None) for b in range(256)
+        ]
+        self._fuse_unk = bool(model.get("fuse_unk"))
+        self._model_unk = model.get("unk_token")
+        self._ignore_merges = bool(model.get("ignore_merges"))
+        self._cache: dict[str, tuple[int, ...]] = {}
+
+        # -- added-token splitter. Alternation sorted longest-first == leftmost-longest.
+        if self._added:
+            alts = sorted(self._added, key=len, reverse=True)
+            self._added_re: Optional[re.Pattern] = re.compile("|".join(re.escape(a) for a in alts))
+        else:
+            self._added_re = None
+
+        # -- Metaspace
+        pt = tokenizer_json["pre_tokenizer"]
+        self._prepend_scheme = pt.get("prepend_scheme", "always" if pt.get("add_prefix_space", True) else "never")
+
+        # -- special tokens (strings), resolved config -> special_tokens_map -> defaults
+        def pick(name: str, default: Optional[str] = None) -> Optional[str]:
+            v = _content(self._config.get(f"{name}_token"))
+            if v is None:
+                v = _content(self._special_map.get(f"{name}_token"))
+            if v is None:
+                v = default
+            return v if (v is None or v in self._vocab) else None
+
+        self.bos_token = pick("bos", "<s>")
+        self.eos_token = pick("eos", "</s>")
+        self.unk_token = pick("unk", self._model_unk)
+        self.pad_token = pick("pad")  # v0.3 ships none; callers fall back to eos
+        self._post_bos = self._post_processor_bos(tokenizer_json.get("post_processor"))
+
+    # ------------------------------------------------------------------ config
+
+    @staticmethod
+    def _validate_pipeline(tj: dict) -> None:
+        model = tj["model"]
+        _expect(model.get("type") == "BPE", f"model.type={model.get('type')!r}")
+        _expect(not model.get("dropout"), "BPE dropout")
+        _expect(not model.get("continuing_subword_prefix"), "continuing_subword_prefix")
+        _expect(not model.get("end_of_word_suffix"), "end_of_word_suffix")
+        _expect(tj.get("normalizer") is None, "a non-null normalizer")
+        pt = tj.get("pre_tokenizer") or {}
+        _expect(pt.get("type") == "Metaspace", f"pre_tokenizer {pt.get('type')!r}")
+        _expect(pt.get("replacement") == _SP, "Metaspace replacement other than U+2581")
+        _expect(pt.get("split") is False, "Metaspace split=true")
+        _expect(
+            pt.get("prepend_scheme", "always") in ("first", "always", "never"),
+            f"Metaspace prepend_scheme {pt.get('prepend_scheme')!r}",
+        )
+        dec = tj.get("decoder") or {}
+        chain = dec.get("decoders") if dec.get("type") == "Sequence" else None
+        _expect(chain is not None and len(chain) == 4, "decoder other than the 4-step Sequence")
+        r, bf, fu, st = chain
+        _expect(
+            r.get("type") == "Replace"
+            and r.get("pattern") == {"String": _SP}
+            and r.get("content") == " ",
+            "decoder[0] other than Replace(U+2581 -> ' ')",
+        )
+        _expect(bf.get("type") == "ByteFallback", "decoder[1] other than ByteFallback")
+        _expect(fu.get("type") == "Fuse", "decoder[2] other than Fuse")
+        _expect(
+            st.get("type") == "Strip"
+            and st.get("content") == " "
+            and st.get("start") == 1
+            and st.get("stop") == 0,
+            "decoder[3] other than Strip(' ', start=1, stop=0)",
+        )
+
+    @staticmethod
+    def _post_processor_bos(pp: Optional[dict]) -> Optional[str]:
+        """The single-sequence template must be exactly [<special>, A]; returns
+        that special token's content (the BOS), or None if there is no post-processor."""
+        if pp is None:
+            return None
+        _expect(pp.get("type") == "TemplateProcessing", f"post_processor {pp.get('type')!r}")
+        single = pp.get("single", [])
+        _expect(
+            len(single) == 2 and "SpecialToken" in single[0] and "Sequence" in single[1],
+            "post_processor single template other than [SpecialToken, Sequence]",
+        )
+        return single[0]["SpecialToken"]["id"]
+
+    # --------------------------------------------------------------- id access
+
+    @property
+    def bos_token_id(self) -> Optional[int]:
+        return self._vocab.get(self.bos_token) if self.bos_token is not None else None
+
+    @property
+    def eos_token_id(self) -> Optional[int]:
+        return self._vocab.get(self.eos_token) if self.eos_token is not None else None
+
+    @property
+    def pad_token_id(self) -> Optional[int]:
+        """None on v0.3 -- callers (align.data.encoding.pad_token_id) fall back to eos."""
+        return self._vocab.get(self.pad_token) if self.pad_token is not None else None
+
+    @property
+    def unk_token_id(self) -> Optional[int]:
+        return self._vocab.get(self.unk_token) if self.unk_token is not None else None
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._base_vocab)
+
+    @property
+    def all_special_tokens(self) -> list[str]:
+        return sorted(self._special_contents, key=lambda t: self._vocab[t])
+
+    @property
+    def all_special_ids(self) -> list[int]:
+        return [self._vocab[t] for t in self.all_special_tokens]
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self._vocab)
+
+    def get_added_vocab(self) -> dict[str, int]:
+        return dict(self._added)
+
+    def __len__(self) -> int:
+        return len(self._vocab)
+
+    def __repr__(self) -> str:
+        return f"Tokenizer(n={len(self)}, merges={len(self._ranks)}, added={len(self._added)})"
+
+    def convert_ids_to_tokens(self, ids: Union[int, Iterable[int]]):
+        """Pieces for ids; None for ids outside the vocab (HF returns None too)."""
+        if isinstance(ids, int):
+            return self._piece(ids)
+        return [self._piece(int(i)) for i in _as_list(ids)]
+
+    def convert_tokens_to_ids(self, tokens: Union[str, Iterable[str]]):
+        """Ids for pieces; unknown pieces map to the unk id (HF behaviour)."""
+        unk = self.unk_token_id
+        if isinstance(tokens, str):
+            return self._vocab.get(tokens, unk)
+        return [self._vocab.get(t, unk) for t in tokens]
+
+    def _piece(self, i: int) -> Optional[str]:
+        if i < 0:
+            raise ValueError(f"negative token id {i}: a label-mask value leaked into a token-id path")
+        return self._id_to_token[i] if i < len(self._id_to_token) else None
+
+    # ------------------------------------------------------------------ encode
+
+    def __call__(self, text: Union[str, Sequence[str]], add_special_tokens: bool = True) -> dict:
+        """{"input_ids": ...}. A list of strings gives a list of id lists.
+        Padding/truncation/tensors are deliberately NOT supported (TypeError on
+        unknown kwargs) -- callers pad their own batches."""
+        if isinstance(text, str):
+            return {"input_ids": self.encode(text, add_special_tokens=add_special_tokens)}
+        return {"input_ids": [self.encode(t, add_special_tokens=add_special_tokens) for t in text]}
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
+        if not isinstance(text, str):
+            raise TypeError(f"encode() takes str, got {type(text).__name__}")
+        try:
+            text.encode("utf-8")
+        except UnicodeEncodeError as e:  # lone surrogate: Rust strings cannot hold one
+            raise ValueError(f"text is not valid Unicode (lone surrogate at index {e.start})") from e
+        ids: list[int] = []
+        if add_special_tokens and self._post_bos is not None:
+            ids.append(self._vocab[self._post_bos])
+        for k, (is_added, seg) in enumerate(self._split_added(text)):
+            if is_added:
+                ids.append(self._added[seg])
+                continue
+            piece = self._metaspace(seg, k == 0)
+            if piece:
+                ids.extend(self._bpe(piece))
+        return ids
+
+    def _split_added(self, text: str):
+        """Yield (is_added, substring). Added tokens are matched on the RAW text,
+        leftmost-longest, and are never fed to BPE (so nothing merges across them).
+        No empty segments are produced (matches HF: "" encodes to no tokens)."""
+        if self._added_re is None or not text:
+            if text:
+                yield False, text
+            return
+        pos = 0
+        for m in self._added_re.finditer(text):
+            if m.start() > pos:
+                yield False, text[pos : m.start()]
+            yield True, m.group(0)
+            pos = m.end()
+        if pos < len(text):
+            yield False, text[pos:]
+
+    def _metaspace(self, seg: str, is_first_split: bool) -> str:
+        """Metaspace(replacement=U+2581, prepend_scheme=..., split=false).
+
+        ' ' -> U+2581 (ONLY the space char; '\\n' and '\\t' stay literal and later
+        byte-fall-back). prepend_scheme="first" prepends U+2581 only when this
+        segment is split #0 of the WHOLE input, i.e. the text does not begin with
+        an added token. Added-token splits are counted (they occupy an index)
+        but skip the pre-tokenizer, so a leading special token CONSUMES the
+        'first' slot and nothing after it is ever prefixed:
+            "hello<s>world"  -> ▁hello <s> world
+            "<s>hello"       -> <s> <s> hello      (bare 'hello', NO ▁ -- measured
+                                                    against the legacy tokenizer)
+        A segment that already starts with U+2581 (leading space / literal U+2581)
+        gets no second prefix. 'always' prefixes every such segment; 'never' none."""
+        s = seg.replace(" ", _SP)
+        scheme = self._prepend_scheme
+        if scheme == "never" or s.startswith(_SP):
+            return s
+        if scheme == "always" or is_first_split:
+            return _SP + s
+        return s
+
+    def _bpe(self, word: str) -> tuple[int, ...]:
+        cached = self._cache.get(word)
+        if cached is not None:
+            return cached
+        vocab = self._base_vocab
+        if self._ignore_merges and word in vocab:
+            res = (vocab[word],)
+        else:
+            res = self._merge(self._initial_symbols(word))
+        if len(word) <= _CACHE_MAX_SEGMENT_CHARS:
+            if len(self._cache) >= _CACHE_MAX_ENTRIES:
+                self._cache.clear()
+            self._cache[word] = res
+        return res
+
+    def _initial_symbols(self, word: str) -> list[str]:
+        """One symbol per char if the char is a vocab piece; else its UTF-8 bytes
+        as <0xNN> pieces (byte_fallback); else <unk> (consecutive unks fuse).
+
+        Mirrors the Rust `BPE::merge_word` loop INCLUDING a quirk: an unknown char
+        only records a PENDING unk, and the pending unk is flushed solely by the
+        next in-vocab char (or end of word). A byte-fallback expansion in between
+        neither flushes nor resets it, so the unk lands AFTER those byte pieces
+        (and, with fuse_unk, unks on both sides of a byte-fallback char fuse into
+        one). Unreachable for v0.3 (all 256 <0xNN> pieces exist, so the unk path
+        never fires), but the loader accepts other byte_fallback tables."""
+        vocab = self._base_vocab
+        syms: list[str] = []
+        unk = self._model_unk
+        pending_unk = False
+        for ch in word:
+            if ch in vocab:
+                if pending_unk:
+                    syms.append(unk)  # type: ignore[arg-type]
+                    pending_unk = False
+                syms.append(ch)
+                continue
+            if self._byte_fallback:
+                pieces = [self._byte_pieces[b] for b in ch.encode("utf-8")]
+                if None not in pieces:
+                    syms.extend(pieces)  # type: ignore[arg-type]
+                    continue  # pending unk is deliberately left untouched (Rust parity)
+            if unk is not None:
+                if pending_unk and not self._fuse_unk:
+                    syms.append(unk)
+                pending_unk = True
+        if pending_unk:
+            syms.append(unk)  # type: ignore[arg-type]
+        return syms
+
+    def _merge(self, syms: list[str]) -> tuple[int, ...]:
+        """Rank-ordered pair merging over a doubly linked list with a lazy-deletion
+        min-heap: O(n log n), so a multi-thousand-char segment (split=false makes
+        every text segment ONE word) is not quadratic. Tie-break is lowest rank,
+        then leftmost position -- identical to the Rust implementation."""
+        n = len(syms)
+        vocab = self._base_vocab
+        if n < 2:
+            return tuple(vocab[s] for s in syms)
+        ranks = self._ranks
+        nxt = list(range(1, n + 1))
+        nxt[-1] = -1
+        prv = list(range(-1, n - 1))
+        heap: list[tuple[int, int, str, str]] = []
+        for i in range(n - 1):
+            r = ranks.get((syms[i], syms[i + 1]))
+            if r is not None:
+                heap.append((r, i, syms[i], syms[i + 1]))
+        heapq.heapify(heap)
+        pop, push = heapq.heappop, heapq.heappush
+        while heap:
+            _, i, left, right = pop(heap)
+            if syms[i] != left:  # stale: left side was merged or absorbed (None)
+                continue
+            j = nxt[i]
+            if j == -1 or syms[j] != right:  # stale: right neighbour changed
+                continue
+            new = left + right
+            syms[i] = new
+            syms[j] = None  # type: ignore[call-overload]
+            nj = nxt[j]
+            nxt[i] = nj
+            if nj != -1:
+                prv[nj] = i
+            p = prv[i]
+            if p != -1:
+                r = ranks.get((syms[p], new))
+                if r is not None:
+                    push(heap, (r, p, syms[p], new))
+            if nj != -1:
+                r = ranks.get((new, syms[nj]))
+                if r is not None:
+                    push(heap, (r, i, new, syms[nj]))
+        out = []
+        i = 0
+        while i != -1:
+            out.append(vocab[syms[i]])
+            i = nxt[i]
+        return tuple(out)
+
+    # ------------------------------------------------------------------ decode
+
+    def decode(self, ids, skip_special_tokens: bool = False) -> str:
+        """ids -> text. Default skip_special_tokens=False matches HF (callers that
+        want specials dropped, e.g. the inference engine, pass True explicitly).
+
+        Ids outside the vocab are silently dropped (HF/Rust behaviour); negative
+        ids raise (they are masks that leaked, never tokens).
+
+        Decoder chain, applied to the list of pieces:
+          Replace(U+2581 -> ' ') per piece;
+          ByteFallback: each maximal run of <0xNN> pieces is UTF-8 decoded as ONE
+            unit; an INVALID run becomes one U+FFFD **per byte** (Rust semantics,
+            not Python errors='replace', which collapses ill-formed subparts);
+          Fuse: concatenate;
+          Strip(' ', start=1): drop at most ONE leading space of the fused text.
+        The Strip is why decode([id_of('▁hello')]) == 'hello', not ' hello': a
+        chunk decoded in isolation loses its leading space. Never decode token
+        deltas and concatenate -- decode the whole sequence (the engine does).
+        """
+        pieces: list[str] = []
+        special = self._special_contents
+        for i in _as_list(ids):
+            t = self._piece(int(i))
+            if t is None:
+                continue
+            if skip_special_tokens and t in special:
+                continue
+            pieces.append(t)
+        return self._decode_pieces(pieces)
+
+    @staticmethod
+    def _decode_pieces(pieces: list[str]) -> str:
+        out: list[str] = []
+        run = bytearray()
+
+        def flush() -> None:
+            if run:
+                try:
+                    out.append(run.decode("utf-8"))
+                except UnicodeDecodeError:
+                    out.append("�" * len(run))
+                run.clear()
+
+        for p in pieces:
+            p = p.replace(_SP, " ")
+            if len(p) == 6 and p.startswith("<0x") and p.endswith(">") and p[3] in _HEX and p[4] in _HEX:
+                run.append(int(p[3:5], 16))
+            else:
+                flush()
+                out.append(p)
+        flush()
+        text = "".join(out)  # Fuse
+        return text[1:] if text.startswith(" ") else text  # Strip(' ', start=1)
+
+    # -------------------------------------------------------------------- save
+
+    def save_pretrained(self, save_directory: str) -> tuple[str, ...]:
+        """Write tokenizer.json, tokenizer_config.json, special_tokens_map.json
+        (+ tokenizer.model if the source had one). Copy-through of the source
+        bytes when loaded from a dir; synthesized (HF-shaped) otherwise. Each file
+        is written to a sibling temp name then os.replace()d, so a kill mid-save
+        leaves the old file or the new file, never a torn one."""
+        os.makedirs(save_directory, exist_ok=True)
+        files = dict(self._files)
+        if "tokenizer.json" not in files:
+            files["tokenizer.json"] = json.dumps(self._tokenizer_json, ensure_ascii=False).encode("utf-8")
+        if "tokenizer_config.json" not in files:
+            files["tokenizer_config.json"] = json.dumps(self._synth_config(), indent=2).encode("utf-8")
+        if "special_tokens_map.json" not in files:
+            files["special_tokens_map.json"] = json.dumps(self._synth_special_map(), indent=2).encode("utf-8")
+        written = []
+        for name, blob in files.items():
+            dst = os.path.join(save_directory, name)
+            tmp = dst + f".tmp.{os.getpid()}"
+            with open(tmp, "wb") as f:
+                f.write(blob)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp, dst)
+            written.append(dst)
+        return tuple(written)
+
+    def _synth_config(self) -> dict:
+        return {
+            "add_bos_token": self._post_bos is not None,
+            "add_eos_token": False,
+            "bos_token": self.bos_token,
+            "eos_token": self.eos_token,
+            "unk_token": self.unk_token,
+            "pad_token": self.pad_token,
+            "clean_up_tokenization_spaces": False,
+            "legacy": False,
+            "model_max_length": int(1e30),
+            "tokenizer_class": "LlamaTokenizer",
+        }
+
+    def _synth_special_map(self) -> dict:
+        d = {}
+        for name in ("bos", "eos", "unk", "pad"):
+            tok = getattr(self, f"{name}_token")
+            if tok is not None:
+                d[f"{name}_token"] = {
+                    "content": tok,
+                    "lstrip": False,
+                    "normalized": False,
+                    "rstrip": False,
+                    "single_word": False,
+                }
+        return d
+
+    # -------------------------------------------------------------------- load
+
+    @classmethod
+    def from_pretrained(cls, path: Union[str, os.PathLike]) -> "Tokenizer":
+        """Load from a directory holding tokenizer.json (a checkpoint dir, or an HF
+        snapshot dir), a path to a tokenizer.json, or an HF cache model dir
+        (`models--org--name`, newest snapshot with a tokenizer.json). LOCAL ONLY:
+        never touches the network, never downloads."""
+        d = _resolve_dir(os.fspath(path))
+        files: dict[str, bytes] = {}
+        for name in _TOKENIZER_FILES + _OPTIONAL_FILES:
+            p = os.path.join(d, name)
+            if os.path.isfile(p):
+                with open(p, "rb") as f:
+                    files[name] = f.read()
+        tj = json.loads(files["tokenizer.json"])
+        cfg = json.loads(files["tokenizer_config.json"]) if "tokenizer_config.json" in files else None
+        smap = json.loads(files["special_tokens_map.json"]) if "special_tokens_map.json" in files else None
+        return cls(tj, cfg, smap, files=files)
+
+
+def _as_list(ids) -> list:
+    """Accept list/tuple/iterable, numpy arrays and torch tensors (via .tolist())."""
+    if hasattr(ids, "tolist"):
+        ids = ids.tolist()
+    if isinstance(ids, int):
+        return [ids]
+    return list(ids)
+
+
+def _resolve_dir(path: str) -> str:
+    if os.path.isfile(path) and os.path.basename(path) == "tokenizer.json":
+        return os.path.dirname(os.path.abspath(path))
+    if os.path.isfile(os.path.join(path, "tokenizer.json")):
+        return path
+    snaps = os.path.join(path, "snapshots")
+    if os.path.isdir(snaps):
+        cands = [
+            os.path.join(snaps, s)
+            for s in os.listdir(snaps)
+            if os.path.isfile(os.path.join(snaps, s, "tokenizer.json"))
+        ]
+        if cands:
+            return max(cands, key=os.path.getmtime)
+    raise FileNotFoundError(
+        f"no tokenizer.json at {path!r} (also tried {snaps!r}). align.tokenizer is "
+        "local-only by design: it never downloads. Pass a checkpoint dir or an HF "
+        "snapshot dir containing tokenizer.json."
+    )
+
+
+def load_tokenizer(path: Union[str, os.PathLike]) -> Tokenizer:
+    """Load a Tokenizer from a local dir/file. See Tokenizer.from_pretrained."""
+    return Tokenizer.from_pretrained(path)

code/align/train/__init__.py

@@ -0,0 +1,16 @@
+"""The distributed training core (section-07).
+
+  * `offload_adamw.py` -- optimizer-state-only CPU offload; params and grads
+    stay on the training device, moments and the fp32 master live in host RAM.
+  * `fsdp.py`          -- FSDP2 wrapping + mixed-precision policy, called from
+    the launch scripts, never from the Trainer.
+  * `trainer.py`       -- the one Trainer for SFT and DPO: micro-step loop,
+    seeded lossless resume, optimizer-step scheduler, checkpoint cadence.
+  * `probes.py`        -- trunk-movement probe, run-health tripwires, the
+    20-step throughput probe.
+  * `sft_loss.py`      -- adapter from a collated batch to section-05's
+    chunked fused CE.
+
+Every module here imports torch at module scope, so nothing above the resource
+guards may import this package.
+"""

code/align/train/adamw8bit.py

@@ -0,0 +1,630 @@
+"""Block-wise 8-bit AdamW in pure torch: the optimizer that lets full-FT 7B launch.
+
+WHY THIS EXISTS (the arithmetic, from .tmp/handoff/review_trainer.md): every
+shipped full-FT config names `optimizer.name = "adamw8bit"`, and the two
+alternatives are both unlaunchable on this box at 7.248B params / 4 ranks:
+
+  * fp32 AdamW on the device needs 14.5GB/rank of moments on top of a 9.6GB
+    activation peak against a 23GB A10G.
+  * `OffloadAdamW` pins the WHOLE rank master (n/world*4 B) while its staging
+    cap is two layer-shards (n/layers/world*4*2 B): required/cap = layers /
+    (world * shards) = 4, so it refuses to construct. Past that, its host RSS
+    projects to ~105GB against ~85GB MemAvailable, and the checkpoint clones
+    master + both moments (12 B/param).
+
+This optimizer keeps both Adam moments as uint8 codes ON THE DEVICE and has no
+master copy at all.
+
+MEMORY BUDGET, per rank-local parameter element (block_size B = 2048):
+
+    exp_avg      1 byte code  + 4/B byte absmax   = 1.00195 B
+    exp_avg_sq   1 byte code  + 4/B byte absmax   = 1.00195 B
+    ------------------------------------------------------------
+    total        2 + 8/B  = 2.0039 B/param         (fp32 AdamW: 8; OffloadAdamW: 12 host)
+
+    7.248e9 params / 4 ranks = 1.812e9 elements/rank  ->  3.63 GB (3.38 GiB)/rank
+    14.5 GB across the job, all on the GPU, ZERO host RSS (which is what
+    `guard.projected_full_ft_rss_gb` already assumes for non-offload).
+
+  Tensors below `min_8bit_size` (4096) elements -- norm scales, biases -- keep
+  fp32 moments (8 B/param); there are ~65 of them in a 32-layer Mistral and
+  they total well under a megabyte. The step additionally holds one CHUNK of
+  fp32 transients (default 2**21 elements, < 100MB), never a whole tensor.
+  `projected_state_bytes` computes the exact figure from element counts and is
+  what the launch banner prints.
+
+THE STEP, per chunk of whole blocks: dequantize both moments -> fp32 AdamW
+update (decoupled weight decay, bias correction; the same op order as
+`torch.optim.AdamW`'s single-tensor path, so the fp32 fallback is BIT-EXACT
+against it) -> requantize -> store. The update uses the freshly computed
+fp32 moments, not the requantized ones: the quantization error only enters
+through the NEXT step's decayed moment, which is the same order bitsandbytes'
+fused kernel uses (quantization is a storage format, not part of the rule).
+
+KNOWN LIMITATION, measured and bounded (not a bug to chase): SECOND-MOMENT
+UNDERFLOW. v is a square, so its dynamic range is the SQUARE of the
+gradient's. The smallest nonzero unsigned code is 3.25e-7 of the block absmax,
+so an element whose gradient is < ~5e-4 of its block's largest one stores v = 0.
+The update never divides by that stored zero (it uses the fresh fp32 moment,
+see above), but the decayed history of v is lost for those elements: right
+after a cold start a tiny-gradient element sees only (1-beta2)*g^2 instead of
+the running average, so its step is up to lr / sqrt(1 - beta2) (4.5 lr at
+beta2 = 0.95) instead of ~lr. Measured on a Gaussian-gradient linear layer: a
+handful of elements (~0.1%) are off by ~1-2 lr on the second step and the mean
+error over all elements is ~1% of the travel -- and that is a COLD-START /
+CHANGING-GRADIENT picture. It does NOT relax for an element whose gradient
+stays below the threshold: v is re-floored to 0 every step, so the over-step is
+PERMANENT, and it GROWS with beta2 (adversarial review, constant g = 1e-4 of the
+block max, 300 steps, travel vs torch.optim.AdamW: 3.9x at beta2 = 0.95, 9.9x at
+0.999 and still rising toward 1/sqrt(1-beta2) = 31.6x; g = 1e-3 of the max is
+1.07x, g = 1e-2 is 1.00x). Clamping every positive v up to the smallest nonzero
+code would turn that into a bounded UNDER-step (0.16x, independent of beta2) at
+the price of freezing those elements; that is a different optimizer, so it is
+the owner's call (`.tmp/review_adamw8bit/probe3.py`), not made here.
+bitsandbytes has the same property
+(same map, same nearest-code rounding); the alternative of storing sqrt(v)
+doubles the relative error of the denominator in exchange for the range, which
+is a different optimizer than the one specified, so it is not done here.
+`tests/test_adamw8bit.py::test_second_moment_underflow_...` pins the bound.
+
+NO MASTER COPY, AND WHAT THAT COSTS. The parameters are bf16 and the update is
+tiny: at lr 1e-5 an update is ~1e-5 against a bf16 weight of ~2e-2 whose
+spacing is 1.2e-4, so round-to-nearest drops almost every update -- a run
+whose loss curve is flat and whose weights "never move" (the same
+silent-no-train class as the ZeRO-3 + Adafactor incident). The fp32 master that
+`OffloadAdamW` pays 4 B/param for exists to prevent exactly this. Here the
+write-back to a bf16 parameter is STOCHASTIC ROUNDING by default: the fp32
+result is rounded up or down with probability proportional to proximity, so the
+parameter moves by the right amount IN EXPECTATION. It is deterministic given
+the seed, and its generator state is checkpointed, so resume stays bit-exact.
+Pass `stochastic_rounding=False` for round-to-nearest (what bnb does).
+
+FSDP2 / DTENSOR (mirrors `OffloadAdamW`, same defect history): under
+`fully_shard` the parameters are DTensors. AdamW is elementwise, so stepping
+`param.to_local()` against the grad's local shard is EXACT; every access goes
+through `_localize` (a storage-sharing view, so in-place updates reach the
+DTensor). Numels are THIS RANK'S, never the global count.
+
+GLOBAL GRAD NORM (mirrors `OffloadAdamW._grad_norm`): the sum of squares is
+all-reduced across ranks before the root, so `last_grad_norm` is the true
+global norm, identical on every rank, and is recorded BEFORE clipping --
+the tripwires that read it must see an explosion, not its clipped image.
+The same all-reduced value is the finiteness gate: a non-finite gradient on
+ANY rank makes the global sum non-finite on EVERY rank, so all ranks skip
+together with one collective instead of two. (In float64 the sum of squares
+of finite fp32 values cannot overflow, so non-finite sum <=> non-finite grad.)
+
+RESUME: `state_dict()` is a snapshot of the uint8 codes, the absmax, the step
+counts, the hyper-parameters and the rounding RNG, all cloned to CPU. Reloading
+copies them back byte for byte, so a resumed run is bit-identical to an
+uninterrupted one. The state is SHARD-shaped: it records `world_size` and
+every tensor's element count and REFUSES a mismatched resume rather than
+aligning rank r's moments to a different slice of the parameter.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Iterable, Optional, Sequence
+
+import torch
+
+from ..utils.dist import rank, world_size
+from .offload_adamw import OffloadAdamW
+from .quant8 import (
+    codebook,
+    dequantize_rows,
+    quantize_rows,
+    zero_code,
+)
+
+logger = logging.getLogger(__name__)
+
+STATE_FORMAT = "adamw8bit-blockwise-v1"
+DEFAULT_BLOCK_SIZE = 2048
+#: bnb's `min_8bit_size`: below this a tensor's moments stay fp32.
+DEFAULT_MIN_8BIT_SIZE = 4096
+#: Elements processed per fused dequant/update/requant pass (rounded down to
+#: whole blocks). Bounds the fp32 transients; see the module docstring.
+DEFAULT_CHUNK_ELEMS = 2**21
+_FP32_MAX = float(torch.finfo(torch.float32).max)
+
+
+def adamw_update_(
+    p: torch.Tensor,
+    g: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    step: int,
+) -> None:
+    """One AdamW update on fp32 tensors, IN PLACE on `p`, `m`, `v`.
+
+    The op sequence is `torch.optim.AdamW._single_tensor_adamw` verbatim
+    (decoupled decay first, `lerp_` for the first moment, `addcmul_` for the
+    second, bias-correction folded into step_size and the denominator). That
+    is what makes "fp32 fallback == torch AdamW" a property of the
+    construction: a reordering that is algebraically identical is NOT
+    bit-identical in fp32, and a test that tolerates 1e-6 would hide the day
+    someone swaps `lerp_` for a hand-written EMA.
+    """
+    if weight_decay != 0:
+        p.mul_(1 - lr * weight_decay)
+    m.lerp_(g, 1 - beta1)
+    v.mul_(beta2).addcmul_(g, g, value=1 - beta2)
+    bias_correction1 = 1 - beta1**step
+    bias_correction2 = 1 - beta2**step
+    step_size = lr / bias_correction1
+    bias_correction2_sqrt = bias_correction2**0.5
+    denom = (v.sqrt() / bias_correction2_sqrt).add_(eps)
+    p.addcdiv_(m, denom, value=-step_size)
+
+
+def stochastic_round_to_bf16(x: torch.Tensor, generator: Optional[torch.Generator]) -> torch.Tensor:
+    """fp32 -> bf16 with stochastic rounding (unbiased: E[out] == x).
+
+    bf16 is exactly the top 16 bits of an fp32. Adding a uniform random
+    integer in [0, 2**16) to the raw bit pattern and truncating the low 16 bits
+    rounds UP with probability (low16 / 2**16) -- i.e. in proportion to how
+    close x is to the next representable value -- and an already-representable
+    x (low16 == 0) is returned unchanged. The add acts on magnitude for both
+    signs, because the sign is the top bit and the carry never reaches it for
+    finite inputs. Truncation back to bf16 is then exact (`.to(bfloat16)` of a
+    value whose low 16 bits are zero).
+    """
+    x = x.contiguous()
+    noise = torch.randint(
+        0, 1 << 16, x.shape, dtype=torch.int32, device=x.device, generator=generator
+    )
+    bits = (x.view(torch.int32) + noise) & -65536
+    return bits.view(torch.float32).to(torch.bfloat16)
+
+
+class BlockwiseAdamW8bit:
+    """AdamW with block-wise 8-bit moments, stepping each rank's LOCAL shard.
+
+    `step()` returns False when the grad-skip gate refused the step (same
+    contract as `OffloadAdamW`), True when it applied.
+
+    Args:
+        params: the parameters (plain tensors or FSDP2 DTensors). Floating
+            point, contiguous locally, no duplicates.
+        lr, betas, eps, weight_decay: AdamW; decay is DECOUPLED (applied to
+            the weight, not folded into the gradient).
+        max_grad_norm: if set, gradients are scaled by
+            `min(1, max_grad_norm / (norm + 1e-6))` -- `clip_grad_norm_`'s rule
+            -- using the GLOBAL norm. `last_grad_norm` is the pre-clip norm.
+            The scale is applied inside the fused update; `p.grad` itself is
+            not rewritten (it is about to be zeroed, and rewriting a bf16 grad
+            in place would only add rounding).
+        block_size: elements per absmax block (default 2048).
+        min_8bit_size: tensors with fewer local elements keep fp32 moments.
+        chunk_elems: fp32-transient bound for the fused update.
+        stochastic_rounding: stochastic (default) vs nearest rounding when the
+            fp32 result is written back to a bf16 parameter. See module doc.
+        seed: base seed of the rounding RNG (offset by rank).
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.Tensor],
+        lr: float,
+        betas: tuple[float, float] = (0.9, 0.95),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        max_grad_norm: Optional[float] = None,
+        block_size: int = DEFAULT_BLOCK_SIZE,
+        min_8bit_size: int = DEFAULT_MIN_8BIT_SIZE,
+        chunk_elems: int = DEFAULT_CHUNK_ELEMS,
+        stochastic_rounding: bool = True,
+        seed: int = 0,
+    ) -> None:
+        self.params: list[torch.Tensor] = list(params)
+        if not self.params:
+            raise ValueError("BlockwiseAdamW8bit got an empty parameter list")
+        if len({id(p) for p in self.params}) != len(self.params):
+            raise ValueError(
+                "BlockwiseAdamW8bit got the same parameter object twice (tied weights?): "
+                "it would be stepped twice per step with one shared state slot"
+            )
+        if not (isinstance(block_size, int) and block_size >= 1):
+            raise ValueError(f"block_size must be a positive int, got {block_size!r}")
+        if min_8bit_size < 0 or chunk_elems < 1:
+            raise ValueError("min_8bit_size must be >= 0 and chunk_elems >= 1")
+        if lr < 0 or eps <= 0 or weight_decay < 0:
+            raise ValueError(f"invalid AdamW hyper-parameters lr={lr} eps={eps} wd={weight_decay}")
+        if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+            raise ValueError(f"betas must be in [0, 1), got {betas}")
+        if max_grad_norm is not None and not max_grad_norm > 0:
+            raise ValueError(f"max_grad_norm must be positive or None, got {max_grad_norm}")
+
+        self.block_size = block_size
+        self.min_8bit_size = min_8bit_size
+        self.chunk_blocks = max(1, chunk_elems // block_size)
+        self.stochastic_rounding = bool(stochastic_rounding)
+        self.seed = int(seed)
+        self.banner_tag = "adamw8bit"
+        self.skipped_steps = 0
+        self.last_grad_norm = float("nan")
+        #: One group, like OffloadAdamW: `Trainer._apply_scheduled_lr` writes
+        #: the scheduled lr into every group, and `step` reads it from here.
+        self.param_groups = [
+            {
+                "params": self.params,
+                "lr": float(lr),
+                "betas": (float(betas[0]), float(betas[1])),
+                "eps": float(eps),
+                "weight_decay": float(weight_decay),
+                "max_grad_norm": None if max_grad_norm is None else float(max_grad_norm),
+            }
+        ]
+        self._generators: dict[str, torch.Generator] = {}
+
+        self._state: list[dict] = []
+        for i, p in enumerate(self.params):
+            if not p.dtype.is_floating_point:
+                raise ValueError(f"param {i} has non-floating dtype {p.dtype}")
+            local = self._localize(p)
+            if not local.is_contiguous():
+                raise ValueError(
+                    f"param {i} (local shape {tuple(local.shape)}) is not contiguous: the "
+                    "blockwise state is laid out over the flat local shard. FSDP2 dim-0 "
+                    "shards are contiguous; a transposed/strided view is not supported."
+                )
+            self._state.append(self._init_slot(local))
+
+        logger.info(
+            "BlockwiseAdamW8bit: %d tensors, %.2fM local elements, state %.3fGiB "
+            "(block %d, min_8bit_size %d, stochastic_rounding=%s)",
+            len(self.params),
+            sum(s["numel"] for s in self._state) / 1e6,
+            self.state_bytes() / 2**30,
+            block_size,
+            min_8bit_size,
+            self.stochastic_rounding,
+        )
+
+    # -- layout --------------------------------------------------------------
+
+    @staticmethod
+    def _localize(t):
+        """This rank's local shard of an FSDP2 DTensor (a storage-sharing
+        view), or `t` itself. Identical contract to OffloadAdamW._localize."""
+        return t.to_local() if hasattr(t, "to_local") else t
+
+    def _is_quantized(self, numel: int) -> bool:
+        return numel >= self.min_8bit_size and numel > 0
+
+    def _init_slot(self, local: torch.Tensor) -> dict:
+        n = local.numel()
+        dev = local.device
+        slot: dict = {"numel": n, "step": 0, "quantized": self._is_quantized(n)}
+        if slot["quantized"]:
+            nb = -(-n // self.block_size)
+            # Codes are allocated in WHOLE blocks (the last block is padded with
+            # the zero code) so every chunk is a clean [blocks, block_size] row
+            # matrix; the pad costs < one block per tensor.
+            slot["m_codes"] = torch.full(
+                (nb, self.block_size), zero_code(True), dtype=torch.uint8, device=dev
+            )
+            slot["v_codes"] = torch.full(
+                (nb, self.block_size), zero_code(False), dtype=torch.uint8, device=dev
+            )
+            slot["m_absmax"] = torch.zeros(nb, dtype=torch.float32, device=dev)
+            slot["v_absmax"] = torch.zeros(nb, dtype=torch.float32, device=dev)
+        else:
+            slot["m"] = torch.zeros(n, dtype=torch.float32, device=dev)
+            slot["v"] = torch.zeros(n, dtype=torch.float32, device=dev)
+        return slot
+
+    @staticmethod
+    def projected_state_bytes(
+        local_numels: Sequence[int],
+        block_size: int = DEFAULT_BLOCK_SIZE,
+        min_8bit_size: int = DEFAULT_MIN_8BIT_SIZE,
+    ) -> int:
+        """Exact optimizer-state bytes for these RANK-LOCAL element counts.
+
+        Large tensors: `2*n + 8*n/block_size` (two uint8 code arrays plus two
+        fp32 absmax arrays), with `n` rounded up to whole blocks because that
+        is how the codes are allocated. Small tensors: 8*n (two fp32 moments).
+        Computable BEFORE construction, so a launch banner or preflight can
+        state the number without allocating it.
+        """
+        total = 0
+        for n in local_numels:
+            if n >= min_8bit_size and n > 0:
+                nb = -(-n // block_size)
+                total += 2 * nb * block_size + 8 * nb
+            else:
+                total += 8 * n
+        return int(total)
+
+    def state_bytes(self) -> int:
+        """Bytes actually held by this optimizer's state tensors on this rank."""
+        total = 0
+        for s in self._state:
+            for k, v in s.items():
+                if isinstance(v, torch.Tensor):
+                    total += v.numel() * v.element_size()
+        return total
+
+    def describe_placement(self) -> dict:
+        """Where the state lives (device/dtype/bytes), for the banner and the
+        placement test -- a 'the state is 8-bit' claim nobody checks is how it
+        quietly stops being true."""
+        ents = []
+        for s in self._state:
+            for k, v in s.items():
+                if isinstance(v, torch.Tensor):
+                    ents.append(
+                        {"key": k, "device": v.device.type, "dtype": str(v.dtype),
+                         "bytes": v.numel() * v.element_size()}
+                    )
+        return {"state": ents, "state_bytes": self.state_bytes()}
+
+    # -- rounding RNG --------------------------------------------------------
+
+    def _generator(self, device: torch.device) -> torch.Generator:
+        key = device.type
+        gen = self._generators.get(key)
+        if gen is None:
+            gen = torch.Generator(device=device)
+            gen.manual_seed(self.seed + rank())
+            self._generators[key] = gen
+        return gen
+
+    def _write_param(self, dst: torch.Tensor, src32: torch.Tensor) -> None:
+        """Write the fp32 result into the (possibly bf16) parameter chunk."""
+        if dst.dtype == torch.float32:
+            if dst.data_ptr() != src32.data_ptr():
+                dst.copy_(src32)
+        elif dst.dtype == torch.bfloat16 and self.stochastic_rounding:
+            dst.copy_(stochastic_round_to_bf16(src32, self._generator(dst.device)))
+        else:
+            dst.copy_(src32)
+
+    # -- the step ------------------------------------------------------------
+
+    @torch.no_grad()
+    def step(self) -> bool:
+        """One optimizer step, or a refusal.
+
+        Order is load-bearing: the (global) norm / finiteness verdict is reached
+        BEFORE any moment is touched, so a skipped step leaves the state
+        byte-identical. A discovered-too-late NaN would already have poisoned
+        the second moment, and every later step is scaled by it.
+        """
+        grads = [None if p.grad is None else self._localize(p.grad) for p in self.params]
+        present = [g for g in grads if g is not None]
+        if not present:
+            raise RuntimeError(
+                "BlockwiseAdamW8bit.step() called with no gradients on any parameter. "
+                "A step that silently does nothing is indistinguishable from a run "
+                "that trained."
+            )
+        for i, (g, s) in enumerate(zip(grads, self._state)):
+            if g is not None and g.numel() != s["numel"]:
+                raise ValueError(
+                    f"param {i}: local grad has {g.numel()} elements but the state was "
+                    f"built for {s['numel']} (grad of a different shard layout?)"
+                )
+
+        norm = OffloadAdamW._grad_norm(present)
+        if not math.isfinite(norm):
+            self.skipped_steps += 1
+            # NaN, not 0.0: a 0.0 in the log reads as 'perfectly converged'.
+            self.last_grad_norm = float("nan")
+            logger.warning(
+                "grad-skip gate: non-finite gradient, skipping optimizer step "
+                "(%d skipped so far). Moments and step counts left untouched.",
+                self.skipped_steps,
+            )
+            return False
+        self.last_grad_norm = norm  # PRE-clip, on purpose
+
+        group = self.param_groups[0]
+        coef = 1.0
+        max_norm = group["max_grad_norm"]
+        if max_norm is not None:
+            coef = min(1.0, max_norm / (norm + 1e-6))
+
+        for p, g, st in zip(self.params, grads, self._state):
+            if g is None or st["numel"] == 0:
+                continue
+            st["step"] += 1
+            local = self._localize(p).detach().view(-1)
+            gflat = g.detach().reshape(-1)
+            if st["quantized"]:
+                self._step_quantized(local, gflat, st, group, coef)
+            else:
+                self._step_fp32(local, gflat, st, group, coef)
+        return True
+
+    def _hyper(self, group: dict, st: dict) -> dict:
+        return dict(
+            lr=group["lr"],
+            beta1=group["betas"][0],
+            beta2=group["betas"][1],
+            eps=group["eps"],
+            weight_decay=group["weight_decay"],
+            step=st["step"],
+        )
+
+    def _step_fp32(self, pflat, gflat, st, group, coef) -> None:
+        g32 = gflat.to(torch.float32)
+        if coef < 1.0:
+            g32 = g32 * coef  # out of place: g32 may alias the live .grad
+        p32 = pflat.to(torch.float32)  # aliases pflat when it is already fp32
+        adamw_update_(p32, g32, st["m"], st["v"], **self._hyper(group, st))
+        self._write_param(pflat, p32)
+
+    def _step_quantized(self, pflat, gflat, st, group, coef) -> None:
+        B = self.block_size
+        n = st["numel"]
+        n_blocks = st["m_codes"].shape[0]
+        cb_m = codebook(True, pflat.device)
+        cb_v = codebook(False, pflat.device)
+        hyper = self._hyper(group, st)
+        for b0 in range(0, n_blocks, self.chunk_blocks):
+            b1 = min(n_blocks, b0 + self.chunk_blocks)
+            lo, hi = b0 * B, min(n, b1 * B)
+            ln = hi - lo
+            g32 = gflat[lo:hi].to(torch.float32)
+            if coef < 1.0:
+                g32 = g32 * coef
+            m_full = dequantize_rows(st["m_codes"][b0:b1], st["m_absmax"][b0:b1], cb_m).view(-1)
+            v_full = dequantize_rows(st["v_codes"][b0:b1], st["v_absmax"][b0:b1], cb_v).view(-1)
+            pchunk = pflat[lo:hi]
+            p32 = pchunk.to(torch.float32)
+            # [:ln] are views of the block-padded buffers: the pad tail is never
+            # touched by the update (it stays exactly 0), then the whole padded
+            # rows are requantized.
+            adamw_update_(p32, g32, m_full[:ln], v_full[:ln], **hyper)
+            self._write_param(pchunk, p32)
+            # THE V-OVERFLOW POISON (adversarial review, 2026-10-12). The finiteness
+            # gate in `step` tests the NORM, which is accumulated in float64 and so
+            # stays finite for a gradient the fp32 SQUARE cannot hold: one element of
+            # 1e21 gives v = (1-beta2)*g^2 = inf. torch.optim.AdamW survives that
+            # (m / (inf + eps) == 0), but here inf becomes the block's absmax, the
+            # requantize divides inf/inf = NaN, the NaN lands in the uint8 codes,
+            # and the block's weights go NaN on the NEXT step, with the gate seeing
+            # only finite gradients. Clamping v to the largest finite fp32 keeps the
+            # update identical (m / 1.8e19 is still 0 for any sane m) and the stored
+            # state finite. Only v can overflow: m is a convex mix of finite values.
+            v_full.clamp_(max=_FP32_MAX)
+            codes, absmax = quantize_rows(m_full.view(-1, B), cb_m)
+            st["m_codes"][b0:b1] = codes
+            st["m_absmax"][b0:b1] = absmax
+            codes, absmax = quantize_rows(v_full.view(-1, B), cb_v)
+            st["v_codes"][b0:b1] = codes
+            st["v_absmax"][b0:b1] = absmax
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        for p in self.params:
+            if set_to_none:
+                p.grad = None
+            elif p.grad is not None:
+                p.grad.zero_()
+
+    # -- resume --------------------------------------------------------------
+
+    def state_dict(self) -> dict:
+        """A DETACHED CPU COPY of everything the next step depends on.
+
+        `.to("cpu", copy=True)`, not `.to("cpu")`: on a CPU tensor the latter
+        returns the SAME object, and a "snapshot" that aliases the live codes
+        is mutated by the very next step -- the shared-state aliasing incident
+        OffloadAdamW documents, here with uint8 codes that no comparison would
+        flag as drifting. Everything is copied so the dict is safe to pickle
+        straight into the checkpoint gather.
+        """
+        slots = []
+        for s in self._state:
+            slots.append(
+                {
+                    k: (v.detach().to("cpu", copy=True) if isinstance(v, torch.Tensor) else v)
+                    for k, v in s.items()
+                }
+            )
+        group = self.param_groups[0]
+        return {
+            "format": STATE_FORMAT,
+            "world_size": world_size(),
+            "block_size": self.block_size,
+            "min_8bit_size": self.min_8bit_size,
+            "stochastic_rounding": self.stochastic_rounding,
+            "param_group": {k: v for k, v in group.items() if k != "params"},
+            "state": slots,
+            "skipped_steps": self.skipped_steps,
+            "last_grad_norm": self.last_grad_norm,
+            "sr_rng": {k: g.get_state().clone() for k, g in self._generators.items()},
+        }
+
+    def load_state_dict(self, sd: dict) -> None:
+        """Restore the state byte for byte. Validates EVERYTHING before writing
+        anything, so a refused load leaves the optimizer exactly as it was."""
+        if sd.get("format") != STATE_FORMAT:
+            raise ValueError(
+                f"optimizer state format {sd.get('format')!r} is not {STATE_FORMAT!r}; "
+                "refusing to interpret it (a 32-bit or bnb state is not loadable here)"
+            )
+        saved_ws, now_ws = sd.get("world_size"), world_size()
+        if saved_ws is not None and saved_ws != now_ws:
+            raise ValueError(
+                f"this optimizer state was saved at world_size={saved_ws} but the current "
+                f"run has world_size={now_ws}. The 8-bit moments are SHARD-shaped: loading "
+                "them under a different world size would align each rank's state to a "
+                "different slice of the parameter, and training would continue plausibly "
+                "on corrupted state. Resume at the original world size, or start fresh."
+            )
+        if sd["block_size"] != self.block_size or sd["min_8bit_size"] != self.min_8bit_size:
+            raise ValueError(
+                f"state was built with block_size={sd['block_size']}, min_8bit_size="
+                f"{sd['min_8bit_size']} but this optimizer has {self.block_size}, "
+                f"{self.min_8bit_size}: the absmax blocks would not line up"
+            )
+        saved = sd["state"]
+        if len(saved) != len(self._state):
+            raise ValueError(
+                f"checkpoint holds {len(saved)} tensors but this optimizer has {len(self._state)}"
+            )
+        for i, (cur, new) in enumerate(zip(self._state, saved)):
+            if cur["numel"] != new["numel"] or cur["quantized"] != new["quantized"]:
+                raise ValueError(
+                    f"tensor {i}: checkpoint has numel={new['numel']} quantized="
+                    f"{new['quantized']}, this optimizer has numel={cur['numel']} quantized="
+                    f"{cur['quantized']} (a different shard layout)"
+                )
+            for k, v in cur.items():
+                if isinstance(v, torch.Tensor):
+                    nv = new.get(k)
+                    if nv is None or nv.shape != v.shape or nv.dtype != v.dtype:
+                        raise ValueError(
+                            f"tensor {i} field {k!r}: checkpoint "
+                            f"{None if nv is None else (tuple(nv.shape), nv.dtype)} vs "
+                            f"expected {(tuple(v.shape), v.dtype)}"
+                        )
+        if bool(sd.get("stochastic_rounding", self.stochastic_rounding)) != self.stochastic_rounding:
+            logger.warning(
+                "resuming with stochastic_rounding=%s but the checkpoint was written with %s: "
+                "the trajectory will differ from an uninterrupted run",
+                self.stochastic_rounding,
+                sd.get("stochastic_rounding"),
+            )
+
+        # The rounding RNG is part of the state being restored. A generator that
+        # already exists here (this optimizer has stepped) and is ABSENT from `sd`
+        # (saved before any bf16 step, or from a fp32 run) would otherwise keep its
+        # advanced stream, and the rewound optimizer would round differently from an
+        # uninterrupted one -- "byte for byte" quietly untrue (adversarial review,
+        # 2026-10-12). Dropping them re-seeds lazily from `seed + rank`, which is
+        # exactly the state a never-stepped optimizer has.
+        self._generators.clear()
+
+        # copy_ INTO the existing tensors: keeps their device, and means loading
+        # one checkpoint dict into two optimizers cannot make them share state.
+        for cur, new in zip(self._state, saved):
+            for k, v in cur.items():
+                if isinstance(v, torch.Tensor):
+                    v.copy_(new[k])
+            cur["step"] = int(new["step"])
+        group = self.param_groups[0]
+        for k, v in sd["param_group"].items():
+            group[k] = tuple(v) if k == "betas" else v
+        self.skipped_steps = int(sd.get("skipped_steps", 0))
+        self.last_grad_norm = float(sd.get("last_grad_norm", float("nan")))
+        for key, st in sd.get("sr_rng", {}).items():
+            first = self.params[0].device
+            gen = self._generator(first if first.type == key else torch.device(key))
+            gen.set_state(st.detach().to("cpu", copy=True))
+
+
+#: The name the factory and docs use.
+AdamW8bit = BlockwiseAdamW8bit

code/align/train/ckpt_format.py

@@ -0,0 +1,124 @@
+"""Weight-file formats the Trainer can write and read: the model-family seam.
+
+Why this exists. `Trainer.save_checkpoint` / `resume_from` called the dense
+Mistral converters (`save_nmo_as_hf`, `hf_to_nmo(expected_n_layers=
+cfg.model.n_layers)`) directly. Those tables know the 7B dense key names and
+nothing about routers or stacked 3-D expert tensors, so an OLMoE run could be
+built and stepped but never checkpointed -- the first save would raise on a
+key it cannot map (or, worse for a model family added later, silently drop
+parameters if a converter ever became lenient). The seam makes the family an
+explicit, injected choice with ONE contract:
+
+    save_weights(align_state_dict, out_dir, dtype)  -> writes model weights +
+                                                       config.json, atomically
+    load_weights(path)                              -> align state dict
+
+and nothing else about the checkpoint changes: optimizer state, `resume.pt`,
+fsync/rename/completeness all stay in the Trainer and are family-independent.
+
+The default is the dense format built from `cfg.model`, so every existing
+caller is unchanged. A family is selected by construction (`Trainer(...,
+weight_format=OlmoeFormat(olmoe_cfg))`), never sniffed from the state dict:
+guessing the family from key names is how a misconfigured run loads the wrong
+converter and half-works.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Mapping, Optional, Protocol
+
+import torch
+
+
+class WeightFormat(Protocol):
+    def save_weights(
+        self, state: Mapping[str, torch.Tensor], out_dir: str, dtype: Optional[torch.dtype]
+    ) -> None: ...
+
+    def load_weights(self, path: str) -> dict[str, torch.Tensor]: ...
+
+    def copy_tokenizer(self, snapshot: str, dst_dir: str) -> list[str]:
+        """Copy the tokenizer files that belong to THIS family into `dst_dir`;
+        return the names copied (empty = weights-only checkpoint)."""
+        ...
+
+
+class DenseFormat:
+    """The original behaviour: dense Mistral-family HF checkpoints."""
+
+    def __init__(self, model_cfg: Any):
+        self.model_cfg = model_cfg
+
+    def save_weights(self, state, out_dir, dtype):
+        from ..model import checkpoint as ckpt
+
+        ckpt.save_nmo_as_hf(dict(state), self.model_cfg, out_dir, dtype=dtype)
+
+    def load_weights(self, path):
+        from ..model import checkpoint as ckpt
+
+        return ckpt.hf_to_nmo(
+            ckpt.load_hf_safetensors(path), expected_n_layers=self.model_cfg.n_layers
+        )
+
+    def copy_tokenizer(self, snapshot, dst_dir):
+        from ..model import checkpoint as ckpt
+
+        return ckpt.copy_tokenizer_files(snapshot, dst_dir)
+
+
+#: Files that make up an HF tokenizer, in either of the layouts OLMoE-class
+#: models ship (BPE `tokenizer.json`, or `vocab.json` + `merges.txt`).
+OLMOE_TOKENIZER_FILES = (
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "vocab.json",
+    "merges.txt",
+    "tokenizer.model",
+    "generation_config.json",
+)
+
+
+class OlmoeFormat:
+    """OLMoE-architecture checkpoints (routers + stacked experts), HF-compatible.
+
+    TOKENIZER. The dense path copies the pinned Mistral v0.3 tokenizer into
+    every checkpoint (same vocab by construction); doing that for an OLMoE
+    checkpoint would stamp a 32768-token tokenizer onto a 50304-token model --
+    loads fine, tokenizes wrong. So this format copies tokenizer files only
+    from an EXPLICIT `tokenizer_dir`, and without one writes a weights-only
+    checkpoint (and says so), never a guess.
+    """
+
+    def __init__(self, olmoe_cfg: Any, tokenizer_dir: Optional[str] = None):
+        self.olmoe_cfg = olmoe_cfg
+        self.tokenizer_dir = tokenizer_dir
+
+    def save_weights(self, state, out_dir, dtype):
+        from ..model.olmoe_checkpoint import save_olmoe_as_hf
+
+        save_olmoe_as_hf(dict(state), self.olmoe_cfg, out_dir, dtype=dtype)
+
+    def load_weights(self, path):
+        from ..model.olmoe_checkpoint import hf_to_align
+        from ..model.checkpoint import load_hf_safetensors
+
+        return hf_to_align(load_hf_safetensors(path), self.olmoe_cfg)
+
+    def copy_tokenizer(self, snapshot, dst_dir):  # noqa: ARG002 -- the dense snapshot is not ours
+        import os
+        import shutil
+
+        if not self.tokenizer_dir:
+            return []
+        copied = []
+        for name in OLMOE_TOKENIZER_FILES:
+            src = os.path.join(self.tokenizer_dir, name)
+            if os.path.isfile(src):
+                shutil.copyfile(src, os.path.join(dst_dir, name))
+                copied.append(name)
+        return copied
+
+
+__all__ = ["WeightFormat", "DenseFormat", "OlmoeFormat"]

code/align/train/dpo_loss.py

@@ -0,0 +1,237 @@
+"""DPO loss adapter: paired batch -> section-05's chunked fused DPO loss.
+
+Mirror of `sft_loss.py`: this module owns plumbing only -- unpacking the paired
+batch, the caller-side label shift, and packaging metrics -- while the kernel in
+`align.kernels.fused_dpo` owns the objective.
+
+The one thing worth stating here rather than in the kernel: the reference
+log-probabilities arrive as PRECOMPUTED CONSTANTS, produced offline by
+section-12 in a single no-grad pass over the SFT checkpoint. There is no
+reference model resident at DPO time, which is what makes 7B DPO fit on 4x A10G
+at all. The kernel detaches them on entry and a test proves no gradient path
+survives; this adapter's job is simply never to build one, so it refuses a
+batch that does not carry them rather than silently substituting zeros -- a
+zero reference turns DPO into plain likelihood maximisation on the chosen side,
+which trains, converges, and is not DPO.
+
+`DpoLossConfig` is the trainer-facing default surface. Note `token_chunk_size`:
+the pre-registered config is micro_batch 1 x seq 1024, i.e. ONE pair per step,
+and at one pair the pair-chunking does nothing at all. Leaving it None there
+would silently restore the [1024, 32768] logit tensor the fused loss exists to
+remove, so the trainer-facing default sets it.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+
+from .trainer import LossOutput
+
+logger = logging.getLogger(__name__)
+
+#: Section-05's kernel module, named so the both-directions test can block the
+#: import and exercise the refusal branch (same pattern as `sft_loss`).
+FUSED_DPO_MODULE = "align.kernels.fused_dpo"
+
+#: Fallback only; the authority is the kernel's constant and section-03's
+#: collator, which both use -100.
+IGNORE_INDEX = -100
+
+
+@dataclass
+class DpoLossConfig:
+    """Trainer-facing DPO loss settings. Section-12 nests this in its stage
+    config; every default that encodes a decision says why."""
+
+    #: Preference sharpness. 0.1 is the pre-registered value for both arms.
+    beta: float = 0.1
+    #: One of align.kernels.fused_dpo.DPO_VARIANTS. `dpop` is the pre-registered
+    #: tripwire RESPONSE for logps/chosen declining while the margin rises; it
+    #: is not the primary configuration and switching to it is a documented
+    #: deviation, not a tuning knob.
+    variant: str = "sigmoid"
+    #: Pairs per chunk.
+    chunk_size: int = 1
+    #: Tokens per inner chunk. NOT None: at micro_batch 1 there is one pair per
+    #: step and pair-chunking alone would materialise the full logit tensor.
+    #: 128 rather than 256, DERIVED from `fused_ce.chunked_peak_bound` at the
+    #: 7B shapes (V 32768, H 4096, 2x1024 rows): 256 predicts a 192 MiB peak
+    #: against a 256 MiB budget (0.75, too tight to survive an allocator
+    #: fluctuation in the M5 smoke), while 128 predicts 128 MiB (0.50). The
+    #: matmul is still large enough to be efficient at 128 x 32768.
+    token_chunk_size: Optional[int] = 128
+    #: NLL auxiliary on the chosen completion. 0.0 is provably inert (the term
+    #: is skipped structurally, not multiplied by zero), so the primary
+    #: pre-registered run is unaffected by this field existing.
+    nll_coeff: float = 0.0
+    #: DPO-Positive penalty weight; read only when variant == "dpop".
+    dpop_lambda: float = 50.0
+
+
+def _load_fused():
+    try:
+        import importlib
+
+        return importlib.import_module(FUSED_DPO_MODULE)
+    except Exception:  # noqa: BLE001
+        return None
+
+
+def fused_dpo_available() -> bool:
+    return _load_fused() is not None
+
+
+def _unpack(batch: Any):
+    """Accept a paired batch object or a mapping. Six tensors are mandatory and
+    none of them has a defensible default."""
+    required = (
+        "input_ids_chosen",
+        "labels_chosen",
+        "input_ids_rejected",
+        "labels_rejected",
+        "ref_logp_chosen",
+        "ref_logp_rejected",
+    )
+    if isinstance(batch, dict):
+        got = {k: batch.get(k) for k in required}
+    else:
+        got = {k: getattr(batch, k, None) for k in required}
+    missing = [k for k, v in got.items() if v is None]
+    if missing:
+        raise TypeError(
+            f"DPO batch is missing {missing}. The reference log-probabilities "
+            "in particular have no default: substituting zeros would turn DPO "
+            "into likelihood maximisation on the chosen completion, which "
+            "trains and converges and is not DPO. Section-12's "
+            "precompute step produces them."
+        )
+    return got
+
+
+def make_dpo_loss_adapter(
+    config: Optional[DpoLossConfig] = None,
+    require_fused: bool = False,
+    decoupled_backward: bool = False,
+    grad_scale: float = 1.0,
+):
+    """Build the DPO loss adapter the Trainer calls once per micro-step.
+
+    `require_fused=True` refuses the reference path, which is what the 7B stage
+    scripts pass: the eager reference materialises BOTH [P, T, 32768] logit
+    tensors and is a memory profile no A10G survives at 7B.
+
+    `decoupled_backward=True` (the 7B stage setting) backwards each side's
+    trunk separately so both checkpointed graphs are never resident at once;
+    `grad_scale` must then be 1/grad_accum, matching the scaling the Trainer
+    applies to the loss it backwards.
+    """
+    cfg = config or DpoLossConfig()
+    module = _load_fused()
+    if module is None:
+        if require_fused:
+            raise RuntimeError(
+                f"the fused chunked DPO loss ({FUSED_DPO_MODULE}."
+                "fused_dpo_loss) is not importable. Refusing to fall back to "
+                "the reference path for a full-scale run: it materialises both "
+                "[P, T, 32768] logit tensors and their gradients."
+            )
+        raise RuntimeError(
+            f"{FUSED_DPO_MODULE} is not importable and there is no reference "
+            "fallback worth offering for DPO at any useful scale."
+        )
+    if cfg.variant not in module.DPO_VARIANTS:
+        raise ValueError(f"unknown DPO variant {cfg.variant!r}; known: {list(module.DPO_VARIANTS)}")
+    ignore_index = module.IGNORE_INDEX
+
+    def _kernel(h_c, h_r, w, got, target_c, target_r):
+        return module.fused_dpo_loss(
+            h_c,
+            h_r,
+            w,
+            target_c,
+            target_r,
+            got["ref_logp_chosen"],
+            got["ref_logp_rejected"],
+            beta=cfg.beta,
+            variant=cfg.variant,
+            chunk_size=cfg.chunk_size,
+            nll_coeff=cfg.nll_coeff,
+            dpop_lambda=cfg.dpop_lambda,
+            ignore_index=ignore_index,
+            token_chunk_size=cfg.token_chunk_size,
+        )
+
+    def _full_w(model):
+        # FSDP2 params are sharded DTensors OUTSIDE a module __call__; passing
+        # one into the fused kernel's mm mixes Tensor and DTensor (the
+        # hook-bypass class, DPO edition). full_tensor() is the DIFFERENTIABLE
+        # all-gather: backward reduce-scatters d_weight onto the shards.
+        w = model.lm_head.weight
+        return w.full_tensor() if hasattr(w, "full_tensor") else w
+
+    def adapter(model, batch) -> LossOutput:
+        got = _unpack(batch)
+        # Caller-shifted, exactly as in sft_loss: position t predicts t+1.
+        target_c = got["labels_chosen"][:, 1:]
+        target_r = got["labels_rejected"][:, 1:]
+
+        if not decoupled_backward:
+            hidden_c = model.forward_hidden_states(input_ids=got["input_ids_chosen"])[:, :-1]
+            hidden_r = model.forward_hidden_states(input_ids=got["input_ids_rejected"])[:, :-1]
+            loss, metrics = _kernel(hidden_c, hidden_r, _full_w(model), got, target_c, target_r)
+            module.require_metrics(metrics)
+            n_tokens = int((target_c != ignore_index).sum() + (target_r != ignore_index).sum())
+            return LossOutput(loss=loss, n_tokens=n_tokens, metrics=dict(metrics))
+
+        # DECOUPLED BACKWARD -- one trunk graph at a time. Holding BOTH
+        # sides' checkpointed 7B graphs plus fp32 shard gradients peaked at
+        # 21.5GB and OOM'd every rank (measured; batch-size independent).
+        # DPO's coupling between the sides is one scalar per pair, so:
+        #   1. no-grad trunk forwards; the kernel runs over LEAF copies of
+        #      the hiddens plus the gathered lm_head, and one backward on
+        #      that small graph yields d_hidden per side analytically and
+        #      accumulates the lm_head shard grad;
+        #   2. each side then re-forwards WITH grad and backwards against
+        #      its saved d_hidden -- never two trunk graphs at once.
+        # The re-forward recomputes hiddens that differ from the no-grad
+        # pass only by kernel-dispatch noise (~1e-7 relative, the same class
+        # the batch-invariance receipt records); act-ckpt already relies on
+        # this determinism. The trainer's `(loss/accum).backward()` contract
+        # is honoured with a detached leaf proxy; the REAL grads are scaled
+        # here by grad_scale (1/accum), so nothing is double-counted.
+        with torch.no_grad():
+            h_c0 = model.forward_hidden_states(input_ids=got["input_ids_chosen"])[:, :-1]
+            h_r0 = model.forward_hidden_states(input_ids=got["input_ids_rejected"])[:, :-1]
+        h_c = h_c0.detach().requires_grad_(True)
+        h_r = h_r0.detach().requires_grad_(True)
+        loss, metrics = _kernel(h_c, h_r, _full_w(model), got, target_c, target_r)
+        module.require_metrics(metrics)
+        (loss * grad_scale).backward()
+        g_c, g_r = h_c.grad, h_r.grad
+        if g_c is None or g_r is None:
+            raise RuntimeError(
+                "decoupled DPO backward produced no hidden gradients; the "
+                "kernel did not differentiate w.r.t. the hidden leaves"
+            )
+        del h_c, h_r, h_c0, h_r0
+
+        hc = model.forward_hidden_states(input_ids=got["input_ids_chosen"])[:, :-1]
+        hc.backward(gradient=g_c)
+        del hc, g_c
+        hr = model.forward_hidden_states(input_ids=got["input_ids_rejected"])[:, :-1]
+        hr.backward(gradient=g_r)
+        del hr, g_r
+
+        n_tokens = int((target_c != ignore_index).sum() + (target_r != ignore_index).sum())
+        # Leaf proxy: the trainer's (loss/accum).backward() runs on THIS and
+        # touches no model gradient; the value is the real loss for the log.
+        proxy = loss.detach().clone().requires_grad_(True)
+        return LossOutput(loss=proxy, n_tokens=n_tokens, metrics=dict(metrics))
+
+    adapter.uses_fused = True
+    adapter.config = cfg
+    return adapter

code/align/train/dpo_stage.py

@@ -0,0 +1,1362 @@
+"""DPO stage: config, entry gates, banner, tripwire response, arm manifest.
+
+Section-12 is the second and last GPU training stage, run as TWO
+pre-registered arms. Everything numerical belongs to other sections -- the
+objective is section-05's `fused_dpo`, the loop is section-07's `Trainer`, the
+tripwire DETECTORS are section-07's `probes.py`, the reference constants are
+`ref_cache.py`. What lives here is the part that decides whether a run is
+allowed to start and what its result is permitted to be.
+
+That sounds like bureaucracy and is not. The predecessor project's expensive
+failures were never wrong arithmetic; they were runs that completed, produced
+publishable-looking numbers, and measured something other than what was
+claimed. So:
+
+  * **Entry gates run in a fixed order and any failure blocks everything
+    after it, including the model load.** Ordered by cost, so a stage that was
+    never going to be valid dies in milliseconds instead of after putting
+    14.5 GB onto a shared no-swap box.
+  * **The tripwire response is pre-registered.** `logps/chosen` declining
+    while the margin rises is the DPO chosen-likelihood-collapse mode, and it
+    is *maximal for Arm B by design*: pi_SFT already assigns the chosen
+    continuations high likelihood, so the optimizer's cheapest margin comes
+    from crushing both logprobs. The response -- rerun with `dpop` or the NLL
+    auxiliary, and record THAT as the arm's result -- was written down before
+    any run, and no third configuration is reachable without a new
+    pre-registration.
+  * **The primary metrics are margin and held-out preference accuracy**, not
+    benchmark accuracy. The pre-registered behavioural expectations are Arm A
+    within +-2 pp of SFT (below floor is the EXPECTED outcome) and Arm B's
+    honest range INCLUDING ZERO. Section-13 owns the paired benchmark
+    evaluation; a stage that graded itself on benchmarks would be choosing its
+    own success criterion after seeing the data.
+"""
+
+from __future__ import annotations
+
+import collections
+import json
+import logging
+import math
+import os
+from dataclasses import dataclass, field, fields, replace
+from typing import Any, Callable, Mapping, Optional, Sequence
+
+logger = logging.getLogger(__name__)
+
+NAN = float("nan")
+
+#: Entry gates, in the order they run. Ordered by COST: the cheap file reads
+#: come first so an invalid stage fails before it touches a GPU. Data, so the
+#: test iterates it instead of copying a list.
+GATE_ORDER: tuple[str, ...] = (
+    "sft_quality_gate",  # section-11 ran and passed its own gate
+    "build_manifest",  # shard counts match the pinned values; Arm A DESCOPED
+    "ref_cache_key",  # the reference constants belong to THIS run
+    "path_bias",  # measured, and its decision honored
+    "throughput_probe",  # 20 steps -> banner wall-clock (never asserted ahead)
+    "rss_preflight",  # host RSS headroom for the full-FT job class
+)
+
+#: Banner keys this stage adds on top of `registry.MANDATORY_BANNER_KEYS`.
+#:
+#: Deliberately NOT registered into the global registry at import time.
+#: `registry.register_mandatory_key` mutates a module-level list shared by
+#: every script in the process; importing this module would then make an
+#: unrelated banner (trace, generate) incomplete. The contract is identical --
+#: a missing key raises, a partial banner never prints -- it is just enforced
+#: without a global side effect.
+DPO_BANNER_KEYS: tuple[str, ...] = (
+    "arm",
+    "loss_type",
+    "beta",
+    "variant",
+    "lr",
+    "epochs",
+    "global_batch",
+    "grad_accum",
+    "expected_pairs",
+    "optimizer_step_formula",
+    "scheduler_horizon_optimizer_steps",
+    "dataset_hashes",
+    "ref_cache_key",
+    "path_bias",
+    "probe_step_time_s",
+    "probe_derived_wall_clock_h",
+    "rss_budget",
+    "attention_backend",
+    "length_drift_baseline",
+    "length_drift_baseline_source",
+    "sft_checkpoint",
+    "pre_registration",
+)
+
+#: The only configurations the pre-registration permits when the
+#: chosen-logprob tripwire fires. Anything else requires a new dated entry in
+#: docs/RESULTS.md written BEFORE the rerun, not after seeing it.
+#:
+#: `dpop_nll_aux` is the default because it is what the frozen
+#: pre-registration literally says: "rerun THAT arm once with dpop + NLL
+#: auxiliary". The section-12 spec phrases the same response as "the `dpop`
+#: variant OR the NLL auxiliary", so the two single-knob forms are kept
+#: available -- but docs/RESULTS.md is the frozen document and the conjunction
+#: is what it froze, so the conjunction is the default. (Discrepancy raised
+#: with the lead rather than resolved silently in either direction.)
+PRE_REGISTERED_RESPONSES: tuple[str, ...] = ("dpop_nll_aux", "dpop", "nll_aux")
+
+#: NLL auxiliary weight used when the response is `nll_aux`. Small: the term
+#: exists to put a floor under the chosen logprob, not to turn DPO into SFT.
+NLL_AUX_COEFF = 0.1
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DpoStageConfig:
+    """Stage settings for one arm. Per-field docstrings, per repo convention:
+    the config file is where someone will try to change these at 2am."""
+
+    #: `arm_a` (argilla, response-level, 2,310 pairs) or `arm_b` (curated
+    #: PRM800K shared-prefix step-local pairs). Selects everything else.
+    arm: str = "arm_a"
+    #: Section-11's SFT checkpoint. Both arms fine-tune from the SAME one, so
+    #: any arm-vs-arm difference is the preference data, not the start point.
+    sft_checkpoint: str = ""
+    #: Section-11's quality-gate receipt. The stage refuses to launch unless
+    #: it exists AND records a pass.
+    sft_quality_gate_path: str = ""
+    #: Pair shards for this arm (train split) and the problem-level held-out
+    #: split used for preference accuracy.
+    shard_dir: str = ""
+    heldout_shard_dir: str = ""
+    #: Reference-logprob caches written by scripts/precompute_ref_logprobs.py.
+    ref_cache_path: str = ""
+    heldout_ref_cache_path: str = ""
+    #: The mandatory pre-launch path-bias measurement's report.
+    path_bias_report_path: str = ""
+    #: Section-03's build manifest. Its hash is part of the ref-cache key, so
+    #: a shard rebuild cannot leave a stale cache looking valid.
+    manifest_path: str = "/apps/development/rlhf/align/data/manifest.json"
+    #: docs/RESULTS.md must already carry this arm's pre-registration
+    #: (expectations + tripwire response) before the run starts.
+    #: The pre-registration is a FROZEN, dated document that lives with the
+    #: predecessor experiment record (read-only; the gate only greps it).
+    #: Copying it would fork the thing that is supposed to be frozen, so the
+    #: default stays on the original until it is migrated deliberately.
+    results_doc: str = "/apps/development/rlhf/nano-mistral-o1/docs/RESULTS.md"
+    run_dir: str = ""
+
+    #: Post-filter pair count for this arm, from the build manifest. Drives
+    #: the optimizer-step horizon, so it is a measured number, never a guess.
+    expected_pairs: int = 2310
+
+    # -- objective ---------------------------------------------------------
+    #: Pre-registered for both arms. `dpop`/`nll_aux` are reachable only as
+    #: the tripwire RESPONSE, and switching is a recorded deviation.
+    variant: str = "sigmoid"
+    #: Preference sharpness. 0.1, both arms, pre-registered.
+    beta: float = 0.1
+    #: NLL auxiliary weight. 0.0 is provably inert in section-05's kernel (the
+    #: term is skipped structurally, not multiplied by zero).
+    nll_coeff: float = 0.0
+    dpop_lambda: float = 50.0
+    #: Which pre-registered response to apply if the chosen-logprob tripwire
+    #: fires. Defaults to the conjunction docs/RESULTS.md froze; the two
+    #: single-knob forms are permitted, anything else raises.
+    tripwire_response: str = "dpop_nll_aux"
+
+    # -- model -------------------------------------------------------------
+    #: Attention backend for the policy model. NEVER handed raw to a
+    #: constructor: `NmoMistralForCausalLM.__init__` routes it through
+    #: `align.model.attention.select_attn_implementation`, which RAISES if the
+    #: request cannot be honoured -- for `triton_fa2` that means section-06's
+    #: equivalence-gate receipt must be valid for THIS kernel source on THIS
+    #: GPU. The pre-registered value for BOTH arms is `triton_fa2`
+    #: (docs/RESULTS.md dated entry, 2026-08-17); the neutral default here is
+    #: sdpa so the pin lives in the arm configs, where changing it is a
+    #: visible diff against the pre-registration. Membership is validated at
+    #: config load: a typo'd backend raises at launch, in milliseconds, not
+    #: as a selector error with four GPUs claimed.
+    attn_implementation: str = "sdpa"
+
+    # -- optimization ------------------------------------------------------
+    #: 5e-7, pre-registered. Two orders below SFT's: DPO moves a model that is
+    #: already where it should be.
+    lr: float = 5e-7
+    #: 2 pre-registered, 3 maximum.
+    epochs: int = 2
+    #: ~32 via gradient accumulation (e.g. 4 ranks x micro 2 x accum 4).
+    global_batch: int = 32
+    #: >= 4 is MANDATORY for full-FT with offloaded AdamW (section-07 refuses
+    #: less): at accum 1 the CPU optimizer step dominates wall-clock.
+    grad_accum: int = 4
+    micro_batch_size: int = 2
+    num_workers: int = 2
+    save_every_micro_steps: int = 100
+    keep_checkpoints: int = 2
+    #: Each rank writes its own optimizer_rank{r}.pt instead of gathering the
+    #: full optimizer state (see TrainConfig.per_rank_optimizer_state). Without
+    #: this field the key was REJECTED by struct mode in DPO TOMLs and, worse,
+    #: scripts/dpo.py built its RunConfig without ever setting it -- so the
+    #: trainer-level flag was unreachable from a DPO run. Default False keeps
+    #: the gathered layout; a resume at a different world_size is refused by
+    #: the trainer when True.
+    per_rank_optimizer_state: bool = False
+
+    #: Global-norm gradient clip (OptimizerConfig.max_grad_norm), forwarded by
+    #: scripts/dpo.py::_run_config. Declared here because the DPO RunConfig is
+    #: built from defaults: before this field the clip could not be stated in a
+    #: DPO TOML at all, only inherited. The recorded grad norm stays PRE-clip.
+    max_grad_norm: Optional[float] = 1.0
+
+    # -- run-health bounds (G0), forwarded to TrainConfig by the stage script's
+    # _run_config: a field that is not copied trains at the default while the
+    # TOML says otherwise. Semantics: align/config.py::TrainConfig.
+    max_wall_seconds: Optional[float] = None
+    max_consecutive_skipped_steps: int = 3
+    max_skipped_fraction: float = 0.05
+
+    # -- measurement -------------------------------------------------------
+    #: Fires the path-bias rerun when max|delta ref-margin| exceeds this
+    #: fraction of the early-margin scale. 0.05 = a 5% perturbation of the
+    #: quantity DPO reads. Chosen as a threshold to MEASURE against, not a
+    #: tolerance to pass: the measured single-GPU-vs-sharded deltas are
+    #: expected orders below it, and if they are not, the cache is rebuilt
+    #: through the sharded forward rather than argued about.
+    path_bias_threshold: float = 0.05
+    #: Denominator of that ratio: the typical magnitude of the implicit margin
+    #: early in the run, where the loss argument sits near zero and a
+    #: reference bias is least attenuated.
+    path_bias_early_margin_scale: float = 0.5
+    #: Held-out preference accuracy cadence, in optimizer steps.
+    heldout_every_optimizer_steps: int = 50
+    #: Pairs the chosen-longer fraction is computed over, as a ROLLING
+    #: WINDOW. Not per micro-batch, which is what the statistic naively is:
+    #: at the pre-registered micro_batch 2 the per-batch fraction can only be
+    #: 0, 0.5 or 1, so against a +-0.10 band it false-fires ~78% of the time.
+    #:
+    #: 512 is DERIVED, not chosen. SE of a fraction near the 0.5489 baseline
+    #: is sqrt(p(1-p)/k), so the band sits at 0.10/SE standard errors:
+    #:   k=2 -> 0.28 SE (78% false-fire)     k=128 -> 2.27 SE (2.3e-02)
+    #:   k=32 -> 1.14 SE (26%)               k=256 -> 3.22 SE (1.3e-03)
+    #:                                       k=512 -> 4.55 SE (5.4e-06)
+    #: At 512 the per-check false-fire rate is 5e-06; over Arm B's ~3,000
+    #: logged micro-steps that is ~1.6e-02 expected spurious alarms for the
+    #: whole run. 256 would give ~4 -- enough to bury a real one.
+    length_window_pairs: int = 512
+    #: Which tripwires this arm ARMS, copied into `RunConfig.tripwires.armed`.
+    #: Both, for both arms: the collapse tripwire is the one with a
+    #: pre-registered response, and the length-drift one guards the 0.5489
+    #: baseline. Section-07 halts at the first log step if an armed tripwire's
+    #: inputs are missing or NaN, so arming here is a commitment that this
+    #: stage's loss adapter emits them -- not a hopeful setting.
+    armed_tripwires: tuple[str, ...] = ("logps_chosen", "length_drift", "grad_norm")
+    #: Length-drift band around the arm's MEASURED chosen-longer baseline
+    #: (Arm A: 0.5489 from the build manifest -- never hardcoded here).
+    length_band: float = 0.10
+    logps_window: int = 20
+
+    def __post_init__(self) -> None:
+        if self.arm not in ("arm_a", "arm_b"):
+            raise ValueError(f"arm={self.arm!r} must be 'arm_a' or 'arm_b'")
+        if self.tripwire_response not in PRE_REGISTERED_RESPONSES:
+            raise ValueError(
+                f"tripwire_response={self.tripwire_response!r} is not "
+                f"pre-registered; permitted: {list(PRE_REGISTERED_RESPONSES)}"
+            )
+        # Imported here, not at module top: align.model.attention pulls torch,
+        # and this dataclass must stay constructible before the resource
+        # guards have configured the process.
+        from ..model.attention import NMO_IMPLEMENTED
+
+        if self.attn_implementation not in NMO_IMPLEMENTED:
+            raise ValueError(
+                f"attn_implementation={self.attn_implementation!r} is not an "
+                f"align attention code path; permitted: {sorted(NMO_IMPLEMENTED)}. "
+                "A typo must raise at config load, never fall through to a "
+                "default backend -- a run that reports one backend while "
+                "executing another poisons every downstream comparison."
+            )
+
+
+def load_stage_config(path: str, overrides: Optional[Sequence[str]] = None) -> DpoStageConfig:
+    """TOML (+ dotted `key=value` overrides) -> DpoStageConfig, struct mode on:
+    a typo'd key raises at launch instead of silently training with the
+    default it failed to override."""
+    from align.config import load_config
+
+    return load_config(DpoStageConfig, path, list(overrides or []))
+
+
+def rotate_stale_checkpoints(ckpt_dir: str) -> Optional[str]:
+    """Move a previous run's checkpoint dir aside, ONCE, for a fresh arm.
+
+    COLLECTIVE: every rank calls it. Only rank 0 touches the filesystem and
+    every rank then waits at a barrier, so all ranks proceed with the same
+    directory state. Run on every rank, the bare `os.rename` let exactly one
+    rank win and the other three die with FileNotFoundError (ranks straddling
+    a second boundary even computed different destinations); the survivors
+    had already entered Trainer construction. Returns the destination on
+    rank 0 when a rotation happened, else None.
+    """
+    import time
+
+    from ..utils.dist import barrier, is_rank0
+
+    dst = None
+    if is_rank0() and os.path.isdir(ckpt_dir) and os.listdir(ckpt_dir):
+        dst = f"{ckpt_dir}_stale_{time.strftime('%Y%m%d-%H%M%S')}"
+        os.rename(ckpt_dir, dst)
+        logger.warning("rotated stale checkpoints to %s", dst)
+    barrier()
+    return dst
+
+
+def probe_micro_cap(micro_step: int, probe_steps: int) -> int:
+    """The `max_micro_steps` cap for the launch probe, on the ABSOLUTE counter.
+
+    The Trainer tests `micro_step >= limit` against the absolute micro-step
+    count, which a resume restores. A bare `probe_steps` cap on a run resumed
+    at micro N >= probe_steps therefore trained exactly one micro-step and
+    left `throughput_probe` one time to aggregate (ValueError), so every
+    supervisor retry died here without a checkpoint. The probe is the first
+    `probe_steps` micro-steps of THIS process, wherever it starts.
+    """
+    return int(micro_step) + int(probe_steps)
+
+
+def reduce_tripwire_alarms(alarms: Sequence[Mapping]) -> list[dict]:
+    """Every rank's tripwire alarms, unioned: the same list on all ranks.
+
+    The pre-registered rerun is a COLLECTIVE decision (a second model load, a
+    new FSDP wrap, a new Trainer). The trainer's alarms are built from
+    rank-local metrics over different pairs, and only HALT verdicts are
+    reduced inside the trainer, so a non-HALT alarm (the chosen-logprob
+    decline) can fire on some ranks only. Acting on the local list sends the
+    firing ranks into the rerun's collectives while the rest write the
+    manifest and exit. COLLECTIVE: every rank must call it. Each alarm gains
+    a `rank` key naming the rank that raised it.
+    """
+    from ..utils.dist import all_gather_objects
+
+    union: list[dict] = []
+    for r, rank_alarms in enumerate(all_gather_objects([dict(a) for a in alarms])):
+        for a in rank_alarms:
+            union.append({**a, "rank": r})
+    return union
+
+
+def resolve_pair_split_dir(shard_dir: str, split: str) -> str:
+    """The directory holding `<split>-NNNN.receipt.json`, found by DISCOVERY.
+
+    THE HAND-COMPOSED-PATH DEFECT, DPO edition. The attempt-1 configs pointed
+    at `data/shards/arm_a-train`, composed from a naming convention; the store
+    actually writes `data/shards/arm_a/train/arm_a-train-0000.receipt.json`.
+    Section-11 hit the identical defect and fixed it with
+    `stage_smoke.resolve_split_dir`, which searches for the receipts the store
+    itself wrote; this delegates to that same function (read-only) rather than
+    growing a second, drifting copy. Accepts either the shard root or the
+    exact split directory -- both contain the split's receipts -- and a
+    missing split reports the splits that ARE present.
+    """
+    from .stage_smoke import resolve_split_dir
+
+    return str(resolve_split_dir(shard_dir, split))
+
+
+def is_hf_checkpoint(path: Optional[str]) -> bool:
+    """A directory is a loadable checkpoint iff it carries `config.json` --
+    the file `AutoConfig.from_pretrained` reads first. A run dir that merely
+    CONTAINS checkpoints is not itself one."""
+    return bool(path) and os.path.isfile(os.path.join(str(path), "config.json"))
+
+
+def resolve_sft_checkpoint(configured: str, receipt_checkpoint: Optional[str]) -> tuple[str, str]:
+    """The policy's starting checkpoint: `(path, source)`, or a raise.
+
+    Section-11's Trainer writes `<run_dir>/checkpoints/micro_NNNNNNNN`
+    (trainer.save_checkpoint) and there is NO stable `final` alias -- the
+    attempt-1 configs pointed at `checkpoints/sft/final`, a path nothing ever
+    writes, which would have failed at model load with the gates already
+    passed. What IS stable is the quality-gate receipt: section-11 records
+    WHICH micro-step checkpoint passed its gate in the receipt's `checkpoint`
+    field, so that is the discovery mechanism. A configured path that is
+    itself a checkpoint wins, as the explicit override for rerunning against
+    a pinned checkpoint.
+    """
+    if is_hf_checkpoint(configured):
+        return configured, "config"
+    if is_hf_checkpoint(receipt_checkpoint):
+        return str(receipt_checkpoint), "sft_quality_gate_receipt"
+    raise FileNotFoundError(
+        f"no SFT checkpoint found. Configured sft_checkpoint={configured!r} "
+        "is not a checkpoint directory (no config.json), and the quality-gate "
+        f"receipt records checkpoint={receipt_checkpoint!r}, which is not one "
+        "either. Section-11's trainer writes "
+        "<run_dir>/checkpoints/micro_NNNNNNNN and its receipt names the one "
+        "that passed the gate; refusing to guess a starting point for a "
+        "full-FT arm."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Step arithmetic
+# ---------------------------------------------------------------------------
+
+
+def optimizer_steps(pairs: int, global_batch: int, epochs: int) -> int:
+    """epochs * ceil(pairs / global_batch).
+
+    Written once, here, because the iterations-vs-steps confusion is a
+    documented predecessor incident: an LR horizon quoted in MICRO-steps
+    schedules the cosine over 4-8x the real run, so the LR never decays and
+    the run ends at nearly its peak rate with no sign anything was wrong.
+    Ceiling, not floor -- a trailing partial batch is still an optimizer step.
+    """
+    if pairs < 1:
+        raise ValueError(f"pairs={pairs} must be >= 1")
+    if global_batch < 1:
+        raise ValueError(f"global_batch={global_batch} must be >= 1")
+    if epochs < 1:
+        raise ValueError(f"epochs={epochs} must be >= 1")
+    return epochs * math.ceil(pairs / global_batch)
+
+
+# ---------------------------------------------------------------------------
+# Entry gates
+# ---------------------------------------------------------------------------
+
+
+def run_entry_gates(
+    cfg: DpoStageConfig, gates: Mapping[str, Callable[[DpoStageConfig], dict]]
+) -> list[dict]:
+    """Run every gate in `GATE_ORDER`, stopping at the first failure.
+
+    Returns one record per gate that ran. A gate absent from `gates` RAISES:
+    an unimplemented gate treated as passed is how a run launches with an
+    unchecked reference cache, and the whole point of the ordered sequence is
+    that reaching gate N means gates 1..N-1 held.
+    """
+    report: list[dict] = []
+    for name in GATE_ORDER:
+        fn = gates.get(name)
+        if fn is None:
+            raise RuntimeError(
+                f"entry gate {name!r} has no implementation. A missing gate is "
+                "not a passed gate; refusing to launch."
+            )
+        result = fn(cfg)
+        record = {"gate": name}
+        if isinstance(result, Mapping):
+            record.update(result)
+        report.append(record)
+        logger.info("entry gate ok: %s", name)
+    return report
+
+
+#: The gate whose pass this stage depends on, named in section-11's own
+#: `stage_gates.GATE_ORDER` rather than spelled here as a string this module
+#: invented.
+SFT_GATE_NAME = "quality_gate"
+
+
+def read_sft_quality_gate(path: str) -> dict:
+    """Section-11's quality-gate verdict: exists, and records a PASS.
+
+    Reads section-11's OWN artifact. `scripts/sft.py` records its gates
+    through `stage_gates.GateLedger`, which writes `stage_gates.json` as
+    `{"eval_source_hash": ..., "gates": {"quality_gate": {"passed": ...}}}`.
+    An earlier draft of this function read a standalone
+    `{"passed": true}` receipt that nothing writes -- the gate would have
+    raised FileNotFoundError with four GPUs already claimed. `path` may name
+    the ledger file or the run directory containing it.
+
+    The comparison itself is NOT restated here. Section-11's
+    `stage_gates.quality_gate` delegates to section-10's `assert_meets_golden`,
+    which compares against the WILSON LOWER BOUND -- the golden anchors are
+    themselves measurements with sampling noise, so demanding the point
+    estimate meet them would fail correct runs about half the time. This stage
+    consumes that verdict; a second copy of the CI reasoning here would drift
+    from it, and the drift would be invisible.
+
+    Three states, and the third is the one that gets skipped. Missing: section
+    11 never ran, or ran and never recorded. Present and failed: obvious.
+    Present and SILENT on this gate: also a refusal, because an artifact that
+    does not answer the question is indistinguishable from one that answers
+    'no', and defaulting to 'yes' would let a partial ledger authorise eight
+    GPU-hours.
+    """
+    from .stage_gates import LEDGER_FILENAME, PENDING
+
+    resolved = path
+    if path and os.path.isdir(path):
+        resolved = os.path.join(path, LEDGER_FILENAME)
+    if not resolved or not os.path.isfile(resolved):
+        raise FileNotFoundError(
+            f"SFT quality gate ledger {resolved!r} is missing. DPO fine-tunes "
+            "the SFT checkpoint, so an SFT run that never passed its own gate "
+            "makes both arms measurements of a broken starting point -- at "
+            "full GPU cost, with nothing downstream able to tell. Point "
+            "sft_quality_gate_path at section-11's run dir or its "
+            f"{LEDGER_FILENAME}."
+        )
+    with open(resolved) as f:
+        blob = json.load(f)
+
+    # Accept the ledger shape, and a bare {"passed": ...} object too: the
+    # latter is what a hand-written override looks like, and refusing it would
+    # push someone toward editing the ledger instead.
+    gates = blob.get("gates", blob if "passed" in blob else {})
+    entry = gates.get(SFT_GATE_NAME, blob if "passed" in blob else PENDING)
+    if not isinstance(entry, dict) or "passed" not in entry:
+        raise RuntimeError(
+            f"{resolved!r} does not record whether the SFT {SFT_GATE_NAME!r} "
+            f"passed (it reads {entry!r}). An artifact that cannot answer the "
+            "question is not a pass."
+        )
+    if not entry["passed"]:
+        raise RuntimeError(
+            f"the SFT quality gate did not pass: "
+            f"{entry.get('detail', '(no detail recorded)')}. Fix section-11 "
+            "before spending this stage's GPU budget."
+        )
+    return {
+        "passed": True,
+        "detail": entry.get("detail", ""),
+        "recorded_at": entry.get("recorded_at"),
+        "eval_source_hash": blob.get("eval_source_hash"),
+        "checkpoint": blob.get("checkpoint", entry.get("checkpoint")),
+        "ledger": resolved,
+    }
+
+
+def check_arm_descoped(arm: str, manifest: Mapping) -> dict:
+    """Honour section-03's Arm A DESCOPED marker (filter yield < 1,500).
+
+    The failure this prevents is quiet and expensive: training Arm A on
+    whatever survived the filter and reporting it under the pre-registered
+    arm's name. A descoped arm is a RESULT -- it gets a manifest saying so --
+    not an absence.
+    """
+    descoped = bool(manifest.get("arm_a_descoped", False)) and arm == "arm_a"
+    return {
+        "arm": arm,
+        "descoped": descoped,
+        "train": not descoped,
+        "detail": (
+            "section-03 marked Arm A descoped (filter yield below the "
+            "1,500 floor); refusing to train it and recording the "
+            "descope as this arm's result"
+            if descoped
+            else "arm is in scope"
+        ),
+    }
+
+
+def build_manifest_gate(cfg: DpoStageConfig, manifest: Mapping) -> dict:
+    """Counts match the pinned values and the build verified its own pins."""
+    if not manifest.get("pins_verified", False):
+        raise RuntimeError(
+            "the build manifest reports pins_verified=false: the corpus was "
+            "built in re-measurement mode and its counts were never checked "
+            "against the pinned values. Sections 11 and 12 refuse to launch "
+            f"against it. Violations: {manifest.get('pin_violations')}"
+        )
+    counts = manifest.get("counts", {})
+    key = "arm_a_pairs" if cfg.arm == "arm_a" else "arm_b_curated"
+    measured = counts.get(key)
+    if measured is None:
+        raise KeyError(
+            f"build manifest carries no {key!r} count, so the optimizer-step "
+            "horizon cannot be derived from a measured pair count."
+        )
+    if int(measured) != int(cfg.expected_pairs):
+        raise RuntimeError(
+            f"{cfg.arm}: manifest reports {measured} pairs but the stage "
+            f"config expects {cfg.expected_pairs}. The horizon, the LR "
+            "schedule and the ref cache were all built for the config number; "
+            "one of the two is stale."
+        )
+    return {"pairs": int(measured), **check_arm_descoped(cfg.arm, manifest)}
+
+
+# ---------------------------------------------------------------------------
+# Banner
+# ---------------------------------------------------------------------------
+
+
+def measure_chosen_longer_fraction(rows: Sequence[Any]) -> float:
+    """Fraction of pairs whose CHOSEN completion is longer, from the rows.
+
+    Needed because section-03's build manifest records Arm A's baseline
+    (0.5489) and not Arm B's, while the length-drift tripwire needs one per
+    arm: `LengthDriftTripwire` refuses a NaN baseline outright, and borrowing
+    Arm A's would arm Arm B against a distribution it was never drawn from --
+    a tripwire that then fires forever or never, both of which are worse than
+    no tripwire because they look like measurements.
+
+    Completion length, not sequence length: Arm B's two sides share their
+    entire prefix, so total length is dominated by a term identical on both
+    sides and would compress the fraction toward 0.5 regardless of the data.
+    """
+    import numpy as np
+
+    longer = 0
+    counted = 0
+    for r in rows:
+        lc = len(np.asarray(r["chosen_input_ids"])) - int(r["chosen_prompt_len"])
+        lr = len(np.asarray(r["rejected_input_ids"])) - int(r["rejected_prompt_len"])
+        if lc <= 0 and lr <= 0:
+            continue  # no completion on either side
+        counted += 1
+        longer += int(lc > lr)
+    return longer / counted if counted else NAN
+
+
+def dpo_banner(
+    *,
+    cfg: DpoStageConfig,
+    manifest: Mapping,
+    ref_cache_key,
+    path_bias: Mapping,
+    probe: Mapping,
+    rss: Mapping,
+    backend: str,
+    sft_gate: Mapping,
+    length_baseline: Optional[tuple[float, str]] = None,
+    pre_registration: str = "docs/RESULTS.md",
+) -> dict:
+    """Everything this run depends on, or a raise.
+
+    Two things are enforced here rather than merely reported. The length-drift
+    baseline is READ FROM THE MANIFEST (Arm A measured 0.5489; Arm B's is
+    different, so a constant would mis-arm the tripwire on one arm forever),
+    and a path-bias decision of 'rerun through the sharded forward' must
+    actually have been honoured -- a measurement taken, recorded, and ignored
+    is worse than one never taken, because it leaves evidence that the
+    question was considered.
+    """
+    baseline_key = f"{cfg.arm}_chosen_longer_fraction"
+    baselines = manifest.get("length_bias_baseline", {})
+    if length_baseline is not None:
+        baseline, baseline_source = length_baseline
+        if math.isnan(baseline):
+            raise ValueError(
+                f"{baseline_key}: the supplied baseline is NaN, which means "
+                "nobody measured it. The length-drift tripwire refuses a NaN "
+                "baseline, so this would disarm it silently."
+            )
+    elif baseline_key in baselines:
+        baseline, baseline_source = baselines[baseline_key], "build_manifest"
+    else:
+        raise KeyError(
+            f"build manifest has no {baseline_key!r} and no measured baseline "
+            "was supplied. The length-drift tripwire needs THIS arm's "
+            "chosen-longer fraction; defaulting to the other arm's would "
+            "mis-arm it permanently. scripts/dpo.py measures it from the "
+            "arm's own shards when the manifest lacks it "
+            "(measure_chosen_longer_fraction)."
+        )
+
+    if path_bias.get("fired") and getattr(ref_cache_key, "forward_impl", None) != "fsdp2_sharded":
+        raise RuntimeError(
+            "the path-bias measurement fired (action="
+            f"{path_bias.get('action')!r}, ratio={path_bias.get('ratio')}) but "
+            f"the reference cache is tagged "
+            f"{getattr(ref_cache_key, 'forward_impl', None)!r}, not "
+            "'fsdp2_sharded'. Re-run the precompute through the sharded "
+            "forward. A decision that was measured, recorded and then ignored "
+            "is worse than one never taken."
+        )
+
+    horizon = optimizer_steps(cfg.expected_pairs, cfg.global_batch, cfg.epochs)
+    return {
+        "arm": cfg.arm,
+        "loss_type": f"dpo_{cfg.variant}",
+        "variant": cfg.variant,
+        "beta": cfg.beta,
+        "lr": cfg.lr,
+        "epochs": cfg.epochs,
+        "global_batch": cfg.global_batch,
+        "grad_accum": cfg.grad_accum,
+        "expected_pairs": cfg.expected_pairs,
+        "optimizer_step_formula": f"{cfg.epochs} x ceil({cfg.expected_pairs}/{cfg.global_batch}) "
+        f"= {horizon}",
+        "scheduler_horizon_optimizer_steps": horizon,
+        "dataset_hashes": dict(manifest.get("hashes", {})),
+        "ref_cache_key": json.loads(ref_cache_key.to_json()),
+        "path_bias": dict(path_bias),
+        "probe_step_time_s": probe.get("step_time_s_median", NAN),
+        "probe_derived_wall_clock_h": probe.get("projected_wall_clock_h", NAN),
+        "rss_budget": dict(rss),
+        "attention_backend": backend,
+        "length_drift_baseline": baseline,
+        "length_drift_baseline_source": baseline_source,
+        "sft_checkpoint": sft_gate.get("checkpoint", cfg.sft_checkpoint),
+        "pre_registration": pre_registration,
+    }
+
+
+def render_dpo_banner(resolved: Mapping, title: str = "align dpo") -> str:
+    """Render, or raise. Never a partial banner.
+
+    Checks this stage's keys AND section-01's global mandatory set, without
+    mutating the global registry -- see DPO_BANNER_KEYS for why.
+    """
+    from ..registry import MANDATORY_BANNER_KEYS
+
+    required = list(DPO_BANNER_KEYS) + [
+        k for k in MANDATORY_BANNER_KEYS if k not in DPO_BANNER_KEYS
+    ]
+    missing = [k for k in DPO_BANNER_KEYS if k not in resolved]
+    unresolved = [k for k in DPO_BANNER_KEYS if k in resolved and resolved[k] is None]
+    if missing or unresolved:
+        raise RuntimeError(
+            "DPO launch banner incomplete -- refusing to start. A run that "
+            f"cannot answer its question does not begin. missing={missing} "
+            f"unresolved(None)={unresolved}. Undefined values must be "
+            "EXPLICIT: float('nan') or 'PENDING', never absent."
+        )
+    width = 78
+    lines = ["=" * width, f"| {title} [{resolved['arm']}]", "=" * width]
+    for key in required:
+        if key in resolved:
+            lines.append(f"| {key:<36} {resolved[key]!r}")
+    for key in sorted(k for k in resolved if k not in required):
+        lines.append(f"| {key:<36} {resolved[key]!r}   (extra)")
+    lines.append("=" * width)
+    return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# Tripwire response
+# ---------------------------------------------------------------------------
+
+
+def apply_tripwire_response(
+    cfg: DpoStageConfig, alarms: Sequence[Mapping]
+) -> tuple[DpoStageConfig, dict]:
+    """Map a fired chosen-logprob tripwire to the PRE-REGISTERED response.
+
+    Returns `(config to rerun the arm with, record)`. The record goes into the
+    arm manifest, so the deviation is part of the result rather than a note
+    somebody remembers.
+
+    The alarm being responded to: `logps/chosen` declining while the margin
+    improves. DPO's objective is satisfied by separating the two logprobs, and
+    crushing the rejected one separates them just as well as raising the
+    chosen one -- while degrading the policy. The mode is MAXIMAL FOR ARM B by
+    design, because pi_SFT already assigns those chosen continuations high
+    likelihood, so there is little headroom above and a great deal below.
+
+    Every permitted response puts a floor under the chosen logprob: `dpop`
+    adds an explicit penalty when it falls below the reference, `nll_aux` adds
+    a likelihood term on the chosen completion, and `dpop_nll_aux` -- the
+    frozen pre-registration's wording -- does both. Beta is NOT changed, and
+    the arm re-runs ONCE: a second rerun would be a third configuration.
+    """
+    # "logps_chosen" is the WIRE NAME the trainer records (trainer.py appends
+    # {"alarm": name} for the armed wire key). This filter matched a name
+    # nobody emits ("logps_chosen_decline") until 2026-08-17 -- a fired
+    # collapse tripwire could never trigger the pre-registered rerun, and the
+    # unit test asserted the filter against its own literal instead of the
+    # trainer's. Silent-failure trap #7's shape: armed-looking, disconnected.
+    fired = [a for a in alarms if a.get("alarm") == "logps_chosen"]
+    if not fired:
+        return cfg, {
+            "fired": False,
+            "rerun_required": False,
+            "pre_registered": True,
+            "configuration": f"{cfg.variant} beta={cfg.beta}",
+            "detail": "no chosen-logprob alarm; primary configuration stands as the arm's result",
+        }
+
+    if cfg.tripwire_response not in PRE_REGISTERED_RESPONSES:
+        raise ValueError(
+            f"tripwire_response={cfg.tripwire_response!r} is outside the "
+            f"pre-registration ({list(PRE_REGISTERED_RESPONSES)}). A third "
+            "configuration requires a NEW pre-registration written before the "
+            "rerun -- choosing one after seeing the data is how a stage picks "
+            "its own success criterion."
+        )
+    if cfg.tripwire_response == "dpop_nll_aux":
+        new = replace(cfg, variant="dpop", nll_coeff=NLL_AUX_COEFF)
+    elif cfg.tripwire_response == "dpop":
+        new = replace(cfg, variant="dpop")
+    else:
+        new = replace(cfg, variant="sigmoid", nll_coeff=NLL_AUX_COEFF)
+    record = {
+        "fired": True,
+        "rerun_required": True,
+        "pre_registered": True,
+        "alarms": [dict(a) for a in fired],
+        "response": cfg.tripwire_response,
+        "configuration": (
+            f"{new.variant} beta={new.beta} nll_coeff={new.nll_coeff} dpop_lambda={new.dpop_lambda}"
+        ),
+        "detail": (
+            "chosen-logprob collapse detected; the arm re-runs with the "
+            "pre-registered response and THAT configuration is recorded "
+            "as the arm's result"
+        ),
+    }
+    logger.error("tripwire response applied: %s", record["configuration"])
+    return new, record
+
+
+# ---------------------------------------------------------------------------
+# Arm result manifest
+# ---------------------------------------------------------------------------
+
+
+def arm_result_manifest(
+    *,
+    cfg: DpoStageConfig,
+    descoped: bool,
+    summary: Mapping,
+    heldout: Mapping,
+    banner: Mapping,
+    checkpoint: Optional[str],
+    tripwire: Optional[Mapping] = None,
+) -> dict:
+    """What section-13 reads to select an arm and run the paired evaluation.
+
+    A descoped arm still produces a manifest. Its accuracy is NaN, not 0.0:
+    zero is a plausible catastrophic result and would read as one, while NaN
+    says nobody measured -- the distinction section-13's arm-selection rule
+    depends on.
+    """
+    return {
+        "arm": cfg.arm,
+        "status": "descoped" if descoped else "trained",
+        "checkpoint": None if descoped else checkpoint,
+        "loss_variant_actually_used": None if descoped else cfg.variant,
+        "beta": cfg.beta,
+        "nll_coeff": cfg.nll_coeff,
+        "lr": cfg.lr,
+        "epochs": cfg.epochs,
+        "global_batch": cfg.global_batch,
+        "optimizer_steps_planned": optimizer_steps(
+            cfg.expected_pairs, cfg.global_batch, cfg.epochs
+        ),
+        "optimizer_steps_run": summary.get("optimizer_step", 0),
+        "heldout_preference_accuracy": heldout.get("accuracy", NAN),
+        "heldout_ci": [heldout.get("ci_low", NAN), heldout.get("ci_high", NAN)],
+        "heldout_n": heldout.get("n", 0),
+        "final_mean_margin": heldout.get("mean_margin", NAN),
+        "tripwire": dict(tripwire) if tripwire else {"fired": False},
+        "tripwire_alarms": list(summary.get("tripwire_alarms", [])),
+        "trunk_report": dict(summary.get("trunk_report", {})),
+        "checkpoints_written": list(summary.get("checkpoints_written", [])),
+        "banner": dict(banner),
+    }
+
+
+def write_arm_result(path: str, manifest: Mapping) -> str:
+    """Write the arm manifest: rank 0 only, atomically.
+
+    COLLECTIVE (every rank calls it; the closing barrier keeps ranks together).
+    Run on all four torchrun ranks with a bare `open(path, "w")`, the ranks
+    truncated and wrote one path within a millisecond, with rank-local content
+    (`checkpoint` is "" off rank 0): the recorded Arm B manifest is
+    unparseable. Section-13 reads this file to select an arm, so it is written
+    once, via tmp + os.replace, and rank 0's `checkpoint` is the real one.
+    """
+    from ..utils.dist import barrier, is_rank0
+    from ..utils.guard import assert_paths_off_root
+    from .ref_cache import _atomic_write
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(path))])
+    if is_rank0():
+
+        def _dump(tmp: str) -> None:
+            with open(tmp, "w") as f:
+                # default=str: the manifest is assembled by the caller from eval
+                # output and this writes an ARM'S RESULT after its full DPO run --
+                # the exact shape (harness-measured values serialized at the end)
+                # that killed SFT attempt 6 with a Tensor in a gate receipt.
+                json.dump(manifest, f, indent=2, default=str)
+                f.flush()
+                os.fsync(f.fileno())
+
+        _atomic_write(path, _dump)
+        logger.info("arm result manifest -> %s", path)
+    barrier()
+    return path
+
+
+# ---------------------------------------------------------------------------
+# Trainer-facing batch + loss adapter
+# ---------------------------------------------------------------------------
+
+
+class DpoBatch:
+    """A `Preference` batch, iterable over its pair ids.
+
+    Section-07's `Trainer` iterates the object it is handed and logs
+    `list(batch)` as that step's identity -- its own tests pass samplers that
+    yield row-index lists, for which `list(batch)` is exactly right. A
+    `Preference` dataclass is not iterable, so a bare one would raise
+    TypeError at the first log step.
+
+    Wrapping supplies the identity the log wants, and supplies a better one
+    than row indices: which PAIRS a step consumed. Row indices are meaningless
+    across a reshuffle or a rebuild; pair ids survive both, so a suspicious
+    step stays reproducible afterwards.
+    """
+
+    __slots__ = ("preference", "pair_ids")
+
+    def __init__(self, preference):
+        self.preference = preference
+        self.pair_ids = [(m or {}).get("pair_id") for m in preference.metadata]
+
+    def __iter__(self):
+        return iter(self.pair_ids)
+
+    def __len__(self) -> int:
+        return len(self.pair_ids)
+
+    def to(self, device) -> "DpoBatch":
+        return DpoBatch(self.preference.to(device))
+
+
+class PreferenceBatchSource:
+    """What section-07's `Trainer` iterates: `set_epoch` / `__iter__` / `__len__`.
+
+    The Trainer walks a batch source directly rather than a `DataLoader`, and
+    that is a deliberate section-07 choice: batch order becomes a pure
+    function of (seed, epoch), which is what makes the seeded-resume contract
+    checkable. It also means there are no dataloader worker processes in this
+    loop at all -- the `<=2 workers per rank` rule that the /dev/shm SIGBUS
+    incident produced has nothing to bind to here, because nothing is forked.
+    Collation happens inline, in the training process.
+
+    That is affordable for DPO specifically: pair shards are memory-mapped and
+    a micro-batch is two short sequences, so collation is microseconds against
+    a 7B forward. It would not be affordable for a tokenization-heavy loop.
+    """
+
+    def __init__(self, rows: Sequence[Any], sampler, pad_token_id: int):
+        self.rows = list(rows)
+        self.sampler = sampler
+        self.pad_token_id = int(pad_token_id)
+
+    def set_epoch(self, epoch: int) -> None:
+        self.sampler.set_epoch(epoch)
+
+    def __len__(self) -> int:
+        return len(self.sampler)
+
+    def __iter__(self):
+        from ..data.loader import collate_preference
+
+        for indices in self.sampler:
+            yield DpoBatch(
+                collate_preference([self.rows[i] for i in indices], pad_token_id=self.pad_token_id)
+            )
+
+
+def make_preference_sampler(
+    rows, *, max_tokens: int, max_rows: int, seed: int = 0, rank=None, world_size=None
+):
+    """Section-03's length-grouped sampler, RANK-SHARDED.
+
+    THE UNSHARDED-SAMPLER INCIDENT (~24 GPU-hours on Arm A). This used to
+    return the bare `LengthGroupedSampler`, so every rank walked the same
+    batch list in the same order and all four trained the SAME pair every
+    micro-step. The effective global batch was micro x accum (8), not
+    micro x accum x world (32) as pre-registered; one "epoch" therefore
+    consumed a quarter of the corpus; the horizon covered half an epoch; and
+    past it the cosine schedule clamps to lr=0, so the run kept stepping on
+    frozen weights for 287 more steps. Nothing failed -- the loss curve was
+    plausible and the manifest reported the intended batch size throughout.
+
+    SHARDING THE BATCHES, NOT THE ROWS, and the choice matters twice over:
+
+      * `PreferenceBatchSource` indexes the FULL row list with whatever
+        indices the sampler yields, so slicing batches keeps every index
+        valid and needs no remapping -- one seam instead of two that can
+        disagree.
+      * `DistributedBatchSampler` truncates the batch list to a multiple of
+        world_size so EVERY RANK RUNS THE SAME NUMBER OF STEPS. A rank that
+        finished early would leave its peers blocked inside a collective,
+        which surfaces hours later as an NCCL timeout rather than as an error
+        where it happened. Row-slicing (`rows[rank::world]` before batching)
+        cannot promise this: the sampler is token-budgeted, so different row
+        slices yield different batch counts.
+
+    The cost is a bounded tail: up to (world_size - 1) batches are dropped per
+    epoch. Section-03 chose that over padding with repeats, which would
+    double-count rows and corrupt the epoch semantics the horizon is derived
+    from.
+
+    `rank`/`world_size` default to the live environment but are explicit
+    parameters so four ranks can be simulated in one CPU process -- which is
+    how this contract is tested, since the incident is a property of the
+    relationship BETWEEN ranks and no single-process run can see it.
+    """
+    from ..data.loader import (
+        DistributedBatchSampler,
+        LengthGroupedSampler,
+        preference_length_key,
+    )
+    from ..utils.dist import rank as env_rank
+    from ..utils.dist import world_size as env_world_size
+
+    lengths = [preference_length_key(r) for r in rows]
+    base = LengthGroupedSampler(lengths, max_tokens=max_tokens, max_rows=max_rows, seed=seed)
+
+    r = env_rank() if rank is None else int(rank)
+    ws = env_world_size() if world_size is None else int(world_size)
+    if ws <= 1:
+        return base
+    if not 0 <= r < ws:
+        raise ValueError(f"rank {r} is outside world_size {ws}")
+    return DistributedBatchSampler(base, rank=r, world_size=ws)
+
+
+def _check_global_batch(
+    global_batch: int, world_size: int, micro_batch_size: int, grad_accum: int
+) -> int:
+    """The arithmetic behind both launch refusals; returns the live global batch."""
+    live = world_size * micro_batch_size * grad_accum
+    if global_batch != live:
+        raise RuntimeError(
+            f"global_batch={global_batch} (the pre-registered value the banner "
+            f"and horizon are quoted in) != world_size {world_size} x micro "
+            f"{micro_batch_size} x grad_accum {grad_accum} = {live}. The run "
+            "would consume a different number of pairs per optimizer step than "
+            "it reports. Launch with the pre-registered rank count or change "
+            "the config deliberately."
+        )
+    return live
+
+
+def live_world_size() -> int:
+    """The data-parallel width this process is ACTUALLY part of.
+
+    The initialised process group wins over the environment: `WORLD_SIZE` is
+    what torchrun exported, the group is what the collectives will average
+    over, and a launch decision belongs to the latter. Falls back to the
+    environment before `init_process_group` (and to 1 outside torchrun).
+    """
+    import torch.distributed as dist
+
+    if dist.is_available() and dist.is_initialized():
+        return int(dist.get_world_size())
+    from ..utils.dist import world_size
+
+    return int(world_size())
+
+
+def assert_launch_topology(cfg: "DpoStageConfig", *, world_size: Optional[int] = None) -> dict:
+    """LAUNCH REFUSAL: world_size x micro x grad_accum must equal the
+    pre-registered `global_batch` (2026-10-12 review).
+
+    Nothing used to check it. `global_batch` is a config number that the
+    banner, `optimizer_steps()` and the arm manifest all quote, while what a
+    run really consumes per optimizer step is `world_size * micro_batch_size *
+    grad_accum`. A 2-rank launch of the 4-rank pre-registration therefore ran
+    16 pairs per step under a banner that said 32 and a quoted horizon of 146
+    (really 290): the experiment recorded as one thing and executed as
+    another, with every number finite. `assert_dpo_horizon_consistent`
+    carries the same check but sits AFTER the 14.5GB model load and was never
+    passed the pre-registered value by the launcher; this one runs before any
+    gate, in milliseconds, fed by the live topology.
+
+    `world_size=None` reads it from the live process group (see
+    `live_world_size`); tests and tools pass it explicitly.
+    """
+    ws = live_world_size() if world_size is None else int(world_size)
+    if ws < 1:
+        raise ValueError(f"world_size={ws} must be >= 1")
+    live = _check_global_batch(cfg.global_batch, ws, cfg.micro_batch_size, cfg.grad_accum)
+    return {
+        "world_size": ws,
+        "micro_batch_size": cfg.micro_batch_size,
+        "grad_accum": cfg.grad_accum,
+        "global_batch": live,
+    }
+
+
+def expected_ref_cache_key(cfg: "DpoStageConfig", manifest: Mapping, sft_checkpoint: str):
+    """The key this run requires of its reference caches.
+
+    Five components are EXPECTATIONS built from this run's own facts (content
+    hash of the SFT weights, tokenizer, shard manifest, max_length, truncation)
+    and one is a RECORD: `forward_impl`, how the cache's own pass ran. The
+    record is read from the cache itself, never from the path-bias report.
+
+    WHY NOT THE REPORT (2026-10-12 review). The launcher used to take
+    forward_impl from the report's `cache_forward_impl`, which
+    precompute_ref_logprobs wrote as "fsdp2_sharded" whenever `--path-bias`
+    was passed -- but a `--path-bias` run does NOT recompute the caches, so the
+    field described the measuring run, not the cache. The documented happy
+    path (single-GPU precompute, then `--path-bias`, decision fired=False)
+    was refused at the cache gate: "Mismatched: forward_impl
+    (cache='single_gpu', run='fsdp2_sharded')". The rule that a FIRED decision
+    must be honoured by a sharded cache is enforced where it belongs, by
+    `dpo_banner`; here the held-out cache must agree with the train cache's
+    tag, and everything else is asserted by `load_ref_cache(expect=...)`.
+    """
+    from .ref_cache import build_ref_cache_key, peek_ref_cache_key
+
+    cache_key = peek_ref_cache_key(cfg.ref_cache_path)
+    forward_impl = cache_key.forward_impl if cache_key is not None else "single_gpu"
+    return build_ref_cache_key(sft_checkpoint, dict(manifest), forward_impl=forward_impl)
+
+
+def load_run_ref_caches(cfg: "DpoStageConfig", manifest: Mapping, sft_checkpoint: str):
+    """Load (train, held-out) reference caches under `expected_ref_cache_key`.
+
+    Both caches are loaded with `expect=`, so a missing file, a stale
+    checkpoint, a rebuilt shard set, a different tokenizer/max_length, or a
+    held-out cache computed through a different forward than the train cache
+    all raise here, at the gate, naming the guilty component.
+    """
+    from .ref_cache import load_ref_cache
+
+    key = expected_ref_cache_key(cfg, manifest, sft_checkpoint)
+    train = load_ref_cache(cfg.ref_cache_path, expect=key)
+    held = load_ref_cache(cfg.heldout_ref_cache_path, expect=key)
+    return train, held
+
+
+def assert_dpo_horizon_consistent(
+    scheduler_total_steps: int,
+    *,
+    n_pairs: int,
+    world_size: int,
+    micro_batch_size: int,
+    grad_accum: int,
+    epochs: int,
+    slack: Optional[int] = None,
+    global_batch: Optional[int] = None,
+) -> dict:
+    """Refuse to launch when the LR horizon disagrees with the corpus.
+
+    `global_batch`, when given, is the PRE-REGISTERED number the banner and the
+    horizon formula (`optimizer_steps`) are quoted in; it must equal what the
+    live topology actually consumes per optimizer step,
+    `world_size * micro_batch_size * grad_accum`. Without this the banner can
+    say global_batch 32 / horizon 146 while a 2-rank launch really runs 16 /
+    290: the scheduler-vs-corpus check below passes (it uses the live world
+    size) and the banner is simply wrong about what ran.
+
+    THE ZOMBIE-RUN REFUSAL. Arm A ran a 146-step horizon over a corpus needing
+    ~289 per-rank steps, so the cosine finished at half an epoch; past it
+    `WarmupCosineSchedule` clamps to exactly 0 (min_lr_frac=0) and the loop
+    kept stepping on frozen weights for 287 more steps. The horizon is a
+    PRE-REGISTERED quantity: if it and the arithmetic disagree, one of them is
+    wrong and no amount of training will reconcile them.
+
+    Computed from the SHARDED per-rank row count, because that is what a rank
+    actually walks. `slack` absorbs the token-budgeted sampler landing a batch
+    or two either side of the row-count arithmetic; it does not absorb a
+    factor.
+    """
+    import math as _math
+
+    if min(n_pairs, world_size, micro_batch_size, grad_accum, epochs) < 1:
+        raise ValueError("every horizon input must be >= 1")
+
+    if global_batch is not None:
+        _check_global_batch(global_batch, world_size, micro_batch_size, grad_accum)
+    per_rank_rows = _math.ceil(n_pairs / world_size)
+    rows_per_step = micro_batch_size * grad_accum
+    expected = epochs * _math.ceil(per_rank_rows / rows_per_step)
+    if slack is None:
+        # PROPORTIONAL slack: the token-budgeted sampler legitimately cuts
+        # batches short of max_rows, so the sampler-true horizon runs a few
+        # PERCENT over the row arithmetic (Arm B: 728 vs 708, 2.8%, which a
+        # fixed slack of 2 wrongly refused). 5% absorbs batch-shortfall while
+        # still catching the zombie class (Arm A was 2.1x, far outside).
+        slack = max(2, _math.ceil(0.05 * expected))
+    report = {
+        "scheduler_total_steps": int(scheduler_total_steps),
+        "expected_steps": expected,
+        "per_rank_rows": per_rank_rows,
+        "rows_per_optimizer_step_per_rank": rows_per_step,
+        "unique_rows_per_optimizer_step": rows_per_step * world_size,
+        "epochs_covered": (scheduler_total_steps / expected * epochs if expected else float("nan")),
+    }
+    if abs(int(scheduler_total_steps) - expected) > slack:
+        covered = report["epochs_covered"]
+        raise RuntimeError(
+            f"DPO horizon is inconsistent with the corpus: the scheduler will "
+            f"run {scheduler_total_steps} optimizer steps but "
+            f"{n_pairs} pairs / {world_size} ranks = {per_rank_rows} rows per "
+            f"rank at {rows_per_step} rows per step needs {expected} "
+            f"(epochs={epochs}, slack={slack}). That horizon covers "
+            f"{covered:.2f} epoch(s).\n\nA horizon shorter than the data is "
+            "not a shorter run -- past it the cosine schedule clamps the "
+            "learning rate to exactly 0 and the loop keeps stepping on frozen "
+            "weights, which is how Arm A spent ~24 GPU-hours training "
+            "nothing. Fix the horizon or the corpus arithmetic before "
+            "launching."
+        )
+    return report
+
+
+#: The canonical name for the implicit reward margin -- section-05's kernel
+#: emits it and section-07's Trainer reads it. Imported from the Trainer
+#: rather than spelled again here: a second copy of the string, asserting what
+#: the Trainer reads, is precisely the drift that disarmed the tripwire the
+#: first time.
+from .trainer import MARGIN_KEY as KERNEL_MARGIN_KEY  # noqa: E402
+
+
+def _arm_tripwire_metrics(metrics: dict, preference, window=None) -> dict:
+    """Assert this adapter supplies what the armed tripwires consume.
+
+    THE DISARMED-TRIPWIRE INCIDENT, in its final form. Section-05's kernel
+    emits the implicit margin as `rewards/margins`; section-07's Trainer used
+    to seed a DIFFERENT key, `margin`, with NaN and feed the tripwire from
+    that. Nothing produced it, `_slope` returns NaN for any window containing
+    NaN, and the trip condition requires a non-NaN slope -- so the
+    chosen-likelihood-collapse tripwire, the one with a pre-registered
+    response and the one maximal for Arm B by design, could never fire while
+    appearing fully wired.
+
+    The cure is section-07's, not this module's: the Trainer now reads
+    `MARGIN_KEY` directly, tripwires declare `INPUT_KEYS`, and
+    `assert_tripwires_armed` halts at the first log step if an armed guard's
+    inputs are missing or NaN. This function no longer bridges names -- an
+    alias here would write a key nothing reads.
+
+    What remains is a PRESENCE check at the seam, which is still worth its
+    line: it catches a section-05 rename here, one function call after the
+    kernel returns, rather than at the first log step with the metric already
+    gone. `chosen_longer_fraction` is supplied for the same reason -- the
+    length-drift tripwire's only input, and this adapter is the only thing
+    that can compute it.
+    """
+    if KERNEL_MARGIN_KEY not in metrics:
+        raise KeyError(
+            f"the DPO kernel did not emit {KERNEL_MARGIN_KEY!r} (present: "
+            f"{sorted(metrics)}). That is the canonical margin key the "
+            "chosen-logprob tripwire consumes; without it the armed tripwire "
+            "self-test halts the run, and this raise names the cause one call "
+            "after the kernel returned instead of at the first log step."
+        )
+    if window is not None:
+        window.extend(_chosen_longer_flags(preference))
+        metrics["chosen_longer_fraction"] = sum(window) / len(window) if window else NAN
+        metrics["chosen_longer_pairs"] = float(len(window))
+    metrics.setdefault("chosen_longer_fraction", preference.stats()["chosen_longer_fraction"])
+    metrics.setdefault("chosen_longer_fraction_batch", preference.stats()["chosen_longer_fraction"])
+    return metrics
+
+
+def _chosen_longer_flags(preference) -> list[int]:
+    """Per-pair `chosen completion is longer`, skipping pairs with no
+    completion on either side.
+
+    Per PAIR, because the length-drift tripwire's input has to be a fraction
+    over a sample whose size we control; the batch-level statistic is that
+    same fraction over whatever happened to be in the micro-batch, which at
+    micro_batch 2 is three possible values.
+
+    Completion tokens, not sequence tokens: Arm B's two sides share their
+    entire prefix, so a sequence-length comparison is dominated by a term
+    identical on both sides and would compress the fraction toward 0.5
+    regardless of the data -- hiding exactly the drift being watched for.
+    """
+    c = preference.chosen_completion_mask.sum(dim=1)
+    r = preference.rejected_completion_mask.sum(dim=1)
+    return [int(ci > ri) for ci, ri in zip(c.tolist(), r.tolist()) if (ci + ri) > 0]
+
+
+def make_dpo_stage_adapter(
+    cache, cfg: Optional[DpoStageConfig] = None, require_fused: bool = False
+):
+    """Trainer loss adapter: attach cached refs by pair_id, then delegate.
+
+    Bridges three things section-05 and section-03 do not agree on, and does
+    it in one place so the disagreement cannot be re-resolved differently
+    somewhere else:
+
+      * NAMES. Section-03's `Preference` carries `chosen_input_ids` /
+        `ref_chosen_logps`; section-05's adapter unpacks `input_ids_chosen` /
+        `ref_logp_chosen`. Neither is wrong; they were written to different
+        drafts.
+      * LABELS. `Preference` carries a completion MASK, `fused_dpo` wants
+        labels with IGNORE_INDEX outside the completion. The conversion is one
+        `masked_fill`, and getting it backwards would score the prompt -- which
+        cancels in the DPO difference while inflating both magnitudes.
+      * REFERENCES. They come from the cache, joined by pair_id, never from
+        batch order and never defaulted.
+
+    The reference tensors enter as plain constants; section-05's kernel
+    detaches them and proves no gradient path survives. This adapter's only
+    obligation is never to build one, which it discharges by reading them out
+    of a dict of floats.
+    """
+    from ..kernels.fused_dpo import IGNORE_INDEX
+    from .dpo_loss import DpoLossConfig, make_dpo_loss_adapter
+    from .ref_cache import attach_refs
+
+    stage = cfg or DpoStageConfig()
+    #: Rolling per-PAIR history feeding the length-drift tripwire. Lives on
+    #: the adapter, not on the batch: the statistic the tripwire needs spans
+    #: batches, and computing it per batch is what makes it unusable.
+    length_window = collections.deque(maxlen=stage.length_window_pairs)
+    inner = make_dpo_loss_adapter(
+        DpoLossConfig(
+            beta=stage.beta,
+            variant=stage.variant,
+            nll_coeff=stage.nll_coeff,
+            dpop_lambda=stage.dpop_lambda,
+        ),
+        require_fused=require_fused,
+        # One trunk graph at a time: both sides' checkpointed 7B graphs plus
+        # fp32 shard grads measured 21.5GB and OOM'd (batch-size independent).
+        decoupled_backward=require_fused,
+        grad_scale=1.0 / max(1, stage.grad_accum),
+    )
+
+    def adapter(model, batch):
+        import torch
+
+        pref = batch.preference if isinstance(batch, DpoBatch) else batch
+        pref = attach_refs(pref, cache)
+        # The WHOLE batch to the compute device, refs included. The model
+        # forward moves input_ids internally, which let CPU labels survive
+        # to the loss's gather ("index is on cpu") on the stage's first
+        # real step -- the SFT script moves its collated batch explicitly,
+        # and this adapter is the DPO stage's equivalent place.
+        #
+        # Follow the MODEL's device, not `cuda.is_available()`: on a box with
+        # visible GPUs a CPU-resident model (the unit-test tiny model) would
+        # otherwise get CUDA batches and die inside the embedding. Under FSDP2
+        # the parameters are DTensors on this rank's cuda device, so the real
+        # path is unchanged.
+        param = next(iter(model.parameters()), None)
+        if param is not None and param.device.type == "cuda":
+            pref = pref.to(torch.device("cuda", torch.cuda.current_device()))
+        elif param is None and torch.cuda.is_available():
+            pref = pref.to(torch.device("cuda", torch.cuda.current_device()))
+        bound = {
+            "input_ids_chosen": pref.chosen_input_ids,
+            "labels_chosen": pref.chosen_input_ids.masked_fill(
+                pref.chosen_completion_mask == 0, IGNORE_INDEX
+            ),
+            "input_ids_rejected": pref.rejected_input_ids,
+            "labels_rejected": pref.rejected_input_ids.masked_fill(
+                pref.rejected_completion_mask == 0, IGNORE_INDEX
+            ),
+            "ref_logp_chosen": pref.ref_chosen_logps,
+            "ref_logp_rejected": pref.ref_rejected_logps,
+        }
+        out = inner(model, bound)
+        _arm_tripwire_metrics(out.metrics, pref, window=length_window)
+        return out
+
+    adapter.uses_fused = getattr(inner, "uses_fused", True)
+    adapter.config = inner.config
+    adapter.stage_config = stage
+    return adapter

code/align/train/fsdp.py

@@ -0,0 +1,124 @@
+"""FSDP2 wrapping and activation checkpointing.
+
+Called from the LAUNCH SCRIPTS, never from the Trainer: the Trainer receives an
+already-wrapped model, which is what keeps it identical between the SFT and DPO
+stages and testable on CPU without a process group.
+
+Two policies here are not configurable, deliberately:
+
+  * `use_reentrant=False`, hard-asserted. Reentrant activation checkpointing
+    silently ZEROES trunk gradients under sharding. The run trains -- loss
+    falls, throughput is normal -- while only the embeddings and the LM head
+    move. There is no knob to turn it back on, because the only reason anyone
+    would set it is not knowing this. The trunk-movement probe in `probes.py`
+    is the runtime backstop; this is the compile-time one.
+  * No `CPUOffloadPolicy`. FSDP's own CPU offload pays per-layer transfers on
+    every forward AND backward; our offload is optimizer-state-only and
+    happens once per optimizer step (see `offload_adamw.py`). Turning both on
+    would move the same bytes twice.
+
+Sharding is PER TRANSFORMER BLOCK, then the root. Per-block units are what
+make the all-gather granularity match the compute granularity: one unit for
+the whole model would gather all 7.25B parameters before the first matmul.
+"""
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import Any, Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+#: bf16 compute, fp32 reduction. bf16 is mandatory (fp16 overflows the
+#: 32768-vocab logits); the fp32 REDUCE matters independently -- gradient
+#: all-reduce in bf16 loses small updates across 4 ranks precisely when they
+#: are averaged, which is every step.
+PARAM_DTYPE = torch.bfloat16
+REDUCE_DTYPE = torch.float32
+
+
+def transformer_blocks(model) -> list:
+    """The per-block FSDP units: every `DecoderBlock` in the stack.
+
+    Found structurally rather than by isinstance so this module does not
+    import the model (and so a wrapped/compiled model still works).
+    """
+    inner = getattr(model, "model", model)
+    layers = getattr(inner, "layers", None)
+    if layers is None:
+        raise ValueError(
+            "no `.layers` found on the model: FSDP2 here shards per "
+            "transformer block, and a model with no block list would be "
+            "wrapped as ONE unit -- gathering all 7.25B parameters before the "
+            "first matmul, which does not fit on a 23GB card."
+        )
+    return list(layers)
+
+
+def mixed_precision_policy():
+    """`MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=fp32)`."""
+    from torch.distributed.fsdp import MixedPrecisionPolicy
+
+    return MixedPrecisionPolicy(param_dtype=PARAM_DTYPE, reduce_dtype=REDUCE_DTYPE)
+
+
+def apply_fsdp(model, mesh: Optional[Any] = None, reshard_after_forward: bool = True):
+    """Shard per transformer block, then the root.
+
+    Order is load-bearing: the root must be sharded LAST, after its children
+    are already units, or the root unit would own their parameters too.
+    """
+    from torch.distributed.fsdp import fully_shard
+
+    policy = mixed_precision_policy()
+    blocks = transformer_blocks(model)
+    for block in blocks:
+        fully_shard(block, mesh=mesh, mp_policy=policy, reshard_after_forward=reshard_after_forward)
+    fully_shard(model, mesh=mesh, mp_policy=policy, reshard_after_forward=reshard_after_forward)
+    logger.info(
+        "FSDP2: %d per-block units + root, param=%s reduce=%s",
+        len(blocks),
+        PARAM_DTYPE,
+        REDUCE_DTYPE,
+    )
+    return model
+
+
+def _checkpointed_forward(block, *args, **kwargs):
+    from torch.utils.checkpoint import checkpoint
+
+    return checkpoint(block._nmo_inner_forward, *args, use_reentrant=False, **kwargs)
+
+
+def apply_activation_checkpointing(model, use_reentrant: bool = False):
+    """Wrap every transformer block in non-reentrant activation checkpointing.
+
+    Patches each block's `forward` rather than wrapping the block in a new
+    module: a wrapper module would insert a level into the module tree and
+    change every parameter name, which would silently break the HF checkpoint
+    mapping and the interp pillar that loads through it.
+    """
+    if use_reentrant:
+        raise ValueError(
+            "use_reentrant=True is refused. Reentrant activation checkpointing "
+            "silently ZEROES trunk gradients under sharding: the run trains "
+            "the embeddings and the LM head, the loss falls, throughput looks "
+            "normal, and the body of the model never moves. This is a "
+            "documented predecessor failure and there is no configuration in "
+            "which we want it."
+        )
+    blocks = transformer_blocks(model)
+    for block in blocks:
+        if hasattr(block, "_nmo_inner_forward"):
+            continue  # idempotent: never double-wrap
+        block._nmo_inner_forward = block.forward
+        block.forward = functools.partial(_checkpointed_forward, block)
+    logger.info("activation checkpointing on %d blocks (use_reentrant=False)", len(blocks))
+    return model
+
+
+def is_checkpointed(block) -> bool:
+    return hasattr(block, "_nmo_inner_forward")

code/align/train/lora_scale.py

@@ -0,0 +1,110 @@
+"""LoRA scale sweep WITHOUT retraining (SOTA plan K4).
+
+A LoRA layer computes `W x + s * B A x` with s = alpha / r. Evaluating the SAME trained adapter at a
+multiple `k` of that delta (k in {0.5, 0.7, 1.0}: "alpha curve") needs no optimisation: set s -> k*s.
+k = 1 is the trained model, k = 0 is the base, and the curve between them tells whether a regression
+is "too much adapter" (a smaller k recovers it) or "the wrong direction" (it does not).
+
+Two entry points, matching the two ways an adapter is evaluated:
+
+  * `scaled_adapter(model, adapter, k)` -- a context manager over a LIVE model with the adapter attached
+    (the un-merged serving option); `sweep(...)` runs an eval callable at each k.
+  * `merged_state_dict_at_scale(model, full_state, adapter, k)` -- a BASE-keyed merged state dict at
+    scale k, for merged exports, with exactly the bf16-cast fidelity report of
+    `align.ppo.lora.export_merged_state_dict` (the delta is k times smaller, so MORE of it can fall into
+    the bf16 dead zone: read `report["cosine"]` before trusting a small-k merged eval).
+
+Refusals: a layer that has the adapter MERGED into its weight cannot be rescaled (the delta is already
+inside W); the context manager raises instead of evaluating a polluted point. The model is never left
+rescaled: the original scaling is restored in `finally`, on error and on KeyboardInterrupt alike.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import math
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from ..ppo.lora import export_merged_state_dict, lora_modules
+
+#: The plan's K4 grid: fractions of the TRAINED adapter delta.
+DEFAULT_SCALES: tuple[float, ...] = (0.5, 0.7, 1.0)
+
+
+def _check_scale(k: float) -> float:
+    k = float(k)
+    if not math.isfinite(k) or k < 0.0:
+        raise ValueError(f"LoRA delta scale must be a finite number >= 0, got {k!r}")
+    return k
+
+
+@contextlib.contextmanager
+def scaled_adapter(model: nn.Module, adapter: str, k: float) -> Iterator[int]:
+    """Within the block, every layer's `adapter` contributes `k` times its trained delta.
+
+    Yields the number of layers rescaled. Raises (before changing anything) if no layer carries the
+    adapter or if any layer has it merged.
+    """
+    k = _check_scale(k)
+    layers = [(p, m) for p, m in lora_modules(model) if adapter in m.adapters]
+    if not layers:
+        raise KeyError(f"no layer carries adapter {adapter!r}; nothing to scale")
+    merged = [p for p, m in layers if m._merged is not None]
+    if merged:
+        raise RuntimeError(
+            f"adapter {adapter!r} is merged into {len(merged)} layer(s) (first: {merged[0]}); the delta is "
+            "already inside the weight and cannot be rescaled. Unmerge first, or use "
+            "merged_state_dict_at_scale on the unmerged model."
+        )
+    original = [(m.adapters[adapter], m.adapters[adapter].scaling) for _, m in layers]
+    try:
+        for ad, s in original:
+            ad.scaling = s * k
+        yield len(layers)
+    finally:
+        for ad, s in original:
+            ad.scaling = s
+
+
+def sweep(
+    model: nn.Module,
+    adapter: str,
+    eval_fn: Callable[[nn.Module], Any],
+    scales: Iterable[float] = DEFAULT_SCALES,
+) -> dict[float, Any]:
+    """`{k: eval_fn(model)}` with the adapter delta scaled to k each time. `eval_fn` sees the live model."""
+    out: dict[float, Any] = {}
+    for k in scales:
+        k = _check_scale(k)
+        if k in out:
+            raise ValueError(f"duplicate scale {k} in the sweep")
+        with scaled_adapter(model, adapter, k):
+            out[k] = eval_fn(model)
+    return out
+
+
+def merged_state_dict_at_scale(
+    model: nn.Module,
+    full_state: dict[str, torch.Tensor],
+    adapter: str,
+    k: float,
+    report: Optional[dict] = None,
+    min_cosine: Optional[float] = None,
+) -> dict[str, torch.Tensor]:
+    """`export_merged_state_dict` with the adapter delta scaled by `k` (the model is not mutated).
+
+    k = 1.0 is bit-identical to `export_merged_state_dict`. k = 0 returns the base weights (every nonzero
+    delta is zero, so `report` has cosine 1.0 and delta_norm 0). `report["scale"]` records k.
+    """
+    k = _check_scale(k)
+    rep: dict = {}
+    with scaled_adapter(model, adapter, k):
+        out = export_merged_state_dict(model, full_state, adapter, report=rep, min_cosine=min_cosine)
+    rep["scale"] = k
+    if report is not None:
+        report.update(rep)
+    return out

code/align/train/moe_module.py

@@ -0,0 +1,691 @@
+"""MoE auxiliary losses on the training seam: `MoETrainModule`.
+
+The model (`align.model.olmoe`) adds NO aux loss: it RETURNS per-layer router
+statistics (`RouterStats`) when asked. This module is where those become a
+training objective, because only the TrainModule sees all three things the
+objective depends on: the loss adapter's CE (an opaque scalar), the
+gradient-accumulation window (`grad_accum`), and the data-parallel group.
+
+THE MATH, and the one approximation in it
+-----------------------------------------
+Switch/OLMoE load balance, per layer: `E * sum_i f_i * P_i` with `f_i` the hard
+top-k hit fraction (no gradient) and `P_i` the mean router probability
+(gradient). Which batch `f_i` and `P_i` are measured over is the whole question
+(Qiu et al., "Demons in the Detail"): per-micro-batch statistics over-constrain
+the router when micro-batches are domain-homogeneous, because a sequence of
+code must be allowed to use the code experts heavily. The loss should balance
+over the GLOBAL batch -- every rank, every micro-batch of the optimizer step.
+
+Exact global `f_i` is not available at micro-batch `m`: the window's later
+micro-batches have not been run, and their backward cannot be deferred without
+retaining every micro-batch's activations. Demons' Algorithm 1 resolves it with
+a BUFFER: keep the window's running (all-reduced) counts, and take `f_i` from
+the buffer as it stands when micro-batch `m` runs. The approximation is
+therefore confined to `f_i` of micro-batches 1..accum-1; at the LAST micro-batch
+`f_i` is exactly the global-batch fraction, and the logged `moe/lb_loss` (computed
+from the full buffer at the optimizer-step boundary) is exact.
+
+SCALING (why `world * m`). The Trainer divides each micro loss by `grad_accum`
+and the gradient reduction averages over `world` ranks, so micro-batch `m` on
+rank `r` reaches the weights with weight `1/(world*accum)`. For the window's
+gradient to equal the gradient of ONE global loss
+`E * sum_i f_i * (sum_{r,m} prob_sum_{r,m,i}) / N`, each term must be multiplied
+by `world * accum / N_window`, and `N_window` (all valid tokens of the window) is
+unknown until the end; estimated as `n_buf * accum / m` it gives
+`world * m / n_buf`. At `m = accum` the estimate is exact. The same factor
+scales the z-loss. In "micro" scope (the per-rank formulation, kept for
+ablations) the factor is 1 and everything is local.
+
+WHAT IS ALL-REDUCED, and when. One packed float64 tensor per micro-batch
+[counts L*E | tokens L | prob_sum L*E | z_sum L | entropy_sum L]. DETACHED data
+only: the gradient of the loss is each rank differentiating its own
+`prob_sum`/`z_sum`, never a gradient crossing ranks here. Every rank calls the
+collective the same number of times (the Trainer runs micro-steps in lockstep),
+and the buffer is reset at `optim_step()`, which the Trainer calls on every
+boundary whether or not the optimizer applied the step.
+
+DEFAULTS. Load balance and z-loss are OFF for post-training stages and
+0.01 / 0.001 for `pretrain` / `proxy` (see `MoEAuxConfig`); telemetry
+(MaxVio, utilization, dead experts, router entropy) is logged either way.
+
+LIMITS (stated, not hidden)
+ * Adapters that run their own backward and return a detached proxy (the DPO
+   decoupled-backward adapter) cannot carry an aux loss: the router graph is
+   consumed by the adapter's backward. With a nonzero weight this raises; with
+   zero weights telemetry still works (it reads detached data).
+ * Not validated under FSDP2 on GPU (no GPU in this lane): the stats ride out
+   of the wrapped forward as plain dataclass fields rather than pytree leaves.
+ * The returned `LossOutput.loss` is the CE alone (comparable with dense runs).
+   If an aux term is non-finite the returned loss is NaN, so the Trainer's
+   non-finite-loss accounting sees it.
+
+EXPERT PARALLELISM (opt-in: `ep_groups`, default None = nothing below applies)
+-----------------------------------------------------------------------------
+`align.model.moe_ep` shards the experts across an EP group. The caller applies it
+(`apply_expert_parallel(model, groups)`) BEFORE the optimizer is built -- the
+optimizer must hold the shard parameters, and sharding after construction would
+leave it stepping tensors that are no longer in the model. This module then does
+the three things only the TrainModule can:
+  * the aux losses need NO change: the router is replicated, `RouterStats` come from
+    local tokens exactly as without EP, and `MoEAux` all-reduces them over the
+    DATA-parallel group (`group`), which contains the EP group.
+  * `optim_step` puts expert-shard gradients on the replicated-gradient scale
+    (`finalize_expert_grads`: they arrive SUMMED over the EP group) and clips by the
+    EP-GLOBAL norm (`clip_grads_`). The wrapped optimizer must therefore NOT clip
+    itself (`max_grad_norm` must be None): its norm would be shard-local, ranks would
+    clip by different coefficients and the replicated parameters would drift apart.
+    Pass the clip value as `ep_max_grad_norm`; `last_grad_norm` is the PRE-clip global.
+  * `full_model_state_dict` / `load_model_state_dict` gather / slice the expert shards
+    (a collective over the EP group), so a checkpoint is the same file a non-EP run writes.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+
+from ..config import MoEAuxConfig
+from ..model import moe
+from .train_module import LossAdapter, LossOutput, TrainModule
+
+logger = logging.getLogger(__name__)
+
+NAN = float("nan")
+
+#: (lb_weight, z_weight) per stage when the config leaves a weight as None.
+#: Stages that train a router FROM SCRATCH balance it; everything else leaves
+#: the pretrained router alone. See MoEAuxConfig for the argument.
+STAGE_AUX_DEFAULTS: dict[str, tuple[float, float]] = {
+    "pretrain": (0.01, 0.001),
+    "proxy": (0.01, 0.001),
+}
+POST_TRAINING_AUX_DEFAULT = (0.0, 0.0)
+LB_SCOPES = ("global", "micro")
+
+
+def resolve_aux_weights(aux: MoEAuxConfig, stage: str) -> tuple[float, float]:
+    """(lb_weight, z_weight) after applying the stage default to any `None`."""
+    d_lb, d_z = STAGE_AUX_DEFAULTS.get(stage, POST_TRAINING_AUX_DEFAULT)
+    lb = d_lb if aux.lb_weight is None else float(aux.lb_weight)
+    z = d_z if aux.z_weight is None else float(aux.z_weight)
+    for name, v in (("lb_weight", lb), ("z_weight", z)):
+        if not (math.isfinite(v) and v >= 0.0):
+            raise ValueError(f"moe_aux.{name}={v!r} must be a finite number >= 0")
+    if aux.lb_scope not in LB_SCOPES:
+        raise ValueError(f"moe_aux.lb_scope={aux.lb_scope!r}; known: {LB_SCOPES}")
+    return lb, z
+
+
+# ---------------------------------------------------------------------------
+# the window buffer (Demons Alg. 1)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class WindowTotals:
+    """Detached sums over (ranks x micro-batches so far): everything the aux
+    losses and the telemetry need, nothing with a gradient."""
+
+    counts: torch.Tensor  # [L, E] float64: top-k hits
+    n_tokens: torch.Tensor  # [L] float64: valid tokens
+    prob_sum: torch.Tensor  # [L, E] float64: sum of softmax probs
+    z_sum: torch.Tensor  # [L] float64: sum of lse^2
+    entropy_sum: torch.Tensor  # [L] float64
+
+    def __add__(self, o: "WindowTotals") -> "WindowTotals":
+        return WindowTotals(
+            self.counts + o.counts,
+            self.n_tokens + o.n_tokens,
+            self.prob_sum + o.prob_sum,
+            self.z_sum + o.z_sum,
+            self.entropy_sum + o.entropy_sum,
+        )
+
+    @classmethod
+    def of(cls, stats: list[moe.RouterStats]) -> "WindowTotals":
+        d = torch.float64
+        return cls(
+            counts=torch.stack([s.counts for s in stats]).to(d),
+            n_tokens=torch.stack([s.n_tokens for s in stats]).to(d),
+            prob_sum=torch.stack([s.prob_sum.detach() for s in stats]).to(d),
+            z_sum=torch.stack([s.z_sum.detach() for s in stats]).to(d),
+            entropy_sum=torch.stack([s.entropy_sum.detach() for s in stats]).to(d),
+        )
+
+    def all_reduce(self, group=None) -> "WindowTotals":
+        """Sum over a process group in ONE collective (packed)."""
+        import torch.distributed as dist
+
+        n_layers, e = self.counts.shape
+        packed = torch.cat(
+            [
+                self.counts.reshape(-1),
+                self.n_tokens,
+                self.prob_sum.reshape(-1),
+                self.z_sum,
+                self.entropy_sum,
+            ]
+        ).contiguous()
+        dist.all_reduce(packed, group=group)
+        le = n_layers * e
+        o = 0
+        counts = packed[o : o + le].reshape(n_layers, e)
+        o += le
+        n = packed[o : o + n_layers]
+        o += n_layers
+        ps = packed[o : o + le].reshape(n_layers, e)
+        o += le
+        z = packed[o : o + n_layers]
+        o += n_layers
+        ent = packed[o : o + n_layers]
+        return WindowTotals(counts, n, ps, z, ent)
+
+
+def merge_call_stats(calls: list[list[moe.RouterStats]]) -> list[moe.RouterStats]:
+    """Combine the stats of several grad-enabled forwards of ONE micro-batch.
+
+    One call (SFT, pretrain) is the identity. Several (a DPO adapter that runs
+    chosen and rejected separately) are summed per layer: they are different
+    token populations of the same optimizer step, so counts, probability sums,
+    z sums and token counts all ADD, and the losses are then over their union.
+    """
+    if not calls:
+        raise ValueError("no router stats captured")
+    if len(calls) == 1:
+        return calls[0]
+    n_layers = len(calls[0])
+    if any(len(c) != n_layers for c in calls):
+        raise ValueError(f"forwards returned different layer counts: {[len(c) for c in calls]}")
+    out = []
+    for layer in range(n_layers):
+        parts = [c[layer] for c in calls]
+        out.append(
+            moe.RouterStats(
+                counts=sum(p.counts for p in parts),
+                prob_sum=sum(p.prob_sum for p in parts),
+                n_tokens=sum(p.n_tokens for p in parts),
+                z_sum=sum(p.z_sum for p in parts),
+                entropy_sum=sum(p.entropy_sum for p in parts),
+            )
+        )
+    return out
+
+
+class MoEAux:
+    """The aux-loss math and the window buffer, independent of any Trainer.
+
+    `micro(stats)` is called once per micro-batch, in lockstep on every rank,
+    and returns `(lb_term, z_term)` ALREADY scaled for the Trainer's
+    `/grad_accum` and the gradient reduction's mean over ranks (see the module
+    docstring). `window_end()` returns the exact global-window telemetry and
+    resets the buffer.
+    """
+
+    def __init__(
+        self,
+        top_k: int,
+        grad_accum: int,
+        *,
+        scope: str = "global",
+        group=None,
+        world_size: Optional[int] = None,
+    ):
+        if scope not in LB_SCOPES:
+            raise ValueError(f"scope {scope!r}; known: {LB_SCOPES}")
+        if grad_accum < 1:
+            raise ValueError(f"grad_accum must be >= 1, got {grad_accum}")
+        self.top_k = int(top_k)
+        self.grad_accum = int(grad_accum)
+        self.scope = scope
+        self.group = group
+        self._world = world_size
+        self.m = 0  # micro-batches seen in the current window
+        self.buf: Optional[WindowTotals] = None
+
+    # -- world ----------------------------------------------------------
+    def world(self) -> int:
+        if self._world is not None:
+            return self._world
+        import torch.distributed as dist
+
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_world_size(self.group)
+        return 1
+
+    def _distributed(self) -> bool:
+        import torch.distributed as dist
+
+        return dist.is_available() and dist.is_initialized() and self.world() > 1
+
+    def reset(self) -> None:
+        self.m = 0
+        self.buf = None
+
+    # -- one micro-batch --------------------------------------------------
+    def micro(
+        self, stats: list[moe.RouterStats], want_lb: bool = True, want_z: bool = True
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        if self.m >= self.grad_accum:
+            raise RuntimeError(
+                f"{self.m + 1} micro-batches in one accumulation window (grad_accum="
+                f"{self.grad_accum}): optim_step()/reset() was not called at the boundary"
+            )
+        self.m += 1
+        local = WindowTotals.of(stats)
+        glob = local.all_reduce(self.group) if self._distributed() else local
+        self.buf = glob if self.buf is None else self.buf + glob
+        buf = self.buf
+        if not bool((buf.n_tokens > 0).all()):
+            raise ValueError(
+                "a MoE layer has seen zero valid tokens in this window: every position was masked. "
+                "The balance and z-losses are undefined over zero tokens (and returning 0.0 would "
+                "report a perfectly balanced router)."
+            )
+        lb = z = None
+        if self.scope == "global":
+            scale = float(self.world() * self.m)
+            if want_lb:
+                lb = scale * moe.load_balancing_loss(
+                    stats,
+                    self.top_k,
+                    global_totals=moe.Totals(counts=buf.counts, n_tokens=buf.n_tokens),
+                )
+            if want_z:
+                # n_tokens is identical across layers (every layer routes the
+                # same valid tokens); the scalar is what router_z_loss takes.
+                z = scale * moe.router_z_loss(stats, loss_div_factor=buf.n_tokens[0])
+        else:
+            if want_lb:
+                lb = moe.load_balancing_loss(stats, self.top_k)
+            if want_z:
+                z = moe.router_z_loss(stats)
+        return lb, z
+
+    # -- the boundary -----------------------------------------------------
+    def window_values(self) -> dict[str, Any]:
+        """Exact global-window quantities (float64 tensors) from the buffer."""
+        buf = self.buf
+        if buf is None:
+            raise RuntimeError("window_values() before any micro-batch")
+        n_layers, e = buf.counts.shape
+        n = buf.n_tokens.unsqueeze(1)
+        f = buf.counts / (self.top_k * n)
+        p = buf.prob_sum / n
+        lb = (e * (f * p).sum(dim=1)).mean()
+        z = (buf.z_sum / buf.n_tokens).mean()
+        ent = buf.entropy_sum / buf.n_tokens
+        per_layer = [moe.expert_load_metrics(buf.counts[i]) for i in range(n_layers)]
+        return {
+            "lb": lb,
+            "z": z,
+            "entropy": ent,
+            "f": f,
+            "max_vio": torch.stack([d["max_vio"] for d in per_layer]),
+            "utilization": torch.stack([d["utilization"] for d in per_layer]),
+            "frac_dead": torch.stack([d["frac_dead"] for d in per_layer]),
+            "load_entropy": torch.stack([d["load_entropy"] for d in per_layer]),
+        }
+
+
+def window_metrics(values: dict[str, Any], per_layer: bool = False) -> dict[str, float]:
+    """The `moe/*` history keys from `MoEAux.window_values()`. ONE host sync."""
+    n_layers = int(values["max_vio"].numel())
+    flat = torch.cat(
+        [
+            values["lb"].reshape(1),
+            values["z"].reshape(1),
+            values["entropy"],
+            values["max_vio"],
+            values["utilization"],
+            values["frac_dead"],
+            values["load_entropy"],
+        ]
+    ).tolist()
+    lb, z = flat[0], flat[1]
+    o = 2
+    ent = flat[o : o + n_layers]
+    o += n_layers
+    vio = flat[o : o + n_layers]
+    o += n_layers
+    util = flat[o : o + n_layers]
+    o += n_layers
+    dead = flat[o : o + n_layers]
+    o += n_layers
+    lent = flat[o : o + n_layers]
+    out = {
+        "moe/lb_loss": lb,
+        "moe/z_loss": z,
+        "moe/max_vio": sum(vio) / n_layers,
+        "moe/max_vio_worst": max(vio),
+        "moe/utilization": sum(util) / n_layers,
+        "moe/utilization_worst": min(util),
+        "moe/frac_dead_worst": max(dead),
+        "moe/load_entropy": sum(lent) / n_layers,
+        "moe/router_entropy": sum(ent) / n_layers,
+    }
+    if per_layer:
+        for i in range(n_layers):
+            out[f"moe/l{i}/max_vio"] = vio[i]
+            out[f"moe/l{i}/utilization"] = util[i]
+    return out
+
+
+#: The keys a MoE run logs, NaN until the first window completes. Fixed set so
+#: every history.jsonl row has the same columns.
+def undefined_window_metrics(n_layers: Optional[int] = None, per_layer: bool = False) -> dict[str, float]:
+    keys = [
+        "moe/lb_loss",
+        "moe/z_loss",
+        "moe/max_vio",
+        "moe/max_vio_worst",
+        "moe/utilization",
+        "moe/utilization_worst",
+        "moe/frac_dead_worst",
+        "moe/load_entropy",
+        "moe/router_entropy",
+    ]
+    out = {k: NAN for k in keys}
+    if per_layer and n_layers:
+        for i in range(n_layers):
+            out[f"moe/l{i}/max_vio"] = NAN
+            out[f"moe/l{i}/utilization"] = NAN
+    return out
+
+
+# ---------------------------------------------------------------------------
+# the model proxy
+# ---------------------------------------------------------------------------
+
+
+class _StatsModel:
+    """What the loss adapter is handed instead of the model.
+
+    Adapters call `model(input_ids=..., hidden_consumer=...)` and expect the
+    loss (or logits) back. The proxy adds `return_router_stats=True` (only when
+    gradients are enabled -- a reference-model forward under `no_grad` has
+    nothing to balance), unwraps the `(out, stats)` tuple, and stashes the
+    stats for the module. Every other attribute is delegated, so adapters that
+    read `model.lm_head` or `model.attn_implementation` are unaffected.
+    """
+
+    def __init__(self, owner: "MoETrainModule"):
+        object.__setattr__(self, "_owner", owner)
+        object.__setattr__(self, "calls", [])
+        object.__setattr__(self, "attention_mask", None)
+
+    def __call__(self, *args, **kwargs):
+        owner = self._owner
+        model = owner.model
+        want = torch.is_grad_enabled()
+        if want:
+            kwargs["return_router_stats"] = True
+            if kwargs.get("attention_mask") is None and self.attention_mask is not None:
+                # The batch's real-token mask: padding must not vote in the
+                # balance loss. The model's attention applies it too (a no-op
+                # for valid tokens under right padding, identical outputs).
+                kwargs["attention_mask"] = self.attention_mask
+        try:
+            out = model(*args, **kwargs)
+        except TypeError as e:
+            if want and "return_router_stats" in str(e):
+                raise RuntimeError(
+                    "MoETrainModule needs a model whose forward accepts return_router_stats=True "
+                    f"(got {type(model).__name__}); a dense model has no router to balance. "
+                    "Set moe_aux weights to 0 and use the plain TrainModule."
+                ) from e
+            raise
+        if not want:
+            return out
+        out, stats = out
+        self.calls.append(stats)
+        return out
+
+    def forward_hidden_states(self, **kwargs):
+        """Route `forward_hidden_states` through THIS proxy's `__call__`.
+
+        The DPO adapter (coupled AND decoupled) never calls `model(...)`: it
+        calls `model.forward_hidden_states(...)`. Without this method the
+        attribute falls through `__getattr__` to the REAL model's bound method,
+        which runs without `return_router_stats`, so the adapter "ran no
+        grad-enabled forward" as far as the module can tell and the first
+        micro-step raises -- with default settings (telemetry on) on a stage
+        whose docstring promises telemetry still works. Mirrors
+        `OlmoeForCausalLM.forward_hidden_states`, which is `self(...,
+        return_hidden_states=True)`.
+        """
+        return self(**kwargs, return_hidden_states=True)
+
+    def __getattr__(self, name):
+        return getattr(self._owner.model, name)
+
+
+# ---------------------------------------------------------------------------
+# the train module
+# ---------------------------------------------------------------------------
+
+
+class MoETrainModule(TrainModule):
+    """`TrainModule` + router aux losses + MoE telemetry.
+
+    `train_batch` does what the base does -- forward through the loss adapter,
+    then `(loss / grad_accum).backward()` -- with `loss = CE + lb_weight * lb +
+    z_weight * z`. See the module docstring for the window buffer and scaling.
+    """
+
+    def __init__(
+        self,
+        model,
+        loss_adapter: LossAdapter,
+        optimizer,
+        grad_accum: int,
+        *,
+        aux: Optional[MoEAuxConfig] = None,
+        stage: str = "sft",
+        top_k: Optional[int] = None,
+        weight_format=None,
+        group=None,
+        world_size: Optional[int] = None,
+        ep_groups=None,
+        ep_max_grad_norm: Optional[float] = None,
+    ):
+        super().__init__(model, loss_adapter, optimizer, grad_accum)
+        self.ep_groups = ep_groups
+        self.ep_max_grad_norm = ep_max_grad_norm
+        self._ep_norm = float("nan")
+        if ep_groups is not None:
+            self._check_expert_parallel(model, optimizer, ep_groups)
+        self.aux_cfg = aux or MoEAuxConfig()
+        self.stage = stage
+        self.lb_weight, self.z_weight = resolve_aux_weights(self.aux_cfg, stage)
+        if top_k is None:
+            top_k = getattr(getattr(model, "cfg", None), "top_k", None)
+        if top_k is None:
+            raise ValueError("top_k not given and model.cfg.top_k not found")
+        self.top_k = int(top_k)
+        self.weight_format = weight_format
+        self.collect = bool(self.aux_cfg.telemetry or self.lb_weight > 0 or self.z_weight > 0)
+        self._aux = MoEAux(
+            self.top_k, grad_accum, scope=self.aux_cfg.lb_scope, group=group, world_size=world_size
+        )
+        self._proxy = _StatsModel(self)
+        self._last_metrics: dict[str, float] = {}
+        logger.info(
+            "MoE aux: stage=%s lb_weight=%g z_weight=%g scope=%s telemetry=%s top_k=%d",
+            stage,
+            self.lb_weight,
+            self.z_weight,
+            self.aux_cfg.lb_scope,
+            self.aux_cfg.telemetry,
+            self.top_k,
+        )
+
+    # -- expert parallelism (opt-in) ---------------------------------------
+    @staticmethod
+    def _check_expert_parallel(model, optimizer, groups) -> None:
+        """Refuse the three silent EP misconfigurations instead of training through them."""
+        from ..model import moe_ep
+
+        blocks = moe_ep._moe_blocks(model)
+        if not blocks or any(b.ep is None for b in blocks):
+            raise ValueError(
+                "ep_groups was given but the model's MoE blocks are not expert-parallel: call "
+                "align.model.moe_ep.apply_expert_parallel(model, groups) BEFORE building the optimizer"
+            )
+        if any(b.ep.groups is not groups for b in blocks):
+            raise ValueError("ep_groups is not the group set the model was sharded with")
+        if getattr(optimizer, "max_grad_norm", None) is not None:
+            raise ValueError(
+                "the optimizer clips by its own (shard-local) gradient norm; under expert parallelism "
+                "that clips every rank by a different coefficient. Build it with max_grad_norm=None and "
+                "pass ep_max_grad_norm to MoETrainModule"
+            )
+
+    def optim_step(self) -> bool:
+        if self.ep_groups is not None:
+            from ..model import moe_ep
+
+            moe_ep.finalize_expert_grads(self.model, self.ep_groups)
+            if self.ep_max_grad_norm is not None:
+                norm = moe_ep.clip_grads_(self.model, self.ep_groups, self.ep_max_grad_norm)
+            else:
+                norm = moe_ep.global_grad_norm(self.model, self.ep_groups)
+            self._ep_norm = float(norm)
+        try:
+            return super().optim_step()
+        finally:
+            self._aux.reset()
+
+    @property
+    def last_grad_norm(self) -> float:
+        return self._ep_norm if self.ep_groups is not None else super().last_grad_norm
+
+    def full_model_state_dict(self):
+        sd = super().full_model_state_dict()
+        if self.ep_groups is not None:
+            from ..model import moe_ep
+
+            # the gather is a collective (every rank calls it); like the base state dict (rank 0 only
+            # when nothing is a DTensor) the result is merged where the base returned something
+            gathered = moe_ep.gather_full_expert_state(self.model, self.ep_groups, offload_to_cpu=True)
+            if sd:
+                for k, v in gathered.items():
+                    sd[k] = v.cpu()
+        return sd
+
+    def load_model_state_dict(self, nmo_state: dict) -> None:
+        if self.ep_groups is not None:
+            from ..model import moe_ep
+
+            nmo_state = {**nmo_state, **moe_ep.local_expert_slices(self.model, self.ep_groups, nmo_state)}
+        super().load_model_state_dict(nmo_state)
+
+    # -- the loop-facing call ---------------------------------------------
+    def train_batch(self, batch) -> LossOutput:
+        if not self.collect:
+            return super().train_batch(batch)
+
+        proxy = self._proxy
+        proxy.calls.clear()
+        mask = getattr(batch, "attention_mask", None)
+        proxy.attention_mask = mask if isinstance(mask, torch.Tensor) else None
+        out = self.loss_adapter(proxy, batch)
+        calls = list(proxy.calls)
+        proxy.calls.clear()
+        if not calls:
+            raise RuntimeError(
+                "the loss adapter ran no grad-enabled forward through the model, so no router "
+                "statistics exist; MoE aux losses/telemetry cannot be computed"
+            )
+        stats = merge_call_stats(calls)
+        want_lb, want_z = self.lb_weight > 0, self.z_weight > 0
+        lb, z = self._aux.micro(stats, want_lb=want_lb, want_z=want_z)
+
+        ce = out.loss
+        total = ce
+        decoupled = ce.grad_fn is None
+        if want_lb or want_z:
+            if decoupled:
+                raise NotImplementedError(
+                    "the loss adapter returned a detached proxy (decoupled backward, e.g. DPO): "
+                    "its backward already consumed the model graph, so a router aux loss cannot "
+                    "be added. Set moe_aux.lb_weight = z_weight = 0 for this stage "
+                    "(telemetry still works)."
+                )
+            if want_lb:
+                total = total + self.lb_weight * lb
+            if want_z:
+                total = total + self.z_weight * z
+        (total / self.grad_accum).backward()
+
+        metrics = dict(out.metrics)
+        if self._aux.m == self.grad_accum:
+            vals = self._aux.window_values()
+            self._last_metrics = window_metrics(vals, per_layer=self.aux_cfg.per_layer_metrics)
+        elif not self._last_metrics:
+            self._last_metrics = undefined_window_metrics(
+                len(stats), per_layer=self.aux_cfg.per_layer_metrics
+            )
+        metrics.update(self._last_metrics)
+
+        reported = ce.detach()
+        if want_lb or want_z:
+            # 0 * finite = 0 (loss stays exactly CE); 0 * NaN/inf = NaN, so a
+            # poisoned aux term cannot hide behind a healthy CE.
+            reported = reported + 0.0 * (total.detach() - ce.detach())
+        return LossOutput(loss=reported, n_tokens=out.n_tokens, metrics=metrics)
+
+    def zero_grads(self) -> None:
+        super().zero_grads()
+        self._aux.reset()
+
+
+def build_moe_train_module(
+    cfg,
+    model,
+    loss_adapter: LossAdapter,
+    optimizer,
+    *,
+    stage: str,
+    olmoe_cfg,
+    group=None,
+    ep_groups=None,
+    ep_max_grad_norm: Optional[float] = None,
+) -> MoETrainModule:
+    """The MoE module for an OLMoE-architecture model, wired from a RunConfig.
+
+    Carries the OLMoE weight format (the Trainer picks it up from the module),
+    `cfg.moe_aux` and `cfg.optimizer.grad_accum`.
+    """
+    from .ckpt_format import OlmoeFormat
+
+    return MoETrainModule(
+        model,
+        loss_adapter,
+        optimizer,
+        cfg.optimizer.grad_accum,
+        aux=cfg.moe_aux,
+        stage=stage,
+        top_k=olmoe_cfg.top_k,
+        weight_format=OlmoeFormat(olmoe_cfg),
+        group=group,
+        ep_groups=ep_groups,
+        ep_max_grad_norm=ep_max_grad_norm,
+    )
+
+
+__all__ = [
+    "STAGE_AUX_DEFAULTS",
+    "POST_TRAINING_AUX_DEFAULT",
+    "resolve_aux_weights",
+    "WindowTotals",
+    "merge_call_stats",
+    "MoEAux",
+    "window_metrics",
+    "undefined_window_metrics",
+    "MoETrainModule",
+    "build_moe_train_module",
+]

code/align/train/offload_adamw.py

@@ -0,0 +1,680 @@
+"""Optimizer-state-only CPU offload for AdamW.
+
+The arithmetic that forces this design (measured on this box, not re-derived):
+full fine-tuning 7.25B params with AdamW needs ~116GB of optimizer state --
+~29GB per GPU even fully sharded across 4 ranks, against a 23GB A10G budget.
+Moving the state to host RAM fits it. Moving anything ELSE to host RAM does
+not pay: FSDP's own `CPUOffloadPolicy` transfers parameters every forward AND
+backward, while this transfers once per optimizer step, which at grad-accum 4
+is a quarter as often and overlappable with the next accumulation window.
+
+Measured cost (tmp/offload_bench.json): torch CPU AdamW (foreach) 2.56s per
+1.81B-param rank shard; PCIe 13.2GB/s solo scaling to 52.7/53.9GB/s aggregate
+across 4 ranks, giving ~0.82s of transfer per rank per optimizer step. At
+grad-accum >= 4 that is an estimated 20-30% step-time overhead; at accum 1 it
+would be 50-100%. This is WHY accum >= 4 is mandatory for full-FT, and the
+Trainer refuses less.
+
+Design note on the update rule: the CPU step is torch's own `AdamW` driven
+over the offloaded fp32 master copies, not a reimplementation. Bitwise
+equivalence with `torch.optim.AdamW` is then a property of the construction
+rather than something a test has to keep rediscovering after every torch
+release -- and the test still catches the failures that actually happen here,
+which are plumbing failures: a dropped weight_decay, a mis-forwarded beta, a
+master copy that round-trips through a narrower dtype.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Iterable, Optional, Sequence
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+#: Optimizer identities this repo KNOWS. Adafactor is deliberately absent
+#: and refused BY NAME -- see `make_optimizer`.
+#:
+#: `adamw8bit` is BUILDABLE again: `align/train/adamw8bit.py` is a torch-only
+#: block-wise 8-bit AdamW (the bitsandbytes dependency stays banned; see
+#: docs/DEPENDENCY_POLICY.md). The distinction between "a typo'd optimizer
+#: name" and "a known optimizer" must still survive in the error text: the
+#: former is `ValueError("unknown optimizer")`.
+KNOWN_OPTIMIZERS = ("offload_adamw", "adamw8bit")
+
+#: The default optimizer in `RunConfig` and the one the pinned-staging budget
+#: and the RSS projection model. Every SHIPPED full-FT config overrides it to
+#: `adamw8bit`, because at 7.25B params / 4 ranks the host cannot hold this
+#: optimizer's state (see adamw8bit.py for the arithmetic).
+OPTIMIZER_OF_RECORD = "offload_adamw"
+
+
+def pinned_staging_bytes(cfg) -> float:
+    """Total pinned staging budget across all ranks, DERIVED from the config.
+
+    One fp32 layer-shard is `n_params / n_layers / world_size` elements at 4
+    bytes; each rank is allowed `pinned_staging_shards_per_rank` of them.
+
+    Pinned memory is unswappable and unreclaimable, and this box has no swap,
+    so this is a hard budget rather than an estimate. It is the same formula
+    `guard.projected_full_ft_rss_gb` uses for the RSS projection -- deliberately
+    the same, so the launch banner and the preflight cannot disagree about how
+    much host RAM the job intends to lock down.
+    """
+    if cfg.model.n_params <= 0:
+        raise ValueError(
+            f"model.n_params={cfg.model.n_params} is not a parameter count, so "
+            "there is no staging budget to derive. The tiny CI config sets it "
+            "to -1 deliberately (the 7B projection does not apply there); a "
+            "negative budget must raise rather than silently become a negative "
+            "cap that every allocation trivially satisfies."
+        )
+    world = cfg.train.world_size
+    layer_shard_bytes = cfg.model.n_params / cfg.model.n_layers / world * 4.0
+    return layer_shard_bytes * cfg.train.pinned_staging_shards_per_rank * world
+
+
+class OffloadAdamW:
+    """AdamW whose state lives on the host; params and grads never move.
+
+    `step()` returns False when the grad-skip gate refused the step.
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.Tensor],
+        lr: float,
+        betas: tuple[float, float] = (0.9, 0.95),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+        pinned_staging_cap_bytes: Optional[float] = None,
+        overlap: bool = True,
+        state_device: str = "cpu",
+        max_grad_norm: Optional[float] = None,
+    ) -> None:
+        self.params: list[torch.Tensor] = list(params)
+        if not self.params:
+            raise ValueError("OffloadAdamW got an empty parameter list")
+        if max_grad_norm is not None and not (float(max_grad_norm) > 0.0):
+            raise ValueError(
+                f"max_grad_norm must be a positive float or None, got {max_grad_norm!r}"
+            )
+        #: GLOBAL-norm clip (`clip_grad_norm_`'s rule, `min(1, c / (norm + 1e-6))`),
+        #: None = no clipping. THE SILENT-NO-CLIP DEFECT (H2 review, 2026-10-13):
+        #: `OptimizerConfig.max_grad_norm` defaults to 1.0 and every shipped TOML
+        #: states it, but this class had no such parameter and `make_optimizer`
+        #: never forwarded the field -- the one optimizer of record for full-FT
+        #: ran UNCLIPPED while `BlockwiseAdamW8bit` honoured the same config.
+        #: `last_grad_norm` stays the PRE-clip norm (the tripwires must see an
+        #: explosion, not its clipped image).
+        self.max_grad_norm = None if max_grad_norm is None else float(max_grad_norm)
+        self.state_device = torch.device(state_device)
+        self.overlap = bool(overlap)
+        self.banner_tag = "offload_adamw"
+
+        # THIS RANK'S bytes, so the staging cap is a per-rank quantity. Under
+        # FSDP2 every parameter is a DTensor whose `.numel()` is the GLOBAL
+        # element count: a world-size-times overcount that would have refused
+        # (or, worse, mis-banner) any sharded job.
+        self.staging_bytes_required = float(
+            sum(self._localize(p).numel() * 4 for p in self.params)
+        )
+        self.pinned_staging_cap_bytes = pinned_staging_cap_bytes
+        if (
+            pinned_staging_cap_bytes is not None
+            and self.staging_bytes_required > pinned_staging_cap_bytes
+        ):
+            raise ValueError(
+                f"pinned staging would need "
+                f"{self.staging_bytes_required / 2**30:.2f}GiB but the cap is "
+                f"{pinned_staging_cap_bytes / 2**30:.2f}GiB. Pinned memory is "
+                "unswappable and this box has no swap: exceeding the budget "
+                "does not degrade, it wedges the host. Lower "
+                "train.pinned_staging_shards_per_rank or shard further."
+            )
+
+        # fp32 master copies on the host. These are the AUTHORITATIVE weights:
+        # the device params are a bf16 view of them, refreshed every step.
+        self._pin = torch.cuda.is_available() and self.state_device.type == "cpu"
+        # THE FSDP2 DTENSOR DEFECT (found in the align port review, 2026-10-12;
+        # the predecessor never hit it because the shard-local 8-bit optimizer
+        # was the one that actually ran on FSDP2 parameters). Under `fully_shard`
+        # `model.parameters()` yields DTensors, and OffloadAdamW died on them
+        # twice: `is_pinned()` has no DTensor sharding rule at construction,
+        # and the grad-norm accumulator mixed a DTensor into a plain tensor at
+        # the first step. Masters, grads and the H2D copy-back therefore all go
+        # through `_localize` -- this rank's shard, a storage-sharing view -- the
+        # same contract `ShardLocalAdamW8bit` documents (AdamW is elementwise,
+        # so stepping the local shard is exact).
+        self._master = [
+            self._host_tensor(self._localize(p).detach().to(torch.float32)) for p in self.params
+        ]
+        self._inner = torch.optim.AdamW(
+            self._master,
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+        )
+
+        self.skipped_steps = 0
+        self.last_grad_norm = float("nan")
+        logger.info(
+            "OffloadAdamW: %d tensors, %.2fGiB fp32 master on %s (pinned=%s), "
+            "staging %.2fGiB, overlap=%s",
+            len(self.params),
+            sum(m.numel() * 4 for m in self._master) / 2**30,
+            self.state_device,
+            self._pin,
+            self.staging_bytes_required / 2**30,
+            self.overlap,
+        )
+
+    # -- host staging -------------------------------------------------------
+
+    @staticmethod
+    def _localize(t):
+        """This rank's local shard of an FSDP2 DTensor (a view sharing the
+        parameter's storage), or `t` itself for an ordinary tensor."""
+        return t.to_local() if hasattr(t, "to_local") else t
+
+    def _host_tensor(self, t: torch.Tensor) -> torch.Tensor:
+        out = t.to(self.state_device, copy=True)
+        if self._pin and not out.is_pinned():
+            out = out.pin_memory()
+        return out
+
+    # -- the step -----------------------------------------------------------
+
+    def step(self) -> bool:
+        """One optimizer step, or a refusal.
+
+        Order matters and is load-bearing: finiteness is checked BEFORE any
+        state is touched. An implementation that discovers the NaN after
+        updating the moments has already poisoned the second-moment estimate,
+        and every subsequent step is scaled by it.
+        """
+        grads = [None if p.grad is None else self._localize(p.grad) for p in self.params]
+        if all(g is None for g in grads):
+            raise RuntimeError(
+                "OffloadAdamW.step() called with no gradients on any parameter. "
+                "A step that silently does nothing is indistinguishable from a "
+                "run that trained."
+            )
+
+        if not self._all_finite(grads):
+            self.skipped_steps += 1
+            # NaN, not 0.0: the norm of a non-finite gradient is undefined, and
+            # a 0.0 in the log reads as 'perfectly converged'.
+            self.last_grad_norm = float("nan")
+            logger.warning(
+                "grad-skip gate: non-finite gradient, skipping optimizer step "
+                "(%d skipped so far). Moments and step counts left untouched.",
+                self.skipped_steps,
+            )
+            return False
+
+        self.last_grad_norm = self._grad_norm(grads)  # PRE-clip, on purpose
+
+        # `_grad_norm` all-reduces the sum of squares, so `coef` is the same on
+        # every rank: each rank scales its own disjoint shard by the GLOBAL factor.
+        coef = 1.0
+        if self.max_grad_norm is not None:
+            coef = min(1.0, self.max_grad_norm / (self.last_grad_norm + 1e-6))
+
+        # D2H: fp32 grads into the host staging that backs the master copies.
+        for master, g in zip(self._master, grads):
+            if g is None:
+                master.grad = None
+                continue
+            master.grad = self._host_tensor(g.detach().to(torch.float32))
+            if coef != 1.0:
+                master.grad.mul_(coef)
+
+        self._inner.step()
+
+        # H2D: updated master back into the device params, in their own dtype.
+        # copy_ keeps the storage identity FSDP's shard views depend on.
+        with torch.no_grad():
+            for p, master in zip(self.params, self._master):
+                local = self._localize(p).detach()
+                local.copy_(master.to(local.dtype))
+        return True
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        for p in self.params:
+            if set_to_none:
+                p.grad = None
+            elif p.grad is not None:
+                p.grad.zero_()
+        for m in self._master:
+            m.grad = None
+
+    def _all_finite(self, grads: Sequence[Optional[torch.Tensor]]) -> bool:
+        """Global finiteness. Under FSDP a rank whose shard is finite must not
+        step while a peer's is not: the peers would diverge silently and only
+        surface as an NCCL mismatch much later, if at all."""
+        local = True
+        for g in grads:
+            if g is not None and not bool(torch.isfinite(g).all()):
+                local = False
+                break
+        return bool(self._all_reduce_flag(local))
+
+    @staticmethod
+    def _collective_device() -> torch.device:
+        """The device a tensor must live on to join a collective on the default
+        process group: this rank's OWN cuda device under NCCL (a bare .cuda()
+        would put every rank's tensor on device 0 and reduce the wrong ones),
+        CPU under gloo. Keyed on the group's backend, not on
+        `torch.cuda.is_available()`: a gloo group on a box that happens to have
+        GPUs would otherwise be handed a cuda tensor it cannot reduce."""
+        import torch.distributed as dist
+
+        if "nccl" in str(dist.get_backend()).lower() and torch.cuda.is_available():
+            from ..utils.dist import local_rank
+
+            return torch.device("cuda", local_rank())
+        return torch.device("cpu")
+
+    @staticmethod
+    def _all_reduce_flag(local: bool) -> bool:
+        import torch.distributed as dist
+
+        if not dist.is_available() or not dist.is_initialized():
+            return local
+        flag = torch.tensor([1.0 if local else 0.0], device=OffloadAdamW._collective_device())
+        dist.all_reduce(flag, op=dist.ReduceOp.MIN)
+        return bool(flag.item() > 0.0)
+
+    @staticmethod
+    def _grad_norm(grads: Sequence[Optional[torch.Tensor]]) -> float:
+        """The GLOBAL L2 norm of the gradient, identical on every rank.
+
+        Accumulate on the GRADS' device, not on CPU: the accumulator used to be
+        a bare `torch.zeros((), dtype=float64)`, a CPU tensor, which broke the
+        moment the shard-local 8-bit optimizer called it with CUDA grads:
+        "Expected all tensors to be on the same device, but found cuda:0 and cpu".
+
+        THE SHARD-LOCAL NORM DEFECT (align port review, 2026-10-12). Under FSDP2
+        each rank holds a DISJOINT 1/world slice of every gradient, so the sum of
+        squares over the local slice is a norm of a quarter of the gradient --
+        and a different number on every rank. `GradNormTripwire` is a HALT wire
+        reading this value: a rank-local norm lets ONE rank cross the ceiling and
+        raise while its peers continue into the next collective (a hang until the
+        NCCL watchdog), and it under-reads a true explosion by ~sqrt(world). The
+        sum of squares is therefore all-reduced (SUM) before the root. Every
+        rank reaches this call together or not at all, because the finiteness
+        gate that precedes it is itself a collective.
+        """
+        import torch.distributed as dist
+
+        present = [g for g in grads if g is not None]
+        distributed = dist.is_available() and dist.is_initialized()
+        if not present and not distributed:
+            return float("nan")
+        device = present[0].device if present else torch.device("cpu")
+        total = torch.zeros((), dtype=torch.float64, device=device)
+        for g in present:
+            total += g.detach().to(torch.float64).pow(2).sum()
+        if distributed:
+            total = total.to(OffloadAdamW._collective_device())
+            dist.all_reduce(total, op=dist.ReduceOp.SUM)
+        return float(total.sqrt())
+
+    # -- introspection ------------------------------------------------------
+
+    def describe_placement(self) -> dict:
+        """Where every tensor actually lives. Consumed by the launch banner and
+        by the device-partition test; a placement claim nobody checks is how
+        'the state is offloaded' becomes untrue after a refactor."""
+
+        def entry(t: torch.Tensor) -> dict:
+            return {
+                "device": t.device.type,
+                "dtype": str(t.dtype),
+                "pinned": bool(t.is_pinned()) if t.device.type == "cpu" else False,
+                "bytes": t.numel() * t.element_size(),
+            }
+
+        state = []
+        for group in self._inner.state.values():
+            for key, v in group.items():
+                if isinstance(v, torch.Tensor):
+                    state.append({"key": key, **entry(v)})
+        return {
+            "state": state,
+            "master": [entry(m) for m in self._master],
+            "params": [entry(self._localize(p)) for p in self.params],
+            "staging_bytes_required": self.staging_bytes_required,
+            "pinned_staging_cap_bytes": self.pinned_staging_cap_bytes,
+        }
+
+    def state_snapshot_for_tests(self) -> dict:
+        """Flat clone of every optimizer-state tensor, including the step
+        counts. Tests assert these are byte-identical across a skipped step."""
+        out = {}
+        for i, p in enumerate(self._master):
+            for key, v in self._inner.state.get(p, {}).items():
+                out[f"{i}.{key}"] = (
+                    v.detach().clone() if isinstance(v, torch.Tensor) else torch.tensor(float(v))
+                )
+        return out
+
+    # -- resume -------------------------------------------------------------
+
+    def state_dict(self) -> dict:
+        """A DETACHED CLONE of every state tensor, on CPU.
+
+        THE SHARED-STATE ALIASING INCIDENT. Two traps stack here, and both are
+        silent:
+
+        1. `torch.optim.Optimizer.state_dict()` returns per-param dicts that
+           SHARE the live state tensors, and `.to("cpu")` on a tensor already
+           on CPU returns the same object -- so a state dict built without
+           clones is a live view of the optimizer, not a snapshot of it.
+        2. Mutating that returned dict in place (to clone the values) rebinds
+           the optimizer's OWN state entries, so the "snapshot" aliases the
+           live state no matter which order you do it in.
+
+        The result of getting this wrong is not a crash: `load_state_dict`
+        hands the loaded optimizer the same `step`/`exp_avg`/`exp_avg_sq`
+        objects the source still holds, both mutate them every step, Adam's
+        step counter advances twice per step, and the two runs diverge while
+        every state comparison between them says they are identical.
+
+        So this builds a NEW nested structure and never touches what the inner
+        optimizer returned.
+        """
+        inner = self._inner.state_dict()
+        return {
+            "param_groups": inner["param_groups"],
+            "state": {
+                idx: {
+                    k: (v.detach().to("cpu").clone() if isinstance(v, torch.Tensor) else v)
+                    for k, v in group.items()
+                }
+                for idx, group in inner["state"].items()
+            },
+            "master": [m.detach().to("cpu").clone() for m in self._master],
+            "skipped_steps": self.skipped_steps,
+        }
+
+    def load_state_dict(self, sd: dict) -> None:
+        """Restore moments AND master weights, then refresh the device params.
+
+        The master is the authoritative weight copy, so a resume that restored
+        only the moments would carry the pre-checkpoint parameters forward --
+        a run that looks like it resumed and is training a different model.
+        """
+        masters = sd.get("master")
+        if masters is None:
+            raise KeyError(
+                "optimizer state_dict has no 'master' entry; refusing to "
+                "resume from a checkpoint that carries moments but not the "
+                "fp32 master weights."
+            )
+        if len(masters) != len(self._master):
+            raise ValueError(
+                f"checkpoint has {len(masters)} master tensors but this "
+                f"optimizer holds {len(self._master)}"
+            )
+        with torch.no_grad():
+            for dst, src in zip(self._master, masters):
+                dst.copy_(src.to(dst.dtype))
+            for p, master in zip(self.params, self._master):
+                local = self._localize(p).detach()
+                local.copy_(master.to(local.dtype))
+        # Clone on the way IN as well: torch stores these objects directly when
+        # no dtype/device cast is needed, so loading one checkpoint dict into
+        # two optimizers would otherwise make them share moments. Symmetric
+        # with the cloning in state_dict(), and for the same reason.
+        inner_sd = {
+            "param_groups": sd["param_groups"],
+            "state": {
+                idx: {
+                    k: (v.detach().clone() if isinstance(v, torch.Tensor) else v)
+                    for k, v in group.items()
+                }
+                for idx, group in sd["state"].items()
+            },
+        }
+        self._inner.load_state_dict(inner_sd)
+        self.skipped_steps = int(sd.get("skipped_steps", 0))
+
+    def reseed_masters_from_params(self) -> None:
+        """Overwrite the fp32 masters with the CURRENT device params (audit U01-2).
+
+        The master is authoritative: every step rewrites the params from it. A resume that loads
+        weights but no optimizer state (a legacy checkpoint, `reshard(drop_optimizer=True)`) would
+        otherwise leave the masters at the LAUNCH weights and the first step would silently revert
+        the restored model. Moments stay fresh; only the master copy is re-seeded.
+        """
+        with torch.no_grad():
+            for p, master in zip(self.params, self._master):
+                master.copy_(self._localize(p).detach().to(torch.float32))
+
+    @property
+    def param_groups(self):
+        """Exposed so a torch LR scheduler can drive this optimizer directly."""
+        return self._inner.param_groups
+
+
+class ShardLocalAdamW8bit:
+    """Compatibility shim: the legacy name for the shard-local 8-bit AdamW.
+
+    The optimizer itself is `align.train.adamw8bit.BlockwiseAdamW8bit` (pure
+    torch; `make_optimizer` returns THAT directly). This class keeps the
+    legacy wrapper shape -- `_inner`, `_pairs`, and a `state_dict` that wraps
+    the inner state with `world_size`/`n_shards` -- because two invariants are
+    regression-tested against that shape WITHOUT constructing the class
+    (tests build a stub via `__new__` and hand it a fake `_inner`):
+
+      * the recursive CPU clone: a one-level clone once passed ~3.2GB of live
+        CUDA moments straight into the checkpoint pickle, and `gather_object`
+        restored them onto the SENDER'S device on rank 0 and OOM'd the DPO
+        run at its first checkpoint;
+      * the world-size refusal: the saved state is SHARD-shaped, and reloading
+        it under a different world size would silently align rank r's moments
+        to a different slice of the parameter.
+
+    WHY SHARD-LOCAL STEPPING IS EXACT, not an approximation: AdamW is
+    ELEMENTWISE. Every update depends only on that element's own gradient and
+    its own two moment scalars, so there is nothing for a shard boundary to
+    cut. (This is exactly what Adafactor cannot claim, and why Adafactor is
+    banned here: its factored second moment is a row/column statistic that
+    DOES cross shard boundaries.) `to_local()` is a VIEW sharing storage with
+    the DTensor, so stepping it updates the parameter -- measured on 4 real
+    ranks: storage pointers stable across calls and steps, and the DTensor
+    observes the update.
+
+    HISTORY: the first version wrapped bitsandbytes, whose fused kernel
+    cannot consume DTensors ("optimizer_update_8bit_blockwise got mixed
+    torch.Tensor and DTensor"), and was then retired with the dependency. It
+    is now a thin delegate to the in-repo implementation, which does the
+    gating, the GLOBAL grad norm and the state handling itself; the shim only
+    mirrors `skipped_steps` / `last_grad_norm` so legacy readers keep working.
+    """
+
+    def __init__(
+        self,
+        params: Iterable[torch.Tensor],
+        lr: float,
+        betas: tuple[float, float] = (0.9, 0.95),
+        eps: float = 1e-8,
+        weight_decay: float = 0.0,
+    ):
+        # Consume the parameter iterable first so a generator-backed caller
+        # (`model.parameters()`) behaves the same however it is used, and
+        # validate the empty case before anything is built.
+        params = list(params)
+        if not params:
+            raise ValueError("ShardLocalAdamW8bit got an empty parameter list")
+        # Lazy: adamw8bit imports this module at top level (it reuses the
+        # collective helpers), so the reverse import cannot be top-level.
+        from .adamw8bit import BlockwiseAdamW8bit
+
+        self._inner = BlockwiseAdamW8bit(
+            params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay
+        )
+        self._pairs = [(p, self._localize(p).detach()) for p in params]
+        self.skipped_steps = 0
+        self.last_grad_norm = float("nan")
+        self.banner_tag = self._inner.banner_tag
+
+    @staticmethod
+    def _localize(t):
+        return t.to_local() if hasattr(t, "to_local") else t
+
+    def step(self) -> bool:
+        """One optimizer step, or a refusal. Same contract as OffloadAdamW.
+
+        Gating, the global grad norm and the update all live in the inner
+        optimizer; doing them here too would pay the collectives twice."""
+        stepped = self._inner.step()
+        self.skipped_steps = self._inner.skipped_steps
+        self.last_grad_norm = self._inner.last_grad_norm
+        return stepped
+
+    def zero_grad(self, set_to_none: bool = True) -> None:
+        for p, local in self._pairs:
+            p.grad = None if set_to_none else p.grad
+            local.grad = None
+
+    def state_dict(self) -> dict:
+        from ..utils.dist import world_size
+
+        inner = self._inner.state_dict()
+
+        # DETACHED CPU CLONES, RECURSIVELY -- same two aliasing traps
+        # OffloadAdamW.state_dict documents, plus the one that survived the
+        # first fix: bnb nests its real tensors under
+        # state[idx]["__bnb_optimizer_quant_state__"], a dict, so a one-level
+        # clone passed ~3.2GB of live CUDA moments straight into the pickle.
+        # gather_object then restored them onto the SENDER'S device on rank 0
+        # and OOM'd the DPO run at its first checkpoint (SFT's end-save
+        # survived the same bug on VRAM headroom alone).
+        def _to_cpu(o):
+            if isinstance(o, torch.Tensor):
+                return o.detach().to("cpu").clone()
+            if isinstance(o, dict):
+                return {k: _to_cpu(v) for k, v in o.items()}
+            if isinstance(o, (list, tuple)):
+                return type(o)(_to_cpu(v) for v in o)
+            return o
+
+        return {
+            "inner": _to_cpu(inner),
+            "world_size": world_size(),
+            "n_shards": len(self._pairs),
+            "skipped_steps": self.skipped_steps,
+        }
+
+    def load_state_dict(self, sd: dict) -> None:
+        from ..utils.dist import world_size
+
+        saved_ws, now_ws = sd.get("world_size"), world_size()
+        if saved_ws is not None and saved_ws != now_ws:
+            raise ValueError(
+                f"this optimizer state was saved at world_size={saved_ws} but "
+                f"the current run has world_size={now_ws}. The 8-bit moments "
+                "are SHARD-shaped: loading them under a different world size "
+                "would align each rank's state to a different slice of the "
+                "parameter, and training would continue plausibly on "
+                "corrupted state. Resume at the original world size, or start "
+                "fresh."
+            )
+        saved_n = sd.get("n_shards")
+        if saved_n is not None and saved_n != len(self._pairs):
+            raise ValueError(
+                f"checkpoint holds {saved_n} shards, this optimizer has {len(self._pairs)}"
+            )
+        self._inner.load_state_dict(sd["inner"])
+        self.skipped_steps = int(sd.get("skipped_steps", 0))
+
+    @property
+    def param_groups(self):
+        return self._inner.param_groups
+
+
+def make_optimizer(cfg, params: Iterable[torch.Tensor]):
+    """Build the optimizer named by the config, and tag it for the banner.
+
+    Which optimizer actually ran is recorded in the banner: a run that used a
+    different optimizer than its config names is not reproducible from that
+    config.
+
+    `adamw8bit` is the pre-registered on-device fallback for the case where
+    host RAM cannot hold the offloaded state (every shipped full-FT config
+    names it). It is the in-repo torch-only `BlockwiseAdamW8bit`; the banner
+    states its state size from the RANK-LOCAL element count
+    (`2*n + 8*n/block` bytes), so the launch log cannot claim a footprint the
+    optimizer does not have. It is never silently swapped for `offload_adamw`:
+    a run whose banner says one optimizer while the config names another is
+    exactly the irreproducibility this function exists to prevent.
+    """
+    name = cfg.optimizer.name
+    lowered = str(name).lower()
+
+    if lowered == "adafactor":
+        raise ValueError(
+            "Adafactor is BANNED in this repo. Its factored second moment "
+            "crosses DTensor shard boundaries: under ZeRO-3 on this box it "
+            "produced a perfect-looking loss curve while moving ~1% of "
+            "tensors -- a run that trained nothing and reported success for "
+            "hours. AdamW only (offload_adamw is the optimizer of record)."
+        )
+    if lowered not in KNOWN_OPTIMIZERS:
+        raise ValueError(
+            f"unknown optimizer {name!r}; known: {list(KNOWN_OPTIMIZERS)}. "
+            "A typo must not fall through to a default -- the whole point of "
+            "naming the optimizer in the config is that the banner can state "
+            "which one ran."
+        )
+
+    o = cfg.optimizer
+    if lowered == "offload_adamw":
+        # The staging cap is a full-FT host-RAM budget. The tiny CPU harness
+        # has no layer shards to size it from (n_params = -1), so it runs
+        # uncapped rather than against a fabricated number.
+        cap = pinned_staging_bytes(cfg) if cfg.train.full_ft else None
+        opt = OffloadAdamW(
+            params,
+            lr=o.lr,
+            betas=tuple(o.betas),
+            eps=o.eps,
+            weight_decay=o.weight_decay,
+            pinned_staging_cap_bytes=cap,
+            max_grad_norm=getattr(o, "max_grad_norm", None),
+        )
+        opt.banner_tag = (
+            f"offload_adamw (fp32 master + moments on host, "
+            f"{opt.staging_bytes_required / 2**30:.2f}GiB staging)"
+        )
+        return opt
+
+    # Only 'adamw8bit' can reach here (unknown names raised above).
+    # Lazy import: adamw8bit imports this module for the collective helpers.
+    from .adamw8bit import BlockwiseAdamW8bit
+
+    opt = BlockwiseAdamW8bit(
+        params,
+        lr=o.lr,
+        betas=tuple(o.betas),
+        eps=o.eps,
+        weight_decay=o.weight_decay,
+        # OptimizerConfig.max_grad_norm (default 1.0), same rule and same pre-clip
+        # `last_grad_norm` record as the offload_adamw branch above.
+        max_grad_norm=getattr(o, "max_grad_norm", None),
+    )
+    n_local = sum(s["numel"] for s in opt._state)
+    state_bytes = BlockwiseAdamW8bit.projected_state_bytes(
+        [s["numel"] for s in opt._state], opt.block_size, opt.min_8bit_size
+    )
+    opt.banner_tag = (
+        f"adamw8bit (block-wise 8-bit moments ON DEVICE, block {opt.block_size}, "
+        f"{state_bytes / 2**30:.2f}GiB state for {n_local / 1e6:.1f}M local params "
+        f"= 2B + 8B/block per param; no master copy; zero host RSS)"
+    )
+    return opt

code/align/train/probes.py

@@ -0,0 +1,627 @@
+"""Probes and tripwires: the checks a loss curve cannot make.
+
+Every mechanism here exists because the predecessor shipped through the
+failure it detects, with healthy-looking metrics the whole way:
+
+  * TRUNK MOVEMENT. Reentrant gradient checkpointing silently zeroes trunk
+    gradients under sharding. The run trains -- loss falls, throughput is
+    normal -- while only the embeddings and the LM head move. Nothing in the
+    loss, the grad norm, or the throughput distinguishes it from a healthy
+    run, so the probe compares actual parameter values and RAISES at the first
+    log step rather than warning at the hundredth.
+  * logps/chosen DECLINE. DPO's objective only asks for a MARGIN, which is
+    equally well served by pushing the rejected logprob down as by pulling the
+    chosen one up. When it takes the second route the model gets better at the
+    loss and worse at the task. Maximal for Arm B by design, since pi_SFT
+    already assigns the chosen sequence high likelihood.
+  * LENGTH DRIFT. Preference optimization has a documented length bias. The
+    alarm is baseline-RELATIVE (Arm A measured 54.9% chosen-longer), because
+    an absolute threshold says nothing about a dataset whose baseline is not
+    50%.
+
+Tripwires are pure functions over metric windows, so they are tested without a
+model. NaN in means NaN out and never a trip: a NaN coerced to 0.0 reads as
+'flat', and flat-chosen-with-rising-margin is itself the alarm condition.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import Any, Optional, Sequence
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+NAN = float("nan")
+
+#: Parameter-name fragments that are NOT trunk. The probe must exclude them:
+#: the incident's signature is "head and embeddings moved, trunk did not", so
+#: a probe that tracked them would call that run healthy.
+NON_TRUNK_FRAGMENTS = ("embed_tokens", "lm_head", "embed_out", "wte", "wpe")
+
+
+class TrunkMovementProbe:
+    """Snapshot trunk parameters, then assert they actually moved."""
+
+    def __init__(self, model, max_tracked: int = 16, seed: int = 0):
+        self._model = model
+        self.tracked_names: list[str] = []
+        self._snapshots: list[torch.Tensor] = []
+
+        # Only TRAINABLE parameters can move. On a LoRA/partially-frozen model a
+        # stride that landed on frozen base weights would report "trunk frozen"
+        # on a perfectly healthy run (found in the PPO port: the probe was
+        # rebuilt on the real model at resume and tracked frozen tensors).
+        candidates = [
+            (n, p)
+            for n, p in model.named_parameters()
+            if p.requires_grad and not any(f in n for f in NON_TRUNK_FRAGMENTS)
+        ]
+        if not candidates:
+            raise ValueError(
+                "TrunkMovementProbe found no trunk parameters to track "
+                f"(every TRAINABLE parameter was absent or matched {NON_TRUNK_FRAGMENTS}). It cannot report "
+                "'no movement' on an empty set -- that is the alarm condition, "
+                "and reporting it from nothing would be a false alarm forever."
+            )
+        # Deterministic even sample across depth, so a probe on a 32-layer
+        # model does not track 16 tensors from layer 0.
+        stride = max(1, len(candidates) // max_tracked)
+        for name, p in candidates[::stride][:max_tracked]:
+            self.tracked_names.append(name)
+            self._snapshots.append(self.snapshot_tensor(p))
+        self.n_tracked = len(self._snapshots)
+
+    def release(self) -> None:
+        """Drop the model reference and the snapshots (U01-11): the probe is dead after the
+        post-train handoff and must not keep the training replica alive."""
+        self._model = None
+        self._snapshots = []
+
+    @staticmethod
+    def snapshot_tensor(t: Any) -> torch.Tensor:
+        """Detached CPU copy of a parameter, reaching through DTensor.
+
+        Under FSDP2 every sharded parameter is a DTensor whose global view is
+        an unmaterialised placeholder; `.to_local()` is the rank's actual
+        shard. Comparing global views would compare something that never
+        changes on this rank.
+        """
+        if hasattr(t, "to_local"):
+            t = t.to_local()
+        if hasattr(t, "detach"):
+            t = t.detach()
+        return t.to("cpu", copy=True).to(torch.float32)
+
+    def check(self, step: int, model=None, raise_on_frozen: bool = True) -> dict:
+        """Compare against the snapshot. Zero movement RAISES."""
+        target = model or self._model
+        if target is None:
+            raise RuntimeError("TrunkMovementProbe.check after release(): the model reference was dropped")
+        current = dict(target.named_parameters())
+        tensors = [current[n] for n in self.tracked_names]
+        deltas = []
+        for snap, t in zip(self._snapshots, tensors):
+            now = self.snapshot_tensor(t)
+            deltas.append((now - snap).abs().max().item())
+
+        moved = sum(1 for d in deltas if d > 0.0)
+        report = {
+            "step": step,
+            "n_tracked": self.n_tracked,
+            "moved_fraction": moved / self.n_tracked if self.n_tracked else NAN,
+            "max_abs_delta": max(deltas) if deltas else NAN,
+            "tracked_names": list(self.tracked_names),
+        }
+        if moved == 0 and raise_on_frozen:
+            raise RuntimeError(
+                f"TRUNK MOVEMENT PROBE FAILED at step {step}: not one of "
+                f"{self.n_tracked} tracked trunk parameters changed by any "
+                "amount since step 0.\n\nThis is the reentrant-checkpointing "
+                "incident: `torch.utils.checkpoint` with use_reentrant=True "
+                "silently zeroes trunk gradients under sharding, and the run "
+                "keeps training the embeddings and the LM head while the body "
+                "stays frozen. Loss still falls. Throughput is normal. Nothing "
+                "else detects it.\n\nCheck that activation checkpointing was "
+                "wrapped with use_reentrant=False, that the trunk parameters "
+                "are in the optimizer's param groups, and that requires_grad "
+                f"is set on them. Tracked: {self.tracked_names[:4]}..."
+            )
+        logger.info(
+            "trunk probe ok at step %d: %d/%d moved, max|d|=%.3e",
+            step,
+            moved,
+            self.n_tracked,
+            report["max_abs_delta"],
+        )
+        return report
+
+
+# ---------------------------------------------------------------------------
+# tripwires
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TripwireResult:
+    tripped: bool
+    response: str
+    value: float = NAN
+    logps_slope: float = NAN
+    margin_slope: float = NAN
+    detail: str = ""
+
+
+def _slope(values: Sequence[float]) -> float:
+    """Least-squares slope over an evenly spaced window. NaN anywhere in the
+    window makes the slope NaN -- a window with missing data has no trend, and
+    pretending otherwise is how a tripwire fires on absent measurements."""
+    n = len(values)
+    if n < 2 or any(math.isnan(v) for v in values):
+        return NAN
+    xs = list(range(n))
+    mx = sum(xs) / n
+    my = sum(values) / n
+    denom = sum((x - mx) ** 2 for x in xs)
+    if denom == 0:
+        return NAN
+    return sum((x - mx) * (y - my) for x, y in zip(xs, values)) / denom
+
+
+class LogpsChosenTripwire:
+    """Fires when PER-TOKEN chosen NLL rises materially while the margin improves.
+
+    R2 RECALIBRATION (2026-08-23 audit; docs/audits/phase1_postmortem.md).
+    The v1 instrument read raw SUMMED logps/chosen -- corr -0.89 with batch
+    completion length -- and applied a zero-threshold slope test over a
+    20-raw-row window (~1.25 optimizer steps). On the primary Arm B run it
+    fired 3,049 times; on a SHUFFLED copy of the same series, 2,987 times:
+    zero discriminating power. Its false alarm at micro-step 43 selected the
+    objective of the recorded (destroyed) rerun. By the correct metric --
+    per-token NLL, DPOP's own published diagnostic -- that primary run showed
+    no collapse at all (0.487 -> 0.548 nats/token across 727 steps).
+
+    v2 semantics: per-token chosen NLL (= -logps/chosen / lengths/chosen),
+    baseline median over the first `window` rows vs trailing median over the
+    last `window` rows; fires only when the trailing median exceeds baseline
+    by `nll_rise_nats` (default 0.15 -- healthy phase-1 drift measured 0.06,
+    Smaug's documented collapse is +1.45) AND the margin median has risen.
+    LATCHED: once fired it stays fired and stops logging. The calibration
+    contract is pinned by tests: it must NOT fire on a shuffled copy of a
+    healthy series, and MUST fire on an injected +0.3 nats/token shift.
+    """
+
+    #: The metric keys this tripwire READS, declared so the wiring can be
+    #: checked instead of assumed (THE DISARMED-TRIPWIRE INCIDENT, below).
+    INPUT_KEYS: tuple[str, ...] = ("logps/chosen", "lengths/chosen", "rewards/margins")
+
+    #: The response is PRE-REGISTERED: decided before the run, so a tripped
+    #: alarm produces a recorded configuration change rather than a debate
+    #: about whether the number is real. Recorded as the arm's configuration.
+    PRE_REGISTERED_RESPONSE = (
+        "rerun this arm with DPOP (chosen-logprob floor) or an NLL auxiliary "
+        "term, and record the switch as the arm's configuration"
+    )
+
+    def __init__(self, window: int = 400, nll_rise_nats: float = 0.15):
+        self.window = int(window)
+        self.nll_rise_nats = float(nll_rise_nats)
+        self._baseline_nll: list[float] = []
+        self._baseline_margin: list[float] = []
+        self._nll: list[float] = []
+        self._margin: list[float] = []
+        self.fired = False
+        self._fired_at_row: Optional[int] = None
+        self._rows = 0
+
+    def state_dict(self) -> dict:
+        """Baselines, window and latch (U01-6): primitives only, so resume.pt stays
+        `weights_only`-loadable. Without it every supervised relaunch re-baselined."""
+        return {
+            "baseline_nll": [float(v) for v in self._baseline_nll],
+            "baseline_margin": [float(v) for v in self._baseline_margin],
+            "nll": [float(v) for v in self._nll],
+            "margin": [float(v) for v in self._margin],
+            "fired": bool(self.fired),
+            "fired_at_row": self._fired_at_row,
+            "rows": int(self._rows),
+        }
+
+    def load_state_dict(self, sd: dict) -> None:
+        self._baseline_nll = [float(v) for v in sd["baseline_nll"]]
+        self._baseline_margin = [float(v) for v in sd["baseline_margin"]]
+        self._nll = [float(v) for v in sd["nll"]]
+        self._margin = [float(v) for v in sd["margin"]]
+        self.fired = bool(sd["fired"])
+        fr = sd.get("fired_at_row")
+        self._fired_at_row = None if fr is None else int(fr)
+        self._rows = int(sd["rows"])
+
+    @staticmethod
+    def _median(values: Sequence[float]) -> float:
+        # NaN-poisoning median: any NaN in the window makes the statistic NaN,
+        # so missing data can never fire (blind != healthy != alarm).
+        if any(math.isnan(v) for v in values):
+            return NAN
+        s = sorted(values)
+        n = len(s)
+        return s[n // 2] if n % 2 else 0.5 * (s[n // 2 - 1] + s[n // 2])
+
+    def update(self, logps_chosen: float, lengths_chosen: float, margin: float) -> TripwireResult:
+        self._rows += 1
+        lp, ln = float(logps_chosen), float(lengths_chosen)
+        nll = -lp / ln if (not math.isnan(lp) and not math.isnan(ln) and ln > 0) else NAN
+        m = float(margin)
+
+        if self.fired:
+            return TripwireResult(
+                True,
+                self.PRE_REGISTERED_RESPONSE,
+                detail=f"latched since row {self._fired_at_row}",
+            )
+        if len(self._baseline_nll) < self.window:
+            self._baseline_nll.append(nll)
+            self._baseline_margin.append(m)
+            return TripwireResult(
+                False, "", detail=f"baselining {len(self._baseline_nll)}/{self.window}"
+            )
+        self._nll.append(nll)
+        self._margin.append(m)
+        self._nll = self._nll[-self.window :]
+        self._margin = self._margin[-self.window :]
+        if len(self._nll) < self.window:
+            return TripwireResult(False, "", detail=f"windowing {len(self._nll)}/{self.window}")
+
+        base_nll = self._median(self._baseline_nll)
+        cur_nll = self._median(self._nll)
+        base_m = self._median(self._baseline_margin)
+        cur_m = self._median(self._margin)
+        rise = cur_nll - base_nll
+        tripped = (
+            not math.isnan(rise)
+            and not math.isnan(cur_m)
+            and not math.isnan(base_m)
+            and rise > self.nll_rise_nats
+            and cur_m > base_m
+        )
+        if tripped:
+            self.fired = True
+            self._fired_at_row = self._rows
+            logger.error(
+                "TRIPWIRE: per-token chosen NLL rose %.3f nats/token over baseline "
+                "(%.3f -> %.3f) while the DPO margin rises (%.4f -> %.4f). The "
+                "objective is degrading chosen likelihood in real terms. "
+                "Pre-registered response: %s",
+                rise,
+                base_nll,
+                cur_nll,
+                base_m,
+                cur_m,
+                self.PRE_REGISTERED_RESPONSE,
+            )
+        return TripwireResult(
+            tripped,
+            self.PRE_REGISTERED_RESPONSE if tripped else "",
+            value=rise,
+            logps_slope=rise,  # field kept for record-schema compat: nll shift
+            margin_slope=(cur_m - base_m) if not math.isnan(cur_m) else NAN,
+        )
+
+
+class GradNormTripwire:
+    """HALTS the run on a gradient-norm explosion. R3 (2026-08-23 audit).
+
+    grad_norm was logged every optimizer step of both destroyed-run attempts
+    and read by NOTHING: the recorded rerun's first step showed 2,851 vs the
+    primary's 22 (~130x), and the terminal blow-up reached 355,830 with the
+    run continuing 221 steps past the break. A ceiling here would have fired
+    at step 1. Unlike the diagnostic tripwires this one HALTS (HALT=True):
+    training past a gradient explosion spends GPU-hours destroying weights.
+    """
+
+    INPUT_KEYS: tuple[str, ...] = ("grad_norm",)
+    PRE_REGISTERED_RESPONSE = (
+        "halt the run; a gradient explosion is a launch/config defect, not a "
+        "tunable -- diagnose before any rerun"
+    )
+    HALT = True
+    #: `grad_norm` is a TRAINER-produced metric that is legitimately NaN on
+    #: every micro-step without an applied optimizer step (the NaN-reporting
+    #: contract). That is a periodic honest state, not a disarmed feed --
+    #: this wire handles NaN as "blind, cannot fire" itself. The arm-time
+    #: self-test therefore still requires the KEY (name drift stays a launch
+    #: failure) but tolerates NaN for wires that declare this.
+    TOLERATES_NAN_INPUT = True
+
+    def __init__(
+        self,
+        absolute_ceiling: float = 500.0,
+        relative_factor: float = 10.0,
+        window: int = 50,
+        warmup: int = 10,
+    ):
+        self.absolute_ceiling = float(absolute_ceiling)
+        self.relative_factor = float(relative_factor)
+        self.window = int(window)
+        self.warmup = int(warmup)
+        self._history: list[float] = []
+        self.fired = False
+
+    def state_dict(self) -> dict:
+        """Trailing history and latch (U01-6); primitives only."""
+        return {"history": [float(v) for v in self._history], "fired": bool(self.fired)}
+
+    def load_state_dict(self, sd: dict) -> None:
+        self._history = [float(v) for v in sd["history"]]
+        self.fired = bool(sd["fired"])
+
+    def update(self, grad_norm: float) -> TripwireResult:
+        g = float(grad_norm)
+        if self.fired:
+            return TripwireResult(True, self.PRE_REGISTERED_RESPONSE, value=g, detail="latched")
+        if math.isnan(g):
+            return TripwireResult(False, "", value=NAN, detail="NaN grad_norm: blind, cannot fire")
+        history = self._history
+        trailing = sorted(history[-self.window :])
+        med = trailing[len(trailing) // 2] if trailing else NAN
+        over_abs = g > self.absolute_ceiling
+        over_rel = (
+            len(history) >= self.warmup
+            and not math.isnan(med)
+            and med > 0
+            and g > self.relative_factor * med
+        )
+        history.append(g)
+        if over_abs or over_rel:
+            self.fired = True
+            logger.error(
+                "TRIPWIRE (HALT): grad_norm %.1f exceeds %s (trailing median %.2f). %s",
+                g,
+                f"absolute ceiling {self.absolute_ceiling}"
+                if over_abs
+                else f"{self.relative_factor}x trailing median",
+                med,
+                self.PRE_REGISTERED_RESPONSE,
+            )
+            return TripwireResult(True, self.PRE_REGISTERED_RESPONSE, value=g)
+        return TripwireResult(False, "", value=g)
+
+
+class LengthDriftTripwire:
+    """Fires when the chosen-longer fraction leaves a band around the measured
+    baseline for the arm -- but only once it has seen enough pairs to mean it.
+
+    THREE DECLARED STATES, and keeping them distinct is the whole design:
+
+      * BLIND     -- input absent or NaN. The guard cannot see AND cannot say
+                     so. Forbidden; `assert_tripwires_armed` halts the run.
+      * WARMING UP -- inputs present and finite, cumulative pairs below
+                     `min_pairs`. The guard is watching and honestly reports
+                     insufficient sample. Valid; the self-test accepts it.
+      * WATCHING  -- armed and evaluating.
+
+    Why `min_pairs` exists (sec09's measurement): at micro_batch 2 the
+    per-batch chosen-longer fraction can only take the values {0, 0.5, 1}, two
+    of which sit outside a +-0.10 band around 0.549 -- so the alarm fired on
+    78% of batches while measuring nothing but the batch size. Their rolling
+    window fixes the estimator; this threshold fixes when it may speak. From
+    their standard-error table: 512 pairs gives 5.4e-06 false-fire per check,
+    256 gives roughly 4 spurious alarms per Arm-B run. An alarm that fires
+    four times a run is an alarm nobody reads the fifth time.
+
+    The baseline is a config field, never a constant here: Arm A measured
+    0.549 and Arm B's is different, so a hardcoded threshold would mis-fire on
+    one of them permanently.
+    """
+
+    INPUT_KEYS: tuple[str, ...] = ("chosen_longer_fraction", "pairs")
+
+    RESPONSE = (
+        "length bias suspected: report length-controlled win-rate "
+        "alongside the raw metric and record the drift band that fired"
+    )
+
+    def __init__(self, baseline: float, band: float = 0.10, min_pairs: int = 512):
+        if math.isnan(baseline):
+            raise ValueError(
+                "LengthDriftTripwire needs a measured baseline; NaN means "
+                "nobody measured the arm and the check would be meaningless"
+            )
+        self.baseline = float(baseline)
+        self.band = float(band)
+        self.min_pairs = int(min_pairs)
+        self._pairs_seen = 0
+
+    def state_dict(self) -> dict:
+        return {"pairs_seen": int(self._pairs_seen)}
+
+    def load_state_dict(self, sd: dict) -> None:
+        self._pairs_seen = int(sd["pairs_seen"])
+
+    def state(self) -> str:
+        """`warming_up` until the cumulative sample reaches min_pairs, then
+        `watching`. Reported in the manifest so 'armed' is never ambiguous."""
+        return "watching" if self._pairs_seen >= self.min_pairs else "warming_up"
+
+    def update(self, chosen_longer_fraction: float, n_pairs: int = 1) -> TripwireResult:
+        self._pairs_seen += int(n_pairs)
+        v = float(chosen_longer_fraction)
+        if math.isnan(v):
+            # No measurement is not a drift and not an all-clear.
+            return TripwireResult(False, "", value=NAN, detail="chosen_longer_fraction undefined")
+        if self._pairs_seen < self.min_pairs:
+            return TripwireResult(
+                False, "", value=v, detail=f"warming up: {self._pairs_seen}/{self.min_pairs} pairs"
+            )
+
+        tripped = abs(v - self.baseline) > self.band
+        if tripped:
+            logger.error(
+                "TRIPWIRE: chosen-longer fraction %.3f is outside %.3f +- %.3f over %d pairs",
+                v,
+                self.baseline,
+                self.band,
+                self._pairs_seen,
+            )
+        return TripwireResult(
+            tripped,
+            self.RESPONSE if tripped else "",
+            value=v,
+            detail=f"watching ({self._pairs_seen} pairs)",
+        )
+
+
+def assert_tripwires_armed(metrics: dict, tripwires: dict) -> None:
+    """Every ARMED tripwire's inputs must be present and finite, or halt.
+
+    THE DISARMED-TRIPWIRE INCIDENT. section-05's kernel emits the implicit
+    margin as `rewards/margins`; section-07's Trainer seeded `margin` with NaN
+    and the tripwire read that. Nothing produced the key, `_slope` returns NaN
+    for any window containing one, and the trip condition requires a non-NaN
+    slope -- so the chosen-likelihood-collapse tripwire could never fire while
+    appearing fully wired: constructed, updated every log step, and listed as
+    armed in every manifest. A guard that cannot fire is worse than no guard,
+    because it is also a claim that the failure mode is being watched.
+
+    TWO DIFFERENT CONTRACTS, and conflating them is what caused this:
+
+      * A REPORTED statistic that is undefined must be NaN. `0.0` would read
+        as a real measurement ("no preference learned"), which is a finding
+        rather than a missing number.
+      * A tripwire INPUT must never be NaN-seeded. An armed guard reading NaN
+        is running blind, and blind is indistinguishable from healthy in every
+        log it writes.
+
+    Run at the first log step of any stage with armed tripwires, so metric-name
+    drift becomes a launch failure instead of a silent disarming.
+
+    A tripwire reporting `warming_up` PASSES this check. Its inputs are
+    present and finite; it simply has not accumulated enough sample to speak
+    yet, and it says so. Treating that as a failure would halt every DPO run
+    on its first batch, and the fix for THAT would be to delete the self-test
+    -- which is exactly how the original disarmed tripwire came about.
+    """
+    problems = []
+    for name, wire in tripwires.items():
+        for key in getattr(wire, "INPUT_KEYS", ()):
+            if key not in metrics:
+                problems.append(
+                    f"{name}: input {key!r} is ABSENT from the logged metrics "
+                    f"(present: {sorted(metrics)[:8]})"
+                )
+            elif (
+                isinstance(metrics[key], float)
+                and math.isnan(metrics[key])
+                and not getattr(wire, "TOLERATES_NAN_INPUT", False)
+            ):
+                problems.append(f"{name}: input {key!r} is NaN -- the guard would run blind")
+    if problems:
+        raise RuntimeError(
+            "ARMED TRIPWIRE SELF-TEST FAILED -- halting before training "
+            "proceeds:\n  - "
+            + "\n  - ".join(problems)
+            + "\n\nAn armed tripwire whose input is missing or NaN cannot "
+            "fire, and will report itself armed in every manifest anyway. "
+            "Either supply the metric (the loss adapter owns this) or do not "
+            "arm the tripwire for this stage. Canonical margin key: "
+            "'rewards/margins' -- the name section-05's kernel emits."
+        )
+
+
+# ---------------------------------------------------------------------------
+# throughput probe + RSS projection
+# ---------------------------------------------------------------------------
+
+
+def _percentile(sorted_vals: Sequence[float], q: float) -> float:
+    if not sorted_vals:
+        return NAN
+    idx = min(len(sorted_vals) - 1, max(0, int(round(q * (len(sorted_vals) - 1)))))
+    return sorted_vals[idx]
+
+
+def aggregate_micro_step_times(micro_times_s: Sequence[float], grad_accum: int) -> list[float]:
+    """Sum consecutive micro-step durations into OPTIMIZER-step durations.
+
+    The unit bug this exists to kill (task #25): both probes measure
+    per-MICRO durations, while `throughput_report` takes per-OPTIMIZER token
+    counts and horizon. Feeding micro times straight in under-projects
+    wall-clock and over-reports tokens/s by ~grad_accum x (8x on the DPO
+    config) -- a banner that costs GPU-hours wrong by an order of magnitude.
+
+    Only COMPLETE groups are kept: a trailing partial group would read as one
+    impossibly fast optimizer step. Raises when no complete group exists --
+    same contract as throughput_report itself (resolved or refused, never
+    quietly wrong); size probe_steps >= grad_accum at the config level.
+    """
+    if grad_accum <= 0:
+        raise ValueError(f"grad_accum={grad_accum} is not a schedule")
+    times = [float(t) for t in micro_times_s]
+    groups = [
+        sum(times[i : i + grad_accum]) for i in range(0, len(times) - grad_accum + 1, grad_accum)
+    ]
+    if not groups:
+        raise ValueError(
+            f"{len(times)} micro-step times cannot fill one optimizer step at "
+            f"grad_accum={grad_accum}; raise probe_steps to at least grad_accum"
+        )
+    return groups
+
+
+def throughput_report(
+    step_times_s: Sequence[float],
+    tokens_per_step: int,
+    optimizer_steps_planned: int,
+    offload_transfer_s: Optional[float] = None,
+    offload_cpu_step_s: Optional[float] = None,
+) -> dict:
+    """Turn measured step times into the banner's throughput entry.
+
+    Refuses to produce a partial entry: a stage that starts without knowing its
+    step time is a stage committing GPU-hours to a schedule nobody costed. When
+    offload is inactive the offload timings are NaN -- explicitly, because a
+    0.0 there would read as 'offload is free'.
+    """
+    times = [float(t) for t in step_times_s]
+    if not times:
+        raise ValueError(
+            "throughput_report got no step times. A probe that cannot fill "
+            "its banner keys must raise: the banner's contract is that every "
+            "value is resolved or explicitly NaN, never quietly absent."
+        )
+    if optimizer_steps_planned <= 0:
+        raise ValueError(
+            f"optimizer_steps_planned={optimizer_steps_planned} is not a "
+            "schedule; the projected wall-clock would be meaningless"
+        )
+
+    mean = sum(times) / len(times)
+    ordered = sorted(times)
+    return {
+        "n_steps": len(times),
+        "mean_step_time_s": mean,
+        "p50_step_time_s": _percentile(ordered, 0.50),
+        "p90_step_time_s": _percentile(ordered, 0.90),
+        "tokens_per_s": tokens_per_step / mean if mean > 0 else NAN,
+        "projected_epoch_s": optimizer_steps_planned * mean,
+        "offload_transfer_s": (NAN if offload_transfer_s is None else float(offload_transfer_s)),
+        "offload_cpu_step_s": (NAN if offload_cpu_step_s is None else float(offload_cpu_step_s)),
+    }
+
+
+def projected_job_rss(cfg) -> dict:
+    """Projected host RSS for this job, for the guard and the launch banner.
+
+    Delegates to `guard.projected_full_ft_rss_gb` rather than restating the
+    arithmetic. Two formulas for the same budget WILL drift, and then the
+    banner advertises a number the preflight is not enforcing -- while the
+    preflight is the only thing standing between this job and the 2026-08-12
+    host crash (2GB free, load 619, no swap, a healthy run killed by an NCCL
+    timeout five hours into a stalled checkpoint).
+    """
+    from ..utils.guard import projected_full_ft_rss_gb
+
+    return projected_full_ft_rss_gb(cfg)

code/align/train/quant8.py

@@ -0,0 +1,193 @@
+"""Block-wise 8-bit dynamic quantization, in plain torch.
+
+This is the numeric core of `adamw8bit.py` and carries NO optimizer logic, so
+the codebook and the quantize/dequantize round trip can be tested (and
+mis-implemented) on their own. The scheme is Dettmers et al. 2022, "8-bit
+Optimizers via Block-wise Quantization" (arXiv:2110.02861), rebuilt from the
+paper's description with torch only (docs/DEPENDENCY_POLICY.md: bitsandbytes
+is not on the allowed list).
+
+WHY BLOCK-WISE: one absmax for a whole 134M-element embedding shard lets a
+single outlier set the scale for every element, and the other 134M-1 values
+land in the bottom few codes. Normalising each block of `block_size` (default
+2048) elements by its OWN absmax confines an outlier's damage to its block and
+makes the quantization independent of how the tensor happens to be sharded
+across ranks -- which matters, because under FSDP2 the blocks are cut from the
+rank-LOCAL shard.
+
+WHY A DYNAMIC (NON-LINEAR) CODEBOOK: Adam's moments span many decades inside
+one block (the second moment is a square of a heavy-tailed gradient). A
+linear 8-bit grid has uniform ABSOLUTE resolution, so everything below
+~absmax/256 collapses to zero -- and a zero second moment makes the update
+`m / (sqrt(0) + eps)`, i.e. m / 1e-8. The dynamic map spends its 256 codes as
+an exponent (a power of ten, 7 of them) plus a linear fraction inside each
+decade: relative resolution of a few percent from absmax down to ~1e-6*absmax.
+
+THE ONE THING THIS MODULE DOES NOT DO is decide when quantization is safe:
+the caller keeps tensors below `min_8bit_size` in fp32 (small tensors such as
+norm scales are both cheap to keep and the most sensitive to rounding).
+"""
+
+from __future__ import annotations
+
+import functools
+
+import torch
+
+#: 8-bit codes. Not configurable: the code type is `uint8` and the lookup
+#: arithmetic assumes exactly 256 entries.
+N_CODES = 256
+
+
+def create_dynamic_map(signed: bool = True, max_exponent_bits: int = 7, total_bits: int = 8):
+    """The sorted 256-entry dynamic codebook, float32, built like bnb's
+    `create_dynamic_map` (the construction the paper describes).
+
+    Layout: for each exponent `i` in `0..max_exponent_bits-1` the decade
+    `10**(i - (max_exponent_bits-1)) * [0.1, 1]` is cut into linear fraction
+    bins and each bin contributes its MIDPOINT (so no code sits on a decade
+    edge and codes are strictly increasing). Decade `i` gets `2**i` fraction
+    values when signed and `2**(i+1)` when unsigned -- the unsigned map spends
+    the sign bit on one more bit of fraction. Exact `0.0` and `1.0` are then
+    appended. Counting, signed: 2 * (1+2+...+64) + 2 = 256; unsigned:
+    (2+4+...+128) + 2 = 256. Both assertions below fire if someone changes a
+    parameter without redoing the arithmetic.
+
+    Signed map range: [-0.99297, 1.0] (there is no exact -1.0; the most
+    negative code is the midpoint of the top bin, 0.7% from -1). That is the
+    same asymmetry bnb has, and it is bounded: the block's largest-magnitude
+    element, if negative, is stored within 0.7%.
+    Unsigned map range: [0.0, 1.0].
+
+    Built in float64 and cast once, so the table is bit-identical across
+    machines and torch versions (linspace/mean in fp32 is not guaranteed to be).
+    """
+    if total_bits != 8:
+        raise ValueError("create_dynamic_map builds 256-entry (8-bit) maps only")
+    non_sign_bits = total_bits - 1
+    if not 1 <= max_exponent_bits <= non_sign_bits:
+        raise ValueError(f"max_exponent_bits must be in [1, {non_sign_bits}]")
+    additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1
+
+    data: list[float] = []
+    for i in range(max_exponent_bits):
+        if signed:
+            fraction_items = 2 ** (i + non_sign_bits - max_exponent_bits) + 1
+        else:
+            fraction_items = 2 ** (i + non_sign_bits - max_exponent_bits + 1) + 1
+        boundaries = torch.linspace(0.1, 1.0, fraction_items, dtype=torch.float64)
+        means = (boundaries[:-1] + boundaries[1:]) / 2.0
+        scale = 10.0 ** (-(max_exponent_bits - 1) + i)
+        data += (scale * means).tolist()
+        if signed:
+            data += (-scale * means).tolist()
+    if additional_items > 0:
+        boundaries = torch.linspace(0.1, 1.0, additional_items + 1, dtype=torch.float64)
+        means = (boundaries[:-1] + boundaries[1:]) / 2.0
+        data += means.tolist()
+        if signed:
+            data += (-means).tolist()
+    data.append(0.0)
+    data.append(1.0)
+    data.sort()
+    if len(data) != N_CODES:
+        raise AssertionError(f"dynamic map has {len(data)} entries, expected {N_CODES}")
+    out = torch.tensor(data, dtype=torch.float64).to(torch.float32)
+    if not bool((out[1:] > out[:-1]).all()):
+        raise AssertionError("dynamic map is not strictly increasing in float32")
+    return out
+
+
+@functools.lru_cache(maxsize=None)
+def _cached_map(signed: bool) -> torch.Tensor:
+    return create_dynamic_map(signed=signed)
+
+
+_DEVICE_CACHE: dict = {}
+
+
+def codebook(signed: bool, device) -> torch.Tensor:
+    """The shared codebook on `device` (cached; never mutate the result)."""
+    device = torch.device(device)
+    key = (bool(signed), device.type, device.index)
+    cb = _DEVICE_CACHE.get(key)
+    if cb is None:
+        cb = _cached_map(bool(signed)).to(device)
+        _DEVICE_CACHE[key] = cb
+    return cb
+
+
+def zero_code(signed: bool) -> int:
+    """Index of the exact-0.0 code (the all-zero state initialises to it)."""
+    return int((_cached_map(bool(signed)) == 0.0).nonzero()[0])
+
+
+def nearest_code(x: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
+    """Nearest-code index of every element of `x` (any shape), as uint8.
+
+    `searchsorted` gives, for each x, the first index `hi` with cb[hi] >= x; the
+    nearest code is then either `hi` or `hi-1`, decided by comparing x against
+    their MIDPOINT (ties go to the lower code). Clamping `hi` to [1, 255]
+    first makes both candidates always exist, and values outside the table
+    (x <= cb[0] or x >= cb[-1]) resolve to the end codes by the same
+    comparison -- no special-casing, so no off-by-one at the edges.
+
+    `out_int32=True`: the index tensor is the largest transient in a chunked
+    step, and int32 halves it. (The final uint8 cast is what is stored.)
+    """
+    flat = x.reshape(-1).contiguous()
+    hi = torch.searchsorted(cb, flat, out_int32=True).clamp_(1, N_CODES - 1)
+    lo = hi - 1
+    mid = (cb[lo] + cb[hi]) * 0.5
+    pick_hi = flat > mid
+    return torch.where(pick_hi, hi, lo).to(torch.uint8).reshape(x.shape)
+
+
+def quantize_rows(x: torch.Tensor, cb: torch.Tensor):
+    """Quantize a `[n_blocks, block_size]` fp32 matrix, one block per row.
+
+    Returns `(codes uint8 [n_blocks, block_size], absmax fp32 [n_blocks])`.
+    An all-zero block has absmax 0; dividing by it would put NaN through the
+    lookup, so the divisor is replaced by 1 there (the block's values are all
+    0, which then quantize to the zero code and dequantize to 0 * absmax = 0).
+    """
+    absmax = x.abs().amax(dim=1)
+    safe = torch.where(absmax > 0, absmax, torch.ones_like(absmax))
+    normalized = x / safe.unsqueeze(1)
+    return nearest_code(normalized, cb), absmax
+
+
+def dequantize_rows(codes: torch.Tensor, absmax: torch.Tensor, cb: torch.Tensor) -> torch.Tensor:
+    """Inverse of `quantize_rows`: fp32 `[n_blocks, block_size]`.
+
+    The index MUST be a signed integer tensor: `cb[codes]` with a uint8 `codes`
+    is interpreted by torch as a (deprecated) boolean-mask index, not a gather.
+    `index_select` with an int32 index has no such reading.
+    """
+    n_blocks, block_size = codes.shape
+    values = cb.index_select(0, codes.reshape(-1).to(torch.int32)).reshape(n_blocks, block_size)
+    return values * absmax.unsqueeze(1)
+
+
+def quantize_blockwise(x: torch.Tensor, block_size: int = 2048, signed: bool = True):
+    """Quantize any tensor (flattened) -> `(codes uint8 [n_blocks*block_size], absmax)`.
+
+    The final partial block is zero-padded; the pad quantizes to the zero code
+    and is dropped by `dequantize_blockwise(..., numel=)`.
+    """
+    flat = x.detach().reshape(-1).to(torch.float32)
+    n = flat.numel()
+    n_blocks = max(1, -(-n // block_size))
+    padded = torch.zeros(n_blocks * block_size, dtype=torch.float32, device=flat.device)
+    padded[:n] = flat
+    codes, absmax = quantize_rows(padded.view(n_blocks, block_size), codebook(signed, flat.device))
+    return codes.reshape(-1), absmax
+
+
+def dequantize_blockwise(
+    codes: torch.Tensor, absmax: torch.Tensor, numel: int, block_size: int = 2048, signed: bool = True
+) -> torch.Tensor:
+    """Inverse of `quantize_blockwise`, trimmed to `numel` elements (fp32, flat)."""
+    rows = codes.reshape(-1, block_size)
+    out = dequantize_rows(rows, absmax, codebook(signed, codes.device))
+    return out.reshape(-1)[:numel]

code/align/train/qwen3_stage.py

@@ -0,0 +1,466 @@
+"""Qwen3 (Goedel-Prover-V2-8B) SFT stage pieces: model build, LoRA train module, weight format.
+
+`scripts/sft.py` dispatches here when `model.arch == "qwen3"`; the dense-Mistral path never
+imports this module. Everything the Trainer needs that is family-specific lives behind the
+two seams the Trainer already has: `weight_format` (what a checkpoint's weight file is) and
+`train_module` (what touches the model).
+
+THE LoRA CHECKPOINT CONTRACT (same as `PPOTrainModule`, for the same reasons):
+
+  * The weight file of every checkpoint is the MERGED bf16 model (base + scaling * B @ A,
+    accumulated in fp32, cast once). Checkpoint saves are how vLLM / transformers / the
+    rollout engine receive the policy, and they load HF base weights, not adapters.
+  * The adapter itself (fp32, exact) rides in the optimizer payload, so a resume is lossless:
+    base is reloaded from the pinned snapshot at launch, the adapter and the 8-bit moments
+    come back from `optimizer_rank{r}.pt`.
+  * `load_weights` on a LoRA resume reads only the safetensors HEADERS (key set + shapes, on
+    the meta device). Materialising the merged 16 GB on each of four ranks would cost 64 GB of
+    host RAM to verify a key set and then be thrown away -- the 2026-08-12 shape.
+
+FSDP2 + frozen base. Per-block `fully_shard` shards the frozen bf16 base (16.4 GB / world) and
+the fp32 adapters together; only adapter gradients are reduce-scattered. Activation
+checkpointing wraps each block. The root unit also owns embed_tokens / lm_head / final norm.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+import os
+from typing import Any, Optional
+
+import torch
+
+from .ckpt_format import WeightFormat  # noqa: F401  (the Protocol this implements)
+from .train_module import TrainModule
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_ARCHS = ("mistral", "qwen3")
+
+
+class Qwen3StageError(RuntimeError):
+    """A qwen3 stage config that cannot be launched as written."""
+
+
+# ---------------------------------------------------------------------------
+# Config validation
+# ---------------------------------------------------------------------------
+
+
+def validate_qwen3_cfg(cfg, qcfg) -> None:
+    """Refuse launches whose declared numbers disagree with the snapshot they will load.
+
+    Derived-vs-hardcoded: `model.n_params` feeds the RSS projection and the banner; the
+    architecture truth is the snapshot's config.json. A TOML that still carries the Mistral
+    defaults (7.25B / 32 layers) would project the wrong host footprint and print a banner
+    describing a different model.
+    """
+    from ..model.qwen3 import count_params
+
+    n = count_params(qcfg)
+    problems = []
+    if cfg.model.n_params != n:
+        problems.append(f"model.n_params={cfg.model.n_params} but the snapshot implies {n}")
+    if cfg.model.n_layers != qcfg.n_layers:
+        problems.append(f"model.n_layers={cfg.model.n_layers} but the snapshot has {qcfg.n_layers}")
+    if cfg.model.vocab_size != qcfg.vocab_size:
+        problems.append(f"model.vocab_size={cfg.model.vocab_size} but the snapshot has {qcfg.vocab_size}")
+    if cfg.model.hidden_size != qcfg.hidden_size:
+        problems.append(f"model.hidden_size={cfg.model.hidden_size} but the snapshot has {qcfg.hidden_size}")
+    if cfg.model.precision != "bf16":
+        problems.append(f"model.precision={cfg.model.precision!r}: bf16 only")
+    if problems:
+        raise Qwen3StageError("qwen3 stage config disagrees with its snapshot: " + "; ".join(problems))
+    lo = cfg.lora
+    if lo.enabled:
+        if lo.dropout != 0.0:
+            raise Qwen3StageError(
+                f"lora.dropout={lo.dropout}: align.ppo.lora has no dropout path (train-time and "
+                "score-time logprobs would differ); set 0.0"
+            )
+        if cfg.train.full_ft:
+            raise Qwen3StageError("lora.enabled with train.full_ft=true: LoRA freezes the base; set train.full_ft=false")
+    elif not cfg.train.full_ft:
+        raise Qwen3StageError("train.full_ft=false requires lora.enabled=true (otherwise nothing trains)")
+
+
+def resolve_qwen3(cfg):
+    """(snapshot_dir, Qwen3Config) from `paths.model_snapshot`, with the attention request applied."""
+    from ..model import checkpoint as ckpt
+    from ..model.qwen3_checkpoint import read_qwen3_config
+
+    snapshot = ckpt.resolve_snapshot_dir(cfg.paths.model_snapshot)
+    qcfg = read_qwen3_config(snapshot).replace(attn_implementation=cfg.model.attn_implementation)
+    validate_qwen3_cfg(cfg, qcfg)
+    return snapshot, qcfg
+
+
+# ---------------------------------------------------------------------------
+# Weight format
+# ---------------------------------------------------------------------------
+
+
+class Qwen3Format:
+    """`WeightFormat` for Qwen3: sharded HF safetensors + config.json + tokenizer files."""
+
+    def __init__(self, qcfg, tokenizer_dir: str, lora: bool):
+        self.qcfg = qcfg
+        self.tokenizer_dir = tokenizer_dir
+        self.lora = bool(lora)
+
+    def save_weights(self, state, out_dir, dtype):
+        from ..model.qwen3_checkpoint import save_hf_qwen3
+
+        save_hf_qwen3(dict(state), out_dir, cfg=self.qcfg, dtype=dtype)
+
+    def load_weights(self, path):
+        from ..model.checkpoint import load_hf_safetensors
+        from ..model.qwen3_checkpoint import hf_to_align
+
+        if self.lora:
+            return hf_to_align(_hf_headers_on_meta(path), self.qcfg)
+        return hf_to_align(load_hf_safetensors(path), self.qcfg)
+
+    def copy_tokenizer(self, snapshot, dst_dir):  # noqa: ARG002 -- the configured dir is authoritative
+        from ..model.qwen3_checkpoint import copy_tokenizer_files
+
+        return copy_tokenizer_files(self.tokenizer_dir, dst_dir)
+
+
+def _hf_headers_on_meta(path: str) -> dict:
+    """{hf key: empty meta tensor of the stored shape/dtype} for a checkpoint dir. Reads the
+    safetensors headers only -- no tensor data is touched."""
+    import json
+
+    from safetensors import safe_open
+
+    from ..model.checkpoint import INDEX_NAME, WEIGHTS_NAME
+
+    idx = os.path.join(path, INDEX_NAME)
+    if os.path.isfile(idx):
+        with open(idx) as f:
+            files = sorted(set(json.load(f)["weight_map"].values()))
+    elif os.path.isfile(os.path.join(path, WEIGHTS_NAME)):
+        files = [WEIGHTS_NAME]
+    else:
+        raise FileNotFoundError(f"no {WEIGHTS_NAME} or {INDEX_NAME} at {path!r}")
+    out: dict[str, torch.Tensor] = {}
+    for fname in files:
+        with safe_open(os.path.join(path, fname), framework="pt") as f:
+            for k in f.keys():
+                sl = f.get_slice(k)
+                dt = _ST_DTYPES[sl.get_dtype()]
+                out[k] = torch.empty(sl.get_shape(), dtype=dt, device="meta")
+    return out
+
+
+_ST_DTYPES = {"BF16": torch.bfloat16, "F32": torch.float32, "F16": torch.float16}
+
+
+# ---------------------------------------------------------------------------
+# Base identity (a LoRA resume must reattach the adapter to the base it was trained on)
+# ---------------------------------------------------------------------------
+
+
+def base_fingerprint(state) -> str:
+    """sha256 over the frozen base's identity: every non-adapter tensor's name/shape/dtype, every
+    1-D tensor in full (norm weights differ between fine-tunes of one architecture), and the first
+    and last 8 rows of every matrix (embeddings, lm_head and each projection). Cheap (~tens of MB
+    hashed for 8B) yet distinguishes two bases with the SAME key set and shapes -- the only
+    difference a LoRA resume, which never loads the checkpoint's weights, could otherwise miss.
+    Needs plain (unsharded, non-meta) tensors: a shard-local hash would differ per rank."""
+    import hashlib
+
+    h = hashlib.sha256()
+    for k in sorted(k for k in state if ".adapters." not in k):
+        t = state[k]
+        if hasattr(t, "to_local") or t.device.type == "meta":
+            raise ValueError(f"cannot fingerprint {k!r}: sharded or meta tensor")
+        t = t.detach()
+        h.update(f"{k}|{tuple(t.shape)}|{t.dtype}".encode())
+        parts = [t] if t.dim() <= 1 else [t[:8], t[-8:]]
+        for part in parts:
+            h.update(part.to("cpu").contiguous().reshape(-1).view(torch.uint8).numpy().tobytes())
+    return h.hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# LoRA train module
+# ---------------------------------------------------------------------------
+
+
+class LoRASFTModule(TrainModule):
+    """`TrainModule` whose checkpoints publish merged weights and persist the adapter."""
+
+    def __init__(self, model, loss_adapter, optimizer, grad_accum: int, adapter: str, weight_format=None):
+        super().__init__(model, loss_adapter, optimizer, grad_accum)
+        from ..ppo.lora import has_lora
+
+        if not has_lora(model):
+            raise Qwen3StageError("LoRASFTModule needs a model with injected LoRA layers")
+        self.adapter = adapter
+        if weight_format is not None:
+            self.weight_format = weight_format
+        # Identity of the frozen base this adapter trains on (recorded in the optimizer payload,
+        # checked on resume). `build_qwen3_stage_model` stamps it BEFORE the FSDP wrap, when the
+        # tensors are still whole; a directly built model is fingerprinted here if it can be.
+        fp = getattr(model, "_base_fingerprint", None)
+        if fp is None:
+            try:
+                fp = base_fingerprint(model.state_dict())
+            except ValueError as e:
+                logger.warning("LoRA base fingerprint unavailable (%s); a resume cannot verify its base", e)
+        self.base_fingerprint: Optional[str] = fp
+
+    def full_model_state_dict(self):
+        """BASE-keyed weights with the adapter FOLDED IN (see the module docstring).
+
+        `get_model_state_dict(full_state_dict=True, cpu_offload=True)` returns the full dict on
+        RANK 0 ONLY; every other rank gets `{}`. The gather is the collective every rank
+        joins; merging an empty dict is a KeyError on ranks 1..N-1 at the FIRST save, which
+        strands rank 0 in the next collective until the watchdog kills the run.
+        """
+        import torch.distributed as dist
+
+        from ..ppo.lora import export_merged_state_dict
+
+        sd = super().full_model_state_dict()
+        if not sd and dist.is_available() and dist.is_initialized() and dist.get_rank() != 0:
+            return sd
+        self.last_export_report: dict = {}  # retained-delta fidelity of the bf16 cast (warned on when poor)
+        return export_merged_state_dict(self.model, sd, self.adapter, report=self.last_export_report)
+
+    def load_model_state_dict(self, nmo_state: dict) -> None:
+        """LoRA resume: the checkpoint's weights are MERGED and are deliberately NOT loaded --
+        the frozen base was reloaded from the snapshot at launch, and loading merged weights
+        into it would double-count the delta once the adapter is restored from the optimizer
+        payload. The key set is still verified so a checkpoint of another architecture cannot
+        resume silently."""
+        base_keys = {k for k in self.model.state_dict() if ".adapters." not in k}
+        missing, unexpected = sorted(base_keys - set(nmo_state)), sorted(set(nmo_state) - base_keys)
+        if missing or unexpected:
+            raise KeyError(f"resume state mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")
+        logger.info("LoRA resume: merged weights verified (%d tensors) and NOT loaded; adapter comes from the optimizer payload", len(base_keys))
+
+    def optim_state_dict(self) -> dict:
+        from ..ppo.lora import lora_state_dict
+
+        lora = lora_state_dict(self.model, self.adapter)
+        # The ZeRO-3 save defect, structurally: a sharded placeholder serialises as an EMPTY
+        # tensor and the file loads without error while being untrained.
+        empty = [k for k, v in lora.items() if v.numel() == 0]
+        if empty:
+            raise RuntimeError(
+                f"{len(empty)} adapter tensors are empty (e.g. {empty[:2]}): a sharded placeholder was "
+                "serialised instead of the gathered parameter. Refusing to write a checkpoint that "
+                "loads without error and is untrained."
+            )
+        return {"optimizer": self.optimizer.state_dict(), "lora": lora, "base_fingerprint": self.base_fingerprint}
+
+    def load_optim_state_dict(self, sd: dict) -> None:
+        from ..ppo.lora import load_lora_state_dict
+
+        if "lora" not in sd or "optimizer" not in sd:
+            raise KeyError(
+                "checkpoint optimizer payload lacks 'lora'/'optimizer'; resuming would restart the "
+                "policy from the bare base while claiming a resume"
+            )
+        # THE UNBOUND-ADAPTER DEFECT. The checkpoint's merged weights are not loaded, so nothing
+        # else ties this adapter to the base reloaded at launch: a repointed model_snapshot or a
+        # second HF-cache snapshot (same keys, same shapes) would resume "fine" on the wrong base.
+        saved = sd.get("base_fingerprint")
+        if saved is None or self.base_fingerprint is None:
+            logger.warning(
+                "LoRA resume cannot verify the base: base fingerprint missing (checkpoint=%s, this launch=%s)",
+                saved is not None, self.base_fingerprint is not None,
+            )
+        elif saved != self.base_fingerprint:
+            raise RuntimeError(
+                "LoRA resume: the base model loaded at launch is not the base this adapter was trained on "
+                f"(fingerprint {self.base_fingerprint[:12]} != checkpoint {saved[:12]}). Check "
+                "paths.model_snapshot and the HF cache for a second snapshot."
+            )
+        self.optimizer.load_state_dict(sd["optimizer"])
+        load_lora_state_dict(self.model, self.adapter, sd["lora"], strict=True)
+
+
+# ---------------------------------------------------------------------------
+# Model build
+# ---------------------------------------------------------------------------
+
+
+def _vm_rss_gb(field: str = "VmRSS") -> float:
+    with open("/proc/self/status") as f:
+        for line in f:
+            if line.startswith(field + ":"):
+                return float(line.split()[1]) / 1024**2
+    return float("nan")
+
+
+def build_qwen3_stage_model(cfg, attn_implementation: Optional[str] = None):
+    """Rank-serial load -> GPU -> (LoRA inject, freeze base) -> activation checkpointing -> FSDP2.
+
+    RANK-SERIAL because each rank materialises the 16.4 GB bf16 state dict on the host before
+    moving it to its card; four at once is 66 GB on a no-swap box. Ranks load one at a time
+    (barrier-ordered) and return freed heap to the OS (`malloc_trim`) before the next starts.
+    """
+    from ..model import qwen3 as q3
+    from ..model.qwen3_checkpoint import hf_to_align
+    from ..model.checkpoint import load_hf_safetensors
+    from ..ppo.lora import LoRAConfig, inject_lora, mark_only_adapter_trainable
+    from ..utils.dist import barrier, local_rank, rank, world_size
+    from . import fsdp
+
+    snapshot, qcfg = resolve_qwen3(cfg)
+    if attn_implementation:
+        qcfg = qcfg.replace(attn_implementation=attn_implementation)
+    device = torch.device("cuda", local_rank())
+    torch.cuda.set_device(device)
+
+    model = None
+    peaks: dict[str, Any] = {}
+    for r in range(world_size()):
+        if rank() == r:
+            sd = hf_to_align(load_hf_safetensors(snapshot, dtype=torch.bfloat16), qcfg)
+            model = q3.Qwen3ForCausalLM.from_state_dict(qcfg, sd, device="cuda")
+            del sd
+            model = model.to(device)
+            peaks["rss_hwm_gb"] = round(_vm_rss_gb("VmHWM"), 2)  # high-water mark of the load, not the residual
+            gc.collect()
+            try:
+                import ctypes
+
+                ctypes.CDLL("libc.so.6").malloc_trim(0)
+            except (OSError, AttributeError):
+                pass
+            peaks["rss_after_trim_gb"] = round(_vm_rss_gb(), 2)
+        barrier()
+
+    if cfg.lora.enabled:
+        model._base_fingerprint = base_fingerprint(model.state_dict())  # whole tensors, before inject/FSDP
+        lc = LoRAConfig(r=cfg.lora.r, alpha=cfg.lora.alpha, target_modules=tuple(cfg.lora.target_modules), rslora=cfg.lora.rslora)
+        # Identical adapter init on every rank: FSDP2 shards what each rank already holds and
+        # does NOT broadcast, so a per-rank random A would be a different adapter per shard.
+        gen = torch.Generator().manual_seed(cfg.train.seed)
+        hit = inject_lora(model, lc, adapter=cfg.lora.adapter, generator=gen)
+        n_trainable = mark_only_adapter_trainable(model, cfg.lora.adapter)
+        expected = len(cfg.lora.target_modules) * qcfg.n_layers
+        if len(hit) != expected:
+            raise Qwen3StageError(f"LoRA hit {len(hit)} layers, expected {expected} ({len(cfg.lora.target_modules)} targets x {qcfg.n_layers} layers)")
+        peaks["lora_trainable_params"] = n_trainable
+    fsdp.apply_activation_checkpointing(model)
+    if world_size() > 1:
+        fsdp.apply_fsdp(model)
+    model._stage_load_stats = peaks  # for the banner / probe records
+    return model
+
+
+def qwen3_trainer_extras(cfg, model, loss_adapter, run_dir: str) -> dict:
+    """The Trainer kwargs a qwen3 stage adds: weight_format, per-rank tracker, and for LoRA the
+    train_module and the probe view.
+
+    PER-RANK TRACKER. The Trainer's default constructs a `Tracker` on EVERY rank, all appending
+    to the same `<run_dir>/history.jsonl`: four interleaved streams with no rank field, so
+    "did each rank see different rows?" cannot be answered from the record (and torn
+    interleaving is one fsync away). Rank 0 keeps `<run_dir>`; rank r writes `<run_dir>/rank{r}`.
+    """
+    from ..utils.dist import rank as _rank
+    from ..utils.tracking import Tracker
+    from .offload_adamw import make_optimizer
+    from .trainer import MANDATORY_METRICS
+
+    snapshot, qcfg = resolve_qwen3(cfg)
+    r = _rank()
+    extras: dict = {
+        "weight_format": Qwen3Format(qcfg, tokenizer_dir=snapshot, lora=cfg.lora.enabled),
+        "tracker": Tracker(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "sft",
+                           mandatory_keys=MANDATORY_METRICS, run_name=cfg.run_name),
+    }
+    if not cfg.lora.enabled:
+        return extras
+    optimizer = make_optimizer(cfg, [p for p in model.parameters() if p.requires_grad])
+    extras["train_module"] = LoRASFTModule(
+        model, loss_adapter, optimizer, cfg.optimizer.grad_accum, cfg.lora.adapter,
+        weight_format=extras["weight_format"],
+    )
+    from ..ppo.trainer import AdapterProbeView
+
+    extras["probe_model"] = AdapterProbeView(model, cfg.lora.adapter)
+    return extras
+
+
+def qwen3_entry_banner(cfg, ledger) -> dict:
+    """The qwen3 counterpart of `stage_gates.run_entry_gates`: what can be verified WITHOUT the
+    Mistral-specific machinery, and an explicit PENDING for what cannot.
+
+    NOT APPLICABLE here, and why (each is a Mistral/PRM800K artifact, not a skipped check):
+      * golden regression + MATH-500/GSM8K quality gate: they certify the Mistral eval stack
+        against a Mistral checkpoint's recorded score. The qwen3 experiment's evaluation is the
+        Lean verifier loop (P2-B/C).
+      * the 10-step 7B smoke / 20-step probe: `build_stage_model` is Mistral-only. The qwen3
+        memory and throughput are measured by scripts/qwen3_sft_probe.py and recorded in
+        docs; the banner says PENDING rather than quoting a number this launch did not measure.
+    KEPT: batch arithmetic, shard receipts (hashes into the banner), the RSS projection, the
+    sampler-derived scheduler horizon, the snapshot-vs-config agreement.
+    """
+    from ..utils.guard import projected_full_ft_rss_gb
+    from .stage_gates import PENDING, assemble_sft_banner, assert_batch_arithmetic_consistent
+    from .stage_smoke import build_training_sampler, resolve_split_dir, sampler_derived_horizon
+
+    snapshot, qcfg = resolve_qwen3(cfg)
+    assert_batch_arithmetic_consistent(cfg)
+    sampler, dataset = build_training_sampler(cfg)
+    horizon = sampler_derived_horizon(cfg, sampler)
+    # THE OVERLONG-ROW DEFECT. The length-grouped sampler budgets tokens, but a row longer than the
+    # budget is not rejected: it becomes a batch of one and trains. seq_len is the length the
+    # memory envelope was MEASURED at (scripts/qwen3_sft_probe.py); a longer row would be the
+    # first thing to run at an untested length, on the last GPU-hour of an epoch.
+    lengths = dataset.lengths()
+    over = int((lengths > cfg.train.seq_len).sum())
+    if over:
+        raise Qwen3StageError(
+            f"{over} of {len(lengths)} rows are longer than train.seq_len={cfg.train.seq_len} "
+            f"(longest {int(lengths.max())}). Filter them at data-build time or raise seq_len to a "
+            "length that has been measured; the sampler would otherwise train each as a lone batch."
+        )
+    split_dir = resolve_split_dir(cfg.data.shard_dir, cfg.data.sft_split)
+    import json
+
+    hashes = {}
+    for rp in sorted(split_dir.glob("*.receipt.json")):
+        with open(rp) as f:
+            r = json.load(f)
+        hashes[r["name"]] = r["content_hash"]
+    if not hashes:
+        raise Qwen3StageError(f"no shard receipts under {split_dir}")
+    banner = assemble_sft_banner(
+        cfg,
+        attention_backend=qcfg.attn_implementation or "auto(sdpa)",
+        optimizer_impl=f"{cfg.optimizer.name} lora={'on r=%d alpha=%g' % (cfg.lora.r, cfg.lora.alpha) if cfg.lora.enabled else 'off'}",
+        dataset_hashes=hashes,
+        ceiling_results={"lean_verifier": PENDING, "note": "qwen3 stage: evaluated by the Lean loop (P2), not MATH/GSM8K"},
+        probe_report=None,
+        rss_projection=projected_full_ft_rss_gb(cfg),
+        mde={"note": PENDING},
+        scheduler_horizon_optimizer_steps=horizon,
+        scheduler_horizon_rows_derived=None,
+    )
+    banner.update({"arch": cfg.model.arch, "snapshot": snapshot, "n_rows": len(dataset), "seq_len": cfg.train.seq_len})
+    # The Mistral gate ledger enforces the golden->smoke->probe order, which does not apply here
+    # (recording a pass for a gate that did not run would be a lie), so the entry is persisted as
+    # its own receipt next to the ledger instead.
+    from ..utils.dist import is_rank0
+
+    if is_rank0():
+        tmp = os.path.join(ledger.run_dir, "qwen3_entry.json.partial")
+        with open(tmp, "w") as f:
+            json.dump(banner, f, indent=1, sort_keys=True, default=str)
+        os.replace(tmp, os.path.join(ledger.run_dir, "qwen3_entry.json"))
+    return banner
+
+
+__all__ = [
+    "SUPPORTED_ARCHS", "Qwen3StageError", "base_fingerprint", "Qwen3Format", "LoRASFTModule", "validate_qwen3_cfg",
+    "resolve_qwen3", "build_qwen3_stage_model", "qwen3_trainer_extras", "qwen3_entry_banner",
+]

code/align/train/ref_cache.py

@@ -0,0 +1,834 @@
+"""Reference log-probabilities: precompute, cache, join, and the path-bias rule.
+
+DPO needs, for every pair, what the FROZEN reference policy thought of the
+chosen and rejected completions. The textbook implementation keeps a second
+7B model resident and forwards it alongside the policy. That does not fit:
+14.5 GB of policy weights plus optimizer shards plus activations already fills
+a 23 GB A10G, and a second copy fills a second one for no gradient.
+
+So the reference is computed once, offline, and consumed as constants. The
+whole arrangement rests on one property -- with dropout disabled, a no-grad
+pass over a fixed dataset yields the same numbers every time -- and on one
+discipline: the cache must be keyed on EVERYTHING that could change those
+numbers, because every failure mode of a stale cache is silent. A cache from
+the wrong checkpoint, the wrong tokenizer, or the pre-rebuild shards still
+produces a finite log-sigmoid of a difference of plausible numbers. The loss
+decreases. The margin it reports is a margin against something that does not
+exist.
+
+The key therefore carries six components, and each one is a way somebody has
+or could break this:
+
+  * `checkpoint_hash` -- the obvious one; a cache from the previous SFT run.
+    Must be a CONTENT hash (`checkpoint_content_hash`): the tensor-inventory
+    hash cannot tell two SFT checkpoints of one architecture apart.
+    (A seventh, DIAGNOSTIC field, `checkpoint_inventory_hash`, carries the
+    inventory hash beside it: it is compared like the rest, but its only job
+    is to make a `checkpoint_hash` mismatch legible -- inventory equal means
+    "same architecture, different weights", the stale-SFT-attempt case.)
+  * `tokenizer_hash`  -- retokenized pairs are different pairs.
+  * `manifest_hash`   -- the one the plan did not list and this project needs
+    most: section-03's shards are being REBUILT (the duplicated `# Answer`
+    marker fix). A rebuild changes every pair's tokens while leaving the
+    checkpoint, the tokenizer and max_length identical, so nothing else in the
+    key moves.
+  * `max_length` and `truncation_mode` -- a completion truncated differently
+    has a different summed logprob, and summed logprobs are what DPO compares.
+  * `forward_impl` -- a single-GPU forward and an FSDP2-sharded forward are
+    not bitwise identical. That is the PATH-BIAS risk, and it is measured
+    (`path_bias_decision`) rather than assumed away, because early in DPO the
+    policy IS the reference and the loss argument sits at zero, which is
+    exactly where a systematic reference bias does the most damage.
+
+EXISTING CACHES MUST BE RE-PRECOMPUTED (2026-10-12). Before this date the
+launchers filled `checkpoint_hash` from the tensor-inventory hash
+(`checkpoint_fingerprint()["manifest_sha256"]`). Every cache written that way
+carries a key that can no longer be reproduced from a checkpoint's weights, so
+it is refused by `load_ref_cache(expect=...)` -- correctly, because it could not
+have told its own checkpoint from any sibling's. The refusal names
+`checkpoint_hash`; the fix is to re-run scripts/precompute_ref_logprobs.py, not
+to loosen the comparison. Both launchers now build the key through
+`build_ref_cache_key`, the single place the six facts are assembled.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import math
+import os
+from dataclasses import asdict, dataclass, field, fields
+from typing import Any, Iterable, Optional, Sequence
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+NAN = float("nan")
+
+#: One entry per pair. Lengths are carried because `sigmoid_norm` divides by
+#: them; a length of 0 (a prompt-only row) must make the normalised quantity
+#: NaN rather than raise a ZeroDivisionError three weeks later.
+ENTRY_FIELDS = ("ref_logp_chosen", "ref_logp_rejected", "len_chosen", "len_rejected")
+
+#: How the pass actually ran. Not a preference -- a record.
+FORWARD_IMPLS = ("single_gpu", "fsdp2_sharded")
+
+
+@dataclass(frozen=True)
+class RefCacheKey:
+    """Everything that could silently change a cached reference logprob.
+
+    Frozen and compared component-by-component so a mismatch names the guilty
+    field. `RuntimeError: ref cache mismatch` with no further detail sends
+    somebody hunting through three sections; `mismatched: manifest_hash
+    (cache=man-abc, run=man-xyz)` sends them to the shard rebuild.
+    """
+
+    checkpoint_hash: str
+    tokenizer_hash: str
+    manifest_hash: str
+    max_length: int
+    truncation_mode: str
+    forward_impl: str
+    #: DIAGNOSTIC companion to `checkpoint_hash` (see the module docstring).
+    #: Empty only on caches written before 2026-10-12, which are refused on
+    #: `checkpoint_hash` anyway.
+    checkpoint_inventory_hash: str = ""
+
+    def __post_init__(self) -> None:
+        if self.forward_impl not in FORWARD_IMPLS:
+            raise ValueError(
+                f"forward_impl={self.forward_impl!r} must be one of "
+                f"{FORWARD_IMPLS}: it records how the pass ACTUALLY ran, and "
+                "an unrecognised tag cannot be compared against the run's."
+            )
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), sort_keys=True)
+
+    @classmethod
+    def from_json(cls, blob: str) -> "RefCacheKey":
+        return cls(**json.loads(blob))
+
+    def mismatches(self, other: "RefCacheKey") -> list[str]:
+        return [f.name for f in fields(self) if getattr(self, f.name) != getattr(other, f.name)]
+
+
+@dataclass
+class RefCache:
+    key: RefCacheKey
+    #: pair_id -> (ref_logp_chosen, ref_logp_rejected, len_chosen, len_rejected)
+    entries: dict[str, tuple[float, float, float, float]] = field(default_factory=dict)
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def require(self, pair_id: str) -> tuple[float, float, float, float]:
+        """Look up a pair, or raise. There is no default: a zero reference
+        turns DPO into likelihood maximisation on the chosen side, which
+        trains and converges and is not DPO."""
+        try:
+            return self.entries[pair_id]
+        except KeyError:
+            raise KeyError(
+                f"pair_id {pair_id!r} is not in the reference cache "
+                f"({len(self.entries)} entries, key {self.key.to_json()}). "
+                "Refusing to substitute a default reference logprob. Either "
+                "the precompute did not cover this split, or these shards are "
+                "not the ones it ran over."
+            ) from None
+
+
+def build_ref_cache_key(
+    sft_checkpoint: str, manifest: dict, *, forward_impl: str
+) -> RefCacheKey:
+    """The ONE place a launcher turns facts into a `RefCacheKey`.
+
+    Both scripts/precompute_ref_logprobs.py (writer) and scripts/dpo.py
+    (reader) call this, so the key they compare cannot drift apart again. The
+    drift it replaces: each script assembled the key inline and BOTH filled
+    `checkpoint_hash` from the tensor-inventory hash (2026-10-12 review), so
+    two different-weight SFT checkpoints produced the same key. Here
+    `checkpoint_hash` is the WEIGHT-BYTES hash (`checkpoint_content_hash`; one
+    streamed read of the checkpoint, ~14.5GB at 7B) and the inventory hash
+    rides along as `checkpoint_inventory_hash`.
+
+    `manifest_hash` covers the shard rebuild nothing else in the key notices;
+    `forward_impl` is a record of how the pass ran and is the caller's to state.
+    """
+    from ..data.receipts import content_hash
+
+    return RefCacheKey(
+        checkpoint_hash=checkpoint_content_hash(sft_checkpoint),
+        tokenizer_hash=manifest["tokenizer_hash"],
+        manifest_hash=content_hash(
+            [manifest["hashes"][k] for k in sorted(manifest.get("hashes", {}))]
+        ),
+        max_length=int(manifest.get("pinned", {}).get("max_len", 2048)),
+        truncation_mode="right",
+        forward_impl=forward_impl,
+        checkpoint_inventory_hash=checkpoint_inventory_hash(sft_checkpoint),
+    )
+
+
+def peek_ref_cache_key(path: str) -> Optional[RefCacheKey]:
+    """The key stored in a cache file's header, or None when there is no file.
+
+    INSPECTION ONLY -- it validates nothing about the cache. It exists so a
+    launcher can read the one component that is a RECORD rather than an
+    expectation (`forward_impl`: how the cache's own pass ran) from the cache
+    itself, build the expected key around it, and then `load_ref_cache(path,
+    expect=...)`, which asserts the other components. Never a substitute for
+    that load.
+    """
+    from safetensors import safe_open
+
+    if not os.path.isfile(path):
+        return None
+    with safe_open(path, framework="pt") as f:
+        meta = f.metadata() or {}
+    if "key" not in meta:
+        return None
+    return RefCacheKey.from_json(meta["key"])
+
+
+#: What `cache_forward_impl` reads when the report is written before any cache
+#: exists. Not a member of FORWARD_IMPLS on purpose: it can never equal a real
+#: cache's tag, so nothing can mistake "no cache yet" for a measured pass.
+NO_CACHE_YET = "no_cache_at_report_time"
+
+
+def path_bias_report_cache_forward_impl(cache_path: str) -> str:
+    """The `forward_impl` a path-bias report should record as `cache_forward_impl`.
+
+    The report's field is about the CACHE the training run will consume; the
+    `--path-bias` run measures it but does not recompute it. So the value is
+    the on-disk train cache's own tag, or `NO_CACHE_YET` when the measurement
+    was taken first. The launcher does not read this field back for its
+    expectation (it reads the cache), so this is a record, not a lever.
+    """
+    key = peek_ref_cache_key(cache_path)
+    return key.forward_impl if key is not None else NO_CACHE_YET
+
+
+def _checkpoint_files(path: str) -> list[str]:
+    """The weight files of a checkpoint, in the order they are hashed.
+
+    A `.safetensors` file, a directory holding `model.safetensors`, or a
+    sharded directory (index + shards, sorted by name).
+    """
+    from ..model.checkpoint import INDEX_NAME, WEIGHTS_NAME
+
+    if os.path.isfile(path):
+        files = [path]
+    else:
+        single = os.path.join(path, WEIGHTS_NAME)
+        index = os.path.join(path, INDEX_NAME)
+        if os.path.isfile(single):
+            files = [single]
+        elif os.path.isfile(index):
+            with open(index) as f:
+                shards = sorted(set(json.load(f)["weight_map"].values()))
+            files = [os.path.join(path, s) for s in shards]
+            missing = [s for s in files if not os.path.isfile(s)]
+            if missing:
+                raise FileNotFoundError(f"{INDEX_NAME} names shard(s) absent from {path!r}: {missing}")
+        else:
+            raise FileNotFoundError(f"no {WEIGHTS_NAME} or {INDEX_NAME} at {path!r}")
+    return files
+
+
+def checkpoint_inventory_hash(path: str) -> str:
+    """Hash of the tensor INVENTORY (name, shape, dtype) read from the safetensors
+    HEADERS only -- no tensor data is materialised.
+
+    This is the DIAGNOSTIC companion to `checkpoint_content_hash`, not a
+    substitute: it is exactly the kind of hash that cannot distinguish two
+    checkpoints of one architecture, which is its use -- when `checkpoint_hash`
+    mismatches and this one does not, the cache is from the same architecture
+    with different weights (the stale-SFT-attempt case). Deliberately NOT
+    `checkpoint_fingerprint()["manifest_sha256"]`: that routes through
+    `load_hf_safetensors`, which reads the whole 14.5GB into RAM on every rank
+    just to print shapes -- four ranks of it at gate 3 on a no-swap box.
+    """
+    from safetensors import safe_open
+
+    inventory = {}
+    for fp in _checkpoint_files(path):
+        with safe_open(fp, framework="pt") as f:
+            for k in f.keys():
+                sl = f.get_slice(k)
+                inventory[k] = [list(sl.get_shape()), str(sl.get_dtype())]
+    return hashlib.sha256(json.dumps(inventory, sort_keys=True).encode()).hexdigest()[:16]
+
+
+def checkpoint_content_hash(path: str) -> str:
+    """Hash of the checkpoint's WEIGHT BYTES: the value `RefCacheKey.checkpoint_hash` must carry.
+
+    WHY THIS EXISTS (review of 2026-10-12). The launchers used to fill
+    `checkpoint_hash` from `checkpoint_fingerprint(path)["manifest_sha256"]`,
+    which hashes only the tensor inventory (name, shape, dtype). Every full-FT
+    SFT checkpoint of one architecture has the SAME inventory, so a cache from
+    a previous SFT attempt, or from micro_00000100 when the gate passed
+    micro_00000200, produced an identical key and loaded without complaint --
+    the exact stale-cache failure this key is documented to defend against
+    ("a cache from the previous SFT run"). The receipt module says so itself
+    (`weights_sha256` is the content hash; `manifest_sha256` is a structure
+    hash), and `test_fingerprint_does_not_call_a_structure_hash_a_content_hash`
+    pins it. Reading ~14.5GB once per launch is the price of the key meaning
+    what it says; it is streamed in 64MB blocks and never held in memory.
+
+    Accepts a `.safetensors` file, a directory holding `model.safetensors`, or
+    a sharded directory (index + shards). Shard names are hashed together with
+    their bytes, in sorted order, so a renamed or reordered shard set changes
+    the hash too.
+    """
+    files = _checkpoint_files(path)
+    h = hashlib.sha256()
+    for fp in files:
+        h.update(os.path.basename(fp).encode() + b"\0")
+        with open(fp, "rb") as fh:
+            for block in iter(lambda: fh.read(64 << 20), b""):
+                h.update(block)
+    return h.hexdigest()[:16]
+
+
+def _atomic_write(path: str, write) -> None:
+    """Run `write(tmp_path)` then `os.replace` onto `path`.
+
+    A cache or report written straight to its final name and killed mid-write
+    (the host has seized mid-checkpoint before) leaves a truncated file at the
+    name the next launch trusts, and REPLACES a previously good one. The temp
+    file sits beside the target (same filesystem, on /apps) so the rename is
+    atomic, and is removed if the write raises.
+    """
+    tmp = f"{path}.tmp.{os.getpid()}"
+    try:
+        write(tmp)
+        os.replace(tmp, path)
+    except BaseException:
+        if os.path.exists(tmp):
+            os.remove(tmp)
+        raise
+
+
+# ---------------------------------------------------------------------------
+# I/O
+# ---------------------------------------------------------------------------
+
+
+def write_ref_cache(path: str, cache: RefCache) -> str:
+    """One safetensors file per (arm, split): tensors for the values, the key
+    and the pair-id order in the metadata.
+
+    One file, not a tensor file plus a JSON sidecar, because a cache split
+    across two files can half-exist -- and a half-existing cache is the
+    resumable-looking partial checkpoint problem in another costume.
+    """
+    from safetensors.torch import save_file
+
+    from ..utils.guard import assert_paths_off_root
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(path))])
+    pair_ids = list(cache.entries)
+    cols = list(zip(*(cache.entries[p] for p in pair_ids))) if pair_ids else [(), (), (), ()]
+    # float64 throughout: the values arrive as fp32 (the log-softmax is upcast
+    # even in a bf16 forward) and float32 -> float64 is exact, so the stored
+    # number IS the computed number. Storing them as fp32 would be equally
+    # exact today and would silently truncate if the compute dtype ever rose.
+    tensors = {
+        name: torch.tensor(col, dtype=torch.float64) for name, col in zip(ENTRY_FIELDS, cols)
+    }
+    from ..data.loader import PREFERENCE_MASK_VERSION
+
+    _atomic_write(
+        path,
+        lambda tmp: save_file(
+            tensors,
+            tmp,
+            metadata={
+                "key": cache.key.to_json(),
+                "pair_ids": json.dumps(pair_ids),
+                "mask_version": PREFERENCE_MASK_VERSION,
+            },
+        ),
+    )
+    logger.info("ref cache written: %s (%d pairs)", path, len(pair_ids))
+    return path
+
+
+def load_ref_cache(path: str, expect: Optional[RefCacheKey] = None) -> RefCache:
+    """Load a cache and, when `expect` is given, ASSERT the key.
+
+    Loading without `expect` is for inspection tools only. Every training path
+    passes it; that assertion is the entire defence against a stale cache.
+    """
+    from safetensors import safe_open
+
+    if not os.path.isfile(path):
+        raise FileNotFoundError(
+            f"reference cache {path!r} does not exist. Run "
+            "scripts/precompute_ref_logprobs.py for this arm first; DPO "
+            "cannot proceed without it and will not invent one."
+        )
+    with safe_open(path, framework="pt") as f:
+        meta = f.metadata() or {}
+        cols = {name: f.get_tensor(name) for name in ENTRY_FIELDS}
+    if "key" not in meta or "pair_ids" not in meta:
+        raise RuntimeError(
+            f"{path!r} carries no cache key. An unkeyed cache cannot be shown "
+            "to belong to this checkpoint, tokenizer and shard build, so it is "
+            "refused rather than trusted."
+        )
+    from ..data.loader import PREFERENCE_MASK_VERSION
+
+    cache_mask = meta.get("mask_version", "prompt_len_v1")
+    if cache_mask != PREFERENCE_MASK_VERSION:
+        raise RuntimeError(
+            f"{path!r} was precomputed under completion-mask semantics "
+            f"{cache_mask!r} but this build uses {PREFERENCE_MASK_VERSION!r} "
+            "(audit R1, 2026-08-23: the boundary moved from prompt_len to "
+            "shared_prefix_len). Reference logprobs are mask-dependent; "
+            "recompute the cache rather than mixing semantics."
+        )
+    key = RefCacheKey.from_json(meta["key"])
+    if expect is not None:
+        bad = key.mismatches(expect)
+        if bad:
+            detail = "; ".join(
+                f"{n} (cache={getattr(key, n)!r}, run={getattr(expect, n)!r})" for n in bad
+            )
+            legacy = ""
+            if "checkpoint_hash" in bad and not key.checkpoint_inventory_hash:
+                legacy = (
+                    " This cache carries no checkpoint_inventory_hash: it predates "
+                    "the 2026-10-12 content-hash key, whose checkpoint_hash was a "
+                    "tensor-inventory hash that cannot identify a checkpoint. It "
+                    "must be re-precomputed."
+                )
+            raise RuntimeError(
+                f"reference cache {path!r} does not match this run. "
+                f"Mismatched: {detail}.{legacy} A stale reference cache does not fail "
+                "loudly at train time -- the loss stays finite, the curve "
+                "stays healthy, and the margin is measured against a policy "
+                "that never existed. Re-run the precompute."
+            )
+    pair_ids = json.loads(meta["pair_ids"])
+    entries = {p: tuple(float(cols[n][i]) for n in ENTRY_FIELDS) for i, p in enumerate(pair_ids)}
+    return RefCache(key=key, entries=entries)
+
+
+# ---------------------------------------------------------------------------
+# Precompute
+# ---------------------------------------------------------------------------
+
+
+def _labels(
+    input_ids: torch.Tensor, completion_mask: torch.Tensor, ignore_index: int
+) -> torch.Tensor:
+    """Completion-only labels: prompt AND pad positions ignored.
+
+    Excluding the prompt is not cosmetic. For Arm B the entire shared prefix
+    is identical across the two sides, so including it adds the same term to
+    both logprobs: it cancels in the DPO difference while inflating both
+    magnitudes, and any length-normalised variant computed from it is wrong.
+    """
+    return input_ids.masked_fill(completion_mask == 0, ignore_index)
+
+
+@torch.no_grad()
+def _side_logprobs(
+    model, input_ids: torch.Tensor, completion_mask: torch.Tensor, ignore_index: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Summed completion logprobs via the SAME routine `fused_dpo` uses.
+
+    Not a reimplementation: `completion_logprobs_eager` is imported from the
+    kernel module, and the caller-side shift here (`logits[:, :-1]` against
+    `labels[:, 1:]`) is the shift the DPO adapter applies at train time. If
+    precompute and consumption disagreed by one position, DPO would start from
+    a nonzero margin against itself and spend the run optimising a constant
+    offset it can never remove.
+    """
+    from ..kernels.fused_dpo import completion_logprobs_eager
+
+    labels = _labels(input_ids, completion_mask, ignore_index)
+    logits = model(input_ids=input_ids)
+    return completion_logprobs_eager(logits[:, :-1], labels[:, 1:], ignore_index)
+
+
+def precompute_ref_logprobs(
+    model,
+    rows: Sequence[Any],
+    *,
+    key: RefCacheKey,
+    pad_token_id: int,
+    batch_size: int = 8,
+    device: Optional[torch.device] = None,
+) -> RefCache:
+    """One no-grad pass over `rows`, returning a keyed cache.
+
+    `rows` are in the shape `collate_preference` consumes (section-03's
+    `preference_row_from_store` produces them); every row must carry a
+    `pair_id`, because pair_id is how refs find their way back to the right
+    training example after any reshuffle.
+
+    NOTE on batch size -- why it is NOT a key component, MEASURED not assumed.
+    Right-padding plus causal masking makes padded columns unreachable, so the
+    value is mathematically independent of how rows are grouped; the fp32
+    reduction ORDER is not. Measured on the tiny CPU harness (40 pairs, batch
+    size 1 vs 2/4/8/20/40): max |delta logp| 3.8e-06 = 9.7e-08 relative, max
+    |delta reference MARGIN| 5.7e-06, mean delta ~1e-07. Symmetric rounding
+    noise, not a systematic offset -- which is the distinction that decides
+    it, since only a systematic offset biases the loss argument. Run through
+    this module's own `path_bias_decision` the ratio is ~1e-05 against a 0.05
+    threshold, so the same rule that governs single-GPU-vs-sharded says
+    negligible here. Bitwise determinism holds at a FIXED batch size, which is
+    the property the caching design actually needs, and
+    `test_batch_size_variation_is_below_the_path_bias_threshold` keeps this
+    paragraph honest.
+    """
+    from ..data.loader import collate_preference
+    from ..kernels.fused_dpo import IGNORE_INDEX
+
+    if getattr(model, "training", False):
+        raise RuntimeError(
+            "the reference model is in train mode. Mistral-v0.3 sets "
+            "attention_dropout 0.0 so this is PROBABLY harmless -- which is "
+            "exactly why nobody would notice, and why 'probably' is not a "
+            "basis for numbers that get cached and consumed as constants for "
+            "the rest of the stage. Call model.eval()."
+        )
+    if batch_size < 1:
+        raise ValueError(f"batch_size={batch_size} must be >= 1")
+
+    rows = list(rows)
+    entries: dict[str, tuple[float, float, float, float]] = {}
+    for start in range(0, len(rows), batch_size):
+        chunk = rows[start : start + batch_size]
+        batch = collate_preference(chunk, pad_token_id=pad_token_id)
+        if device is not None:
+            batch = batch.to(device)
+        lp_c, len_c = _side_logprobs(
+            model, batch.chosen_input_ids, batch.chosen_completion_mask, IGNORE_INDEX
+        )
+        lp_r, len_r = _side_logprobs(
+            model, batch.rejected_input_ids, batch.rejected_completion_mask, IGNORE_INDEX
+        )
+        for i, meta in enumerate(batch.metadata):
+            pid = meta.get("pair_id")
+            if pid is None:
+                raise ValueError(
+                    "a preference row carries no pair_id. Refs are joined to "
+                    "training rows BY pair_id; a row without one could only be "
+                    "joined by position, which mis-pairs under any reshuffle."
+                )
+            if pid in entries:
+                raise ValueError(
+                    f"duplicate pair_id {pid!r}: the cache is a mapping, so a "
+                    "duplicate would silently overwrite one pair's reference "
+                    "with another's."
+                )
+            entries[pid] = (float(lp_c[i]), float(lp_r[i]), float(len_c[i]), float(len_r[i]))
+    logger.info("precomputed %d reference pairs (%s)", len(entries), key.forward_impl)
+    return RefCache(key=key, entries=entries)
+
+
+# ---------------------------------------------------------------------------
+# Joining refs to batches
+# ---------------------------------------------------------------------------
+
+
+def attach_refs(batch, cache: RefCache):
+    """Fill a `Preference` batch's ref tensors FROM THE CACHE, by pair_id.
+
+    By pair_id and never by position. The section-03 loader is length-grouped
+    and shuffled per epoch, so batch order is not cache order and never will
+    be; a positional join would attach a perfectly plausible reference to the
+    wrong row, produce a finite loss, and point the gradient somewhere
+    arbitrary. Nothing downstream can see it.
+    """
+    pair_ids = []
+    for i, meta in enumerate(batch.metadata):
+        pid = (meta or {}).get("pair_id")
+        if pid is None:
+            raise ValueError(
+                f"row {i} of this batch carries no pair_id, so its reference "
+                "logprobs cannot be looked up. Falling back to batch order is "
+                "the silent-substitution incident this join exists to prevent."
+            )
+        pair_ids.append(pid)
+
+    vals = [cache.require(p) for p in pair_ids]
+    dtype = torch.float32
+    batch.ref_chosen_logps = torch.tensor([v[0] for v in vals], dtype=dtype)
+    batch.ref_rejected_logps = torch.tensor([v[1] for v in vals], dtype=dtype)
+    return batch
+
+
+# ---------------------------------------------------------------------------
+# Held-out preference accuracy
+# ---------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def heldout_preference_accuracy(
+    model, batches: Iterable[Any], cache: RefCache, beta: float, z: float = 1.96
+) -> dict:
+    """Fraction of held-out pairs the policy's IMPLICIT reward prefers.
+
+    The implicit reward difference is beta * [(logp_c - logp_r) - (ref_c -
+    ref_r)]; the pair counts as correct when it is strictly positive. This is
+    the stage's primary health metric together with the margin, and it is
+    deliberately NOT benchmark accuracy: the pre-registered expectations for
+    both arms are about preference fit, and section-13 owns the paired
+    benchmark evaluation.
+
+    Empty or degenerate held-out sets report NaN, never 0. Note the
+    distinction the test pins: at initialisation the policy IS the reference,
+    every implicit margin is exactly 0.0, and nothing is STRICTLY preferred --
+    accuracy 0.0 over a real denominator. That is an honest measurement, and
+    it is a different fact from NaN over a denominator of zero.
+    """
+    from ..eval.metrics import wilson_interval
+    from ..kernels.fused_dpo import IGNORE_INDEX
+
+    was_training = getattr(model, "training", False)
+    model.eval()
+    n = 0
+    correct = 0
+    margin_sum = 0.0
+    try:
+        for batch in batches:
+            if int(batch.chosen_input_ids.shape[0]) == 0:
+                continue
+            lp_c, _ = _side_logprobs(
+                model, batch.chosen_input_ids, batch.chosen_completion_mask, IGNORE_INDEX
+            )
+            lp_r, _ = _side_logprobs(
+                model, batch.rejected_input_ids, batch.rejected_completion_mask, IGNORE_INDEX
+            )
+            for i, meta in enumerate(batch.metadata):
+                pid = (meta or {}).get("pair_id")
+                if pid is None:
+                    raise ValueError("held-out row carries no pair_id")
+                ref_c, ref_r, _, _ = cache.require(pid)
+                m = beta * ((float(lp_c[i]) - float(lp_r[i])) - (ref_c - ref_r))
+                margin_sum += m
+                correct += int(m > 0.0)
+                n += 1
+    finally:
+        if was_training:
+            model.train()
+
+    if n == 0:
+        return {
+            "accuracy": NAN,
+            "ci_low": NAN,
+            "ci_high": NAN,
+            "mean_margin": NAN,
+            "n": 0,
+            "correct": 0,
+            "beta": beta,
+        }
+    lo, hi = wilson_interval(correct, n, z)
+    return {
+        "accuracy": correct / n,
+        "ci_low": lo,
+        "ci_high": hi,
+        "mean_margin": margin_sum / n,
+        "n": n,
+        "correct": correct,
+        "beta": beta,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Path bias
+# ---------------------------------------------------------------------------
+
+
+def path_bias_decision(
+    deltas: Sequence[float], *, beta: float, early_margin_scale: float, threshold: float
+) -> dict:
+    """Decide whether a single-GPU reference cache is usable under FSDP2.
+
+    `deltas` are per-pair differences of the REFERENCE MARGIN (ref_c - ref_r)
+    between the single-GPU forward and the sharded forward. The two are not
+    bitwise identical -- different reduction trees, different collective
+    orders -- and the question is whether the difference matters.
+
+    It matters more here than anywhere else in the project. DPO's loss
+    argument is beta * (pi_margin - ref_margin). Early in the run the policy
+    IS the reference, so that argument sits at zero, which is where the
+    log-sigmoid's gradient is largest and where a systematic offset in
+    ref_margin is least attenuated. A bias that would be invisible against a
+    large margin is fully visible against a zero one.
+
+    The rule compares max|delta| to the early-margin scale. Beta multiplies
+    the perturbation and the reference scale alike, so it CANCELS out of the
+    ratio; it stays in the signature because the receipt must record the beta
+    the decision was made under, and because a reader who does not see the
+    cancellation stated will assume it was forgotten.
+
+    NaN is not small. A NaN delta means a pair could not be compared, and an
+    uncomparable pair forces the firing path rather than comparing False and
+    vanishing. So does an empty measurement: not measuring is not passing.
+    """
+    if not (early_margin_scale > 0) or math.isnan(early_margin_scale):
+        raise ValueError(
+            f"early_margin_scale={early_margin_scale!r} must be a positive "
+            "measured scale; it is the denominator of the decision ratio and "
+            "a zero or NaN one makes the decision undefined rather than lenient"
+        )
+    if not (threshold > 0):
+        raise ValueError(f"threshold={threshold!r} must be positive")
+
+    vals = [float(d) for d in deltas]
+    n_nan = sum(1 for d in vals if math.isnan(d))
+    if not vals or n_nan:
+        detail = (
+            "no deltas were measured; not measuring is not passing"
+            if not vals
+            else f"{n_nan} of {len(vals)} deltas are NaN and a NaN delta is not a small delta"
+        )
+        return {
+            "fired": True,
+            "action": "rerun_precompute_sharded",
+            "ratio": NAN,
+            "max_abs_delta": NAN,
+            "mean_delta": NAN,
+            "median_abs_delta": NAN,
+            "p95_abs_delta": NAN,
+            "n": len(vals),
+            "n_nan": n_nan,
+            "beta": beta,
+            "early_margin_scale": early_margin_scale,
+            "threshold": threshold,
+            "detail": detail,
+        }
+
+    absd = sorted(abs(d) for d in vals)
+    max_abs = absd[-1]
+    ratio = max_abs / early_margin_scale
+    fired = ratio > threshold
+    return {
+        "fired": fired,
+        "action": ("rerun_precompute_sharded" if fired else "keep_single_gpu_cache"),
+        "ratio": ratio,
+        "max_abs_delta": max_abs,
+        "mean_delta": sum(vals) / len(vals),
+        "median_abs_delta": _quantile(absd, 0.5),
+        "p95_abs_delta": _quantile(absd, 0.95),
+        "n": len(vals),
+        "n_nan": 0,
+        "beta": beta,
+        "early_margin_scale": early_margin_scale,
+        "threshold": threshold,
+        "detail": (
+            f"max|delta| {max_abs:.3e} is "
+            f"{'above' if fired else 'within'} {threshold:g} x the "
+            f"early-margin scale {early_margin_scale:g}"
+        ),
+    }
+
+
+def _quantile(sorted_vals: Sequence[float], q: float) -> float:
+    if not sorted_vals:
+        return NAN
+    idx = min(len(sorted_vals) - 1, max(0, int(round(q * (len(sorted_vals) - 1)))))
+    return float(sorted_vals[idx])
+
+
+def batch_invariance_report(
+    model,
+    rows: Sequence[Any],
+    *,
+    key: RefCacheKey,
+    pad_token_id: int,
+    beta: float,
+    early_margin_scale: float,
+    threshold: float,
+    batch_sizes: Sequence[int] = (2, 4, 8, 16),
+) -> dict:
+    """Batched-vs-single-row precompute deltas, as a receipt.
+
+    A SECOND path-bias question, distinct from the single-GPU-vs-sharded one
+    and cheaper to answer: does it matter how many pairs share a forward?
+    Right-padding plus causal masking makes the value grouping-independent in
+    exact arithmetic; the fp32 reduction order is not, so the honest answer is
+    a measurement.
+
+    It decides one thing: whether `batch_size` belongs in the cache key. It is
+    judged by the SAME rule as the sharded comparison (`path_bias_decision`),
+    because it is the same kind of error entering the same place -- a
+    perturbation of the reference margin, read at the point where the loss
+    argument sits near zero.
+
+    Reported per batch size, plus the worst case across all of them. A
+    systematic MEAN delta matters more than a large max: symmetric rounding
+    averages out of the gradient, a consistent offset does not.
+    """
+    single = precompute_ref_logprobs(model, rows, key=key, pad_token_id=pad_token_id, batch_size=1)
+    per_size = {}
+    worst = None
+    for bs in batch_sizes:
+        got = precompute_ref_logprobs(
+            model, rows, key=key, pad_token_id=pad_token_id, batch_size=bs
+        )
+        deltas = [
+            (got.entries[p][0] - got.entries[p][1]) - (single.entries[p][0] - single.entries[p][1])
+            for p in single.entries
+        ]
+        d = path_bias_decision(
+            deltas, beta=beta, early_margin_scale=early_margin_scale, threshold=threshold
+        )
+        per_size[str(bs)] = d
+        if worst is None or (not math.isnan(d["ratio"]) and d["ratio"] > worst["ratio"]):
+            worst = d
+    return {
+        "comparison": "batched_vs_single_row",
+        "n_examples": len(list(rows)),
+        "batch_sizes": list(batch_sizes),
+        "per_batch_size": per_size,
+        "worst": worst,
+        "belongs_in_cache_key": bool(worst and worst["fired"]),
+        "detail": (
+            "if belongs_in_cache_key is true, batch_size must join the "
+            "six key components; measured false means regrouping is "
+            "below the threshold that governs the sharded comparison"
+        ),
+    }
+
+
+def write_path_bias_report(path: str, decision: dict, *, n_examples: int, forward_impl: str) -> str:
+    """The decision is written, not remembered: precompute and training are
+    separate heavy jobs (serialized by the flock, never concurrent) in
+    different processes, and the banner reads this file."""
+    from ..utils.guard import assert_paths_off_root
+
+    assert_paths_off_root([os.path.dirname(os.path.abspath(path))])
+    payload = dict(decision)
+    payload["n_examples"] = int(n_examples)
+    payload["cache_forward_impl"] = forward_impl
+    def _dump(tmp: str) -> None:
+        with open(tmp, "w") as f:
+            json.dump(payload, f, indent=2)
+
+    _atomic_write(path, _dump)
+    logger.info("path-bias report -> %s (fired=%s)", path, decision["fired"])
+    return path
+
+
+def read_path_bias_report(path: str) -> dict:
+    if not os.path.isfile(path):
+        raise FileNotFoundError(
+            f"path-bias report {path!r} is missing. The measurement is "
+            "MANDATORY before a DPO launch: a sharded forward that disagrees "
+            "with the cached single-GPU reference biases the loss argument "
+            "exactly where it is most sensitive. Run "
+            "scripts/precompute_ref_logprobs.py --path-bias."
+        )
+    with open(path) as f:
+        return json.load(f)

code/align/train/sft_loss.py

@@ -0,0 +1,150 @@
+"""SFT loss adapter: collated batch -> section-05's chunked fused CE.
+
+LOSS NORMALISATION (G3 audit U03-2). DEFAULT UNCHANGED; the global token mean is OPT-IN via
+`train.loss_normalization = "global_token_mean"` (Trainer window look-ahead +
+`TrainModule.train_batch(window_tokens=...)`; this adapter supplies `count_tokens`).
+The paragraph below describes the default, "micro_mean".
+The returned loss is the TOKEN MEAN OF ONE MICRO-BATCH, and `n_tokens` is only
+reported. Under grad-accum A and world W the window gradient is therefore the
+mean of A*W micro-batch means -- not the gradient of the global per-token mean
+(sum CE / sum n_tokens): a target token carries weight 1/(A*W*n_micro), so with
+token-budget batches whose target counts differ (280 vs 6) the short batch's
+tokens weigh ~47x more. The claim "matches one large batch" holds only for
+equal-n micro-batches. Making it a global token mean needs the window's token
+total (a look-ahead over the window's batches plus a cross-rank reduce) and the
+adapter returning CE_sum; that changes the objective every recorded SFT run
+was trained under, so it is a protocol decision, not a bug fix made here.
+
+The adapter owns plumbing only -- label shifting, masking, token accounting --
+and section-05 owns the kernel. Keeping the shift on THIS side (caller-shifted
+views) is the arbitrated contract: `fused_linear_cross_entropy` takes hidden
+states and labels that already line up, so there is exactly one place in the
+repo that decides what predicts what.
+
+Why the fused path matters, so the reason survives the handoff: at seq 1024 and
+micro-batch 8 the `[B*T, 32768]` logit tensor is 1.0GB in bf16 and its gradient
+is another 1.0GB. Materialising both spends 2GB of a 23GB card on a tensor
+nobody reads. The fused loss consumes HIDDEN STATES and applies the LM head per
+chunk, which is why `NmoMistralForCausalLM.forward_hidden_states` exists as a
+separate entry point.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+
+from .trainer import LossOutput
+
+logger = logging.getLogger(__name__)
+
+#: Section-05's kernel module. Named here so the both-directions test can block
+#: the import and exercise the refusal branch.
+FUSED_CE_MODULE = "align.kernels.fused_ce"
+
+#: Fallback only. The authority is section-05's constant (imported below when
+#: available) and section-03's collator, which both use -100; this literal
+#: exists so the module still loads if the kernel is missing.
+IGNORE_INDEX = -100
+
+
+def _load_fused():
+    """Return (fused_fn, ignore_index) or (None, IGNORE_INDEX)."""
+    try:
+        import importlib
+
+        mod = importlib.import_module(FUSED_CE_MODULE)
+        return mod.fused_linear_cross_entropy, mod.IGNORE_INDEX
+    except Exception:  # noqa: BLE001
+        return None, IGNORE_INDEX
+
+
+def fused_ce_available() -> bool:
+    return _load_fused()[0] is not None
+
+
+def make_sft_loss_adapter(chunk_size: Optional[int] = None, require_fused: bool = False):
+    """Build the SFT loss adapter the Trainer calls once per micro-step.
+
+    `require_fused=True` refuses to run on the reference path -- which is what
+    the 7B stage scripts pass, so a GPU run cannot quietly spend 2GB of card on
+    a materialised logit tensor because section-05's import broke.
+    """
+    fused_fn, ignore_index = _load_fused()
+    if require_fused and fused_fn is None:
+        raise RuntimeError(
+            f"the fused chunked cross-entropy ({FUSED_CE_MODULE}."
+            "fused_linear_cross_entropy) is not importable. Refusing to fall "
+            "back to the reference path for a full-scale run: it materialises "
+            "the [B*T, 32768] logit tensor AND its gradient, ~2GB of a 23GB "
+            "card at seq 1024."
+        )
+    if fused_fn is None:
+        logger.warning(
+            "%s unavailable; using the reference path. Correct, and "
+            "memory-hungry -- small-model use only.",
+            FUSED_CE_MODULE,
+        )
+
+    kwargs = {} if chunk_size is None else {"chunk_size": chunk_size}
+
+    def adapter(model, batch) -> LossOutput:
+        input_ids, labels = _unpack(batch)
+        # Caller-shifted: position t predicts token t+1. Done once, here.
+        target = labels[:, 1:]
+        n_tokens = int((target != ignore_index).sum())
+        if n_tokens == 0:
+            raise ValueError(
+                "SFT batch has no unmasked target tokens: every position is "
+                f"ignore_index ({ignore_index}). The mean cross-entropy over "
+                "zero positions is undefined, and returning 0.0 would report a "
+                "perfectly fit batch. Check prompt_len against sequence length "
+                "in the collator."
+            )
+
+        if fused_fn is not None:
+            # The loss runs INSIDE the model's forward via `hidden_consumer`,
+            # so section-05 still owns the lm_head matmul while FSDP2 has both
+            # the hidden states and lm_head.weight gathered. Computing it
+            # after the forward returned would see `lm_head.weight` re-sharded
+            # -- the mixed-DTensor failure that killed the first 4-rank smoke.
+            def _consume(hidden, lm_head_weight):
+                return fused_fn(
+                    hidden[:, :-1], lm_head_weight, target, ignore_index=ignore_index, **kwargs
+                )
+
+            loss = model(input_ids=input_ids, hidden_consumer=_consume)
+        else:
+            logits = model(input_ids=input_ids)
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.shape[-1]).float(),
+                target.reshape(-1),
+                ignore_index=ignore_index,
+            )
+
+        return LossOutput(loss=loss, n_tokens=n_tokens, metrics={"tokens_in_loss": float(n_tokens)})
+
+    def count_tokens(batch) -> int:
+        """Unmasked target tokens of `batch` (what `n_tokens` will be), without a forward pass.
+        The Trainer sums this over a look-ahead window for the global token mean."""
+        _, labels = _unpack(batch)
+        return int((labels[:, 1:] != ignore_index).sum())
+
+    adapter.uses_fused = fused_fn is not None
+    adapter.count_tokens = count_tokens
+    return adapter
+
+
+def _unpack(batch: Any) -> tuple[torch.Tensor, torch.Tensor]:
+    """Accept either a `Collated` from section-03 or a plain (ids, labels)."""
+    if hasattr(batch, "input_ids") and hasattr(batch, "labels"):
+        return batch.input_ids, batch.labels
+    if isinstance(batch, (tuple, list)) and len(batch) == 2:
+        return batch[0], batch[1]
+    raise TypeError(
+        f"SFT loss adapter got {type(batch).__name__}; expected a Collated "
+        "batch (input_ids, labels) from align.data.loader.collate_sft"
+    )

code/align/train/sharded_ckpt.py

@@ -0,0 +1,1285 @@
+"""Sharded streaming checkpoints: every rank writes ONLY the weights it owns.
+
+OPT-IN (`train.checkpoint_format = "sharded"`; the default `"gathered"` is
+untouched). Runtime deps: torch + stdlib. The converter half (`export_hf`,
+`reshard`) never touches a GPU and never imports transformers/safetensors.
+
+THE PROBLEM, MEASURED (docs/MEASUREMENTS.md, Track M, OLMoE-1B-7B, 4 ranks)
+    * save   228-241 s, rank-0 RSS peak +28 GB. `get_model_state_dict(
+      full_state_dict=True, cpu_offload=True)` is a full all-gather into rank 0's
+      host RAM (1x weights), and the HF converter then builds a second dict of
+      cloned expert slices plus `.contiguous()` copies (the other 1x): 2x the
+      weights on ONE rank of a no-swap, shared box -- the 2026-08-12 OOM shape.
+      Ranks 1..N-1 sit in a barrier for the whole four minutes.
+    * resume peaks at 90 GiB SUMMED RSS across ranks: `resume_from` reads the
+      whole `model.safetensors` on EVERY rank and then slices out the 1/N it
+      owns.
+OLMo-core avoids both with distributed checkpointing (per-rank files, no full
+gather). This is the same idea at the size this repo needs.
+
+LAYOUT (a checkpoint directory in sharded mode)
+    model_rank00000.safetensors ...   one per rank; that rank's LOCAL shards,
+                                      bf16, standard safetensors (any reader
+                                      opens them), written tensor by tensor
+    optimizer_rank{r}.pt              unchanged: the existing per-rank optimizer
+                                      files (sharded mode implies them)
+    shard_manifest.json               global shapes, placements, per-rank row
+                                      ranges, world size, per-file bytes and
+                                      sha256, the HF export recipe
+    resume.pt                         unchanged loop state (no weights)
+    tokenizer files                   as in the gathered layout
+The manifest is written last, inside the `.partial` staging directory, and the
+directory is published by the existing fsync-then-rename protocol. A directory
+without a complete manifest and every rank file the manifest names is not a
+checkpoint (`sharded_problem`, called from `trainer.checkpoint_problem`).
+
+WHAT EACH PIECE PREVENTS (incident-first, repo style)
+
+THE WRONG-SHARD DEFECT. A rank's file holds rows [off, off+n) of each tensor.
+Loaded under a different world size, rank r would receive rows cut for another
+partitioning and training would continue plausibly on scrambled weights --
+nothing crashes, the loss is merely worse. Resume therefore reads the manifest
+FIRST and REFUSES a different world size with a message naming the offline
+reshard (`scripts/ckpt_convert.py reshard`); it never "tries" a mismatch. The
+per-row ranges recorded in the manifest are also compared against the LIVE
+model's own sharding tensor by tensor, so a model rebuilt with a different mesh
+or an uneven-sharding change is caught even at the same world size.
+
+THE SILENT-CORRUPTION DEFECT. Per-file sha256 is recorded at save and checked
+by each rank, on its OWN file, before any weight is copied. The verdict is then
+reduced across ranks (`any_rank`): a bit-rotted file on rank 2 alone would
+otherwise make rank 2 raise while ranks 0/1/3 entered the next collective and
+hung until the watchdog -- the lone-rank shape every other decision in the
+trainer avoids. The save side uses the same reduction: a rank whose disk fills
+mid-write must stop its peers at once, not leave them in a gather for the
+watchdog's 30 minutes.
+
+THE TWO-COPIES-ON-RANK-0 DEFECT (the measured 28 GB). There is no gather. Each
+rank's peak is ONE tensor (`emit` copies one local shard to host, writes it,
+drops it). The export side is byte-copy `pread` -> `write` in 8 MiB pieces: for
+1:1 families it never builds a tensor at all.
+
+THE PARTIAL-WRITE DEFECT (same as the gathered format). Rank files are written
+to `*.tmp`, fsynced, renamed into the `.partial` staging dir; the staging dir
+is fsync_tree'd and renamed to `micro_N` only after rank 0 has the manifest of
+every rank. A kill anywhere leaves a `.partial` that `_discard_partials`
+removes at the next start, never a `micro_N` that looks resumable.
+
+WHAT THIS DOES NOT DO (read before relying on it)
+    * It is not loadable by `eval_checkpoint.py` / the rollout weight sync /
+      `ref_cache` until exported: those read `model.safetensors`. Export is
+      `ckpt_convert.py export` (streaming, bounded memory).
+    * Placements: `Shard(0)` and `Replicate()` on a 1-D mesh (what
+      `fully_shard` produces). Anything else (HSDP, TP, Shard(d>0)) is refused
+      at SAVE time, not discovered at resume.
+    * A train module that changes what "the weights" mean (LoRA folds an
+      adapter into merged weights; PPO has its own gather) must provide
+      `sharded_model_state()` / `load_sharded_model_state()` or the Trainer
+      refuses `"sharded"` at construction. Silent fallback to the base-class
+      behaviour would checkpoint the UNMERGED base under a merged name.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import os
+import re
+import shutil
+import struct
+from dataclasses import dataclass
+from typing import Any, Callable, Iterable, Mapping, Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+SHARD_MANIFEST_NAME = "shard_manifest.json"
+FORMAT_NAME = "align-sharded-ckpt"
+FORMAT_VERSION = 1
+CHECKPOINT_FORMATS = ("gathered", "sharded")
+
+#: Bytes moved per pread/write while exporting or resharding. Peak RSS of the
+#: converter is a small multiple of this plus one optimizer tensor, whatever the
+#: model size.
+COPY_CHUNK = 8 * 1024 * 1024
+
+#: Files in a sharded checkpoint directory that are NOT tokenizer/config
+#: payload (the exporter copies everything else next to the weights).
+_OWN_FILE_RE = re.compile(
+    r"^(model_rank\d+\.safetensors(\.tmp)?|optimizer_rank\d+\.pt(\.tmp)?|resume\.pt|"
+    + r"integrity\.json|"
+    + re.escape(SHARD_MANIFEST_NAME)
+    + r")$"
+)
+
+
+class ShardedCheckpointError(RuntimeError):
+    """A sharded checkpoint that is malformed, mismatched, or corrupt."""
+
+
+class CheckpointIntegrityError(ShardedCheckpointError):
+    """A payload file's sha256 differs from the one recorded at save time."""
+
+
+class ShardedWorldSizeError(ShardedCheckpointError):
+    """Resume under a world size other than the one that saved the shards."""
+
+
+# ---------------------------------------------------------------------------
+# dtypes
+# ---------------------------------------------------------------------------
+
+_ST_TO_TORCH = {
+    "BOOL": torch.bool,
+    "U8": torch.uint8,
+    "I8": torch.int8,
+    "I16": torch.int16,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+    "I32": torch.int32,
+    "F32": torch.float32,
+    "F64": torch.float64,
+    "I64": torch.int64,
+}
+_TORCH_TO_ST = {v: k for k, v in _ST_TO_TORCH.items()}
+#: safetensors serialises tensors sorted by (dtype DESC, name); this is the
+#: dtype order of its enum. Byte-identical output to `save_file` is a tested
+#: property (it is what lets "export equals gathered export" be a file hash).
+_ST_ORDER = ["BOOL", "U8", "I8", "I16", "F16", "BF16", "I32", "F32", "F64", "I64"]
+_ST_ITEMSIZE = {"BOOL": 1, "U8": 1, "I8": 1, "I16": 2, "F16": 2, "BF16": 2, "I32": 4, "F32": 4, "F64": 8, "I64": 8}
+
+
+def _prod(shape: Iterable[int]) -> int:
+    n = 1
+    for s in shape:
+        n *= int(s)
+    return n
+
+
+def chunk_range(n: int, world_size: int, rank: int) -> tuple[int, int]:
+    """(row offset, row count) of `rank`'s `Shard(0)` slice of `n` rows.
+
+    `torch.chunk` semantics, which is what `DTensor` `Shard(0)` uses: ceil(n/ws)
+    rows per rank, the tail ranks shorter or EMPTY (5 rows over 4 ranks is
+    2,2,1,0). Evenly dividing would put row 4 nowhere. One function for save,
+    load, reshard and export so they cannot disagree about who owns a row.
+    """
+    c = -(-int(n) // int(world_size))
+    off = min(int(rank) * c, int(n))
+    return off, min(c, int(n) - off)
+
+
+# ---------------------------------------------------------------------------
+# safetensors, streaming (writer) and random-access (reader)
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class StEntry:
+    """One tensor to be written: `emit(f)` must write exactly `nbytes` bytes."""
+
+    name: str
+    st_dtype: str
+    shape: tuple
+    emit: Callable[[Any], None]
+
+    @property
+    def nbytes(self) -> int:
+        return _prod(self.shape) * _ST_ITEMSIZE[self.st_dtype]
+
+
+def _st_sort_key(e: StEntry):
+    return (-_ST_ORDER.index(e.st_dtype), e.name)
+
+
+def safetensors_header(entries: list[StEntry], metadata: Optional[Mapping[str, str]] = None) -> tuple[bytes, list[StEntry]]:
+    """(8-byte length prefix + JSON header padded to 8 bytes, entries in file order)."""
+    ordered = sorted(entries, key=_st_sort_key)
+    header: dict[str, Any] = {}
+    if metadata:
+        header["__metadata__"] = dict(metadata)
+    off = 0
+    for e in ordered:
+        header[e.name] = {"dtype": e.st_dtype, "shape": [int(s) for s in e.shape], "data_offsets": [off, off + e.nbytes]}
+        off += e.nbytes
+    raw = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    raw += b" " * ((-len(raw)) % 8)
+    return struct.pack("<Q", len(raw)) + raw, ordered
+
+
+class _HashingWriter:
+    def __init__(self, f):
+        self.f = f
+        self.sha = hashlib.sha256()
+        self.n = 0
+
+    def write(self, b) -> int:
+        self.sha.update(b)
+        self.n += len(b)
+        return self.f.write(b)
+
+
+def write_safetensors_streaming(
+    path: str, entries: list[StEntry], metadata: Optional[Mapping[str, str]] = None
+) -> dict:
+    """Write `path` tensor by tensor, fsync it, return {"bytes", "sha256"}.
+
+    Peak memory is whatever ONE `emit` holds. `safetensors.save_file` needs the
+    whole dict resident, which is the 2x-weights transient this module exists
+    to remove.
+    """
+    head, ordered = safetensors_header(entries, metadata)
+    with open(path, "wb") as raw:
+        w = _HashingWriter(raw)
+        w.write(head)
+        for e in ordered:
+            before = w.n
+            e.emit(w)
+            if w.n - before != e.nbytes:
+                raise ShardedCheckpointError(
+                    f"{e.name}: emitted {w.n - before} bytes, header promised {e.nbytes}"
+                )
+        raw.flush()
+        os.fsync(raw.fileno())
+    return {"bytes": w.n, "sha256": w.sha.hexdigest()}
+
+
+def tensor_bytes_emitter(t: torch.Tensor) -> Callable[[Any], None]:
+    def emit(f):
+        c = t.detach().contiguous().reshape(-1)
+        if c.numel():
+            f.write(memoryview(c.view(torch.uint8).numpy()))
+
+    return emit
+
+
+class ShardFile:
+    """Random-access reader over one safetensors file (header parsed once)."""
+
+    def __init__(self, path: str):
+        self.path = path
+        self.fd = os.open(path, os.O_RDONLY)
+        try:
+            size = os.fstat(self.fd).st_size
+            if size < 8:
+                raise ShardedCheckpointError(f"{os.path.basename(path)} is {size} bytes (no header)")
+            (n,) = struct.unpack("<Q", os.pread(self.fd, 8, 0))
+            if n > size - 8 or n > 256 * 2**20:
+                raise ShardedCheckpointError(f"{os.path.basename(path)} header claims {n} bytes of {size}")
+            header = json.loads(os.pread(self.fd, n, 8))
+            self.data_start = 8 + n
+            self.tensors: dict[str, tuple[str, tuple, int, int]] = {}
+            end = 0
+            for k, v in header.items():
+                if k == "__metadata__":
+                    continue
+                a, b = v["data_offsets"]
+                self.tensors[k] = (v["dtype"], tuple(v["shape"]), int(a), int(b))
+                end = max(end, int(b))
+            if size < self.data_start + end:
+                raise ShardedCheckpointError(
+                    f"{os.path.basename(path)} truncated: {size} bytes, header needs {self.data_start + end}"
+                )
+        except BaseException:
+            os.close(self.fd)
+            raise
+
+    def close(self) -> None:
+        if self.fd >= 0:
+            os.close(self.fd)
+            self.fd = -1
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def copy_range(self, name: str, lo: int, hi: int, sink) -> None:
+        """Stream bytes [lo, hi) of tensor `name` into `sink.write`."""
+        _dt, _shape, a, b = self.tensors[name]
+        if not (0 <= lo <= hi <= b - a):
+            raise ShardedCheckpointError(f"{name}: byte range [{lo},{hi}) outside the {b - a}-byte tensor")
+        pos = self.data_start + a + lo
+        left = hi - lo
+        while left:
+            buf = os.pread(self.fd, min(COPY_CHUNK, left), pos)
+            if not buf:
+                raise ShardedCheckpointError(f"{self.path}: short read at {pos}")
+            sink.write(buf)
+            pos += len(buf)
+            left -= len(buf)
+
+    def tensor(self, name: str) -> torch.Tensor:
+        dt, shape, a, b = self.tensors[name]
+        torch_dt = _ST_TO_TORCH[dt]
+        if b == a:
+            return torch.empty(shape, dtype=torch_dt)
+        buf = bytearray(b - a)
+        mv = memoryview(buf)
+        pos, got = self.data_start + a, 0
+        while got < len(buf):
+            n = os.preadv(self.fd, [mv[got : got + COPY_CHUNK]], pos + got)
+            if n <= 0:
+                raise ShardedCheckpointError(f"{self.path}: short read of {name}")
+            got += n
+        return torch.frombuffer(buf, dtype=torch_dt).reshape(shape)
+
+
+def file_sha256(path: str) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        while True:
+            b = f.read(COPY_CHUNK)
+            if not b:
+                break
+            h.update(b)
+    return h.hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# payload integrity sidecar (audit U01-3)
+# ---------------------------------------------------------------------------
+
+#: `<ckpt>/integrity.json`: {"version": 1, "sha256": {file name: hex}} for the payloads the
+#: weights manifest does not already cover (resume.pt, optimizer_rank*.pt, and the gathered
+#: layout's model.safetensors). A torch.save zip passes the structural check after a same-size
+#: bit flip in a stored tensor, so structure alone cannot catch silent media corruption.
+INTEGRITY_NAME = "integrity.json"
+
+
+def write_integrity(out_dir: str, sha256: Mapping[str, str]) -> None:
+    path = os.path.join(out_dir, INTEGRITY_NAME)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(json.dumps({"version": 1, "sha256": dict(sha256)}, indent=1, sort_keys=True) + "\n")
+        f.flush()
+        os.fsync(f.fileno())
+
+
+def read_integrity(path: str) -> Optional[dict]:
+    """{file name: sha256} or None for a checkpoint that predates the sidecar. A sidecar that
+    exists but cannot be parsed is an error, never 'legacy' (fail closed)."""
+    fp = os.path.join(path, INTEGRITY_NAME)
+    if not os.path.isfile(fp):
+        return None
+    try:
+        with open(fp, encoding="utf-8") as f:
+            blob = json.load(f)
+        out = {str(k): str(v) for k, v in blob["sha256"].items()}
+    except (OSError, ValueError, KeyError, AttributeError, TypeError) as e:
+        raise CheckpointIntegrityError(f"{fp} is unreadable ({type(e).__name__}: {e})") from e
+    return out
+
+
+def verify_integrity(path: str, names: Iterable[str], recorded: Optional[dict] = None) -> Optional[str]:
+    """None iff every name in `names` that the sidecar records matches its sha256.
+
+    A name the sidecar does not list is skipped (a legacy or partial sidecar); a listed file that
+    is missing or differs is a problem.
+    """
+    recorded = read_integrity(path) if recorded is None else recorded
+    if recorded is None:
+        return None
+    for n in names:
+        want = recorded.get(n)
+        if want is None:
+            continue
+        fp = os.path.join(path, n)
+        if not os.path.isfile(fp):
+            return f"{n} is missing (integrity.json records it)"
+        got = file_sha256(fp)
+        if got != want:
+            return f"{n} sha256 {got[:12]}.. != recorded {want[:12]}.. (bit rot or a partial copy)"
+    return None
+
+
+# ---------------------------------------------------------------------------
+# describing a live (possibly DTensor) tensor
+# ---------------------------------------------------------------------------
+
+
+def _local_and_desc(t: torch.Tensor, rank: int, world_size: int) -> tuple[torch.Tensor, dict]:
+    """(this rank's local tensor, descriptor) for a plain tensor or a DTensor.
+
+    Descriptor: global `shape`, `placement` ("shard0" | "replicate"), and for
+    shard0 this rank's `row` = [offset, count]. Everything else is refused HERE
+    (save time): a layout the loader cannot place must not produce a file.
+    """
+    if hasattr(t, "to_local") and hasattr(t, "placements"):
+        pl = tuple(t.placements)
+        mesh = t.device_mesh
+        if mesh.ndim != 1 or mesh.size() != world_size:
+            raise ShardedCheckpointError(
+                f"unsupported mesh {tuple(mesh.shape)} for world_size={world_size}: sharded checkpoints "
+                "support a 1-D mesh over all ranks (what fully_shard builds). HSDP/TP meshes are not "
+                "implemented; use checkpoint_format='gathered'."
+            )
+        if mesh.get_local_rank() != rank:
+            raise ShardedCheckpointError(f"mesh local rank {mesh.get_local_rank()} != process rank {rank}")
+        local = t.to_local()
+        shape = tuple(int(s) for s in t.shape)
+        from torch.distributed.tensor import Replicate, Shard
+
+        if len(pl) == 1 and isinstance(pl[0], Shard) and pl[0].dim == 0:
+            off, n = chunk_range(shape[0], world_size, rank)
+            if tuple(local.shape) != (n,) + shape[1:]:
+                raise ShardedCheckpointError(
+                    f"local shard {tuple(local.shape)} != chunk_range-derived {(n,) + shape[1:]} "
+                    f"(global {shape}, rank {rank}/{world_size}): uneven/padded layout not understood"
+                )
+            return local, {"shape": list(shape), "placement": "shard0", "row": [off, n]}
+        if len(pl) == 1 and isinstance(pl[0], Replicate):
+            return local, {"shape": list(shape), "placement": "replicate", "row": None}
+        raise ShardedCheckpointError(f"unsupported placements {pl}: only Shard(0) and Replicate() are implemented")
+    return t, {"shape": [int(s) for s in t.shape], "placement": "replicate", "row": None}
+
+
+def _store_dtype(t: torch.Tensor, store_dtype: torch.dtype) -> str:
+    dt = store_dtype if t.is_floating_point() else t.dtype
+    try:
+        return _TORCH_TO_ST[dt]
+    except KeyError as e:
+        raise ShardedCheckpointError(f"unsupported tensor dtype {dt}") from e
+
+
+def rank_file_name(rank: int) -> str:
+    return f"model_rank{rank:05d}.safetensors"
+
+
+# ---------------------------------------------------------------------------
+# save (one rank)
+# ---------------------------------------------------------------------------
+
+
+def save_rank_shard(
+    state: Mapping[str, torch.Tensor],
+    out_dir: str,
+    *,
+    rank: int,
+    world_size: int,
+    store_dtype: torch.dtype = torch.bfloat16,
+) -> dict:
+    """Write THIS rank's local shards to `out_dir/model_rank{r}.safetensors`.
+
+    Atomic per file (tmp + fsync + rename). Returns this rank's info for the
+    manifest: {"file": {name, bytes, sha256}, "tensors": {name: descriptor}}.
+    Peak host memory is one tensor.
+    """
+    entries: list[StEntry] = []
+    tensors: dict[str, dict] = {}
+    for name in sorted(state):
+        local, desc = _local_and_desc(state[name], rank, world_size)
+        st_dt = _store_dtype(local, store_dtype)
+        desc["dtype"] = st_dt
+        tensors[name] = desc
+
+        def emit(f, local=local, st_dt=st_dt):
+            host = local.detach().to(device="cpu", dtype=_ST_TO_TORCH[st_dt])
+            tensor_bytes_emitter(host)(f)
+
+        entries.append(StEntry(name, st_dt, tuple(local.shape), emit))
+    fname = rank_file_name(rank)
+    final = os.path.join(out_dir, fname)
+    tmp = final + ".tmp"
+    try:
+        meta = write_safetensors_streaming(tmp, entries, {"format": "pt"})
+        os.replace(tmp, final)
+    except BaseException:
+        for p in (tmp, final):
+            try:
+                os.remove(p)
+            except OSError:
+                pass
+        raise
+    return {"file": {"name": fname, **meta}, "tensors": tensors, "order": list(state)}
+
+
+def build_manifest(
+    infos: list[dict],
+    *,
+    param_order: Optional[list[str]] = None,
+    export: Optional[dict] = None,
+    extra: Optional[dict] = None,
+) -> dict:
+    """Merge every rank's info (rank-ordered) into the manifest. Cross-checks
+    that all ranks saw the same tensors and that shard0 rows tile [0, n)."""
+    ws = len(infos)
+    names = sorted(infos[0]["tensors"])
+    tensors: dict[str, dict] = {}
+    for n in names:
+        d0 = infos[0]["tensors"][n]
+        rows = []
+        for r, info in enumerate(infos):
+            if sorted(info["tensors"]) != names:
+                raise ShardedCheckpointError(f"rank {r} wrote a different tensor set than rank 0")
+            d = info["tensors"][n]
+            for k in ("shape", "placement", "dtype"):
+                if d[k] != d0[k]:
+                    raise ShardedCheckpointError(f"{n}: rank {r} disagrees with rank 0 on {k}: {d[k]} vs {d0[k]}")
+            rows.append(d["row"])
+        entry = {"shape": d0["shape"], "dtype": d0["dtype"], "placement": d0["placement"]}
+        if d0["placement"] == "shard0":
+            pos = 0
+            for off, cnt in rows:
+                if off != min(pos, d0["shape"][0]) and cnt:
+                    raise ShardedCheckpointError(f"{n}: shard rows do not tile ({rows})")
+                pos += cnt
+            if pos != d0["shape"][0]:
+                raise ShardedCheckpointError(f"{n}: shard rows sum to {pos}, global rows {d0['shape'][0]}")
+            entry["rows"] = rows
+        tensors[n] = entry
+    m: dict[str, Any] = {
+        "format": FORMAT_NAME,
+        "version": FORMAT_VERSION,
+        "world_size": ws,
+        "files": {str(r): infos[r]["file"] for r in range(ws)},
+        "tensors": tensors,
+    }
+    if infos[0].get("order") is not None:
+        # The model's own state-dict order: shard planning (qwen3's 5 GB files)
+        # is greedy over insertion order, and the manifest JSON is key-sorted,
+        # so without this the exported file split would differ from the
+        # gathered writer's.
+        m["key_order"] = list(infos[0]["order"])
+    if param_order is not None:
+        m["param_order"] = list(param_order)
+    if export is not None:
+        m["export"] = export
+    if extra:
+        m.update(extra)
+    return m
+
+
+def write_manifest(out_dir: str, manifest: dict) -> None:
+    path = os.path.join(out_dir, SHARD_MANIFEST_NAME)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(json.dumps(manifest, indent=1, sort_keys=True) + "\n")
+        f.flush()
+        os.fsync(f.fileno())
+
+
+# ---------------------------------------------------------------------------
+# read side: manifest, completeness, load
+# ---------------------------------------------------------------------------
+
+
+def is_sharded(path: str) -> bool:
+    return os.path.isfile(os.path.join(path, SHARD_MANIFEST_NAME))
+
+
+def read_manifest(path: str) -> dict:
+    """Parse + structurally validate; raises `ShardedCheckpointError`."""
+    mp = os.path.join(path, SHARD_MANIFEST_NAME)
+    try:
+        with open(mp, encoding="utf-8") as f:
+            m = json.load(f)
+    except (OSError, ValueError) as e:
+        raise ShardedCheckpointError(f"unreadable {SHARD_MANIFEST_NAME} ({e})") from e
+    if not isinstance(m, dict) or m.get("format") != FORMAT_NAME:
+        raise ShardedCheckpointError(f"{SHARD_MANIFEST_NAME} is not an {FORMAT_NAME} manifest")
+    if m.get("version") != FORMAT_VERSION:
+        raise ShardedCheckpointError(f"manifest version {m.get('version')} != supported {FORMAT_VERSION}")
+    ws = m.get("world_size")
+    if not isinstance(ws, int) or ws < 1:
+        raise ShardedCheckpointError(f"manifest world_size {ws!r} is not a positive int")
+    files, tensors = m.get("files"), m.get("tensors")
+    if not isinstance(files, dict) or sorted(files, key=int) != [str(r) for r in range(ws)]:
+        raise ShardedCheckpointError("manifest does not name exactly one file per rank")
+    for r in range(ws):
+        fe = files[str(r)]
+        if not (isinstance(fe, dict) and {"name", "bytes", "sha256"} <= fe.keys()):
+            raise ShardedCheckpointError(f"manifest file entry for rank {r} is malformed")
+    # An EMPTY tensor set is legal: a module whose weights are reconstructible
+    # elsewhere (a frozen base + adapter) may declare `sharded_model_state() == {}`
+    # and checkpoint only its optimizer/adapter state.
+    if not isinstance(tensors, dict):
+        raise ShardedCheckpointError("manifest tensors entry is not an object")
+    for n, e in tensors.items():
+        if e.get("placement") not in ("shard0", "replicate") or e.get("dtype") not in _ST_ITEMSIZE:
+            raise ShardedCheckpointError(f"{n}: malformed tensor entry")
+        if e["placement"] == "shard0" and len(e.get("rows", [])) != ws:
+            raise ShardedCheckpointError(f"{n}: shard0 entry has {len(e.get('rows', []))} row ranges, world {ws}")
+    return m
+
+
+def sharded_problem(path: str, st_header_problem: Callable[[str], Optional[str]]) -> Optional[str]:
+    """None iff every rank file the manifest names is present, the size it
+    recorded, and a structurally complete safetensors file. CHEAP (no payload
+    reads): sha256 is verified on the rank that loads the file."""
+    try:
+        m = read_manifest(path)
+    except ShardedCheckpointError as e:
+        return str(e)
+    for r in range(m["world_size"]):
+        fe = m["files"][str(r)]
+        fp = os.path.join(path, fe["name"])
+        if not os.path.isfile(fp):
+            return f"missing shard file {fe['name']}"
+        if os.path.getsize(fp) != fe["bytes"]:
+            return f"{fe['name']} is {os.path.getsize(fp)} bytes, manifest recorded {fe['bytes']}"
+        problem = st_header_problem(fp)
+        if problem:
+            return problem
+    return None
+
+
+def verify_rank_file(path: str, rank: int) -> Optional[str]:
+    """sha256 of THIS rank's file against the manifest; None if it matches."""
+    m = read_manifest(path)
+    fe = m["files"][str(rank)]
+    got = file_sha256(os.path.join(path, fe["name"]))
+    if got != fe["sha256"]:
+        return f"{fe['name']} sha256 {got[:12]}.. != manifest {fe['sha256'][:12]}.. (bit rot or a partial copy)"
+    return None
+
+
+def check_world_size(path: str, manifest: dict, world_size: int) -> None:
+    if manifest["world_size"] != world_size:
+        raise ShardedWorldSizeError(
+            f"checkpoint {path} was saved with world_size={manifest['world_size']} but this run has "
+            f"world_size={world_size}. Sharded weights are rank-local row slices and cannot be loaded "
+            "under another partitioning. Resume with the same number of ranks, or reshard offline: "
+            f"python scripts/ckpt_convert.py reshard {path} <dst> --world-size {world_size}"
+        )
+
+
+def load_rank_shard(
+    state: Mapping[str, torch.Tensor], path: str, *, rank: int, world_size: int, verify_hash: bool = False
+) -> int:
+    """Copy this rank's file into the LIVE tensors of `state` (a model state
+    dict whose DTensors alias the parameters). Reads only this rank's file, one
+    tensor at a time. Returns the number of tensors loaded."""
+    m = read_manifest(path)
+    check_world_size(path, m, world_size)
+    mt = m["tensors"]
+    missing, unexpected = sorted(set(mt) - set(state)), sorted(set(state) - set(mt))
+    if missing or unexpected:
+        raise KeyError(f"resume state mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")
+    if verify_hash:
+        bad = verify_rank_file(path, rank)
+        if bad:
+            raise ShardedCheckpointError(bad)
+    # validate EVERYTHING before the first copy_: a mismatch on tensor 300 must
+    # not leave a model half-overwritten.
+    plan = []
+    for name in sorted(mt):
+        local, desc = _local_and_desc(state[name], rank, world_size)
+        e = mt[name]
+        if desc["shape"] != e["shape"] or desc["placement"] != e["placement"]:
+            raise ShardedCheckpointError(
+                f"{name}: live {desc['placement']} {desc['shape']} != checkpoint {e['placement']} {e['shape']}"
+            )
+        if e["placement"] == "shard0" and list(e["rows"][rank]) != list(desc["row"]):
+            raise ShardedCheckpointError(
+                f"{name}: live rows {desc['row']} != checkpoint rows {e['rows'][rank]} for rank {rank}"
+            )
+        if local.is_floating_point() != _ST_TO_TORCH[e["dtype"]].is_floating_point:
+            raise ShardedCheckpointError(f"{name}: dtype class differs (live {local.dtype}, ckpt {e['dtype']})")
+        plan.append((name, local))
+    with ShardFile(os.path.join(path, m["files"][str(rank)]["name"])) as sf:
+        missing_in_file = [n for n, _ in plan if n not in sf.tensors]
+        if missing_in_file:
+            raise ShardedCheckpointError(f"{sf.path} lacks tensors {missing_in_file[:5]}")
+        with torch.no_grad():
+            for name, local in plan:
+                t = sf.tensor(name)
+                if tuple(t.shape) != tuple(local.shape):
+                    raise ShardedCheckpointError(f"{name}: file shape {tuple(t.shape)} != live {tuple(local.shape)}")
+                local.copy_(t.to(device=local.device))
+    return len(plan)
+
+
+class _ListSink:
+    def __init__(self):
+        self.parts: list[bytes] = []
+
+    def write(self, b):
+        self.parts.append(bytes(b))
+        return len(b)
+
+
+class GlobalReader:
+    """Byte-range reads over the GLOBAL (row-major) tensors of a sharded checkpoint."""
+
+    def __init__(self, path: str, manifest: Optional[dict] = None):
+        self.path = path
+        self.m = manifest or read_manifest(path)
+        self.files = {
+            r: ShardFile(os.path.join(path, self.m["files"][str(r)]["name"])) for r in range(self.m["world_size"])
+        }
+
+    def close(self):
+        for f in self.files.values():
+            f.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *a):
+        self.close()
+
+    def info(self, name: str) -> tuple[str, tuple]:
+        e = self.m["tensors"][name]
+        return e["dtype"], tuple(e["shape"])
+
+    def nbytes(self, name: str) -> int:
+        dt, shape = self.info(name)
+        return _prod(shape) * _ST_ITEMSIZE[dt]
+
+    def copy_range(self, name: str, lo: int, hi: int, sink) -> None:
+        """Stream global bytes [lo, hi) of tensor `name` into `sink`."""
+        e = self.m["tensors"][name]
+        if not (0 <= lo <= hi <= self.nbytes(name)):
+            raise ShardedCheckpointError(f"{name}: byte range [{lo},{hi}) outside tensor")
+        if e["placement"] == "replicate":
+            self.files[0].copy_range(name, lo, hi, sink)
+            return
+        dt, shape = self.info(name)
+        rb = _prod(shape[1:]) * _ST_ITEMSIZE[dt]
+        for r, (off, cnt) in enumerate(e["rows"]):
+            a, b = max(lo, off * rb), min(hi, (off + cnt) * rb)
+            if a < b:
+                self.files[r].copy_range(name, a - off * rb, b - off * rb, sink)
+
+    def tensor(self, name: str) -> torch.Tensor:
+        dt, shape = self.info(name)
+        sink = _ListSink()
+        self.copy_range(name, 0, self.nbytes(name), sink)
+        buf = bytearray(b"".join(sink.parts))
+        if not buf:
+            return torch.empty(shape, dtype=_ST_TO_TORCH[dt])
+        return torch.frombuffer(buf, dtype=_ST_TO_TORCH[dt]).reshape(shape)
+
+
+def load_full_state_dict(path: str) -> dict[str, torch.Tensor]:
+    """The WHOLE state dict, assembled in host memory. For tests and small
+    models only -- the production path is `export_hf`, which never does this."""
+    with GlobalReader(path) as g:
+        return {n: g.tensor(n) for n in g.m["tensors"]}
+
+
+# ---------------------------------------------------------------------------
+# describing the weight format (so the checkpoint is self-exporting)
+# ---------------------------------------------------------------------------
+
+
+def describe_weight_format(weight_format: Any, dtype: torch.dtype = torch.bfloat16) -> Optional[dict]:
+    """{"family", "hf_config"} for the formats we can export; None otherwise
+    (the checkpoint still saves and resumes, `export` then refuses with a clear
+    message instead of guessing a config)."""
+    # Subclasses count: scripts wrap OlmoeFormat to serialise its loads
+    # (olmoe_sft_smoke.SerialWeightFormat) and must stay exportable.
+    names = {c.__name__ for c in type(weight_format).__mro__}
+    name = next((n for n in ("DenseFormat", "OlmoeFormat", "Qwen3Format") if n in names), type(weight_format).__name__)
+    try:
+        if name == "DenseFormat":
+            from ..model import hf_config as hfc
+
+            return {"family": "dense", "hf_config": hfc.to_hf_config_dict(weight_format.model_cfg, dtype=dtype)}
+        if name == "OlmoeFormat":
+            from ..model.olmoe_checkpoint import config_to_hf_dict
+
+            return {"family": "olmoe", "hf_config": config_to_hf_dict(weight_format.olmoe_cfg, dtype=dtype)}
+        if name == "Qwen3Format":
+            return {"family": "qwen3", "hf_config": weight_format.qcfg.to_hf_config(dtype=dtype)}
+    except Exception as e:  # noqa: BLE001 -- a config we cannot describe must not block saving weights
+        logger.warning("cannot describe weight format %s for export (%s: %s)", name, type(e).__name__, e)
+    return None
+
+
+# ---------------------------------------------------------------------------
+# HF export plan + streaming export
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class HFEntry:
+    hf_name: str
+    st_dtype: str
+    shape: tuple
+    src: str  # align key
+    lo: int  # global byte range in the source tensor
+    hi: int
+
+
+def _meta(manifest: dict) -> dict[str, torch.Tensor]:
+    order = manifest.get("key_order") or list(manifest["tensors"])
+    return {
+        n: torch.empty(tuple(manifest["tensors"][n]["shape"]), dtype=_ST_TO_TORCH[manifest["tensors"][n]["dtype"]], device="meta")
+        for n in order
+    }
+
+
+def build_hf_plan(manifest: dict) -> list[HFEntry]:
+    """Plan the HF tensor list from the manifest ALONE (no data). 1:1 families
+    run the family's own converter on META tensors and map outputs back by
+    object identity, so the name table is the converter's, not a copy of it.
+    OLMoE's stacked experts become per-expert byte ranges of the fused tensor."""
+    if not manifest["tensors"]:
+        raise ShardedCheckpointError("this checkpoint holds no weight tensors (adapter/optimizer-only); nothing to export")
+    ex = manifest.get("export")
+    if not ex:
+        raise ShardedCheckpointError(
+            "this checkpoint carries no export recipe (its weight format is not Dense/Olmoe/Qwen3); "
+            "export cannot guess a config"
+        )
+    fam, hfc_d = ex["family"], ex["hf_config"]
+    meta = _meta(manifest)
+    ident = {id(v): k for k, v in meta.items()}
+    entries: list[HFEntry] = []
+
+    def one_to_one(out: Mapping[str, torch.Tensor]) -> None:
+        for hf_name, v in out.items():
+            src = ident.get(id(v))
+            if src is None:
+                raise ShardedCheckpointError(f"{hf_name}: not a 1:1 image of an align tensor")
+            e = manifest["tensors"][src]
+            entries.append(HFEntry(hf_name, e["dtype"], tuple(e["shape"]), src, 0, _prod(e["shape"]) * _ST_ITEMSIZE[e["dtype"]]))
+
+    if fam == "dense":
+        from ..model import checkpoint as ckpt
+        from ..model import hf_config as hfc
+
+        cfg = hfc.from_hf_config_dict(hfc_d)
+        one_to_one(ckpt.nmo_to_hf(meta, expected_n_layers=cfg.n_layers))
+    elif fam == "qwen3":
+        from ..model.qwen3 import Qwen3Config
+        from ..model.qwen3_checkpoint import align_to_hf
+
+        one_to_one(align_to_hf(meta, Qwen3Config.from_hf_config(hfc_d)))
+    elif fam == "olmoe":
+        from ..model import olmoe_checkpoint as oc
+
+        cfg = oc.config_from_hf_dict(hfc_d)
+        real = oc.align_to_hf(meta, cfg)  # the converter's own key set and shapes: the oracle
+        inter = cfg.intermediate_size
+        fused = set()
+        for i in range(cfg.n_layers):
+            gu_k, dn_k = (k.format(i=i) for k in oc._AL_FUSED)
+            fused.update((gu_k, dn_k))
+            gu, dn = manifest["tensors"][gu_k], manifest["tensors"][dn_k]
+            es = _ST_ITEMSIZE[gu["dtype"]]
+            e_n, two_i, d = gu["shape"]
+            for e in range(e_n):
+                base = e * two_i * d * es
+                half = inter * d * es
+                entries.append(HFEntry(oc._HF_EXPERT.format(i=i, e=e, p="gate"), gu["dtype"], (inter, d), gu_k, base, base + half))
+                entries.append(HFEntry(oc._HF_EXPERT.format(i=i, e=e, p="up"), gu["dtype"], (inter, d), gu_k, base + half, base + 2 * half))
+                dbase = e * dn["shape"][1] * dn["shape"][2] * es
+                entries.append(HFEntry(oc._HF_EXPERT.format(i=i, e=e, p="down"), dn["dtype"], tuple(dn["shape"][1:]), dn_k, dbase, dbase + _prod(dn["shape"][1:]) * es))
+        rest = {k: v for k, v in real.items() if id(v) in ident and ident[id(v)] not in fused}
+        one_to_one(rest)
+        got = {e.hf_name: e.shape for e in entries}
+        want = {k: tuple(v.shape) for k, v in real.items()}
+        if got != want:
+            raise ShardedCheckpointError(
+                f"olmoe export plan disagrees with align_to_hf: only-plan={sorted(set(got) - set(want))[:3]} "
+                f"only-converter={sorted(set(want) - set(got))[:3]} "
+                f"shape-diff={[k for k in got if k in want and got[k] != want[k]][:3]}"
+            )
+    else:
+        raise ShardedCheckpointError(f"unknown export family {fam!r}")
+    return entries
+
+
+def export_hf(src: str, dst: str, *, verify: bool = True) -> dict:
+    """Sharded checkpoint -> HF-layout directory, tensor by tensor.
+
+    Byte-identical to what the gathered path writes (`save_nmo_as_hf` /
+    `save_olmoe_as_hf` / `save_hf_qwen3`) -- tested by file hash -- with peak RSS
+    of one `COPY_CHUNK` instead of 2x the weights. `dst` is staged as
+    `dst.partial`, fsynced and renamed: an interrupted export leaves nothing at
+    `dst`.
+    """
+    if os.path.exists(dst):
+        raise FileExistsError(f"{dst} exists; export never overwrites")
+    m = read_manifest(src)
+    if verify:
+        for r in range(m["world_size"]):
+            bad = verify_rank_file(src, r)
+            if bad:
+                raise ShardedCheckpointError(bad)
+    plan = build_hf_plan(m)
+    fam = m["export"]["family"]
+    staging = dst.rstrip("/") + ".partial"
+    if os.path.isdir(staging):
+        shutil.rmtree(staging)
+    os.makedirs(staging)
+    try:
+        with GlobalReader(src, m) as g:
+            def make_entry(h: HFEntry) -> StEntry:
+                def emit(f, h=h):
+                    g.copy_range(h.src, h.lo, h.hi, f)
+
+                return StEntry(h.hf_name, h.st_dtype, h.shape, emit)
+
+            by_name = {h.hf_name: h for h in plan}
+            if fam == "qwen3":
+                from ..model.qwen3_checkpoint import MAX_SHARD_BYTES, _plan_shards
+
+                sizes = {h.hf_name: _prod(h.shape) * _ST_ITEMSIZE[h.st_dtype] for h in plan}
+                groups = _plan_shards(sizes, MAX_SHARD_BYTES)
+                weight_map: dict[str, str] = {}
+                for n, keys in enumerate(groups, start=1):
+                    fname = f"model-{n:05d}-of-{len(groups):05d}.safetensors"
+                    write_safetensors_streaming(
+                        os.path.join(staging, fname), [make_entry(by_name[k]) for k in keys], {"format": "pt"}
+                    )
+                    weight_map.update({k: fname for k in keys})
+                with open(os.path.join(staging, "model.safetensors.index.json"), "w") as fh:
+                    json.dump({"metadata": {"total_size": sum(sizes.values())}, "weight_map": weight_map}, fh, indent=2, sort_keys=True)
+                    fh.flush()
+                    os.fsync(fh.fileno())
+            else:
+                write_safetensors_streaming(
+                    os.path.join(staging, "model.safetensors"), [make_entry(h) for h in plan], {"format": "pt"}
+                )
+        with open(os.path.join(staging, "config.json"), "w", encoding="utf-8") as fh:
+            fh.write(json.dumps(m["export"]["hf_config"], indent=2, sort_keys=True) + "\n")
+            fh.flush()
+            os.fsync(fh.fileno())
+        for fn in sorted(os.listdir(src)):
+            if not _OWN_FILE_RE.match(fn) and os.path.isfile(os.path.join(src, fn)) and fn != "config.json":
+                shutil.copyfile(os.path.join(src, fn), os.path.join(staging, fn))
+        _fsync_tree(staging)
+        os.rename(staging, dst)
+        _fsync_dir(os.path.dirname(os.path.abspath(dst)))
+    except BaseException:
+        shutil.rmtree(staging, ignore_errors=True)
+        raise
+    return {"tensors": len(plan), "family": fam, "dst": dst}
+
+
+def _fsync_dir(p: str) -> None:
+    fd = os.open(p, os.O_RDONLY)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _fsync_tree(root: str) -> None:
+    dirs = []
+    for dp, _dn, fns in os.walk(root):
+        dirs.append(dp)
+        for fn in fns:
+            fd = os.open(os.path.join(dp, fn), os.O_RDONLY)
+            try:
+                os.fsync(fd)
+            finally:
+                os.close(fd)
+    for d in reversed(dirs):
+        _fsync_dir(d)
+
+
+# ---------------------------------------------------------------------------
+# offline reshard
+# ---------------------------------------------------------------------------
+
+
+class ReshardUnsupported(ShardedCheckpointError):
+    """Optimizer state whose layout cannot be re-cut without knowing the optimizer."""
+
+
+def _check_opt_leaf(vals: list, shapes_old: list[tuple], name: str, where: str) -> None:
+    """Every old rank's piece must have exactly the LOCAL shard shape of the parameter it is
+    filed under (shape metadata only: reads nothing from an mmap'd file)."""
+    for r, (v, want) in enumerate(zip(vals, shapes_old)):
+        if tuple(v.shape) != tuple(want):
+            raise ReshardUnsupported(
+                f"{where}: tensor shape {tuple(v.shape)} on old rank {r} != the local shard {tuple(want)} of "
+                f"parameter {name!r}. This optimizer state is not elementwise-per-shard (quantised/blocked "
+                "moments?) or its parameter order differs from the model's; it cannot be re-cut safely."
+            )
+
+
+def _cut_opt_leaf(vals: list, old_rows: Optional[list], new_ws: int, new_rank: int) -> torch.Tensor:
+    """NEW rank `new_rank`'s piece of one per-param optimizer tensor, built from only the
+    old pieces that overlap its rows (no whole-tensor concatenation, no other new rank)."""
+    if old_rows is None:  # replicated parameter
+        return vals[0].clone()
+    total = sum(c for _o, c in old_rows)
+    off, n = chunk_range(total, new_ws, new_rank)
+    parts = []
+    for v, (qo, qc) in zip(vals, old_rows):
+        a, b = max(off, qo), min(off + n, qo + qc)
+        if a < b:
+            parts.append(v[a - qo : b - qo])
+    if not parts:
+        return torch.empty((0,) + tuple(vals[0].shape[1:]), dtype=vals[0].dtype)
+    return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0].clone()
+
+
+def iter_reshard_optimizer(old: list[dict], manifest: dict, new_ws: int):
+    """Yield the re-cut optimizer state dict of NEW rank 0, 1, ... one at a time.
+
+    THE ALL-RANKS-AT-ONCE DEFECT this shape prevents. The first version returned a
+    list holding every new rank's state, so its peak was the WHOLE optimizer
+    (12 B/param for fp32-master AdamW: 83 GB for a 7B model) while the CLI
+    estimated `optimizer_bytes / new_ws` and the 45 GB floor check was sized on
+    that. Measured on 4 old -> 4 new ranks, 50 M params: anonymous-RSS growth 576
+    MiB against an estimate of 144 MiB, i.e. `new_ws` times the estimate. On this
+    no-swap shared box a reshard that passes the preflight at 21 GB and takes 83
+    GB is the 2026-08-12 OOM. Here only the old pieces that overlap the rank being
+    built are read (mmap'd), and the caller writes and drops each yielded dict
+    before asking for the next, so the peak is one new rank's state.
+
+    Every shape/consistency refusal is raised BEFORE the first yield, so a
+    resharder that cannot prove it understands the state writes no file at all.
+    """
+    order = manifest.get("param_order")
+    if not order:
+        raise ReshardUnsupported("manifest has no param_order: cannot map optimizer slots to parameters")
+    ws = manifest["world_size"]
+    if len(old) != ws:
+        raise ReshardUnsupported(f"{len(old)} optimizer file(s) for a {ws}-rank checkpoint")
+    if "world_size" in old[0]:
+        raise ReshardUnsupported("optimizer state is stamped with its world_size (quantised shard-local state)")
+
+    def entry(pname: str) -> dict:
+        e = manifest["tensors"].get(pname)
+        if e is None:
+            raise ReshardUnsupported(f"param_order names {pname!r}, which is not a tensor of this checkpoint")
+        return e
+
+    def local_shape(pname: str, r: int) -> tuple:
+        e = entry(pname)
+        if e["placement"] == "replicate":
+            return tuple(e["shape"])
+        return (e["rows"][r][1],) + tuple(e["shape"][1:])
+
+    def rows_of(pname: str) -> Optional[list]:
+        e = entry(pname)
+        return None if e["placement"] == "replicate" else [tuple(x) for x in e["rows"]]
+
+    def agree(vals: list, where: str):
+        v0 = vals[0]
+        for v in vals[1:]:
+            same = torch.equal(v, v0) if isinstance(v0, torch.Tensor) else v == v0
+            if not same:
+                raise ReshardUnsupported(f"{where}: differs across old ranks ({v0!r} vs {v!r})")
+        return v0
+
+    # -- pass 1: validate everything, build a per-key plan (no tensor payload read) -----------
+    plan: list[tuple] = []
+    for key in old[0]:
+        vals = [o[key] for o in old]
+        if key == "state":
+            leaves = []  # (idx, k, kind, payload)
+            for idx in vals[0]:
+                if not 0 <= int(idx) < len(order):
+                    raise ReshardUnsupported(f"state slot {idx} is outside param_order ({len(order)} entries)")
+                pname = order[int(idx)]
+                for k in vals[0][idx]:
+                    leaf = [v[idx][k] for v in vals]
+                    where = f"state[{idx}][{k}]"
+                    if isinstance(leaf[0], torch.Tensor) and leaf[0].ndim >= 1:
+                        _check_opt_leaf(leaf, [local_shape(pname, r) for r in range(ws)], pname, where)
+                        leaves.append((idx, k, "cut", (leaf, rows_of(pname))))
+                    else:
+                        leaves.append((idx, k, "same", agree(leaf, where)))
+            plan.append((key, "state", leaves))
+        elif key == "master":
+            if len(vals[0]) != len(order) or any(len(v) != len(order) for v in vals):
+                raise ReshardUnsupported(
+                    f"master has {[len(v) for v in vals]} entries per old rank, param_order {len(order)}: the "
+                    "optimizer was not built over the model's trainable parameters in order"
+                )
+            masters = []
+            for i, pname in enumerate(order):
+                leaf = [v[i] for v in vals]
+                _check_opt_leaf(leaf, [local_shape(pname, r) for r in range(ws)], pname, f"master[{i}]")
+                masters.append((leaf, rows_of(pname)))
+            plan.append((key, "master", masters))
+        else:
+            if _has_tensor(vals[0], min_ndim=1):
+                raise ReshardUnsupported(f"optimizer key {key!r} holds sharded tensors this resharder does not understand")
+            plan.append((key, "plain", _agree_tree(vals, key)))
+
+    # -- pass 2: build ONE new rank at a time ---------------------------------------------------
+    for r in range(new_ws):
+        out: dict = {}
+        for key, kind, payload in plan:
+            if kind == "state":
+                st: dict = {}
+                for idx, k, how, p in payload:
+                    st.setdefault(idx, {})[k] = _cut_opt_leaf(p[0], p[1], new_ws, r) if how == "cut" else p
+                out[key] = st
+            elif kind == "master":
+                out[key] = [_cut_opt_leaf(leaf, rows, new_ws, r) for leaf, rows in payload]
+            else:
+                out[key] = payload
+        yield out
+
+
+def reshard_optimizer(old: list[dict], manifest: dict, new_ws: int) -> list[dict]:
+    """Re-cut per-rank optimizer state dicts for `new_ws`, all at once (tests, small models).
+
+    Production resharding uses `iter_reshard_optimizer` and writes each rank
+    before building the next; this wrapper holds every new rank in memory.
+
+    Handles the layout `torch.optim.AdamW` / `OffloadAdamW` produce:
+    `state[idx][k]` tensors and a `master` list, each shaped like the LOCAL
+    shard of parameter `param_order[idx]`. Scalars (`step`) and non-tensor
+    leaves must agree across old ranks. ANYTHING else that holds a tensor with
+    ndim >= 1 (adamw8bit's quantised blocks, a `world_size`-stamped state)
+    raises `ReshardUnsupported`: re-cutting what we do not understand is how a
+    resume loads plausible, wrong moments.
+    """
+    return list(iter_reshard_optimizer(old, manifest, new_ws))
+
+
+def _has_tensor(o: Any, min_ndim: int) -> bool:
+    if isinstance(o, torch.Tensor):
+        return o.ndim >= min_ndim
+    if isinstance(o, dict):
+        return any(_has_tensor(v, min_ndim) for v in o.values())
+    if isinstance(o, (list, tuple)):
+        return any(_has_tensor(v, min_ndim) for v in o)
+    return False
+
+
+def _agree_tree(vals: list, where: str):
+    v0 = vals[0]
+    for v in vals[1:]:
+        if repr(v) != repr(v0):
+            raise ReshardUnsupported(f"{where}: differs across old ranks")
+    return v0
+
+
+def reshard(src: str, dst: str, new_ws: int, *, drop_optimizer: bool = False) -> dict:
+    """Offline reshard of a sharded checkpoint to `new_ws` ranks.
+
+    Model shards are re-cut by byte range (bounded memory). Optimizer files are
+    re-cut by `reshard_optimizer` (peak RSS = one NEW rank's optimizer state --
+    at 7B fp32-master AdamW that is 12 B/param/new_ws; run it on an idle box).
+    With `drop_optimizer=True` the optimizer files are omitted and the loop
+    state is marked, so resume starts FRESH moments: a recorded deviation, and
+    NOT lossless when the optimizer holds the fp32 master.
+    `resume.pt` is carried over with its world-size stamp rewritten and the
+    per-rank RNG list dropped (the trainer restarts RNG from the seed with a
+    warning; carrying N_old streams into N_new ranks would be wrong).
+    """
+    if new_ws < 1:
+        raise ValueError("new_ws must be >= 1")
+    m = read_manifest(src)
+    for r in range(m["world_size"]):
+        bad = verify_rank_file(src, r)
+        if bad:
+            raise ShardedCheckpointError(bad)
+    if os.path.exists(dst):
+        raise FileExistsError(f"{dst} exists; reshard never overwrites")
+    staging = dst.rstrip("/") + ".partial"
+    if os.path.isdir(staging):
+        shutil.rmtree(staging)
+    os.makedirs(staging)
+    try:
+        infos = []
+        with GlobalReader(src, m) as g:
+            for r in range(new_ws):
+                entries, tdesc = [], {}
+                for name in sorted(m["tensors"]):
+                    e = m["tensors"][name]
+                    dt, shape = e["dtype"], tuple(e["shape"])
+                    if e["placement"] == "shard0":
+                        off, cnt = chunk_range(shape[0], new_ws, r)
+                        lshape = (cnt,) + shape[1:]
+                        rb = _prod(shape[1:]) * _ST_ITEMSIZE[dt]
+                        lo, hi = off * rb, (off + cnt) * rb
+                        tdesc[name] = {"shape": list(shape), "placement": "shard0", "row": [off, cnt], "dtype": dt}
+                    else:
+                        lshape, lo, hi = shape, 0, g.nbytes(name)
+                        tdesc[name] = {"shape": list(shape), "placement": "replicate", "row": None, "dtype": dt}
+
+                    def emit(f, name=name, lo=lo, hi=hi):
+                        g.copy_range(name, lo, hi, f)
+
+                    entries.append(StEntry(name, dt, lshape, emit))
+                fname = rank_file_name(r)
+                meta = write_safetensors_streaming(os.path.join(staging, fname), entries, {"format": "pt"})
+                infos.append({"file": {"name": fname, **meta}, "tensors": tdesc, "order": m.get("key_order")})
+        new_manifest = build_manifest(
+            infos,
+            param_order=m.get("param_order"),
+            export=m.get("export"),
+            extra={"resharded_from": {"world_size": m["world_size"], "source": os.path.basename(src.rstrip("/"))}},
+        )
+        write_manifest(staging, new_manifest)
+
+        state = torch.load(os.path.join(src, "resume.pt"), weights_only=True, map_location="cpu")
+        state.pop("rng_per_rank", None)
+        if drop_optimizer:
+            state["optimizer_state"] = {"layout": "none", "world_size": new_ws}
+        else:
+            state["optimizer_state"] = {"layout": "per_rank_files", "world_size": new_ws}
+            olds = [
+                torch.load(os.path.join(src, f"optimizer_rank{r}.pt"), weights_only=True, map_location="cpu", mmap=True)
+                for r in range(m["world_size"])
+            ]
+            for r, sd in enumerate(iter_reshard_optimizer(olds, m, new_ws)):
+                p = os.path.join(staging, f"optimizer_rank{r}.pt")
+                with open(p, "wb") as f:
+                    torch.save(sd, f)
+                    f.flush()
+                    os.fsync(f.fileno())
+                del sd  # one new rank's state alive at a time
+            del olds
+        with open(os.path.join(staging, "resume.pt"), "wb") as f:
+            torch.save(state, f)
+            f.flush()
+            os.fsync(f.fileno())
+        write_integrity(
+            staging,
+            {
+                fn: file_sha256(os.path.join(staging, fn))
+                for fn in sorted(os.listdir(staging))
+                if fn == "resume.pt" or re.fullmatch(r"optimizer_rank\d+\.pt", fn)
+            },
+        )
+        for fn in sorted(os.listdir(src)):
+            if not _OWN_FILE_RE.match(fn) and os.path.isfile(os.path.join(src, fn)):
+                shutil.copyfile(os.path.join(src, fn), os.path.join(staging, fn))
+        _fsync_tree(staging)
+        os.rename(staging, dst)
+        _fsync_dir(os.path.dirname(os.path.abspath(dst)))
+    except BaseException:
+        shutil.rmtree(staging, ignore_errors=True)
+        raise
+    return {"dst": dst, "world_size": new_ws, "from": m["world_size"], "optimizer": "dropped" if drop_optimizer else "resharded"}
+
+
+def peak_rss_kb() -> int:
+    """VmHWM of this process in kB (psutil is not a dependency)."""
+    with open("/proc/self/status") as f:
+        for line in f:
+            if line.startswith("VmHWM:"):
+                return int(line.split()[1])
+    return -1
+
+
+__all__ = [
+    "CHECKPOINT_FORMATS",
+    "SHARD_MANIFEST_NAME",
+    "ShardedCheckpointError",
+    "ShardedWorldSizeError",
+    "ReshardUnsupported",
+    "chunk_range",
+    "save_rank_shard",
+    "build_manifest",
+    "write_manifest",
+    "read_manifest",
+    "is_sharded",
+    "sharded_problem",
+    "verify_rank_file",
+    "load_rank_shard",
+    "load_full_state_dict",
+    "describe_weight_format",
+    "build_hf_plan",
+    "export_hf",
+    "reshard",
+    "reshard_optimizer",
+    "iter_reshard_optimizer",
+    "peak_rss_kb",
+]

code/align/train/stage_gates.py

@@ -0,0 +1,1002 @@
+"""Stage gate sequencing and decisions (section-11).
+
+This is the first stage that spends real GPU-hours, so the job of this module
+is to make wasting them structurally impossible rather than merely unlikely.
+Everything here is a pure decision over recorded facts; the stage script drives
+it and owns the side effects.
+
+WHY THE ORDER IS FIXED, and why it is enforced instead of documented:
+
+  golden_regression -> smoke_7b -> probe_20step -> rss_preflight
+                    -> train -> eval -> quality_gate
+
+  * The golden regression runs before ANY new training because it validates
+    the eval stack itself against numbers that are already known. A defective
+    harness silently corrupts every downstream measurement, and a full-epoch
+    run that under-performs a 0.35-epoch run then looks like a modeling
+    result instead of the stack bug it is.
+  * The 7B smoke runs before the probe because the chunked-loss peak-memory
+    assertion is only meaningful at 7B on the real 4-GPU topology. A
+    unit-scale memory check cannot catch what a 32768-vocab logit tensor does
+    to a 23GB card.
+  * The probe runs before the preflight and the launch because a measured step
+    time that invalidates the schedule must stop the run BEFORE it starts, not
+    three hours in.
+
+MODULE-LEVEL IMPORTS ARE DELIBERATELY TORCH-FREE. `scripts/sft.py` imports this
+module above `apply_resource_guards()`, and torch reads
+PYTORCH_CUDA_ALLOC_CONF, the cache paths and the thread caps AT IMPORT -- so
+anything that pulls torch in transitively (probes, trainer, model) is imported
+lazily inside the function that needs it. A test asserts this.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import math
+import os
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, Optional
+
+from ..registry import MANDATORY_BANNER_KEYS, assert_banner_complete
+
+logger = logging.getLogger(__name__)
+
+NAN = float("nan")
+
+#: Explicit "not measured". Never absent, never a plausible substitute.
+PENDING = "PENDING"
+
+#: The fixed stage sequence. Data, not prose: the ledger iterates it.
+GATE_ORDER: tuple[str, ...] = (
+    "golden_regression",
+    "smoke_7b",
+    "probe_20step",
+    "rss_preflight",
+    "train",
+    "eval",
+    "quality_gate",
+)
+
+#: Gates whose result depends on the eval stack. A change under `align/eval/`
+#: invalidates exactly these -- re-running the 7B smoke or the probe because
+#: someone edited a verifier would burn GPU-hours for nothing.
+EVAL_DEPENDENT_GATES: frozenset[str] = frozenset({"golden_regression", "eval", "quality_gate"})
+
+LEDGER_FILENAME = "stage_gates.json"
+
+
+def _is_rank0() -> bool:
+    """Rank from the environment, so this works before (and without) any
+    process group -- same convention as the heavy-job flock."""
+    return int(os.environ.get("RANK", "0")) == 0
+
+
+class GateOrderError(RuntimeError):
+    """Raised when a gate runs before its predecessors have passed."""
+
+
+@dataclass
+class GateDecision:
+    passed: bool
+    halt_message: str = ""
+    detail: dict = None  # type: ignore[assignment]
+
+
+# ---------------------------------------------------------------------------
+# ledger
+# ---------------------------------------------------------------------------
+
+
+def eval_source_hash(repo_root: Optional[str] = None) -> str:
+    """Hash the `align/eval/` source tree.
+
+    The golden regression certifies the harness, so its pass is only valid for
+    the harness that produced it. Keying the ledger entry on this hash is what
+    turns "we ran the golden regression once" into "we ran it against THIS
+    code".
+    """
+    root = repo_root or os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    eval_dir = os.path.join(root, "align", "eval")
+    h = hashlib.sha256()
+    for dirpath, _dirs, files in sorted(os.walk(eval_dir)):
+        for fn in sorted(files):
+            if not fn.endswith(".py"):
+                continue
+            path = os.path.join(dirpath, fn)
+            h.update(os.path.relpath(path, eval_dir).encode())
+            with open(path, "rb") as f:
+                h.update(f.read())
+    return h.hexdigest()[:16]
+
+
+class GateLedger:
+    """Records gate results, enforces the order, survives a restart."""
+
+    def __init__(self, run_dir: str, eval_source_hash: Optional[str] = None):
+        self.run_dir = run_dir
+        os.makedirs(run_dir, exist_ok=True)
+        self.path = os.path.join(run_dir, LEDGER_FILENAME)
+        self.eval_source_hash = eval_source_hash
+        self._entries: dict[str, dict] = {}
+        self._load()
+
+    # -- persistence --------------------------------------------------------
+
+    def _load(self) -> None:
+        if not os.path.isfile(self.path):
+            return
+        try:
+            with open(self.path) as f:
+                blob = json.load(f)
+        except (OSError, json.JSONDecodeError):
+            logger.warning("unreadable gate ledger at %s; starting fresh", self.path)
+            return
+
+        recorded_hash = blob.get("eval_source_hash")
+        self._entries = dict(blob.get("gates", {}))
+        if (
+            self.eval_source_hash is not None
+            and recorded_hash is not None
+            and recorded_hash != self.eval_source_hash
+        ):
+            dropped = [g for g in EVAL_DEPENDENT_GATES if g in self._entries]
+            for g in dropped:
+                del self._entries[g]
+            if dropped:
+                logger.warning(
+                    "align/eval/ changed since this ledger was written (%s -> %s); "
+                    "invalidated %s. The golden regression certifies the eval "
+                    "stack, so its pass does not carry across a change to it.",
+                    recorded_hash,
+                    self.eval_source_hash,
+                    sorted(dropped),
+                )
+
+    def _save(self) -> None:
+        """Persist. RANK 0 ONLY -- see `record`."""
+        if not _is_rank0():
+            return
+        tmp = self.path + ".partial"
+        with open(tmp, "w") as f:
+            json.dump(
+                {"eval_source_hash": self.eval_source_hash, "gates": self._entries},
+                f,
+                indent=1,
+                sort_keys=True,
+                # Runs on EVERY record: the site with the most chances to meet
+                # a non-serializable detail (attempt-6 class).
+                default=str,
+            )
+        os.replace(tmp, self.path)
+
+    # -- recording ----------------------------------------------------------
+
+    def record(self, gate: str, passed: bool, detail: str = "") -> dict:
+        if gate not in GATE_ORDER:
+            raise GateOrderError(f"unknown gate {gate!r}; the sequence is {list(GATE_ORDER)}")
+        idx = GATE_ORDER.index(gate)
+        for earlier in GATE_ORDER[:idx]:
+            if not self.is_passed(earlier):
+                raise GateOrderError(
+                    f"refusing to record {gate!r}: {earlier!r} has not passed "
+                    f"(it reads {self.status()[earlier]!r}). The stage sequence "
+                    f"is fixed at {list(GATE_ORDER)} and each gate exists to "
+                    "stop the next one from wasting GPU-hours. Run the earlier "
+                    "gate, or fix why it failed."
+                )
+        entry = {
+            "passed": bool(passed),
+            "detail": detail,
+            "recorded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        }
+        # EVERY rank updates its own in-memory view; only RANK 0 persists.
+        #
+        # THE LEDGER RENAME RACE (killed attempt 3, past the smoke): all four
+        # ranks must run the gates -- the 7B smoke needs the full FSDP
+        # topology -- so all four called this, all four wrote the SAME
+        # `.partial`, the first rename consumed it and the other three died on
+        # FileNotFoundError. A gate is a property of the JOB, not of a rank.
+        #
+        # NO BARRIER, deliberately, and this is the safer half of the choice.
+        # Every rank keeps an identical in-memory ledger (the order checks and
+        # the passed flag are deterministic from local state), so no rank ever
+        # needs to READ the file mid-run -- which means no collective here,
+        # which means no way for a rank-0 raise to strand its peers waiting on
+        # one. A barrier would have re-introduced the NCCL-timeout shape to
+        # fix a filesystem race. Gate failures are themselves deterministic
+        # across ranks (same config, same files, host-wide MemAvailable), so
+        # all ranks raise together and torchrun propagates it.
+        self._entries[gate] = entry
+        self._save()
+        logger.info("gate %s: %s (%s)", gate, "PASS" if passed else "FAIL", detail or "-")
+        return entry
+
+    # -- reading ------------------------------------------------------------
+
+    def is_passed(self, gate: str) -> bool:
+        return self._entries.get(gate, {}).get("passed") is True
+
+    def status(self) -> dict:
+        """Every gate present; unrun gates read PENDING, never absent."""
+        return {g: (self._entries[g] if g in self._entries else PENDING) for g in GATE_ORDER}
+
+    def next_gate(self) -> Optional[str]:
+        for gate in GATE_ORDER:
+            if not self.is_passed(gate):
+                return gate
+        return None
+
+    def require(self, gate: str) -> None:
+        if not self.is_passed(gate):
+            raise GateOrderError(
+                f"gate {gate!r} has not passed (reads {self.status()[gate]!r}); "
+                "refusing to continue"
+            )
+
+
+# ---------------------------------------------------------------------------
+# individual gates
+# ---------------------------------------------------------------------------
+
+
+def _mem_available_gb() -> float:
+    with open("/proc/meminfo") as f:
+        for line in f:
+            if line.startswith("MemAvailable:"):
+                return float(line.split()[1]) / 1024**2
+    raise RuntimeError("/proc/meminfo has no MemAvailable line")
+
+
+def rss_preflight(cfg, mem_available_gb: Optional[float] = None) -> dict:
+    """`MemAvailable >= projected_job_RSS + headroom`, or refuse to start.
+
+    The generic 45GB launch floor is insufficient for this job class and that
+    insufficiency IS the 2026-08-12 incident: the offloaded optimizer state
+    (~87GB) arrives over the first optimizer steps, AFTER any launch-time free
+    check has passed. On a box with no swap there is no overflow -- exhausting
+    RAM wedges the host in D-state and takes unrelated runs down with it.
+    """
+    from .probes import projected_job_rss  # imports torch; kept lazy
+
+    projection = projected_job_rss(cfg)
+    projected = projection["total_gb"]
+    headroom = cfg.gates.rss_headroom_gb
+    required = projected + headroom
+    available = _mem_available_gb() if mem_available_gb is None else float(mem_available_gb)
+
+    result = {
+        "passed": available >= required,
+        "mem_available_gb": available,
+        "projected_job_rss_gb": projected,
+        "headroom_gb": headroom,
+        "required_gb": required,
+        "breakdown": projection,
+    }
+    if not result["passed"]:
+        raise RuntimeError(
+            f"RSS PREFLIGHT FAILED: MemAvailable {available:.1f}GB < projected "
+            f"job RSS {projected:.1f}GB + {headroom:.0f}GB headroom = "
+            f"{required:.1f}GB.\nBreakdown: {projection}\nThis box has NO SWAP, "
+            "so there is no overflow: the offloaded optimizer state lands after "
+            "launch and would wedge the host mid-run (2026-08-12: 2GB free, "
+            "load 619, a healthy run killed by a 5-hour NCCL stall). Wait for "
+            "the host to quiesce, or switch optimizer.name to adamw8bit, which "
+            "drops the host-state term entirely."
+        )
+    logger.info("RSS preflight ok: %.1fGB available >= %.1fGB required", available, required)
+    return result
+
+
+def schedule_check(probe_report: dict, cfg, optimizer_steps_planned: int) -> dict:
+    """Project epoch wall-clock from the measured step time, or refuse.
+
+    A measured step time that invalidates the schedule must stop the run before
+    it starts. Three hours into an epoch is too late to discover that the
+    epoch was never going to fit the window.
+    """
+    step_time = float(probe_report.get("mean_step_time_s", NAN))
+    if math.isnan(step_time):
+        raise RuntimeError(
+            "schedule check refused: the probe did not measure a step time "
+            "(mean_step_time_s is NaN). Every comparison against NaN is False, "
+            "so a naive budget check would read 'within budget' and launch. "
+            "Re-run the 20-step probe."
+        )
+    if optimizer_steps_planned <= 0:
+        raise ValueError(f"optimizer_steps_planned={optimizer_steps_planned} is not a schedule")
+
+    projected = step_time * optimizer_steps_planned
+    allowed = cfg.gates.epoch_budget_s * cfg.gates.schedule_budget_factor
+    result = {
+        "passed": projected <= allowed,
+        "mean_step_time_s": step_time,
+        "optimizer_steps_planned": optimizer_steps_planned,
+        "projected_epoch_s": projected,
+        "budget_s": cfg.gates.epoch_budget_s,
+        "allowed_s": allowed,
+    }
+    if not result["passed"]:
+        raise RuntimeError(
+            f"SCHEDULE CHECK FAILED: measured step time {step_time:.2f}s x "
+            f"{optimizer_steps_planned} optimizer steps projects an epoch "
+            f"wall-clock of {projected:.0f}s, beyond the allowed "
+            f"{allowed:.0f}s (budget {cfg.gates.epoch_budget_s:.0f}s x factor "
+            f"{cfg.gates.schedule_budget_factor}). Refusing to launch: the "
+            "GPU-hours are spent whether or not the epoch finishes. Either "
+            "raise gates.epoch_budget_s deliberately, or find the step-time "
+            "regression first."
+        )
+    logger.info("schedule ok: %.0fs projected <= %.0fs allowed", projected, allowed)
+    return result
+
+
+def quality_gate(report: dict, cfg) -> GateDecision:
+    """Does the new SFT checkpoint reach the golden anchors?
+
+    Delegates the accuracy-vs-CI comparison to section-10's
+    `assert_meets_golden` rather than restating it: that function compares
+    against the WILSON LOWER BOUND (the golden numbers are themselves
+    measurements with sampling noise, so demanding the point estimate meet
+    them would fail correct runs about half the time), and two copies of that
+    reasoning would drift.
+
+    NaN halts. That is the extraction-ceiling incident in gate form: a PERFECT
+    policy scored 0/100 for two full PPO runs because extraction was broken,
+    and an undefined accuracy that compared False would read as 'above golden'.
+    """
+    from ..eval.harness import assert_meets_golden
+
+    golden = {"math500": cfg.gates.golden_math500, "gsm8k": cfg.gates.golden_gsm8k}
+    try:
+        assert_meets_golden(report, golden)
+    except RuntimeError as exc:
+        return GateDecision(
+            passed=False,
+            halt_message=(
+                f"{exc}\n\nSTAGE HALTED. Section-12 (DPO) must not start on a "
+                "failed SFT gate. Bisect in this order before touching any "
+                "hyperparameter -- DATA (is the corpus the rebuilt one, and "
+                "does a decoded sample match the training format?), then LOSS "
+                "(does fused_ce still match the eager reference?), then "
+                "TRAINER (did the trunk-movement probe pass, and did the "
+                "schedule run in optimizer steps?), then EVAL (does the golden "
+                "regression still reproduce 0.09/0.258?). A full-epoch run "
+                "below a 0.35-epoch anchor is near-certain evidence of a "
+                "new-stack defect, not a modeling result."
+            ),
+            detail={"golden": golden, "report": report.get("per_set", {})},
+        )
+    return GateDecision(passed=True, detail={"golden": golden})
+
+
+def write_quality_gate_receipt(cfg, decision: "GateDecision", checkpoint: str = "") -> str:
+    """Write the receipt section-12 reads before it will start.
+
+    The contract is narrow on purpose: a JSON object with a BOOLEAN `passed`.
+    Section-12's `read_sft_quality_gate` refuses anything else -- absent,
+    PENDING, or a truthy string all read as "not passed", because the whole
+    point is that a DPO run cannot begin on an SFT stage that did not clear
+    its own gate.
+    """
+    path = cfg.paths.sft_quality_gate_path
+    if not _is_rank0():
+        # Same rename race as the ledger: one receipt, written by the job.
+        return path
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    payload = {
+        "passed": bool(decision.passed),
+        # `detail` is the key SECTION-12's reader quotes in its refusal
+        # message. Emitting only `halt_message` would leave their error
+        # reading "(no detail recorded)" -- technically a correct refusal,
+        # and useless to whoever has to act on it.
+        "detail": (
+            decision.halt_message.splitlines()[0]
+            if decision.halt_message
+            else "quality gate passed"
+        ),
+        "halt_message": decision.halt_message,
+        "checkpoint": checkpoint,
+        "golden": (decision.detail or {}).get("golden", {}),
+        "recorded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+    }
+    tmp = path + ".partial"
+    with open(tmp, "w") as f:
+        # default=str: this receipt is written AFTER the full epoch and both
+        # evals -- a non-serializable value here would discard the run's
+        # verdict at maximum cost (the smoke-receipt Tensor crash, attempt 6).
+        json.dump(payload, f, indent=1, sort_keys=True, default=str)
+    os.replace(tmp, path)
+    logger.info("SFT quality-gate receipt written: %s (passed=%s)", path, payload["passed"])
+    return path
+
+
+def assert_build_is_current(manifest_path: str, cfg) -> dict:
+    """Refuse to train on a data build that predates the last semantic fix.
+
+    SECTION-03 REOPENED: every SFT target rendered `# Answer` TWICE -- the raw
+    PRM800K final step carries the marker inline and the renderer prepended it
+    again. Extraction survived by taking the last match, so nothing failed
+    loudly; the cost is training on a format the extraction ceiling never
+    exercised.
+
+    Read honestly from the manifest, with no hardcoded dates: `pins_verified`
+    must be True, `pin_violations` must be empty, and `build_version` must
+    reach the configured minimum. A MISSING `build_version` raises -- an
+    unstamped build is not evidence of a current one.
+    """
+    if not os.path.isfile(manifest_path):
+        raise RuntimeError(
+            f"no data build manifest at {manifest_path}. The stage consumes "
+            "shard hashes from the build artifact and never rebuilds data; "
+            "run scripts/build_data.py first."
+        )
+
+    # Section-03 drops a DEFECTIVE marker beside a build it has condemned.
+    # Checked FIRST and independently of build_version: a human-written "do
+    # not train on this" is the most direct signal there is, and honouring it
+    # costs one stat() while ignoring it costs a 4-hour epoch.
+    defective = os.path.join(os.path.dirname(os.path.abspath(manifest_path)), "DEFECTIVE")
+    if os.path.isfile(defective):
+        with open(defective) as f:
+            reason = f.read().strip()
+        raise RuntimeError(
+            f"the data build at {os.path.dirname(defective)} is marked "
+            f"DEFECTIVE by section-03:\n\n{reason}\n\nRefusing to train on "
+            "condemned shards."
+        )
+
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+
+    if manifest.get("pins_verified") is not True:
+        raise RuntimeError(
+            f"{manifest_path}: pins_verified is "
+            f"{manifest.get('pins_verified')!r}, not True. The corpus counts "
+            "were not corroborated against their pinned values, so the stage "
+            "cannot say what it would be training on."
+        )
+    violations = manifest.get("pin_violations") or []
+    if violations:
+        raise RuntimeError(
+            f"{manifest_path}: pins_verified is True but pin_violations is "
+            f"non-empty: {violations}. Trust the violations, not the flag."
+        )
+
+    version = manifest.get("build_version")
+    minimum = cfg.gates.min_build_version
+    if version is None:
+        raise RuntimeError(
+            f"{manifest_path} has no `build_version` field, so the stage "
+            "cannot tell whether this corpus predates the doubled-`# Answer`-"
+            "marker fix. An unstamped build is NOT assumed current: section-03 "
+            "stamps build_version (>= 2 after the fix; 1 is everything before "
+            "it) and this stage requires at least "
+            f"{minimum}. Rebuild the corpus, or add the field if the rebuild "
+            "already happened."
+        )
+    if int(version) < minimum:
+        raise RuntimeError(
+            f"{manifest_path}: build_version {version} < required {minimum}. "
+            "This corpus predates the doubled-`# Answer`-marker fix -- every "
+            "SFT target renders the marker twice, which trains a format the "
+            "extraction ceiling never exercised. Rebuild before spending "
+            "GPU-hours on it."
+        )
+    logger.info("data build ok: build_version=%s, pins verified", version)
+    return manifest
+
+
+def assert_distinct_devices() -> dict:
+    """Every rank must own a DIFFERENT physical GPU. Microseconds, at startup.
+
+    THE DUPLICATE-GPU DEFECT. torchrun starts four processes and assigns no
+    devices; nothing bound a rank to its LOCAL_RANK, so every `.cuda()` landed
+    on device 0 and NCCL died with "Duplicate GPU detected : rank 2 and rank 0
+    both on CUDA device 1b0" -> ncclInvalidUsage. The symptom appeared deep in
+    a 7B forward+backward and read like a sharding or dtype problem; the cause
+    was four processes sharing one card.
+
+    Two checks, and the second is the one that matters. `current_device() ==
+    local_rank()` catches the missing bind. But the REAL invariant is that the
+    four devices are physically distinct -- a correct-looking local index can
+    still collide under an unexpected CUDA_VISIBLE_DEVICES -- so the ranks
+    all-gather their device UUIDs and assert uniqueness. UUID rather than
+    pci_bus_id: `get_device_properties(i).pci_bus_id` returns the local index
+    (0 here), which is exactly the number that would agree while the hardware
+    collides.
+
+    Raises with the full rank -> device map, so the failure names itself
+    instead of surfacing 400 lines later as an NCCL error code.
+    """
+    import torch
+
+    from ..utils.dist import local_rank, world_size
+
+    if not torch.cuda.is_available():
+        return {"checked": False, "reason": "no CUDA device visible"}
+
+    current, lr = torch.cuda.current_device(), local_rank()
+    uuid = str(getattr(torch.cuda.get_device_properties(current), "uuid", current))
+    ws = world_size()
+
+    import torch.distributed as dist
+
+    distributed = ws > 1 and dist.is_available() and dist.is_initialized()
+
+    # THE COLLECTIVE COMES FIRST, and my own both-directions test is why. My
+    # first version raised the local mismatch BEFORE gathering: under a real
+    # collision ranks 1-3 raised immediately while rank 0 (whose
+    # current_device DID equal its local_rank) entered the all-gather alone
+    # and waited forever for peers that had already died. A guard against a
+    # hang must not be able to hang. So every rank enters the collective
+    # unconditionally, then all of them evaluate the SAME gathered map and
+    # raise identically.
+    gathered: list = [(0, lr, current, uuid)]
+    if distributed:
+        # GLOO, not the default NCCL group. NCCL cannot complete a collective
+        # on ranks that share a device -- it raises ncclInvalidUsage -- so a
+        # guard that gathers over NCCL reports the very failure it exists to
+        # name, in NCCL's words instead of ours. gloo is CPU-side and does not
+        # care about GPU binding, so the duplicate map survives to the error
+        # message. Verified: with all four ranks forced onto device 0 this
+        # prints the rank->device map; over NCCL it printed "invalid usage".
+        group = getattr(assert_distinct_devices, "_guard_group", None)
+        if group is None:
+            group = dist.new_group(backend="gloo")
+            assert_distinct_devices._guard_group = group
+        gathered = [None] * ws
+        dist.all_gather_object(gathered, (dist.get_rank(), lr, current, uuid), group=group)
+
+    device_map = {r: {"local_rank": l, "current_device": c, "uuid": u} for r, l, c, u in gathered}
+
+    mismatched = {r: v for r, v in device_map.items() if v["current_device"] != v["local_rank"]}
+    if mismatched:
+        raise RuntimeError(
+            f"rank binding is wrong on {sorted(mismatched)}: "
+            f"{mismatched}. Call align.utils.dist.init_distributed() BEFORE any "
+            "CUDA op -- it does init_process_group + set_device(local_rank) -- "
+            "and place tensors with torch.device('cuda', local_rank()), never "
+            "a bare .cuda()."
+        )
+
+    uuids = [v["uuid"] for v in device_map.values()]
+    if len(set(uuids)) != len(uuids):
+        collisions: dict = {}
+        for r, v in device_map.items():
+            collisions.setdefault(v["uuid"], []).append(r)
+        dupes = {u: rs for u, rs in collisions.items() if len(rs) > 1}
+        raise RuntimeError(
+            f"DUPLICATE GPU: {len(uuids)} ranks are sharing "
+            f"{len(set(uuids))} physical device(s). Ranks per device: {dupes}. "
+            f"Full map: {device_map}. Every rank must bind its own card before "
+            "any CUDA allocation; this is the ncclInvalidUsage that surfaces "
+            "later as a mysterious failure inside the model's forward."
+        )
+    logger.info("device binding ok: %d rank(s) on %d distinct GPU(s)", len(uuids), len(set(uuids)))
+    return {"checked": True, "world_size": len(uuids), "device_map": device_map}
+
+
+def derived_global_batch_rows(cfg) -> int:
+    """Rows per OPTIMIZER step, derived from the numbers that actually
+    produce them: grad_accum x micro_batch_size x world_size."""
+    return cfg.optimizer.grad_accum * cfg.train.micro_batch_size * cfg.train.world_size
+
+
+def assert_batch_arithmetic_consistent(cfg) -> dict:
+    """`global_batch_rows` must equal accum x micro x world, or refuse.
+
+    THE TWO-HORIZONS DEFECT (found pre-launch). configs/sft.toml carried
+    grad_accum 4 x micro_batch 3 x world 4 = 48 rows per optimizer step while
+    declaring global_batch_rows = 24. Nothing reconciled them, and the two
+    numbers feed DIFFERENT consumers: the banner's reported horizon is
+    ceil(rows / global_batch_rows), while the Trainer's LR schedule is
+    ceil(len(sampler) x epochs / accum) and never reads global_batch_rows at
+    all. An LR schedule that ends at half or a third of its reported horizon
+    is the iterations-vs-steps incident wearing a new coat -- the run finishes
+    at a learning rate nobody chose, and the banner records the intended one.
+
+    This makes the config internally consistent. It does NOT by itself make
+    the banner agree with the Trainer -- see `assert_horizon_agrees`, because
+    the length-grouped sampler batches by TOKEN budget, so its batch count is
+    not rows/global_batch under any value.
+    """
+    derived = derived_global_batch_rows(cfg)
+    declared = cfg.train.global_batch_rows
+    if declared != derived:
+        raise RuntimeError(
+            f"batch arithmetic is inconsistent: train.global_batch_rows="
+            f"{declared} but grad_accum({cfg.optimizer.grad_accum}) x "
+            f"micro_batch_size({cfg.train.micro_batch_size}) x world_size("
+            f"{cfg.train.world_size}) = {derived}. One of these is wrong and "
+            "the stage cannot tell which: they drive DIFFERENT consumers "
+            "(the banner's reported horizon vs the optimizer step the trainer "
+            "actually takes), so a mismatch means the LR schedule ends "
+            f"{max(declared, derived) / min(declared, derived):.1f}x away from "
+            "the horizon the banner reports. Set global_batch_rows to "
+            f"{derived}, or change accum/micro/world so the product is "
+            f"{declared}."
+        )
+    return {"global_batch_rows": declared, "derived": derived}
+
+
+def assert_horizon_agrees(sampler_horizon: int, cfg, tolerance: float = 0.0) -> dict:
+    """The banner's horizon and the Trainer's must be the SAME number.
+
+    The Trainer schedules on `ceil(len(sampler) x epochs / accum)`; the banner
+    reports `ceil(rows / global_batch_rows) x epochs`. Because the sampler
+    batches by token budget rather than row count, these agree only by
+    coincidence. Rather than pick one silently, this refuses on disagreement
+    and prints both -- an LR horizon is a pre-registered quantity and choosing
+    it is a human decision, not a rounding convention.
+    """
+    from .stage_smoke import _sft_rows
+
+    config_horizon = scheduler_horizon_optimizer_steps(
+        rows=_sft_rows(cfg), global_batch=cfg.train.global_batch_rows, epochs=cfg.train.epochs
+    )
+    if sampler_horizon != config_horizon:
+        spread = max(sampler_horizon, config_horizon) / max(1, min(sampler_horizon, config_horizon))
+        raise RuntimeError(
+            f"HORIZON MISMATCH: the trainer will schedule over "
+            f"{sampler_horizon} optimizer steps (len(sampler) x epochs / "
+            f"accum) but the config projects {config_horizon} "
+            f"(rows / global_batch_rows x epochs) -- {spread:.2f}x apart. The "
+            "LR schedule is a pre-registered quantity; the stage refuses to "
+            "guess which horizon was intended. Reconcile global_batch_rows "
+            "with the sampler's token budget, or set the horizon explicitly."
+        )
+    return {"horizon_optimizer_steps": sampler_horizon}
+
+
+def stamp_golden_report(path: str, repo_root: Optional[str] = None) -> str:
+    """Stamp the golden report with the hash of the eval source that PRODUCED it (audit U01-9).
+
+    Called by `scripts/eval.py --mode=golden` right after the run (also when the gate raised: the
+    failing report is evidence too). `record_golden_from_report` refuses a report whose stamp is
+    not the current eval source, so editing align/eval/ and re-recording from the old report
+    cannot re-certify the new code. Atomic rewrite; returns the stamped hash.
+    """
+    h = eval_source_hash(repo_root)
+    with open(path, "rb") as f:
+        report = json.loads(f.read())
+    report["eval_source_hash"] = h
+    tmp = path + ".partial"
+    with open(tmp, "w") as f:
+        json.dump(report, f, indent=1, default=str)
+        f.flush()
+        os.fsync(f.fileno())
+    os.replace(tmp, path)
+    return h
+
+
+def record_golden_from_report(
+    ledger: GateLedger, report_path: Optional[str] = None, cfg=None
+) -> dict:
+    """Bridge the golden-regression RUN to the gate LEDGER, by re-deciding.
+
+    THE MISSING-BRIDGE DEFECT (pre-launch): `scripts/eval.py --mode=golden`
+    writes its report and passes or raises IN PROCESS; nothing ever wrote the
+    verdict into the ledger. `require_golden_receipt` reads
+    `ledger.is_passed("golden_regression")`, so after a PASSING golden run the
+    SFT stage would still have refused on PENDING.
+
+    VERIFY FROM EVIDENCE, NOT FROM A FLAG. This does not take "the run
+    exited 0" as the receipt -- it loads the persisted report and re-runs the
+    verdict through section-10's own `assert_meets_golden`, so there is ONE
+    authority for what meeting the golden numbers means. A report that no
+    longer passes cannot leave a stale PASS in the ledger.
+
+    Three outcomes, and the middle one is the one that gets skipped:
+      * report passes  -> record PASSED, with the path and sha256 in the
+        detail so the receipt names the exact bytes it certified;
+      * report FAILS   -> record FAILED with the failure text. Failing is not
+        not-yet-run: a recorded failure BLOCKS every later gate, whereas
+        leaving it unrecorded would let a rerun quietly try again;
+      * report absent  -> record NOTHING, leaving the gate PENDING. The run
+        has not happened; inventing a verdict from a missing file is the
+        failure this whole gate exists to prevent.
+
+    sha256 is computed inline rather than via `receipts.sha256_file` to keep
+    this module importable above the resource guards -- see the module header.
+    """
+    from ..eval.golden import GOLDEN, GOLDEN_N, assert_meets_golden
+
+    path = report_path or (cfg.paths.golden_report_path if cfg is not None else None)
+    if not path:
+        raise ValueError("record_golden_from_report needs a report path or a cfg")
+
+    if not os.path.isfile(path):
+        logger.warning(
+            "golden report %s is absent; leaving golden_regression PENDING. "
+            "Run `scripts/eval.py --mode=golden` before the stage.",
+            path,
+        )
+        return {"recorded": False, "passed": None, "path": path, "reason": "report absent"}
+
+    with open(path, "rb") as f:
+        raw = f.read()
+    digest = hashlib.sha256(raw).hexdigest()
+    report = json.loads(raw)
+
+    # U01-9: the ledger drops the golden pass when align/eval/ changes (`GateLedger._load`), but a
+    # report produced by the OLD source would otherwise be re-recorded as a fresh PASS here. The
+    # report must name the source that produced it, and it must be the source we are running.
+    stamped = report.get("eval_source_hash") if isinstance(report, dict) else None
+    current = ledger.eval_source_hash or eval_source_hash()
+    if stamped is None:
+        logger.warning(
+            "golden report %s carries no eval_source_hash (it predates producer stamping): "
+            "cannot verify it was produced by the current align/eval/ source",
+            path,
+        )
+    elif stamped != current:
+        logger.error(
+            "golden report %s was produced by eval source %s but align/eval/ is now %s; refusing to "
+            "re-certify it. Re-run `scripts/eval.py --mode=golden`.",
+            path,
+            stamped,
+            current,
+        )
+        return {
+            "recorded": False,
+            "passed": None,
+            "path": path,
+            "reason": f"golden report produced by a different eval source ({stamped} != {current})",
+        }
+
+    try:
+        assert_meets_golden(report, GOLDEN, expected_n=GOLDEN_N)
+    except (RuntimeError, AssertionError) as exc:
+        detail = json.dumps({"report": path, "sha256": digest, "failure": str(exc)[:600]})
+        ledger.record("golden_regression", passed=False, detail=detail)
+        logger.error("golden regression FAILED per %s: %s", path, exc)
+        return {
+            "recorded": True,
+            "passed": False,
+            "path": path,
+            "sha256": digest,
+            "failure": str(exc),
+        }
+
+    detail = json.dumps(
+        {
+            "report": path,
+            "sha256": digest,
+            "golden": GOLDEN,
+            "expected_n": GOLDEN_N,
+            "per_set": report.get("per_set", {}),
+        },
+        default=str,
+    )
+    ledger.record("golden_regression", passed=True, detail=detail)
+    logger.info("golden regression PASSED per %s (sha256 %s)", path, digest[:12])
+    return {"recorded": True, "passed": True, "path": path, "sha256": digest}
+
+
+def require_golden_receipt(ledger: GateLedger, cfg) -> None:
+    """The golden regression must have PASSED, not merely been attempted.
+
+    PENDING is never passed. This is the precondition that keeps a defective
+    eval stack from certifying everything downstream of it.
+    """
+    if not cfg.gates.require_golden_receipt:
+        logger.warning(
+            "gates.require_golden_receipt is False -- training will "
+            "start without a validated eval stack"
+        )
+        return
+    if not ledger.is_passed("golden_regression"):
+        raise GateOrderError(
+            "the golden-regression gate has not passed (it reads "
+            f"{ledger.status()['golden_regression']!r}). Refusing to train: "
+            "that run is what validates the eval stack -- verifier, prompt "
+            "template, extraction, inference path -- against numbers that are "
+            "already known. Without it, every accuracy this stage produces is "
+            "unfalsifiable. Run `scripts/eval.py --mode=golden` first."
+        )
+
+
+# ---------------------------------------------------------------------------
+# banner
+# ---------------------------------------------------------------------------
+
+
+def assemble_sft_banner(
+    cfg,
+    attention_backend: str,
+    optimizer_impl: str,
+    dataset_hashes: Optional[dict],
+    ceiling_results: Optional[dict],
+    probe_report: Optional[dict],
+    rss_projection: Optional[dict],
+    mde: Optional[dict],
+    scheduler_horizon_optimizer_steps: Optional[int],
+    scheduler_horizon_rows_derived: Optional[int] = None,
+) -> dict:
+    """Every mandatory banner value, resolved or explicitly PENDING/NaN.
+
+    `beta` is NaN for SFT rather than 0.0: 0.0 is a legitimate DPO beta, so a
+    zero here would be indistinguishable from a real setting.
+
+    THE HORIZON IS THE TRAINER'S, NOT THE CONFIG'S. `scheduler_horizon_
+    optimizer_steps` must be the sampler-derived number the scheduler will
+    actually run; the rows/global_batch projection rides beside it under a
+    LABELED secondary key. A banner that records a horizon the scheduler does
+    not run is the config-translation value-drift class in its purest form --
+    the record looks right and the run is something else. A mismatch is LOGGED
+    loudly here rather than raised; the hard refusal
+    (`assert_horizon_agrees`) is wired after the first SFT run so it cannot
+    become a fresh launch blocker.
+    """
+    banner = {
+        "attention_backend": attention_backend,
+        "loss_type": "sft_ce",
+        "beta": NAN,
+        "lr": cfg.optimizer.lr,
+        "scheduler_horizon_optimizer_steps": (
+            scheduler_horizon_optimizer_steps
+            if scheduler_horizon_optimizer_steps is not None
+            else PENDING
+        ),
+        "dataset_hashes": dataset_hashes,
+        "ceiling_results": ceiling_results,
+        "rss_budget": rss_projection,
+        "pinned_staging_gb": (rss_projection or {}).get("pinned_staging_gb", PENDING),
+        "probe_step_time_s": (probe_report or {}).get("mean_step_time_s", PENDING),
+        "mde": mde,
+        "optimizer_impl": optimizer_impl,
+        "grad_accum": cfg.optimizer.grad_accum,
+        "probe_p90_step_time_s": (probe_report or {}).get("p90_step_time_s", PENDING),
+        "scheduler_horizon_rows_derived": (
+            scheduler_horizon_rows_derived
+            if scheduler_horizon_rows_derived is not None
+            else PENDING
+        ),
+    }
+    if (
+        scheduler_horizon_rows_derived is not None
+        and scheduler_horizon_optimizer_steps is not None
+        and scheduler_horizon_rows_derived != scheduler_horizon_optimizer_steps
+    ):
+        logger.warning(
+            "HORIZON DIVERGENCE: the scheduler will run %s optimizer steps "
+            "(sampler-derived, AUTHORITATIVE) but rows/global_batch_rows "
+            "projects %s -- %.2fx apart. The banner records the former "
+            "because that is what runs. Reconcile global_batch_rows with the "
+            "sampler's token budget; the hard refusal lands after this run.",
+            scheduler_horizon_optimizer_steps,
+            scheduler_horizon_rows_derived,
+            max(scheduler_horizon_optimizer_steps, scheduler_horizon_rows_derived)
+            / max(1, min(scheduler_horizon_optimizer_steps, scheduler_horizon_rows_derived)),
+        )
+    # Raises rather than printing a partial banner: a run that cannot answer
+    # its own question does not begin.
+    assert_banner_complete(banner)
+    return banner
+
+
+def assert_probe_resolved(banner: dict) -> None:
+    """A stage that REQUIRES the probe refuses to launch on PENDING.
+
+    Banner completeness and banner resolution are different properties: every
+    key being present is not the same as every value having been measured.
+    """
+    pending = [
+        k
+        for k in ("probe_step_time_s", "rss_budget", "scheduler_horizon_optimizer_steps")
+        if banner.get(k) == PENDING or banner.get(k) is None
+    ]
+    if pending:
+        raise RuntimeError(
+            f"refusing to launch: {pending} still read PENDING. The 20-step "
+            "probe and the RSS projection must be MEASURED before the epoch "
+            "starts -- their whole purpose is to stop a run that cannot fit "
+            "its schedule or its host RAM."
+        )
+    step_time = banner["probe_step_time_s"]
+    if isinstance(step_time, float) and math.isnan(step_time):
+        raise RuntimeError(
+            "probe_step_time_s is NaN: the probe ran but measured nothing. "
+            "Refusing to launch on an unmeasured schedule."
+        )
+
+
+# ---------------------------------------------------------------------------
+# CLI contract + stage drivers (thin; scripts/sft.py owns the side effects)
+# ---------------------------------------------------------------------------
+
+
+def parse_args(argv: Optional[list[str]] = None):
+    """`torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml
+    --optimizer.lr=1e-5` -- first positional is the TOML, the rest are dotted
+    overrides handed to section-01's loader (struct mode on, so a typo'd key
+    raises at launch instead of training with the default it failed to
+    override).
+
+    Parsed by hand rather than with argparse: `--optimizer.lr=1e-5` is not an
+    argparse option, it is a dotted-override entry (`section.field=value`), and argparse rejects
+    every one of them as unrecognized.
+    """
+    import sys as _sys
+
+    argv = list(_sys.argv[1:] if argv is None else argv)
+    if "-h" in argv or "--help" in argv:
+        print(__doc__)
+        raise SystemExit(0)
+
+    positional = [a for a in argv if not a.startswith("-")]
+    overrides = [a.lstrip("-") for a in argv if a.startswith("-")]
+    if not positional:
+        raise SystemExit(
+            "usage: sft.py <config.toml> [--section.field=value ...]\n"
+            "The stage TOML is required: a stage that picks its own defaults "
+            "is a stage nobody can reproduce."
+        )
+    if len(positional) > 1:
+        raise SystemExit(
+            f"expected exactly one config TOML, got {positional}. Extra "
+            "positional arguments are silently ignored by most launchers, "
+            "which is how a run ends up using the wrong config."
+        )
+    bad = [o for o in overrides if "=" not in o]
+    if bad:
+        raise SystemExit(f"overrides must be dotted key=value pairs; got {bad}")
+    return type("Args", (), {"config": positional[0], "overrides": overrides})()
+
+
+def run_entry_gates(
+    cfg,
+    ledger: GateLedger,
+    smoke_steps: int = 10,
+    probe_steps: int = 20,
+    mem_available_gb: Optional[float] = None,
+    optimizer_steps_planned: Optional[int] = None,
+) -> dict:
+    """The M5 entry sequence, in order: 7B smoke -> 20-step probe -> RSS
+    preflight. GPU-only; imported lazily so this module stays torch-free."""
+    from .stage_smoke import run_smoke_and_probe  # GPU-only helpers
+
+    require_golden_receipt(ledger, cfg)
+    assert_batch_arithmetic_consistent(cfg)
+
+    smoke = run_smoke_and_probe(cfg, smoke_steps=smoke_steps, probe_steps=probe_steps)
+    ledger.record("smoke_7b", passed=True, detail=json.dumps(smoke["smoke"], default=str))
+    ledger.record("probe_20step", passed=True, detail=json.dumps(smoke["probe"], default=str))
+
+    steps = optimizer_steps_planned or smoke["optimizer_steps_planned"]
+    schedule_check(smoke["probe"], cfg, optimizer_steps_planned=steps)
+    rss = rss_preflight(cfg, mem_available_gb=mem_available_gb)
+    ledger.record(
+        "rss_preflight",
+        passed=True,
+        detail=json.dumps({k: v for k, v in rss.items() if k != "breakdown"}),
+    )
+
+    return assemble_sft_banner(
+        cfg,
+        attention_backend=smoke["attention_backend"],
+        optimizer_impl=smoke["optimizer_impl"],
+        dataset_hashes=smoke["dataset_hashes"],
+        ceiling_results=smoke["ceiling_results"],
+        probe_report=smoke["probe"],
+        rss_projection=rss["breakdown"],
+        mde=smoke["mde"],
+        scheduler_horizon_optimizer_steps=steps,
+        scheduler_horizon_rows_derived=smoke.get("optimizer_steps_rows_derived"),
+    )
+
+
+def evaluate_and_gate(cfg, ledger: GateLedger, checkpoint: str) -> GateDecision:
+    """Run the post-training eval and apply the quality gate, recording both."""
+    from .stage_smoke import run_stage_eval  # GPU-only helper
+
+    report = run_stage_eval(cfg, checkpoint)
+    ledger.record("eval", passed=True, detail=json.dumps(report.get("per_set", {}), default=str))
+    decision = quality_gate(report, cfg)
+    ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400])
+    write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint)
+    return decision

code/align/train/stage_smoke.py

@@ -0,0 +1,530 @@
+"""GPU-only stage helpers: the 7B smoke, the throughput probe, the stage eval.
+
+Separated from `stage_gates.py` so that module stays importable above the
+resource guards. Everything here touches torch, the 7B weights, or the eval
+harness, so nothing here runs in the CPU suite -- the gate DECISIONS these
+feed are what the CPU tests cover.
+
+The peak-memory assertion lives here rather than in the unit tests because it
+is only meaningful at 7B on the real topology: the tensor it proves we avoid
+is `[B*T, 32768]`, which at seq 1024 x micro-batch 8 is 1.0GB in bf16 plus
+another 1.0GB of gradient. A unit-scale check cannot see that.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _load_snapshot_state(snapshot: str) -> dict:
+    """HF-layout state dict from a snapshot dir, single-file OR sharded, cast
+    to bf16 shard-by-shard.
+
+    The pinned Mistral-7B-v0.3 snapshot is THREE shards plus
+    `model.safetensors.index.json`; `transformers.from_pretrained` used to
+    resolve that. The layout logic (and the index-vs-shards completeness
+    check) now lives in `ckpt.load_hf_safetensors`, so every caller pointed at
+    a base snapshot gets the same behaviour; this wrapper only pins the dtype.
+    """
+    from ..model import checkpoint as ckpt
+
+    return ckpt.load_hf_safetensors(snapshot, dtype=torch.bfloat16)
+
+
+def build_stage_model(
+    cfg, attn_implementation: Optional[str] = None, weights_dir: Optional[str] = None
+):
+    """The real 7B model, FSDP2-sharded per block with activation
+    checkpointing. Wrapping happens HERE (the stage), never inside the
+    Trainer -- the Trainer receives an already-wrapped model, which is what
+    keeps it identical between the SFT and DPO stages.
+
+    `weights_dir` names an HF-layout checkpoint whose weights replace the
+    pinned snapshot's (the architecture still comes from the snapshot's
+    config.json, which a full-FT checkpoint shares). Without it the model is
+    the PRETRAINED BASE -- right for training, wrong for evaluating a trained
+    checkpoint (G3 audit U01-4: the stage-11 gate scored the base)."""
+    from ..model import checkpoint as ckpt
+    from ..model.hf_config import read_config
+    from ..model.mistral import NmoMistralForCausalLM
+    from . import fsdp
+
+    # Local only -- never let a stage reach the hub and start a 15GB download
+    # onto a shared box mid-launch. `resolve_snapshot_dir` accepts the
+    # configured `snapshots/` parent and resolves the one hash dir under it.
+    snapshot = ckpt.resolve_snapshot_dir(cfg.paths.model_snapshot)
+    # The pre-port path went through transformers' `from_pretrained` (and its
+    # MistralConfig defaults); `read_config` states every architecture field
+    # explicitly and refuses a config.json that omits one.
+    model_cfg = read_config(snapshot)
+    model_cfg.attn_implementation = attn_implementation or cfg.model.attn_implementation
+    # Shards are read straight from safetensors (the snapshot is bf16 on disk;
+    # the cast is a no-op there and a safety net for an fp32 checkpoint).
+    hf_state = _load_snapshot_state(weights_dir if weights_dir is not None else snapshot)
+    model = NmoMistralForCausalLM.from_state_dict(
+        model_cfg,
+        ckpt.hf_to_nmo(hf_state),
+        attn_implementation=model_cfg.attn_implementation,
+    )
+    del hf_state
+
+    # Onto THIS rank's device before sharding. The weights arrive on CPU from
+    # safetensors, and `fully_shard` on CPU parameters shards the wrong
+    # thing; the batches are already `.cuda()`. Uses the local rank explicitly
+    # rather than a bare `.cuda()`, which is device 0 for every rank.
+    from ..utils.dist import local_rank
+
+    device = torch.device("cuda", local_rank()) if torch.cuda.is_available() else None
+    if device is not None:
+        torch.cuda.set_device(device)
+        model = model.to(device)
+
+    fsdp.apply_activation_checkpointing(model)
+    fsdp.apply_fsdp(model)
+    return model
+
+
+def assert_chunked_loss_memory(cfg, model, batch) -> dict:
+    """The chunked-loss peak-memory assertion, at 7B scale.
+
+    MEASURES THE LOSS, NOT THE STEP -- and the first 7B run is why. I had
+    wrapped the whole adapter forward+backward, so the measured peak was
+    6103 MiB against a [1002, 32768] logit tensor of 125 MiB: a 40x "failure"
+    that was really the 7B model's own activations and FSDP all-gathers being
+    attributed to the loss. The gate asserts that CHUNKING controls the loss's
+    memory; feeding it the model's memory makes it unsatisfiable by any
+    implementation, which is the same unsatisfiable-gate trap section-05
+    already hit once with the lm_head gradient.
+
+    So the hidden states are produced OUTSIDE the measured region and
+    detached, and the lm_head weight is materialised once via `full_tensor()`
+    (it is a DTensor under FSDP2 and the fused loss needs the whole thing).
+    What remains inside the measurement is exactly the chunked loss and its
+    backward.
+
+    The full step's peak is still REPORTED, as `full_step_peak_bytes` -- it is
+    the number that tells us what a rank actually costs -- it just is not what
+    this gate asserts on.
+    """
+    from ..kernels.fused_ce import (
+        IGNORE_INDEX,
+        assert_chunking_saves,
+        assert_peak_memory_below_full_logits,
+        fused_linear_cross_entropy,
+        lm_head_grad_allowance,
+    )
+    from .sft_loss import make_sft_loss_adapter
+
+    input_ids, labels = batch
+    target = labels[:, 1:]
+
+    # 1. The real step, executed and measured for the record (this is also the
+    #    end-to-end wiring check: FSDP2 hooks, checkpointing, fused path).
+    torch.cuda.reset_peak_memory_stats()
+    out = make_sft_loss_adapter(require_fused=True)(model, (input_ids, labels))
+    out.loss.backward()
+    torch.cuda.synchronize()
+    full_step_peak = int(torch.cuda.max_memory_allocated())
+    model.zero_grad(set_to_none=True)
+    loss_value = float(out.loss.detach())
+    del out
+
+    # 2. The gate: the chunked loss in isolation, over detached hidden states.
+    with torch.no_grad():
+        hidden = model(input_ids=input_ids, return_hidden_states=True)
+    h = hidden[:, :-1].detach().clone().requires_grad_(True)
+    del hidden
+    # n_rows DERIVED FROM THE TENSOR, never from the config. The measured
+    # batch is 1002 rows where seq_len x micro_batch would predict 3072 --
+    # the sampler is token-budgeted -- and a budget computed from the config
+    # would silently be 3x the shape actually being measured.
+    n_rows = int(h.reshape(-1, h.shape[-1]).shape[0])
+    w = model.lm_head.weight
+    w_full = w.full_tensor() if hasattr(w, "full_tensor") else w
+    w_full = w_full.detach().clone().requires_grad_(True)
+
+    def _loss_only():
+        loss = fused_linear_cross_entropy(h, w_full, target, ignore_index=IGNORE_INDEX)
+        loss.backward()
+        return loss
+
+    report = assert_peak_memory_below_full_logits(
+        _loss_only,
+        n_rows,
+        cfg.model.vocab_size,
+        device=h.device,
+        # DTYPE-AWARE, and section-05's gate now refuses the alternative. The
+        # fp32 default (8 bytes/element) came to 1024 MiB, which EXCEEDED the
+        # entire 775.8 MiB measured peak -- so `attributable` clamped to zero
+        # and the gate passed no matter what the loss did. Under FSDP2
+        # mixed-precision the lm_head is bf16 with fp32 gradient accumulation:
+        # 6 bytes/element, 768 MiB. That leaves a real 7.8 MiB attributable
+        # against a 125.2 MiB budget -- a measurement instead of a rubber stamp.
+        # `param_dtype` is the lm_head's own dtype under FSDP2
+        # mixed-precision (bf16); `accum_dtype` is the fp32 the fused kernel
+        # accumulates d_weight in. Their sum -- 6 bytes/element, 768 MiB -- is
+        # what the module provably holds at the instant d_weight is cast to
+        # the parameter dtype, which is the peak.
+        unavoidable_bytes=lm_head_grad_allowance(
+            cfg.model.vocab_size,
+            cfg.model.hidden_size,
+            param_dtype=torch.bfloat16,
+            accum_dtype=torch.float32,
+        ),
+        label="section-11 M5 7B smoke",
+    )
+
+    # DIFFERENTIAL gate: chunked vs unchunked at the SAME shape, so every
+    # confounder cancels and no allowance is needed. An implementation that
+    # materialised the full logit tensor would show ~0 saving. This is the
+    # check that cannot be made vacuous by a mis-specified allowance, which is
+    # the failure the previous version of this gate actually had.
+    def _make_fn(chunk_size: int):
+        def _run():
+            h.grad = None
+            w_full.grad = None
+            fused_linear_cross_entropy(
+                h, w_full, target, chunk_size=chunk_size, ignore_index=IGNORE_INDEX
+            ).backward()
+
+        return _run
+
+    report["chunking_saves"] = assert_chunking_saves(
+        _make_fn,
+        n_rows=n_rows,
+        vocab=cfg.model.vocab_size,
+        chunk_rows=128,
+        device=h.device,
+        label="section-11 M5 7B smoke (differential)",
+    )
+
+    report["full_step_peak_bytes"] = full_step_peak
+    report["step_loss"] = loss_value
+    logger.info(
+        "7B smoke: loss=%.4f  full-step peak=%.0fMiB  loss-only peak=%.0fMiB",
+        loss_value,
+        full_step_peak / 2**20,
+        report.get("peak_bytes", 0) / 2**20,
+    )
+    return report
+
+
+def resolve_split_dir(shard_dir, split: str):
+    """Find the directory holding `<split>-NNNN.receipt.json`, by DISCOVERY.
+
+    THE HAND-COMPOSED-PATH DEFECT (caught at launch). I wrote
+    `ShardDataset(cfg.data.shard_dir, names=[cfg.data.sft_split])`, expecting a
+    flat `data/shards/sft-train.receipt.json`. Section-03's real layout is
+    `data/shards/sft/train/sft-train-0000.receipt.json`. Resolved by searching
+    for the receipts the store itself wrote rather than rebuilding
+    `<root>/<group>/<split>` from a naming convention, so a layout change moves
+    this by itself and a missing split reports the splits that ARE present.
+    """
+    from pathlib import Path
+
+    root = Path(shard_dir)
+    if not root.is_dir():
+        raise FileNotFoundError(
+            f"shard root {root} does not exist; the stage consumes shards "
+            "built by scripts/build_data.py and never builds them itself"
+        )
+
+    found: dict[str, "Path"] = {}
+    for receipt in root.rglob("*.receipt.json"):
+        prefix = receipt.name.split(".receipt.json")[0].rsplit("-", 1)[0]
+        found.setdefault(prefix, receipt.parent)
+    if split not in found:
+        raise FileNotFoundError(
+            f"no shards for split {split!r} under {root}. Splits present: "
+            f"{sorted(found)}. (The stage never composes this path from a "
+            "naming convention -- it looks for the receipts the store wrote.)"
+        )
+    return found[split]
+
+
+def open_split(cfg, split: Optional[str] = None):
+    """The ONE way this stage opens shards. Both the smoke batches and the
+    training sampler go through here, so they cannot disagree about layout."""
+    from ..data.store import ShardDataset
+
+    name = split or cfg.data.sft_split
+    return ShardDataset(resolve_split_dir(cfg.data.shard_dir, name))
+
+
+def build_training_sampler(cfg, rank: int | None = None, world_size: int | None = None):
+    """The sampler the Trainer will actually iterate -- THIS RANK's share.
+
+    Built HERE and reused by scripts/sft.py and scripts/pretrain.py so the
+    stage cannot time, project, and then train on three different batch
+    streams. The LR horizon is derived from THIS object's length, which is
+    therefore the per-rank micro-step count the Trainer really runs.
+
+    Under world_size > 1 the token-budgeted batch list is sharded by
+    `DistributedBatchSampler` (the same seam dpo_stage uses). Before this,
+    every rank iterated the full unsharded list: a 4-rank run averaged four
+    IDENTICAL gradients -- the throughput of one GPU at the cost of four --
+    and scheduled on a horizon 4x the per-rank work. Nothing failed; the loss
+    curve was that of a smaller-batch run.
+
+    `rank`/`world_size` default to the live environment and are explicit so
+    the inter-rank contract can be tested in one CPU process.
+    """
+    from ..data.loader import DistributedBatchSampler, LengthGroupedSampler
+    from ..utils.dist import rank as env_rank
+    from ..utils.dist import world_size as env_world_size
+
+    dataset = open_split(cfg)
+    lengths = [len(dataset[i].input_ids) for i in range(len(dataset))]
+    base = LengthGroupedSampler(
+        lengths, max_tokens=cfg.train.seq_len * cfg.train.micro_batch_size, seed=cfg.train.seed
+    )
+    r = env_rank() if rank is None else int(rank)
+    ws = env_world_size() if world_size is None else int(world_size)
+    if ws <= 1:
+        return base, dataset
+    if not 0 <= r < ws:
+        raise ValueError(f"rank {r} is outside world_size {ws}")
+    return DistributedBatchSampler(base, rank=r, world_size=ws), dataset
+
+
+def sampler_derived_horizon(cfg, sampler) -> int:
+    """The horizon the TRAINER will schedule on: the same arithmetic
+    `Trainer.__init__` uses, `ceil(len(sampler) x epochs / accum)`.
+
+    This is the authoritative number. The rows/global_batch figure is a
+    projection over a ROW count, but the sampler batches by TOKEN budget, so
+    the two agree only by coincidence -- and the banner must record the
+    schedule that runs, not the one that was estimated.
+    """
+    import math as _math
+
+    total_micro = len(sampler) * cfg.train.epochs
+    return max(1, _math.ceil(total_micro / cfg.optimizer.grad_accum))
+
+
+def run_smoke_and_probe(cfg, smoke_steps: int = 10, probe_steps: int = 20) -> dict:
+    """10-step 7B smoke (with the peak-memory assertion) then the 20-step
+    throughput probe. Returns everything the launch banner needs."""
+    from ..data.loader import assert_worker_count
+    from .offload_adamw import make_optimizer
+    from .probes import aggregate_micro_step_times, throughput_report
+    from .sft_loss import make_sft_loss_adapter
+    from .trainer import scheduler_horizon_optimizer_steps
+
+    assert_worker_count(cfg.data.num_workers)
+    model = build_stage_model(cfg)
+    optimizer = make_optimizer(cfg, [p for p in model.parameters() if p.requires_grad])
+    adapter = make_sft_loss_adapter(require_fused=True)
+
+    batches = _stage_batches(cfg, n=smoke_steps + probe_steps)
+    memory_report = assert_chunked_loss_memory(cfg, model, batches[0])
+
+    step_times: list[float] = []
+    for i, batch in enumerate(batches):
+        t0 = time.perf_counter()
+        out = adapter(model, batch)
+        (out.loss / cfg.optimizer.grad_accum).backward()
+        if (i + 1) % cfg.optimizer.grad_accum == 0:
+            optimizer.step()
+            optimizer.zero_grad()
+        torch.cuda.synchronize()
+        elapsed = time.perf_counter() - t0
+        if i >= smoke_steps:  # the smoke steps are warmup
+            step_times.append(elapsed)
+
+    # AUTHORITATIVE: what the Trainer will actually schedule on.
+    sampler, _dataset = build_training_sampler(cfg)
+    planned = sampler_derived_horizon(cfg, sampler)
+    # SECONDARY, labeled: the rows/global_batch projection, kept beside it so
+    # a divergence is visible in the record rather than silently resolved.
+    rows_derived = scheduler_horizon_optimizer_steps(
+        rows=_sft_rows(cfg), global_batch=cfg.train.global_batch_rows, epochs=cfg.train.epochs
+    )
+    probe = throughput_report(
+        # micro -> optimizer-step durations (task #25): tokens_per_step and
+        # `planned` below are per-optimizer-step units.
+        step_times_s=aggregate_micro_step_times(step_times, cfg.optimizer.grad_accum),
+        tokens_per_step=cfg.train.seq_len
+        * cfg.train.micro_batch_size
+        * cfg.optimizer.grad_accum
+        * cfg.train.world_size,
+        optimizer_steps_planned=planned,
+        offload_transfer_s=None,
+        offload_cpu_step_s=None,
+    )
+
+    return {
+        "smoke": {"steps": smoke_steps, **memory_report},
+        "probe": probe,
+        "optimizer_steps_planned": planned,
+        "optimizer_steps_rows_derived": rows_derived,
+        "attention_backend": model.attn_implementation,
+        "optimizer_impl": getattr(optimizer, "banner_tag", cfg.optimizer.name),
+        "dataset_hashes": _manifest_hashes(cfg),
+        "ceiling_results": _ceiling_results(cfg),
+        "mde": _mde(cfg),
+    }
+
+
+def run_stage_eval(cfg, checkpoint: str) -> dict:
+    """Evaluate a checkpoint on MATH-500 + GSM8K under section-10's harness."""
+    from ..eval.harness import EvalConfig, load_eval_set, run_eval
+    from ..eval.verifier import OutcomeVerifier
+
+    verifier = OutcomeVerifier()
+    reports = {}
+    for name, path, limit in (
+        ("math500", cfg.paths.math500_path, cfg.eval.math500_n),
+        ("gsm8k", cfg.paths.gsm8k_path, cfg.eval.gsm8k_n),
+    ):
+        problems = load_eval_set(path, limit=limit)
+        report = run_eval(
+            problems, _stage_generate_fn(cfg, checkpoint), verifier, EvalConfig(), model_name=name
+        )
+        # EvalReport has no `.results`; per-set scores live in `per_set`,
+        # keyed by each problem's source (one key for a single-set call).
+        # The old attribute error was latent -- this path never ran to
+        # completion before 2026-08-18.
+        reports[name] = report.per_set[name]
+    return {
+        "per_set": reports,
+        "config": {
+            "mde": _mde(cfg),
+            "n_problems": sum(r.get("n", 0) for r in reports.values()),
+            "eval_set_hash": _manifest_hashes(cfg),
+        },
+    }
+
+
+# -- small helpers the stage script fills in from the build artifact ---------
+
+
+def _stage_batches(cfg, n: int):
+    from ..data.loader import collate_sft
+    from ..utils.dist import local_rank
+
+    dataset = open_split(cfg)
+    mb = cfg.train.micro_batch_size
+    out = []
+    for i in range(n):
+        rows = [dataset[(i * mb + j) % len(dataset)] for j in range(mb)]
+        batch = collate_sft(rows, pad_token_id=cfg.data.pad_token_id)
+        # Explicit device: a bare .cuda() is device 0 for every rank unless
+        # set_device already ran, and relying on that ordering is what the
+        # duplicate-GPU guard exists to stop us doing.
+        dev = torch.device("cuda", local_rank())
+        out.append((batch.input_ids.to(dev), batch.labels.to(dev)))
+    return out
+
+
+def _sft_rows(cfg) -> int:
+    """Corpus size from the BUILD ARTIFACT, never a literal.
+
+    The plan says "~11.8K"; the build measured 12,410 (walk- and max_len-
+    dependent). The schedule horizon must follow the corpus that actually
+    exists, or the LR decays against a row count nobody trained on.
+    """
+    import json
+
+    with open(cfg.data.manifest_path) as f:
+        counts = json.load(f).get("counts", {})
+    rows = counts.get("sft_rows")
+    if not rows:
+        raise RuntimeError(
+            f"{cfg.data.manifest_path} has no counts.sft_rows; the scheduler "
+            "horizon cannot be derived and must not be guessed."
+        )
+    return int(rows)
+
+
+def _manifest_hashes(cfg) -> dict:
+    import json
+
+    with open(cfg.data.manifest_path) as f:
+        return json.load(f).get("hashes", {})
+
+
+def _ceiling_results(cfg) -> dict:
+    """The extraction ceiling, re-run against THIS stage's eval sets.
+
+    `preflight_extraction` returns NaN with a loud warning when no source
+    exists, and a NaN here is what the banner shows and the gate refuses --
+    which is the point: an outcome signal nobody ceiling-tested is how a
+    PERFECT policy scored 0/100 for two full PPO runs.
+    """
+    from ..eval.roundtrip import preflight_extraction
+
+    return {
+        "extraction_roundtrip": preflight_extraction(
+            [cfg.paths.math500_path, cfg.paths.gsm8k_path], eos="</s>", label="section-11 SFT stage"
+        )
+    }
+
+
+def _mde(cfg) -> dict:
+    from ..eval.metrics import min_detectable_effect
+
+    return {
+        "math500": min_detectable_effect(cfg.eval.math500_n),
+        "gsm8k": min_detectable_effect(cfg.eval.gsm8k_n),
+    }
+
+
+def _stage_generate_fn(cfg, checkpoint: str):
+    """Bind section-09's Engine to the harness's `generate_fn(prompts, cfg,
+    seeds)` contract, over the weights IN `checkpoint`.
+
+    THE BASE-MODEL-GATE DEFECT (G3 audit U01-4). This used to build the model
+    from `cfg.paths.model_snapshot` and use `checkpoint` only to find
+    tokenizer.json, so the post-epoch quality gate (and the receipt section-12
+    reads) scored the untouched base model; and its signature
+    `(prompts, seeds=None, **kw)` could not even accept the harness's three
+    positional arguments. Generation is the golden gate's greedy path
+    (`golden_generate_fn`), so the numbers are comparable to the anchors the
+    gate compares against.
+
+    PAIRED SEEDS. The Engine owns ONE sampling generator per instance, so
+    per-problem seeds cannot be threaded through it. Under greedy decoding
+    (the eval default) the seeds are irrelevant and dropping them is exact; a
+    SAMPLED eval would silently decorrelate the arms and invalidate every
+    McNemar comparison, so it is refused instead.
+    """
+    import os
+
+    from ..eval.golden import golden_generate_fn
+    from ..model.checkpoint import resolve_snapshot_dir
+    from ..tokenizer import load_tokenizer
+
+    # Weights-only checkpoints (pre tokenizer-copy fix) have no tokenizer.json
+    # -- that crash cost stage-11's first post-epoch eval. Vocab never changes
+    # during full-FT, so the pinned snapshot's tokenizer IS the checkpoint's.
+    # (`load_tokenizer` reads tokenizer.json directly; there is no slow
+    # sentencepiece-conversion path left to fall into.)
+    tok_src = checkpoint
+    if not os.path.isfile(os.path.join(checkpoint, "tokenizer.json")):
+        tok_src = resolve_snapshot_dir(cfg.paths.model_snapshot)
+        logger.warning(
+            "checkpoint %s has no tokenizer.json; loading tokenizer from the pinned snapshot %s",
+            checkpoint,
+            tok_src,
+        )
+    tokenizer = load_tokenizer(tok_src)
+    model = build_stage_model(cfg, weights_dir=checkpoint)
+    greedy_fn = golden_generate_fn(model, tokenizer)
+
+    def generate_fn(prompts, eval_cfg, seeds):
+        if eval_cfg.temperature != 0.0:
+            raise ValueError(
+                f"stage eval is greedy-only: temperature={eval_cfg.temperature} needs the "
+                "paired per-problem seeds, which the Engine cannot thread (one generator "
+                "per engine); dropping them would invalidate every paired comparison"
+            )
+        return greedy_fn(prompts, eval_cfg, seeds)
+
+    return generate_fn

code/align/train/train_module.py

@@ -0,0 +1,283 @@
+"""The TrainModule: model + loss adapter + optimizer behind one seam.
+
+This is the OLMo-core separation (olmo_core/train/train_module/train_module.py)
+ported as a PATTERN, not as code: the Trainer owns the loop, the schedule, the
+checkpoint orchestration, the instruments and history.jsonl; the TrainModule
+owns everything that touches the model -- forward, backward, the gated
+optimizer step, and the model/optimizer state movement that FSDP2 makes
+non-trivial. The Trainer never calls `.backward()` or `.step()` itself and
+never inspects a DTensor; the TrainModule never opens a file, never logs a
+metric, and never decides WHEN anything happens.
+
+Why the seam matters here and not just aesthetically: stage-b GRPO (M-C/M-D)
+needs a second module -- same Trainer loop, same checkpoint orchestration,
+same tripwires-and-history machinery, different train_batch (rollout batch,
+vendored GRPO loss, rho weights). Under the old single-class Trainer that
+would have meant a second trainer or a stage flag forest; under this seam it
+is one new TrainModule subclass.
+
+Two contracts are load-bearing and inherited verbatim from the old Trainer:
+
+  * `train_batch` performs EXACTLY `(out.loss / grad_accum).backward()`. The
+    DPO decoupled-backward adapter returns a detached leaf proxy sized so
+    that this call applies the trainer's accumulation contract to gradients
+    the adapter already planted -- any cleverness here (scaling elsewhere,
+    calling backward twice, retaining graphs) silently double-scales DPO
+    gradients.
+  * The FSDP2 state movement (full gather on save, CPU-slice per-tensor load
+    on resume) keeps its exact production-proven form; both carry incident
+    history in their docstrings.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Optional, Protocol
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LossOutput:
+    """What every loss adapter returns.
+
+    `metrics` carries stage-specific values for the log (DPO's margin,
+    logps/chosen, logps/rejected). Undefined entries must be NaN, never 0.0 --
+    a 0.0 margin reads as 'no preference learned', which is a finding, not a
+    missing measurement.
+    """
+
+    loss: torch.Tensor
+    n_tokens: int
+    metrics: dict[str, float] = field(default_factory=dict)
+
+
+class LossAdapter(Protocol):
+    """The interface the TrainModule consumes. Section-05 owns the two real
+    implementations; neither the Trainer nor the TrainModule imports either.
+
+    SFT:  batch -> chunked fused CE over hidden states, so the
+          [B*T, 32768] logit tensor never materialises.
+    DPO:  batch (with precomputed reference logprobs as constants) ->
+          fused DPO loss, reporting margin and both logps in `metrics`.
+    """
+
+    def __call__(self, model: Any, batch: Any) -> LossOutput: ...
+
+
+LOSS_NORMALIZATIONS = ("micro_mean", "global_token_mean")
+
+
+def reduce_window_tokens(local_tokens: int) -> int:
+    """SUM of `local_tokens` over every rank (gloo group; identity on one process).
+
+    The denominator of the global token mean: the number of target tokens in the whole accumulation
+    window of the WHOLE world. A collective -- every rank must call it at the same program point.
+    """
+    import torch.distributed as dist
+
+    if not (dist.is_available() and dist.is_initialized()):
+        return int(local_tokens)
+    from ..utils.dist import gloo_group
+
+    g = gloo_group()
+    if g is None:
+        return int(local_tokens)
+    t = torch.tensor([int(local_tokens)], dtype=torch.int64)
+    dist.all_reduce(t, op=dist.ReduceOp.SUM, group=g)
+    return int(t.item())
+
+
+def _dp_world_size() -> int:
+    import torch.distributed as dist
+
+    return dist.get_world_size() if (dist.is_available() and dist.is_initialized()) else 1
+
+
+class TrainModule:
+    """Owns the model, the loss adapter, and the optimizer.
+
+    The default module wraps the existing SFT/DPO path unchanged; stage-b
+    GRPO subclasses `train_batch`.
+    """
+
+    def __init__(self, model, loss_adapter: LossAdapter, optimizer, grad_accum: int):
+        if grad_accum < 1:
+            raise ValueError(f"grad_accum must be >= 1, got {grad_accum}")
+        self.model = model
+        self.loss_adapter = loss_adapter
+        self.optimizer = optimizer
+        self.grad_accum = int(grad_accum)
+
+    # -- the two loop-facing calls -------------------------------------------
+
+    def train_batch(self, batch, window_tokens: Optional[int] = None) -> LossOutput:
+        """Forward + backward for one micro-batch.
+
+        Scales by accumulation so the gradient matches one large batch rather
+        than the sum of `grad_accum` batches. This exact call is the contract
+        the DPO decoupled-backward proxy honors -- see the module docstring.
+
+        `window_tokens` (None = `loss_normalization="micro_mean"`, the original
+        behaviour, byte-for-byte the line above) switches to the GLOBAL TOKEN
+        MEAN (audit U03-2): the adapter's micro-batch mean is turned back into
+        the micro-batch's summed CE (`loss * n_tokens`) and divided by the
+        target-token count of the whole window across all ranks. FSDP2/DDP then
+        AVERAGE the per-rank gradients over the world, so the local term is
+        multiplied by the world size to make that average a SUM:
+        grad = (1/N) * sum over every rank and micro-batch of grad(sum CE).
+        """
+        out = self.loss_adapter(self.model, batch)
+        if window_tokens is None:
+            (out.loss / self.grad_accum).backward()
+        else:
+            if window_tokens < 1:
+                raise ValueError(f"window_tokens must be >= 1, got {window_tokens}")
+            scale = _dp_world_size() * float(out.n_tokens) / float(window_tokens)
+            (out.loss * scale).backward()
+        return out
+
+    def optim_step(self) -> bool:
+        """Gated optimizer step + zero_grad; True iff the step was applied.
+
+        The gate (non-finite grads skip the step) lives in the optimizer
+        wrapper; this module only reports what happened so the Trainer can
+        keep its rule that the schedule counts steps that ACTUALLY happened.
+        """
+        stepped = self.optimizer.step()
+        self.optimizer.zero_grad()
+        return stepped
+
+    def zero_grads(self) -> None:
+        self.optimizer.zero_grad()
+
+    def set_train(self) -> None:
+        self.model.train()
+
+    # -- optimizer telemetry the Trainer logs --------------------------------
+
+    @property
+    def last_grad_norm(self) -> float:
+        return self.optimizer.last_grad_norm
+
+    @property
+    def skipped_steps(self) -> int:
+        return getattr(self.optimizer, "skipped_steps", 0)
+
+    # -- state movement -------------------------------------------------------
+
+    def full_model_state_dict(self):
+        """An UNSHARDED, CPU state dict that safetensors can actually write.
+
+        THE DTENSOR CHECKPOINT DEFECT (found before it could fire, one gate
+        after the ledger race). Under FSDP2 `model.state_dict()` returns
+        DTensors -- global shape, local shard -- and safetensors dies on them
+        with "Attempted to access the data pointer on an invalid python
+        storage". Measured on the tiny model at 4 ranks: a [64,64] weight
+        comes back as a DTensor holding a [16,64] shard. Nothing in the CPU
+        suite can see this; a single process has no sharding.
+
+        `get_model_state_dict(full_state_dict=True, cpu_offload=True)` is the
+        canonical gather. It costs a transient full copy in host RAM per rank,
+        which is why the write is rank-0-only rather than the gather being
+        skipped.
+        """
+        import torch.distributed as dist
+
+        if not (dist.is_available() and dist.is_initialized()):
+            return self.model.state_dict()
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+
+        return get_model_state_dict(
+            self.model, options=StateDictOptions(full_state_dict=True, cpu_offload=True)
+        )
+
+    def load_model_state_dict(self, nmo_state: dict) -> None:
+        """Load full CPU tensors into a possibly-FSDP2-sharded model.
+
+        FSDP2 params are DTensors: a plain load_state_dict copies full
+        tensors into shards and throws "mixed torch.Tensor and DTensor"
+        -- first exercised IN PRODUCTION at the stage-11 tail resume,
+        since every earlier launch resumed from nothing. And
+        set_model_state_dict(full_state_dict=True) is the wrong mirror
+        here, and so is per-tensor distribute_tensor over full tensors:
+        both put each WHOLE tensor on GPU before sharding, and the
+        varying-shape transients churn reserved memory toward
+        full-model size per rank (measured: rank 0 at 14GB against a
+        3.4GB shard; the 4-rank resume OOM'd on a card also carrying
+        peers' NCCL IPC mappings). Every rank already read the FULL
+        state from disk, so: slice on CPU, ship ONLY the local shard.
+        """
+        if any(hasattr(p, "to_local") for p in self.model.parameters()):
+            from torch.distributed.tensor import Shard, distribute_tensor
+
+            params = dict(self.model.named_parameters())
+            missing = sorted(set(params) - set(nmo_state))
+            unexpected = sorted(set(nmo_state) - set(params))
+            if missing or unexpected:
+                raise KeyError(
+                    f"resume state mismatch: missing={missing[:5]} unexpected={unexpected[:5]}"
+                )
+            with torch.no_grad():
+                for pname, p in params.items():
+                    full = nmo_state[pname].to(p.dtype)
+                    local = p.to_local()
+                    if tuple(p.placements) == (Shard(0),):
+                        ws = p.device_mesh.size()
+                        rk = p.device_mesh.get_local_rank()
+                        chunk = full.chunk(ws, dim=0)[rk]
+                        if chunk.shape != local.shape:
+                            raise RuntimeError(
+                                f"{pname}: local shard {tuple(local.shape)} != "
+                                f"chunk {tuple(chunk.shape)} (uneven sharding?)"
+                            )
+                        local.copy_(chunk.to(local.device))
+                    else:
+                        p.copy_(
+                            distribute_tensor(full, p.device_mesh, p.placements, src_data_rank=None)
+                        )
+        else:
+            self.model.load_state_dict(nmo_state, strict=True)
+
+    def optim_state_dict(self) -> dict:
+        return self.optimizer.state_dict()
+
+    def load_optim_state_dict(self, sd: dict) -> None:
+        self.optimizer.load_state_dict(sd)
+
+    def reseed_optimizer_from_params(self) -> bool:
+        """After a resume that restored weights but NO optimizer state: let an optimizer that owns an
+        fp32 master copy (OffloadAdamW) re-seed it from the restored params. True iff it did."""
+        fn = getattr(self.optimizer, "reseed_masters_from_params", None)
+        if fn is None:
+            return False
+        fn()
+        return True
+
+    # -- module-owned state that must survive a kill --------------------------
+
+    def extra_state(self) -> dict:
+        """Module-private resume state, saved under `resume.pt["extra_state"]`.
+
+        The Trainer owns weights, optimizer, schedule and stream position; it
+        cannot know about state a SUBCLASS accumulates across steps (the GRPO
+        collapse tripwire's baseline window and latch). That state used to die
+        with the process: a resume re-measured the baseline on the post-resume
+        policy, so a collapse that began before the kill was invisible for
+        ~min_steps_before_fire steps. Subclasses override this pair; the default
+        is an empty dict so the key always exists and `load_extra_state({})`
+        is a no-op, never a KeyError on an older checkpoint's neighbour.
+
+        Contract: only tensors, dicts, lists and primitives. `resume.pt` is
+        read back with `weights_only=True`, so anything else would make the
+        checkpoint unloadable -- fail at SAVE time, not at the resume you
+        needed it for.
+        """
+        return {}
+
+    def load_extra_state(self, sd: dict) -> None:
+        """Inverse of `extra_state`. Called BEFORE the scheduler/stream restore."""
+        return None

code/align/utils/dist.py

@@ -0,0 +1,391 @@
+"""Thin distributed helpers -- env-var based so they work before
+init_process_group, plus the seeded-determinism policy and the
+wait-condition helper.
+
+No framework: torchrun sets RANK/WORLD_SIZE/LOCAL_RANK; we read them.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+
+
+def rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+def local_rank() -> int:
+    return int(os.environ.get("LOCAL_RANK", "0"))
+
+
+def world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def is_rank0() -> bool:
+    return rank() == 0
+
+
+#: Default collective timeout. torch's own default is 10 min for NCCL (30 min
+#: for gloo), which is shorter than a 14.5GB checkpoint write on this box
+#: while peers wait in the post-save barrier; 30 min is long enough for that
+#: and short enough that a genuinely wedged rank is reported in the same
+#: sitting rather than after the 5-hour stall of the 2026-08-12 incident.
+DEFAULT_PG_TIMEOUT_S = 1800.0
+
+PG_TIMEOUT_ENV = "ALIGN_PG_TIMEOUT_S"
+
+
+def pg_timeout_s(environ=None) -> float:
+    """The process-group timeout in seconds: `ALIGN_PG_TIMEOUT_S` or 30 min.
+
+    Parsed strictly. A typo ("30m") that silently fell back to the default
+    would leave an operator believing a 60-second watchdog was armed; a
+    non-positive value would make every collective time out instantly (or,
+    in torch, mean "no timeout"). Both raise instead.
+    """
+    environ = os.environ if environ is None else environ
+    raw = environ.get(PG_TIMEOUT_ENV)
+    if raw is None or raw.strip() == "":
+        return DEFAULT_PG_TIMEOUT_S
+    try:
+        val = float(raw)
+    except ValueError:
+        raise ValueError(f"{PG_TIMEOUT_ENV}={raw!r} is not a number of seconds") from None
+    if not (val > 0.0) or val == float("inf"):
+        raise ValueError(f"{PG_TIMEOUT_ENV}={raw!r} must be a finite number of seconds > 0")
+    return val
+
+
+#: The timeout `init_distributed` armed the default process group with. The
+#: gloo CONTROL group (`gloo_group`) must use the same number -- see there.
+_armed_timeout_s: float | None = None
+
+
+#: The flight-recorder buffer-size variable was renamed TORCH_NCCL_TRACE_BUFFER_SIZE ->
+#: TORCH_FR_BUFFER_SIZE when the recorder moved out of ProcessGroupNCCL (torch 2.11:
+#: FlightRecorder.hpp:111 looks up {"TORCH_FR_BUFFER_SIZE", "TORCH_NCCL_TRACE_BUFFER_SIZE"}
+#: in that order). Both are read by 2.11; only the new one is the supported name.
+_FR_BUFFER_NEW = "TORCH_FR_BUFFER_SIZE"
+_FR_BUFFER_LEGACY = "TORCH_NCCL_TRACE_BUFFER_SIZE"
+_FR_RENAME_VERSION = (2, 11)
+
+
+def _torch_version_tuple(version: str | None = None) -> tuple[int, int]:
+    """(major, minor) of the installed torch, or of an explicit version string."""
+    import re
+
+    if version is None:
+        import torch
+
+        version = torch.__version__
+    m = re.match(r"^(\d+)\.(\d+)", str(version))
+    if not m:
+        raise ValueError(f"cannot parse a torch version from {version!r}")
+    return int(m.group(1)), int(m.group(2))
+
+
+def configure_nccl_diagnostics(
+    run_dir: str | None = None, environ=None, torch_version: str | None = None
+) -> dict:
+    """Turn the PyTorch NCCL flight recorder on, with its dump under the run dir.
+
+    Must run BEFORE `init_process_group`: the recorder and the watchdog read
+    these variables when the NCCL process group is constructed. Defaults only
+    (`setdefault`): an operator's explicit export always wins. Returns what it
+    set, for the launch banner.
+
+    Variable names verified against the installed torch 2.11.0 (headers under
+    torch/include/torch/csrc/distributed/c10d/ and the strings in
+    libtorch_cpu.so / libtorch_cuda.so):
+
+      * TORCH_FR_BUFFER_SIZE (torch >= 2.11, FlightRecorder.hpp:111) or, below 2.11,
+        TORCH_NCCL_TRACE_BUFFER_SIZE (ProcessGroupNCCL.hpp:114; still read as the
+        fallback alias on 2.11). Ring-buffer entries;
+        the 2.11 default is 2000, which one FSDP2 step with 32 layers can wrap, so we
+        ask for 20000. >0 enables the recorder.
+      * TORCH_NCCL_DUMP_ON_TIMEOUT    (ProcessGroupNCCL.hpp:73). Dump the buffer when
+        the watchdog times out. Per the header it must be set together with
+        TORCH_NCCL_ENABLE_MONITORING=1 and a buffer size > 0; we do NOT set the
+        monitoring variable (its abort-on-stuck-heartbeat side effect is a separate
+        decision) and have not re-verified its 2.11 default, so on a build where it
+        is off the dump is simply not triggered.
+      * TORCH_FR_DUMP_TEMP_FILE (new name) / TORCH_NCCL_DEBUG_INFO_TEMP_FILE (legacy
+        name; both strings present in libtorch_cpu.so, neither in a header because the
+        dump writer is compiled). The value is a FILE PREFIX; torch appends the
+        global rank. Without it the dump goes to /tmp/nccl_trace_rank_<r>, which this
+        box forbids (/tmp is noexec and not durable). We set both names so either
+        generation of the reader finds it.
+
+    `run_dir=None` falls back to `<cwd>/nccl_trace` rather than the /tmp default.
+
+    THE DEPRECATED-NAME DEFECT (H2 review, 2026-10-13): on torch >= 2.11 the buffer
+    size is `TORCH_FR_BUFFER_SIZE`; the old `TORCH_NCCL_TRACE_BUFFER_SIZE` is only a
+    fallback alias, so exporting it alone arms the recorder through a name the
+    next release is free to drop -- and the recorder going dark is invisible until
+    the hang it was meant to explain. The legacy name is set ONLY below 2.11, where
+    it is the only one the build knows. An operator's explicit export of EITHER name
+    wins: the new name outranks the old in torch's lookup, so setting it over an
+    operator's legacy export would silently override their value.
+    """
+    environ = os.environ if environ is None else environ
+    new_torch = _torch_version_tuple(torch_version) >= _FR_RENAME_VERSION
+    buffer_name = _FR_BUFFER_NEW if new_torch else _FR_BUFFER_LEGACY
+    operator_set_buffer = _FR_BUFFER_NEW in environ or _FR_BUFFER_LEGACY in environ
+    trace_dir = os.path.join(run_dir or os.getcwd(), "nccl_trace")
+    os.makedirs(trace_dir, exist_ok=True)
+    prefix = os.path.join(trace_dir, "rank_")
+    wanted = {
+        buffer_name: "20000",
+        "TORCH_NCCL_DUMP_ON_TIMEOUT": "1",
+        "TORCH_FR_DUMP_TEMP_FILE": prefix,
+        "TORCH_NCCL_DEBUG_INFO_TEMP_FILE": prefix,
+    }
+    applied = {}
+    for k, v in wanted.items():
+        if k == buffer_name and operator_set_buffer:
+            continue
+        if k not in environ:
+            environ[k] = v
+            applied[k] = v
+    return applied
+
+
+def init_distributed(
+    backend: str = "nccl", run_dir: str | None = None, timeout_s: float | None = None
+) -> None:
+    """Idempotent process-group init for torchrun launches; no-op on a
+    single process.
+
+    Always passes an EXPLICIT timeout (`ALIGN_PG_TIMEOUT_S`, default 30 min)
+    instead of inheriting torch's backend-dependent default, and, for NCCL,
+    arms the flight recorder first (see `configure_nccl_diagnostics`) so a hang
+    leaves a per-rank trace of the last collectives instead of only a
+    watchdog SIGABRT.
+    """
+    if world_size() == 1:
+        return
+    import datetime
+
+    import torch.distributed as dist
+
+    if not dist.is_initialized():
+        if backend == "nccl":
+            configure_nccl_diagnostics(run_dir)
+        global _armed_timeout_s
+        t = pg_timeout_s() if timeout_s is None else float(timeout_s)
+        _armed_timeout_s = t
+        dist.init_process_group(backend=backend, timeout=datetime.timedelta(seconds=t))
+        import torch
+
+        # Only NCCL needs a device bound per rank; a gloo group on a box whose
+        # GPUs are hidden (the CPU test children) must not touch CUDA.
+        if backend == "nccl" and torch.cuda.is_available():
+            torch.cuda.set_device(local_rank())
+
+
+def barrier() -> None:
+    """No-op-safe barrier: single process or uninitialized group passes."""
+    if world_size() == 1:
+        return
+    import torch.distributed as dist
+
+    if dist.is_initialized():
+        dist.barrier()
+
+
+def seed_everything(seed: int, rank_offset: bool = True) -> int:
+    """The seeded-determinism policy.
+
+    `rank_offset=True` gives each rank a distinct-but-derived stream
+    (seed + rank) -- required for dataloader shuffling; checkpoint-resume
+    tests (section-07) and crosscoder streaming (section-08) rely on this
+    exact derivation to reproduce batch order after a restart.
+    """
+    s = seed + (rank() if rank_offset else 0)
+    random.seed(s)
+    try:
+        import numpy as np
+
+        np.random.seed(s % 2**32)
+    except ImportError:
+        pass
+    import sys
+
+    if "torch" in sys.modules:  # never import torch just to seed it
+        import torch
+
+        torch.manual_seed(s)
+    return s
+
+
+def find_pids(interpreter_substr: str, script_path: str, exclude_self: bool = True) -> list[int]:
+    """Find processes by INTERPRETER + FULL SCRIPT PATH, never bare name.
+
+    Incident: a stage-4 launcher waited on `pgrep -f <bare name>`, matched
+    its own command line, and idled 4 GPUs for 3 hours. Matching both the
+    interpreter and the absolute script path (and excluding ourselves and
+    our parent shell) cannot self-match.
+    """
+    me = os.getpid()
+    parent = os.getppid()
+    hits = []
+    for pid_dir in os.listdir("/proc"):
+        if not pid_dir.isdigit():
+            continue
+        pid = int(pid_dir)
+        if exclude_self and pid in (me, parent):
+            continue
+        try:
+            with open(f"/proc/{pid}/cmdline", "rb") as f:
+                argv = f.read().decode(errors="replace").split("\0")
+        except OSError:
+            continue
+        if len(argv) < 2:
+            continue
+        if interpreter_substr in os.path.basename(argv[0]) and any(a == script_path for a in argv):
+            hits.append(pid)
+    return hits
+
+
+def gloo_group():
+    """Cached CPU-side subgroup for object collectives and control-plane
+    gathers. None on a single process. gloo, never NCCL: object collectives
+    stage through pickled CPU tensors, and NCCL binds them to GPU state that
+    the control plane must not depend on.
+
+    THE UNARMED-WATCHDOG DEFECT. `dist.new_group(backend="gloo")` with no
+    `timeout` does NOT inherit the default group's timeout: torch fills in its
+    own per-backend default (`_get_default_timeout`, 30 min for gloo) whatever
+    `init_process_group` was given. Every stop / HALT / skip-cap decision and
+    every checkpoint gather runs on THIS group, so an operator who exported
+    `ALIGN_PG_TIMEOUT_S=120` to be told about a wedged rank within two minutes
+    still waited 30 on the optimizer-step boundary all-reduce -- the exact
+    shape of the 2026-08-12 stall, behind a knob that read as armed. The group
+    is therefore built with the timeout `init_distributed` armed the default
+    group with (or, when a caller initialised the group itself, the env value).
+    """
+    import datetime
+
+    import torch.distributed as dist
+
+    if world_size() == 1 or not dist.is_initialized():
+        return None
+    g = getattr(gloo_group, "_group", None)
+    if g is None:
+        t = _armed_timeout_s if _armed_timeout_s is not None else pg_timeout_s()
+        g = dist.new_group(backend="gloo", timeout=datetime.timedelta(seconds=t))
+        gloo_group._group = g
+    return g
+
+
+def all_reduce_max(values: list[int]) -> list[int]:
+    """Element-wise MAX of a short int vector across ranks, on the gloo group.
+
+    The control-plane primitive: "did ANY rank ...?" decisions (stop signal,
+    wall-clock, HALT, bad step) are made by reducing a 0/1 flag with MAX so that
+    every rank acts on the same answer at the same program point. gloo, not
+    NCCL, for the same reason as `gloo_group` -- and so a decision never
+    queues behind (or deadlocks with) in-flight CUDA collectives. A single
+    process returns the input unchanged.
+    """
+    import torch
+    import torch.distributed as dist
+
+    g = gloo_group()
+    if g is None:
+        return [int(v) for v in values]
+    t = torch.tensor([int(v) for v in values], dtype=torch.int64)
+    dist.all_reduce(t, op=dist.ReduceOp.MAX, group=g)
+    return [int(v) for v in t.tolist()]
+
+
+def any_rank(flag: bool) -> bool:
+    """True iff `flag` is True on at least one rank (identical on all ranks)."""
+    return bool(all_reduce_max([1 if flag else 0])[0])
+
+
+def all_gather_objects(obj) -> list:
+    """`obj` from every rank, rank-ordered, on the gloo group ([obj] on one process)."""
+    import torch.distributed as dist
+
+    g = gloo_group()
+    if g is None:
+        return [obj]
+    out: list = [None] * dist.get_world_size()
+    dist.all_gather_object(out, obj, group=g)
+    return out
+
+
+def capture_rng_state() -> dict:
+    """This process's RNG state: torch CPU, torch CUDA (current device only),
+    python `random`, numpy. Every leaf is a tensor, list, int, float, str or
+    None, because `resume.pt` is read back with `weights_only=True`.
+
+    WHY. `seed_everything` seeds the streams at launch, but a resumed process
+    re-seeds from scratch and so replays the first draws of the original run:
+    dropout masks, sampling noise and any shuffle that is not a pure function
+    of (seed, epoch) restart instead of continuing. The run is then not the run
+    its uninterrupted twin would have been. State is per rank (the streams are
+    seeded seed + rank).
+
+    CUDA: only the CURRENT device, and only if CUDA is already initialised.
+    `torch.cuda.get_rng_state_all()` would create a context on every visible
+    GPU from every rank -- under torchrun that is world_size x n_gpus contexts
+    on a box with 23GB cards.
+    """
+    import torch
+
+    py = random.getstate()
+    out: dict = {
+        "torch": torch.get_rng_state(),
+        "python": [py[0], list(py[1]), py[2]],
+    }
+    try:
+        import numpy as np
+
+        ns = np.random.get_state()
+        out["numpy"] = {
+            "name": str(ns[0]),
+            "keys": [int(k) for k in ns[1]],
+            "pos": int(ns[2]),
+            "has_gauss": int(ns[3]),
+            "cached_gaussian": float(ns[4]),
+        }
+    except ImportError:  # pragma: no cover -- numpy is a hard dependency
+        pass
+    if torch.cuda.is_available() and torch.cuda.is_initialized():
+        dev = torch.cuda.current_device()
+        out["cuda"] = {str(dev): torch.cuda.get_rng_state(dev)}
+    return out
+
+
+def restore_rng_state(sd: dict) -> None:
+    """Inverse of `capture_rng_state`. Missing sections are skipped (an older
+    checkpoint or a CPU-only save resumed on GPU)."""
+    import torch
+
+    if "python" in sd:
+        ver, keys, gauss = sd["python"]
+        random.setstate((int(ver), tuple(int(k) for k in keys), gauss))
+    if "numpy" in sd:
+        import numpy as np
+
+        n = sd["numpy"]
+        np.random.set_state(
+            (
+                n["name"],
+                np.array(n["keys"], dtype=np.uint32),
+                int(n["pos"]),
+                int(n["has_gauss"]),
+                float(n["cached_gaussian"]),
+            )
+        )
+    if "torch" in sd:
+        torch.set_rng_state(sd["torch"].cpu())
+    if "cuda" in sd and torch.cuda.is_available():
+        for idx, state in sd["cuda"].items():
+            if int(idx) < torch.cuda.device_count():
+                # lazily applied if CUDA is not initialised yet
+                torch.cuda.set_rng_state(state.cpu(), device=int(idx))