diff --git a/docs/design/iter-ldi0-01-20260717.md b/docs/design/iter-ldi0-01-20260717.md new file mode 100644 index 00000000..bed480f5 --- /dev/null +++ b/docs/design/iter-ldi0-01-20260717.md @@ -0,0 +1,71 @@ +# LDI0-01 — local-decision evidence contract recommit + +Date: 2026-07-17 +Status: **documentation and source-inventory change only; no code, no checkpoint, +no model-quality or ship claim.** + +## Scope + +LDI0-01 makes the local-decision-interventions research synthesis and architectural +boundaries canonical in the repository so future agents do not rebuild a parallel +harness, rerun the falsified E249–E284 experiments as if open, or treat a local +metric gain as promotion evidence. No event-schema or trainer code, adapter, +checkpoint, or model-card quality update is included. + +Reviewed commit: `c7cb099` (main at review time). + +## Source inventory: before / after + +| | Before | After | +| --- | --- | --- | +| Academic works | 25 | 34 | +| Implementation / doc sources | 8 | 8 | +| Total rows | 33 | 42 | + +Nine works were added (all `implementation_status = Adjacent`, none reimplemented), +with title / authors / date verified against arXiv: + +| Work | arXiv | Category | Relationship | +| --- | --- | --- | --- | +| LoRA | 2106.09685 | adapter_actuator | Adjacent | +| AdaLoRA | 2303.10512 | adapter_actuator | Adjacent | +| PCGrad (Gradient Surgery) | 2001.06782 | multi_objective_balance | Adjacent | +| MGDA (MTL as MOO) | 1810.04650 | multi_objective_balance | Adjacent | +| PICARD | 2109.05093 | constrained_decoding | Adjacent | +| Grammar-Aligned Decoding | 2405.21047 | constrained_decoding | Adjacent | +| Min-p sampling | 2407.01082 | constrained_decoding | Adjacent | +| DeepSeekMath / GRPO | 2402.03300 | verifiable_training | Adjacent | +| TAB-PO | 2603.00025 | local_preference | Adjacent | + +## Duplicate / alternate-URL handling + +- Each new `uri` is a fresh arXiv `abs` URL, checked against the existing 33 rows; + none collides. The adapter family already present (DoRA, PiSSA, MoLoRA, PermDoRA, + and the PEFT LoRA doc row) is distinct from the LoRA / AdaLoRA papers, and + DeepSeekMath / GRPO (2402.03300) is distinct from DeepSeek-R1 (2501.12948). +- No new paper duplicates an existing work under an alternate URL; the repo + convention keeps arXiv as `uri` and any secondary link in `metadata.alternate_urls`. +- `test_local_decision_source_manifest_is_complete` asserts unique `source_id` and + unique `uri` across all 42 rows, required provenance metadata on every row, and the + `{Adapted, Adjacent}` status vocabulary (nothing claimed Faithful). + +## Documents updated + +- `src/slm_training/resources/autoresearch/local-decision-sources.json` — nine rows + plus an honest `source_scope`. +- `docs/design/local-decision-interventions.md` — separation-of-concerns invariants, + owners plus the no-second-stack rule, the E249–E284 falsification chain and current + blocker, the `DecisionEventV2` first-contract note, and a dated Measured-status note. +- `docs/design/research-lineage.md` — refreshed local-decision entry; new works + labeled Adjacent. +- `docs/design/quality-experiment-matrix.md` — an LDI0 index (no claimed rows, no + reserved E-IDs). +- `tests/test_autoresearch/test_harness.py` — updated counts and stronger assertions. + +## Honesty statement + +No model or ship claim. This iteration is documentation and a source inventory. The +E249–E284 chain remains the authoritative measured record; the current blocker — +stable grammar-state support is necessary but not sufficient for the objective — is +recorded, not resolved. `python -m scripts.repo_policy` passes and the focused +source-manifest test is green. diff --git a/docs/design/iter-ldi0-02-decision-events-v2-20260717.md b/docs/design/iter-ldi0-02-decision-events-v2-20260717.md new file mode 100644 index 00000000..6a846626 --- /dev/null +++ b/docs/design/iter-ldi0-02-decision-events-v2-20260717.md @@ -0,0 +1,82 @@ +# LDI0-02 — DecisionEventV2 (state / action-evidence / objective-view) + +Date: 2026-07-17 +Status: **schema + materializers + V1 migration, with tests; no event mining, +no trainer change, no checkpoint, and no model-quality or ship claim.** + +## What and why + +`DecisionEventV2` (`src/slm_training/harnesses/preference/decision_events_v2.py`) +replaces V1's frozen `good_token_ids` / `bad_token_ids` payload as the *primary* +semantic corpus contract with a three-part model that separates concerns V1 +conflated — the fix for the E284 blocker (*stable grammar-state support does not +imply objective/action-partition support*): + +* **`DecisionStateV2`** — the exact, replayable model decision state plus immutable + runtime identities. Its `state_id` is a canonical `content_sha` over state and + identity **only** (never labels, rollout outcomes, or candidate order), and + generalizes V1's TwoTower-only `canvas_ids`+`position` shape to `twotower` and + `causal` (prefix) tracks. +* **`ActionOutcomeV2`** — append-only, per-candidate verifier evidence. The complete + ordered **G0–G12** gate vector is preserved verbatim (validated, never + set-collapsed). Merge deduplicates by content identity, not array order. +* **`ObjectiveView`** — a pure, versioned materialization of one state's action + table. Five materializers: `pareto`, `thresholded`, `single_best_worst`, + `set_valued` (all `trainable=True`), and `constraint_shadow` + (`trainable=False` — a semantic trainer must refuse legality-only evidence). + +`migrate_v1_event` performs a one-way, deterministic, idempotent V1→V2 migration +that marks partial evidence **incomplete** and never fabricates rollout/verifier +vectors. V1 corpora and their loaders are untouched. + +## Fixture-corpus demonstration + +The executable version of this walk-through is +`tests/test_harnesses/preference/test_decision_events_v2.py` (26 tests). One exact +state with two independently-observed actions: + +1. **One exact state** — a `DecisionStateV2` for a `component` decision with + `legal_action_ids = (4, 9, 10)`. Its `state_id` is fixed by the state/identity + fields; constructing it twice yields the same id + (`test_two_label_samples_of_one_state_share_one_state_id`), and it is unchanged + by how action evidence is later ordered or augmented + (`test_state_id_is_independent_of_action_evidence_order_and_augmentation`). +2. **Multiple independently replayed actions** — action `4` observed with reward + `0.9`, action `9` with reward `0.1`, each carrying its own complete G0–G12 + vector; a re-observation of identical evidence collapses under the append-only + merge, new evidence appends + (`test_append_only_merge_dedups_by_content_not_order`). +3. **Two materializers, different views from the same evidence** — `pareto` labels + `good={4}`, `bad={9}`, `unobserved={10}` and is trainable; `constraint_shadow` + labels `good={}`, `bad={}`, `ambiguous={4,9}` and is **not** trainable + (`test_pareto_view_is_trainable_and_partitions_actions`, + `test_constraint_shadow_view_is_not_trainable`). +4. **Stable state fingerprint, changing objective fingerprint** — the manifest + (`decision_state_manifest`) fingerprints states, action evidence, and objective + views separately and order-independently + (`test_manifest_fingerprints_are_separate_and_order_independent`): re-ordering + the evidence rows leaves the outcome fingerprint unchanged, while the state, + outcome, and view fingerprints differ from one another. + +## Validation (fail-closed) + +Enforced and tested: unknown record fields rejected; a tampered `state_id` +rejected; `split` must derive from `group_id`; a TwoTower state requires +`canvas_ids`, a causal state requires `context_ids`; every present verifier vector +must carry the complete ordered G0–G12 gate set with valid statuses; objective-view +partitions must be disjoint; a materializer rejects an outcome whose `state_id` +does not match the state; atomic, duplicate-safe JSONL writes. + +## Honest scope / caveats + +- **No model or ship claim.** This is a corpus/record schema plus pure + materializers and a migration. No event is mined, no trainer consumes V2 yet, no + checkpoint or eval is produced. +- **Named independent-judge rationale** beyond the G11 gate status, and a + full model-driven **logit-replay** integration fixture (versus the schema-level + causal/TwoTower state construction covered here), are the honest follow-ons; the + state records already carry the prefix/canvas + policy/tokenizer/decode/verifier + identities required to replay. +- Verified with `pytest tests/test_harnesses/preference/test_decision_events_v2.py` + (20) and `tests/test_harnesses/preference/test_local_decisions.py` (V1 unchanged, + 11); `ruff` and `python -m scripts.repo_policy` clean. diff --git a/docs/design/iter-ldi0-03-objective-support-admission-20260718.md b/docs/design/iter-ldi0-03-objective-support-admission-20260718.md new file mode 100644 index 00000000..458fbe8e --- /dev/null +++ b/docs/design/iter-ldi0-03-objective-support-admission-20260718.md @@ -0,0 +1,96 @@ +# LDI0-03 — objective-support admission and bounded diagnostics + +Date: 2026-07-18 +Status: **objective-support admission gate, bounded-diagnostic runner, the Tier-1 +objective-geometry pass, the Tier-2 refuse-full-parameter interface, and a V1 +objective-signature trainer-entry refusal all landed with tests. The V2 +`admit_semantic_corpus` CLI wiring (`--require-admission` / `--materializer`) remains +follow-on. No training, checkpoint, model-quality, or ship claim.** + +## Why this exists + +E283 repaired every held-out grammar-**state** support signature, yet E284 still +found **35 held-out objective conflicts**: the FTPO objective depends on the +sampled bad-action set, which is *not* part of the state-support signature, so a +corpus can pass state support while its objective (action-partition) support is +deficient. E285 and E286 then tried a more exact profile but blew the experiment +envelope and produced **invalid evidence**: + +- **E285** — no cumulative wall-clock deadline; ran past 25 minutes, operator-stopped, + **no report** → invalid evidence. +- **E286** — chunked batched vector-Jacobian products; killed at 283.62 s by the + five-minute envelope, **no report**; the implementation was removed rather than + retained on a fast unit test. + +Both are marked invalid historical evidence here; no metric, comparison, or training +decision may be inferred from them. + +## Support criteria compared + +| Criterion | Keys on | Catches E284? | +| --- | --- | --- | +| V1 state support (`decision_support_signature`) | decision kind + legal tokens + **good** tokens | No — excludes the sampled negatives | +| V1 objective signature (`decision_signature`) | + **bad** tokens | Diagnoses it, but not wired as an admission gate | +| **V2 objective-view support** (`objective_view_support`, this issue) | materializer id/config hash + **good/bad** partition, split-aware | **Yes** — held-out objective signatures without train support fail admission | + +The V2 gate is the fix: `admit_semantic_corpus` fails closed before training when a +corpus (a) contains a non-trainable view (e.g. the `constraint_shadow` diagnostic), +(b) carries a materializer that does not match the requested objective, or (c) lacks +train support for any held-out objective signature — naming what is missing so +support can be repaired (never by copying held-out programs). + +## Bounded runtime (the E285/E286 fix) + +`decision_diagnostics.run_bounded_stages` runs diagnostic stages under one cumulative +`DiagnosticBudget` (default and hard cap five minutes), read through a `time` +reference so a deterministic fake clock drives the tests. On expiry the run is +`expired` with **`result: None`** — no partial result is ever presented as a +diagnostic result — and a full-parameter Tier-2 request is refused as +`not_authorized` rather than replaying the invalid E285 full-parameter profile. +`write_diagnostic_report` persists reports atomically (mkstemp + fsync + os.replace). + +## Landed in this iteration + +- `src/slm_training/harnesses/preference/decision_events_v2.py`: + `objective_view_signature`, `objective_view_support`, `admit_semantic_corpus`. +- `src/slm_training/harnesses/preference/decision_diagnostics.py`: `DiagnosticBudget`, + `Deadline`, `run_bounded_stages`, `not_authorized_report`, `write_diagnostic_report`, + and — on that bounded runner — `tier1_objective_geometry` and + `tier2_subspace_gradients`. +- `tier1_objective_geometry` reads already-materialized objective views (`getattr` on + `good_action_ids` / `bad_action_ids`, so it stays decoupled from + `decision_events_v2`) and reports, per corpus, the count of states with an + objective **contradiction** (an action scored good by one view and bad by another — + the logit-space shadow of the E284 conflict) and the mean pairwise Jaccard overlap + of the per-state good-action sets. It computes no gradient and trains nothing. +- `tier2_subspace_gradients` refuses an empty / `None` (full-parameter) request as + `not_authorized` rather than replaying the invalid E285 full-parameter profile; for + an explicit adapter subset it records a bounded plan only, with the gradient + computation deferred to a model stage this module never runs. +- A V1 objective-signature pre-flight refusal now guards `train_local_from_paths` + (`require_objective_support=True`): `objective_signature_support` (keyed on the + good+bad `decision_signature`) fails the entry before the checkpoint is loaded when + a held-out objective signature lacks train support — the E284 blocker. +- Tests: the E284 pattern (passes state support, fails objective support), + constraint-shadow and materializer-mismatch refusals + (`tests/test_harnesses/preference/test_decision_events_v2.py`), the deterministic + fake-clock deadline / no-result-on-expiry / not-authorized cases plus the Tier-1 + contradiction/agreement and Tier-2 refuse/plan cases + (`tests/test_harnesses/preference/test_decision_diagnostics.py`), and the trainer + objective-support refusal (`tests/test_harnesses/preference/test_local_decisions.py`). + `ruff` and `python -m scripts.repo_policy` clean. + +## Honest remaining scope + +- Wiring the **V2** `admit_semantic_corpus` gate into `train_local_from_paths` (and a + `--require-admission` / `--materializer` CLI flag) so the semantic trainer refuses a + non-admitted corpus — a materializer mismatch or a non-trainable view — before the + first optimizer step, mirroring the existing `constraint_shadow` refusal. The V1 + objective-signature refusal already landed here covers the sampled-negative (E284) + case; the V2 gate additionally covers materializer identity. +- The Tier-2 gradient computation itself (the deferred model stage) and the richer + Tier-1 logit-space content (parent legal-space good/bad mass and margins, dominance + under raw vs unit-normalized scaling, held-out signature distances) that build on the + bounded runner landed here. +- These are the next commits. This iteration adds no token/component special cases, + runs no training, and makes no model-quality claim. diff --git a/docs/design/iter-ldi1-01-causal-capture-20260718.md b/docs/design/iter-ldi1-01-causal-capture-20260718.md new file mode 100644 index 00000000..87714e47 --- /dev/null +++ b/docs/design/iter-ldi1-01-causal-capture-20260718.md @@ -0,0 +1,105 @@ +# LDI1-01 — exact causal decision-state capture + +Date: 2026-07-18 +Status: **capture core, DecisionEventV2 emission, TraceStore persistence with a +fail-closed loader, and the plug-in traced-decode + forced-action replay landed with +tests and a fixture-grade evidence artifact. No FTPO trainer, no judge, no adapter +matrix, no training, and no model-quality or ship claim.** + +## Why this exists + +The causal OpenUI plug-in (`models/causal_lm_openui.py`) decodes under a hard grammar +constraint via `model.generate(..., prefix_allowed_tokens_fn=...)`, which hides the +per-step raw logits and the raw-vs-constrained selection. LDI1-01 needs every supervised +causal decision recoverable from **exact prefix token ids and model logits** — never from +decoded strings or later retokenization — and emitted as evidence compatible with the +shared preference harness's `DecisionEventV2` (LDI0-02, SLM-116). This iteration adds that +capture path; it deliberately trains nothing and runs no judge (those are LDI1-02/1-03). + +## What landed + +### Torch-free capture core — `models/causal_trace.py` + +The capture is split into a torch-free core driven by two injected callables so the whole +algorithm is deterministic and unit-testable without torch, a real tokenizer, or the +grammar: + +- `forward_logits(prefix_ids) -> Sequence[float]` — abstracts `model(input_ids).logits[:, -1, :]`; +- `allowed_ids(prefix_ids) -> Sequence[int]` — abstracts the plug-in's `_allowed_ids`. + +`capture_raw_steps(...)` runs a greedy constrained decode and, at each step, computes the +raw (pre-mask) argmax and the constrained selection (greedy over the legal set) from the +**same** logits, so a `constraint_shadow` — a raw winner outside the legal set overridden +by a legal selection — is exact. It returns a `CaptureResult` with the policy-retained +`RawStepObservation`s, the full emitted `generated_token_ids`, and an honest `stop_reason` +(`eos` / `no_legal_continuation` / `max_new_tokens`). Bounded selection policies +(`every`, `constraint_shadow_only`, `margin_threshold`, `sampled_positions`, `named_roles`) +choose which decisions are retained but never change which token is emitted. Forced +(single-legal) steps are recorded as deductions, not decisions, so `decision_index` counts +only real choices. + +### DecisionEventV2 emission + +`emit_causal_decision(obs, identity)` materializes a causal `DecisionStateV2` whose +`context_ids` are the exact integer prefix and whose `grammar_state_hash` is the +content-addressed legal set. Because `DecisionStateV2` carries a single +`policy_checkpoint_sha` with no dedicated adapter field, `fold_policy_identity(base, +adapter)` folds both the base checkpoint and the active adapter into it — so an +adapter-enabled and adapter-disabled capture over the same prefix receive **different** +state identities. A `constraint_shadow` step additionally yields a legality-only +`ActionOutcomeV2` (empty reward/verifier vectors) and a **non-trainable** view via +`materialize_constraint_shadow`; `admit_semantic_corpus` refuses it, so a legality +diagnostic can never supervise a semantic objective. + +### Persistence + fail-closed loader + +`CausalTraceWriter` appends each decision to the shared `TraceStore` (no second trace +format) with a distinct `kind="causal_decision"`, lifting identity hashes to the row top +level, and tracks a reproducible manifest (identities, state/shadow counts, unique legal +sets, duplicate-set reuse, bytes/state). `load_causal_decision_states(store, *, +expected_checkpoint_sha, expected_tokenizer_sha)` fails closed before returning any state +when the checkpoint or tokenizer does not match — mirroring `local_train._validate_identity` +— and `DecisionStateV2.from_dict` re-verifies each state id (tamper check). + +### Plug-in wiring — `models/causal_lm_openui.py` + +`generate_constrained_traced(...)` drives the per-step loop, storing `context_ids` as the +full prefix (prompt + generated suffix) so a consumer can replay +`model(context_ids).logits[:, -1, :]` exactly; `replay_causal_action(state, +forced_action_id, continuation_seed, ...)` applies the forced action to the exact stored +prefix, continues under the deterministic constrained policy, and returns a **pre-judge** +`GeneratedOutcome` for the shared counterfactual owner (it runs no judge itself). +`generate_constrained` (trace-off) is unchanged. An injectable `allowed_ids_fn` documents +the grammar seam and keeps the torch loop testable. + +## Tests and evidence + +- `tests/test_models/test_causal_trace.py` (torch-free): shadow fires only on illegal + raw + legal selection; shadow view non-trainable and admission-refused; forced-vs-decision + indexing; EOS only after the prefix validates; integer prefixes are state authority; + bounded policies; content-addressed legal sets; store round-trip; fail-closed load; + pre-judge candidate shape. +- `tests/test_models/test_causal_trace_plugin.py` (torch): the real torch forward through + the capture loop; **stored logits replay within tolerance**; tracing does not change + emitted tokens; reproducible forced-action replay; adapter-identity folding. +- `tests/test_models/test_causal_trace_fixture.py` + `tests/test_models/fixtures/ldi1/`: + a committed, deterministic, torch-free trace showing exact prefix replay, raw/legal + telemetry, one non-admittable constraint shadow, and one forced legal counterfactual + replay outcome (canonicalization deferred to the strict validator — the fixture does + not assert catalog-validity, which is environment/pack-dependent) — no semantic label, + no model-quality claim. + +Regenerate the fixture with `capture_raw_steps` over the logit/legal map in the fixture +`README.md`, a fixed `CausalTraceIdentity`, and a `TraceStore` opened with explicit +`run_id`/`trace_id`/`span_id` (drop the store's timestamped `manifest.json`; keep the +reproducible `causal_trace_manifest.json`). `ruff` and `python -m scripts.repo_policy` clean. + +## Honest remaining scope + +- Efficiency: `forward_logits` recomputes a full forward per step (no KV cache); capture + correctness, not decode throughput, is the goal here. +- The semantic action evidence (running `replay_causal_action` outcomes through the + counterfactual owner's judge and Pareto partition) is exercised by the owner, not this + module — LDI1-02 (causal PEFT FTPO) and LDI1-03 (the matched matrix) are the follow-ons. +- This iteration adds no token/component special cases, runs no training, and makes no + model-quality claim. diff --git a/docs/design/iter-ldi2-01-twotower-adapter-20260718.md b/docs/design/iter-ldi2-01-twotower-adapter-20260718.md new file mode 100644 index 00000000..3a856055 --- /dev/null +++ b/docs/design/iter-ldi2-01-twotower-adapter-20260718.md @@ -0,0 +1,74 @@ +# LDI2-01 — removable low-rank adapter backend for TwoTower + +Date: 2026-07-18 +Status: **adapter primitive, deterministic attachment, adapter-only training, one-way +merge, and a removable save/load directory (fail-closed on identity) landed with tests. +The base-`.pt` checkpoint interplay, the config/CLI round-trip, and the `artifact_identity` +base-vs-merged distinction are the remaining follow-on. No training run, no promoted +checkpoint, no quality claim.** + +## Why this exists + +LDI2 needs a small, removable, lineage-safe actuator on selected TwoTower denoiser +projections so a local intervention can be trained and measured against the untouched +full-update parent — without importing Hugging Face PEFT. The backend is standard +low-rank delta (not DoRA/PiSSA/AdaLoRA, which are not implemented here). + +## What landed + +- `models/adapters/spec.py` — `TwoTowerAdapterSpec`: torch-free, frozen, versioned, + validated config that round-trips through `to_dict`/`from_dict` (rejects unknown + fields), bound to a base by `base_compatibility_fingerprint` / `base_checkpoint_sha` / + `tokenizer_sha`. +- `models/adapters/low_rank.py` — `LowRankAdapter`: wraps a frozen `nn.Linear` with + `y = W x + (alpha/rank)·B(A(dropout(x)))`, `B` zero-initialized (a fresh adapter is + output-identical to the parent bit-for-bit), enable/disable, and a one-way + `merged_linear()` on a copy that never mutates the parent. +- `models/adapters/twotower_adapter.py` — deterministic target resolution mapping spec + target names (`attn_q/k/v/out`, `cross_attn_*`, `mlp_in/out`) onto the denoiser block + linears, failing closed on unknown targets, non-linear matches, and out-of-range + layers; attaches wrappers in place under a forked RNG (attachment never shifts the + training RNG). The context tower is never adapted. +- `TwoTowerModel` methods: `attach_adapter` (base-fingerprint fail-closed + freeze every + non-adapter parameter so the existing `trainable_parameters()` yields only adapter + tensors), `enable_adapter` / `disable_adapter`, `adapter_parameters`, `has_adapter`, + `active_adapter_identity` (content digest of adapter tensors), and `merge_adapter_copy` + (wrapper-free copy equal to the adapter-enabled map; original untouched). + +## Acceptance covered by tests + +Fresh enabled adapter equals the parent bit-for-bit at every adapted site; disabled +adapter equals the parent after weights change; only adapter tensors receive gradients; +requested targets resolve deterministically and missing/unsupported/out-of-range targets +raise actionable errors; base-fingerprint mismatch fails closed; attachment does not +shift training RNG; adapter-only `trainable_parameters`; merged copy matches the enabled +output and leaves the original adapter intact. Existing `test_twotower.py` unchanged. + +The removable adapter directory landed: `save_adapter(path, provenance=...)` writes +`adapter_config.json` (the spec), `adapter_model.pt` (only the lora tensors — the base +checkpoint is never duplicated), and `adapter_manifest.json` (resolved module map, +parameter names/shapes, trainable parameter count, adapter bytes, base/tokenizer +identity). `load_adapter(path, trainable=...)` validates the whole artifact **before** +mutating the model — tokenizer identity, config/manifest agreement on base fingerprint, +and the exact tensor key set and shapes against the manifest — so a mismatched or +truncated adapter never leaves the model half-attached; only then does it attach (which +fails closed on a base compatibility-fingerprint mismatch) and copy the tensors. + +`base_checkpoint_sha` is recorded for **provenance only**: attachment is gated on the +architecture-level `compatibility_fingerprint` plus `tokenizer_sha`, which do not +distinguish two different checkpoints of the *same* architecture. Rejecting an adapter +trained against a different base checkpoint requires the model to carry its loaded +checkpoint identity (it does not today — a `from_records` model has none), so that +enforcement is deferred with the checkpoint-identity work below and is not yet claimed. + +## Honest remaining scope + +- The base-`.pt` checkpoint interplay is **not** a one-line allow-list: attaching an + adapter renames the wrapped base weight key (`…q_proj.weight` → `…q_proj.base.weight`, + plus `lora_A`/`lora_B`), so `_state_dict_for_checkpoint` / `_load_checkpoint_state` + need to strip the wrapper structure (not just allow-list `adapter.` keys) for a + base checkpoint to round-trip through an adapter-capable model. This needs its own + careful pass. +- Config/CLI round-trip (`ModelBuildConfig` → attach) and the `artifact_identity` / + `compatibility_fingerprint` base-vs-merged distinction. +- These are the next commits. This iteration runs no training and makes no quality claim. diff --git a/docs/design/quality-experiment-matrix.md b/docs/design/quality-experiment-matrix.md index 2de9aa3f..104a1b06 100644 --- a/docs/design/quality-experiment-matrix.md +++ b/docs/design/quality-experiment-matrix.md @@ -1718,3 +1718,21 @@ self-distillation, trajectory RL) are in | E63 | Gate calibration | ECE / selective accuracy / abstention on `FastPathGate` | proposed | | E64 | Trajectory-aligned RL | MDPO/d1-style on intermediate MaskGIT states | proposed | | E65 | Schema generalization | Held-out schemas / rename / `toy-layout` transfer | proposed | + +## LDI (local decision interventions) index + +LDI0-01 recommits the local-decision architecture and its 34-work source inventory +([`local-decision-interventions.md`](local-decision-interventions.md), +[`local-decision-sources.json`](../../src/slm_training/resources/autoresearch/local-decision-sources.json)). +The runnable local-decision experiments already live in the **V10** section above +(E248 control and E249 measured; E250–E254 unrun / fail-closed). LDI reuses that +campaign, the existing preference harness, and the append-only decode traces, and +introduces no second trainer or orchestration stack. + +No LDI matrix rows are claimed here and no E-IDs are reserved: future LDI +experiments draw globally unique E-IDs from the existing allocation process and are +registered as ordinary E rows when they run. + +| Namespace | Scope | Status | +| --- | --- | --- | +| LDI0 | Evidence contract and bounded diagnostics — source inventory, separation-of-concerns invariants, E249–E284 falsification chain | inventory registered; no matrix rows claimed | diff --git a/docs/design/verified-scope-solver.md b/docs/design/verified-scope-solver.md index 6b2f7f32..dfc11b03 100644 --- a/docs/design/verified-scope-solver.md +++ b/docs/design/verified-scope-solver.md @@ -1,12 +1,16 @@ # VSS0-01 (SLM-57): the bounded verified-scope-solver contract -**Status:** specification only. This document defines the guarantee boundary for -the *Verified Scope Solving & Hybrid Realization* project. It adds no solver -code, no runtime dependency, no experiment, and no checkpoint, and makes no model -or ship claim. Existing checkpoints and decode behavior are unchanged until a -later VSS issue enables new flags behind a feature gate. - -**Code:** none (contract/spec only). +**Status:** contract defined (VSS0-01), with the torch-free finite-domain state +subset now implemented. This document defines the guarantee boundary for the +*Verified Scope Solving & Hybrid Realization* project; the immutable finite-domain +state schema and the completion-forest adapter are implemented in +`src/slm_training/dsl/solver/` (see "Implemented state schema" below). No decode-path +integration, runtime dependency in the generation path, experiment, or checkpoint is +added, and no model or ship claim is made. Existing checkpoints and decode behavior +are unchanged until a later VSS issue enables new flags behind a feature gate. + +**Code:** the finite-domain state schema (`state.py`) and completion-forest adapter +(`adapters.py`) in `src/slm_training/dsl/solver/`, torch-free; no decode integration yet. **Reader contract:** `docs/design/` is the source of truth for coding agents and experiment reviewers; the Linear project document is planning context only. The @@ -266,7 +270,7 @@ the closure. | `ScopeContract` / `ScopeKind` (`data/progspec/scopes.py:33`/`:18`) | AST scopes (`COMPONENT_CALL`/`STATEMENT`/`CHILD_LIST`) plus a def/use binder overlay (`definitions`/`uses`/`visible_binders`). | **Extended.** Scope def/use edges feed capsule SCC construction; scopes are **not** assumed independent. | | `dependency_closed_failure_cone` (`data/progspec/scopes.py:113`) | Despite the name, computes the least-common-ancestor of failing AST paths — **not** a dependency closure or SCC. | **Not reused for capsules.** Capsule SCCs are a new, distinct construction; this contract does not duplicate or overload this helper. | | `ChoiceTokenizer` / `ChoiceDecodeState` (`models/choice_tokenizer.py:172`/`:608`) | Grammar-closed choice IR; length-aware legality (`allowed_ids`, `exhaustive_allowed_ids`, `minimal_completion_length`). | **Retained.** The choice IR is the late-realization and verification surface. Not duplicated. | -| `HoleId`, finite `domain`, `verification capsule`, `proof certificate` | Absent today (confirmed greenfield). | **New** dataclasses/semantics introduced by later VSS issues under this contract. | +| `HoleId`, finite `domain`, `verification capsule`, `proof certificate` | Absent today (confirmed greenfield). | **New.** The Torch-free state subset (`HoleId`, `SolverBounds`, `SupportVerdict`, `DomainValue`, `HoleDomain`, `FiniteDomainState`) plus the completion-forest adapter are implemented by VSS0-03 (SLM-59) under [`src/slm_training/dsl/solver/`](../../src/slm_training/dsl/solver/state.py); see [Implemented state schema](#implemented-state-schema-vss0-03--slm-59). Verification capsules and proof certificates remain later-issue work. | ## Constraint evidence (VSS0-02) @@ -366,6 +370,73 @@ the X16-X21 convention ("lineage labels, not reproduced results"). rollback are the closure/deduction model; LDT architecture, alpha supervision, and training remain future work. +## Implemented state schema (VSS0-03 / SLM-59) + +The Torch-free state subset of this contract is implemented under +[`src/slm_training/dsl/solver/`](../../src/slm_training/dsl/solver/state.py) +(`state.py`, `adapters.py`, `__init__.py`). The package imports no `torch` and +performs no model inference; the compiler-forest adapter is loaded lazily via +`__getattr__` so `import slm_training.dsl.solver` stays light. Every +mutation-like operation returns a new validated state. + +**Dataclasses** (all `frozen=True`, JSON-safe): + +- `HoleId(namespace: str, path: tuple[str | int, ...], kind: str)` — element types + survive the JSON round-trip; a `canonical_key` (canonical JSON) provides a + type-safe total order for sorting/lookup instead of native comparison. +- `SolverBounds(max_tokens, max_nodes, max_depth, max_backtracks, + max_verifier_calls)` — non-negative ints; negatives are rejected. +- `SupportVerdict(str, Enum)` — `supported` / `unsupported` / `unknown`. No API + translates `UNKNOWN` to `UNSUPPORTED`. +- `DomainValue(tag, token_ids=(), kind=None, attributes=())` — a deterministic, + tagged JSON value; identity is by value. `token_ids` carries a *full* compiler + path (not only its first token) so grammar-forced suffixes stay + distinguishable. `attributes` is reserved (canonical, unique-key scalars) for + future structured action values. +- `HoleDomain(hole_id, values, metadata=())` — values are deduplicated + (duplicates rejected) and canonically ordered on construction; `is_empty` is + bottom, `is_singleton` marks a resolved hole. `metadata` carries JSON scalars + such as the compiler `coverage` guarantee — never soft scores. +- `FiniteDomainState(problem_id, pack_id, constraint_version, bounds, holes, + decision_level=0, parent_fingerprint=None)` — holes deduplicated and + canonically ordered; `is_bottom`, `is_structurally_solved`, `domain()`, + `refine()`, `with_decision()`, `meet()`, `summary()`, and `to_dict`/`from_dict`. + +**State fingerprint and its exclusions.** `fingerprint` is a SHA-256 over the +canonical JSON (`sort_keys=True`, compact separators) of the *hard state and +configuration only*: `problem_id`, `pack_id`, `constraint_version`, `bounds`, and +the canonically ordered `holes` (each hole's id, values, and metadata). It is +order-insensitive because holes and their values are sorted on construction. It +**excludes** `decision_level` and `parent_fingerprint` (reversible-search trail +lineage), so a state and a same-domain decision child remain interchangeable for +replay; and by construction it excludes logits, soft scores, timestamps, process +IDs, and mutable caches. Soft candidate scores are never stored on the state — a +separate ranker keys scores by `(state_fingerprint, hole_id, value)`. + +**Monotone transitions.** `refine(hole_id, retained_values, *, +certificate_ref=None)` reduces one hole to a subset (added values and unknown +holes are rejected) and leaves the trail lineage untouched; `certificate_ref` is +a forward-compatibility hook and is **not** persisted here (certificate +persistence/replay is VSS1-04), keeping it out of the fingerprint. +`with_decision(hole_id, retained_values)` applies the same monotone reduction but +increments `decision_level` and records the parent fingerprint, claiming no proof +(no certificate). `meet(other)` is the greatest lower bound over a shared +problem/pack/constraint/bounds identity: it intersects the domains of shared +holes (an empty intersection yields bottom) and keeps holes present in only one +state; mismatched identity is rejected. + +**Completion-forest adapter.** `completion_forest_state(*, prefix_ids, forest, +pack_id, constraint_version, bounds)` projects a `CompletionForest` into a single +"next semantic decision" hole keyed by prefix length and prefix fingerprint. Each +candidate carries the full `CompletionPath.token_ids` plus `kind`, and the +forest's `coverage` guarantee is carried verbatim in the hole metadata. An empty +forest projects to bottom; a singleton path projects to a structurally solved +state **for that projection only** — not a globally verified program and no +`SUPPORTED` verdict. The adapter is not on the default decode path, so existing +compiler/model behaviour is unchanged. `CompletionForestProjection` is the +reference `FiniteDomainProjection` seam that future topology-node domains +implement. + ## Measured status 2026-07-17 — specification-only change (VSS0-01 / SLM-57). No solver code, no @@ -384,3 +455,18 @@ synced — no quality or ship claim. Verified by `tests/test_models/test_compiler_decode.py`, `tests/test_models/test_choice_tokenizer.py`, and `python -m scripts.repo_policy`. + +2026-07-17 — VSS0-03 / SLM-59 implements the Torch-free finite-domain state +subset above (`src/slm_training/dsl/solver/`) with `tests/test_dsl/test_solver_state.py` +(28 tests: order-insensitive fingerprints, hard-state/version fingerprint +sensitivity, JSON round-trip stability, bottom/solved semantics, monotone +refinement and candidate-expansion rejection, meet intersection and +mismatched-identity rejection, decision lineage, the empty/singleton/multi-path/ +forced-suffix/partial-coverage adapter fixtures, and a subprocess-verified +Torch-free import). `python -m pytest tests/test_dsl/test_solver_state.py +tests/test_models/test_compiler_decode.py -q` passes (63) and `python -m +scripts.repo_policy` passes. **No model or ship claim**: this is a +model-independent data structure and a projection adapter that the default +decode path does not invoke; no checkpoint, experiment, or eval run is involved. +Exact closure, proof certificates, capsule solving, and learned guidance remain +later VSS issues. diff --git a/src/slm_training/harnesses/preference/decision_diagnostics.py b/src/slm_training/harnesses/preference/decision_diagnostics.py new file mode 100644 index 00000000..f7a00d61 --- /dev/null +++ b/src/slm_training/harnesses/preference/decision_diagnostics.py @@ -0,0 +1,235 @@ +"""Bounded objective-geometry diagnostics for local decisions (LDI0-03). + +A single cumulative wall-time budget (default and hard cap five minutes) is shared +across every diagnostic stage, mirroring ``CampaignBudget`` / +``autoresearch.engine.execute_commands``. On expiry the run reports ``expired`` and +produces **no result artifact** — the E285/E286 lesson that runtime expiry is a +stopped run, never a result. A future Tier-2 trainable-subspace pass refuses an +unauthorized full-parameter request (``not_authorized``) rather than replaying the +invalid E285 full-parameter profile. + +This module writes no model, checkpoint, or training result. Elapsed time is read +through the module-level ``time`` reference so a deterministic fake clock can be +injected in tests (``monkeypatch.setattr(decision_diagnostics.time, "monotonic", +...)``). +""" + +from __future__ import annotations + +import json +import os +import tempfile +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +DEFAULT_WALL_MINUTES = 5.0 +MAX_WALL_MINUTES = 5.0 +DiagnosticStatus = Literal["completed", "expired", "not_authorized"] + +__all__ = [ + "DEFAULT_WALL_MINUTES", + "MAX_WALL_MINUTES", + "Deadline", + "DiagnosticBudget", + "not_authorized_report", + "run_bounded_stages", + "tier1_objective_geometry", + "tier2_subspace_gradients", + "write_diagnostic_report", +] + + +@dataclass(frozen=True) +class DiagnosticBudget: + """A cumulative wall-time budget, defaulted and hard-capped at five minutes.""" + + max_wall_minutes: float = DEFAULT_WALL_MINUTES + + def __post_init__(self) -> None: + if not 0.0 < float(self.max_wall_minutes) <= MAX_WALL_MINUTES: + raise ValueError("max_wall_minutes must be in (0, 5]") + + @property + def seconds(self) -> float: + return float(self.max_wall_minutes) * 60.0 + + +class Deadline: + """A monotonic cumulative deadline shared across diagnostic stages.""" + + def __init__(self, budget: DiagnosticBudget) -> None: + self._budget = budget + self._start = time.monotonic() + self._deadline = self._start + budget.seconds + + def remaining(self) -> float: + return max(0.0, self._deadline - time.monotonic()) + + def expired(self) -> bool: + return time.monotonic() >= self._deadline + + def elapsed(self) -> float: + return time.monotonic() - self._start + + +def run_bounded_stages( + stages: Sequence[tuple[str, Callable[[], Any]]], + *, + budget: DiagnosticBudget | None = None, +) -> dict[str, Any]: + """Run named stages under one cumulative deadline. + + Returns a report whose ``status`` is ``completed`` or ``expired``. On expiry + ``result`` is ``None`` — no partial result is ever represented as a diagnostic + result — and ``stage_telemetry`` records exactly which stages ran before the + stop, so a rejected stage is visible rather than silently retained. + """ + budget = budget or DiagnosticBudget() + deadline = Deadline(budget) + telemetry: list[dict[str, Any]] = [] + results: dict[str, Any] = {} + for name, stage in stages: + if deadline.expired(): + return _report("expired", budget, deadline, telemetry, None) + started = time.monotonic() + output = stage() + telemetry.append({"stage": name, "seconds": time.monotonic() - started}) + results[name] = output + if deadline.expired(): + return _report("expired", budget, deadline, telemetry, None) + return _report("completed", budget, deadline, telemetry, results) + + +def not_authorized_report(reason: str, *, budget: DiagnosticBudget | None = None) -> dict[str, Any]: + """Report a refused stage (e.g. a full-parameter Tier-2 request) with no result.""" + budget = budget or DiagnosticBudget() + return { + "status": "not_authorized", + "max_wall_minutes": budget.max_wall_minutes, + "elapsed_wall_seconds": 0.0, + "stage_telemetry": [], + "reason": reason, + "result": None, + } + + +def _report( + status: DiagnosticStatus, + budget: DiagnosticBudget, + deadline: Deadline, + telemetry: list[dict[str, Any]], + result: dict[str, Any] | None, +) -> dict[str, Any]: + return { + "status": status, + "max_wall_minutes": budget.max_wall_minutes, + "elapsed_wall_seconds": deadline.elapsed(), + "stage_telemetry": telemetry, + "result": result, + } + + +def write_diagnostic_report(path: Path | str, report: dict[str, Any]) -> None: + """Atomically write a diagnostic report (mkstemp + fsync + os.replace). + + An ``expired`` / ``not_authorized`` report carries ``result: None`` — the + artifact records the stop without a diagnostic result. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) + tmp = Path(raw) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True, allow_nan=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + finally: + tmp.unlink(missing_ok=True) + + +def tier1_objective_geometry( + per_state_views: Sequence[Sequence[Any]], + *, + budget: DiagnosticBudget | None = None, +) -> dict[str, Any]: + """Read-only objective geometry over already-materialized objective views. + + ``per_state_views[i]`` is the sequence of ``ObjectiveView`` objects for one state + (one view per materializer / objective). The pass runs inside the bounded runner, + so it obeys the same cumulative deadline as every other diagnostic stage. It reads + ``good_action_ids`` / ``bad_action_ids`` off each view (via ``getattr`` so it stays + decoupled from ``decision_events_v2``) and reports, per corpus: + + - ``objective_contradictions`` — states where some action is scored *good* by one + view and *bad* by another. A contradiction means the objective is not well-posed + on that state; it is the logit-space shadow of the E284 objective conflict. + - ``mean_good_set_overlap`` — mean pairwise Jaccard of the per-state good-action + sets, a coarse measure of how much the views agree on what is preferred + (``None`` when no state carries two or more views). + + This computes no gradient, trains nothing, writes no model, and makes no + model-quality claim — it is geometry over inputs the caller already materialized. + """ + + def _analyze() -> dict[str, Any]: + contradictions = 0 + overlaps: list[float] = [] + for views in per_state_views: + good_sets = [set(getattr(view, "good_action_ids", ())) for view in views] + bad_sets = [set(getattr(view, "bad_action_ids", ())) for view in views] + good_union: set[int] = set().union(*good_sets) if good_sets else set() + bad_union: set[int] = set().union(*bad_sets) if bad_sets else set() + if good_union & bad_union: + contradictions += 1 + for i in range(len(good_sets)): + for j in range(i + 1, len(good_sets)): + left, right = good_sets[i], good_sets[j] + union = left | right + overlaps.append(len(left & right) / len(union) if union else 1.0) + return { + "states": len(per_state_views), + "objective_contradictions": contradictions, + "mean_good_set_overlap": ( + sum(overlaps) / len(overlaps) if overlaps else None + ), + } + + return run_bounded_stages([("objective_geometry", _analyze)], budget=budget) + + +def tier2_subspace_gradients( + *, + trainable_parameter_subset: Sequence[str] | None, + budget: DiagnosticBudget | None = None, +) -> dict[str, Any]: + """Tier-2 adapter-subspace gradient interface (refuses full-parameter). + + A Tier-2 gradient pass is authorized only over an explicit trainable-parameter + subset (named adapter tensors). An empty or ``None`` subset is a full-parameter + request, refused as ``not_authorized`` rather than replaying the invalid E285 + full-parameter profile that blew the runtime envelope. When a subset is supplied + the runner records a bounded *plan* only — the gradient computation is deferred to + a model stage this module never runs, so no gradient is computed and no model is + written here. + """ + subset = tuple(trainable_parameter_subset or ()) + if not subset: + return not_authorized_report( + "full-parameter Tier-2 is not authorized; provide an explicit adapter subset", + budget=budget, + ) + + def _plan() -> dict[str, Any]: + return { + "parameter_subset_size": len(subset), + "parameter_subset": list(subset), + "gradients": "deferred_to_model_stage", + } + + return run_bounded_stages([("subspace_plan", _plan)], budget=budget) diff --git a/src/slm_training/harnesses/preference/decision_events_v2.py b/src/slm_training/harnesses/preference/decision_events_v2.py new file mode 100644 index 00000000..85440795 --- /dev/null +++ b/src/slm_training/harnesses/preference/decision_events_v2.py @@ -0,0 +1,874 @@ +"""DecisionEventV2: a state / action-evidence / objective-view model. + +Extends (does not replace) :mod:`slm_training.harnesses.preference.local_decisions`. +See ``docs/design/local-decision-interventions.md``. The V2 contract separates +three concerns that ``DecisionEventV1`` conflated, which is the fix for the E284 +blocker (*stable grammar-state support does not imply objective/action-partition +support*): + +* :class:`DecisionStateV2` — the exact, replayable model decision state plus + immutable runtime identities. Its ``state_id`` is a canonical hash of state and + identity **only** — never of sampled good/bad labels, rollout outcomes, or + candidate order — so two independent action-label samples of one exact state + produce the *same* state id and merge into one state row. +* :class:`ActionOutcomeV2` — append-only, per-candidate verifier evidence. The + complete ordered G0-G12 gate vector is preserved verbatim (never set-collapsed). +* :class:`ObjectiveView` — a pure, versioned materialization of one state's action + table into a trainable objective. A ``constraint_shadow`` view is explicitly + ``trainable=False`` so semantic trainers cannot consume legality-only evidence. + +V1 corpora remain readable; :func:`migrate_v1_event` performs a one-way, lossless, +evidence-honest migration (it marks partial evidence incomplete and never +fabricates rollout/verifier outcomes). This module writes no model, checkpoint, or +training result. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Literal + +from slm_training.data.verify import Gate +from slm_training.harnesses.preference.counterfactuals import ( + SEMANTIC_VERIFIER_V1, + _METRICS, +) +from slm_training.harnesses.preference.local_decisions import ( + DecisionEventV1, + Split, + split_for_group, +) +from slm_training.lineage.records import content_sha + +SCHEMA_VERSION = 2 +Architecture = Literal["twotower", "causal"] +GateStatusStr = Literal["pass", "fail", "skip"] + +# The complete ordered G0-G12 gate identity a verifier vector must carry. +GATE_ORDER: tuple[str, ...] = tuple(gate.value for gate in Gate) +_GATE_STATUSES = frozenset({"pass", "fail", "skip"}) +_CONSTRAINT_SHADOW_VERIFIER = "constraint_shadow_legality_only" + +__all__ = [ + "SCHEMA_VERSION", + "GATE_ORDER", + "ActionOutcomeV2", + "DecisionStateV2", + "ObjectiveView", + "admit_semantic_corpus", + "append_action_outcomes", + "materialize_constraint_shadow", + "materialize_pareto", + "materialize_set_valued", + "materialize_single_best_worst", + "materialize_thresholded", + "migrate_v1_event", + "objective_view_signature", + "objective_view_support", + "verifier_bundle_hash", +] + + +def _ids(values: Iterable[int]) -> tuple[int, ...]: + """Sorted, deduplicated int tuple (matches the V1 legal/candidate convention).""" + return tuple(sorted({int(value) for value in values})) + + +def _int_tuple(values: Iterable[int]) -> tuple[int, ...]: + """Order-preserving int tuple (for a causal prefix, where order is identity).""" + return tuple(int(value) for value in values) + + +def verifier_bundle_hash(name: str, metrics: Sequence[str]) -> str: + """Stable identity for a verifier bundle (name + ordered metric list).""" + return content_sha({"name": str(name), "metrics": [str(metric) for metric in metrics]}) + + +# --------------------------------------------------------------------------- # +# Decision state. +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class DecisionStateV2: + """One exact, replayable model decision state. + + ``state_id`` is derived on construction from the state and identity fields + only; it is deliberately independent of any action evidence. + """ + + group_id: str + architecture: Architecture + context_text: str + decision_position: int + legal_action_ids: tuple[int, ...] + decision_kind: str + abstract_state_role: str + grammar_state_hash: str + policy_checkpoint_sha: str + tokenizer_sha: str + decode_config_hash: str + verifier_bundle_hash: str + split: Split + context_ids: tuple[int, ...] | None = None + canvas_ids: tuple[int, ...] | None = None + generation_step: int | None = None + version: int = SCHEMA_VERSION + state_id: str = "" + + def __post_init__(self) -> None: + if self.version != SCHEMA_VERSION: + raise ValueError("unsupported decision state version") + if self.architecture not in ("twotower", "causal"): + raise ValueError(f"unknown architecture {self.architecture!r}") + object.__setattr__(self, "legal_action_ids", _ids(self.legal_action_ids)) + if not self.legal_action_ids: + raise ValueError("legal_action_ids must be non-empty") + if self.canvas_ids is not None: + object.__setattr__(self, "canvas_ids", _int_tuple(self.canvas_ids)) + if self.context_ids is not None: + object.__setattr__(self, "context_ids", _int_tuple(self.context_ids)) + if not self.group_id or not self.context_text: + raise ValueError("group_id and context_text are required") + if self.architecture == "twotower": + if not self.canvas_ids: + raise ValueError("twotower state requires canvas_ids") + if not 0 <= self.decision_position < len(self.canvas_ids): + raise ValueError("decision_position must address canvas_ids") + else: # causal + if self.context_ids is None: + raise ValueError("causal state requires context_ids (the prefix)") + if self.decision_position < 0: + raise ValueError("decision_position must be non-negative") + if self.generation_step is not None and self.generation_step < 0: + raise ValueError("generation_step must be non-negative") + for name in ( + "grammar_state_hash", + "policy_checkpoint_sha", + "tokenizer_sha", + "decode_config_hash", + "verifier_bundle_hash", + "abstract_state_role", + "decision_kind", + ): + if not getattr(self, name): + raise ValueError(f"state identity field {name!r} must be non-empty") + if self.split != split_for_group(self.group_id): + raise ValueError("split must be derived from group_id") + object.__setattr__(self, "state_id", self._compute_state_id()) + + def _compute_state_id(self) -> str: + # Exact model state + immutable runtime identity only. No labels, no + # rollout outcomes, no candidate order -> reordering/augmenting action + # evidence never changes the state id. + return content_sha( + { + "group_id": self.group_id, + "architecture": self.architecture, + "context_text": self.context_text, + "context_ids": list(self.context_ids) if self.context_ids is not None else None, + "canvas_ids": list(self.canvas_ids) if self.canvas_ids is not None else None, + "decision_position": self.decision_position, + "generation_step": self.generation_step, + "legal_action_ids": list(self.legal_action_ids), + "decision_kind": self.decision_kind, + "abstract_state_role": self.abstract_state_role, + "grammar_state_hash": self.grammar_state_hash, + "policy_checkpoint_sha": self.policy_checkpoint_sha, + "tokenizer_sha": self.tokenizer_sha, + "decode_config_hash": self.decode_config_hash, + "verifier_bundle_hash": self.verifier_bundle_hash, + } + ) + + def to_dict(self) -> dict[str, Any]: + data = asdict(self) + for name in ("legal_action_ids", "context_ids", "canvas_ids"): + value = data[name] + data[name] = list(value) if value is not None else None + return data + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> DecisionStateV2: + fields = set(cls.__dataclass_fields__) + unknown = set(value) - fields + if unknown: + raise ValueError(f"unknown decision state fields: {sorted(unknown)}") + payload = {key: value[key] for key in value if key != "state_id"} + state = cls(**payload) + recorded = value.get("state_id") + if recorded and recorded != state.state_id: + raise ValueError("decision state_id does not match its state (tampered)") + return state + + +# --------------------------------------------------------------------------- # +# Action evidence. +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class ActionOutcomeV2: + """Append-only verifier evidence for one legal candidate at one state. + + ``verifier_vectors`` is a tuple of complete ordered G0-G12 vectors (one per + rollout); each inner vector is a tuple of ``(gate, status)`` pairs in G0..G12 + order. It may be empty when evidence is incomplete (e.g. migrated from V1), + but any present vector must carry the full ordered gate set. + """ + + state_id: str + action_id: int + legal: bool + rollout_policy_sha: str + continuation_seeds: tuple[int, ...] = () + outcome_hashes: tuple[str, ...] = () + verifier_vectors: tuple[tuple[tuple[str, str], ...], ...] = () + reward_vectors: tuple[tuple[tuple[str, float], ...], ...] = () + mean_value: float | None = None + confidence_interval: tuple[float, float] | None = None + evidence_ids: tuple[str, ...] = () + evidence_confidence: float = 0.0 + version: int = SCHEMA_VERSION + + def __post_init__(self) -> None: + if self.version != SCHEMA_VERSION: + raise ValueError("unsupported action outcome version") + if not self.state_id: + raise ValueError("action outcome requires a state_id") + if not self.rollout_policy_sha: + raise ValueError("action outcome requires a rollout_policy_sha") + object.__setattr__(self, "action_id", int(self.action_id)) + object.__setattr__(self, "continuation_seeds", tuple(int(s) for s in self.continuation_seeds)) + object.__setattr__(self, "outcome_hashes", tuple(str(h) for h in self.outcome_hashes)) + object.__setattr__( + self, "verifier_vectors", tuple(_normalize_gate_vector(v) for v in self.verifier_vectors) + ) + object.__setattr__( + self, "reward_vectors", tuple(_normalize_reward_vector(v) for v in self.reward_vectors) + ) + object.__setattr__(self, "evidence_ids", tuple(str(e) for e in self.evidence_ids)) + if not 0.0 <= float(self.evidence_confidence) <= 1.0: + raise ValueError("evidence_confidence must be in [0, 1]") + if self.confidence_interval is not None: + low, high = self.confidence_interval + if low > high: + raise ValueError("confidence_interval must be ordered (low <= high)") + object.__setattr__(self, "confidence_interval", (float(low), float(high))) + if self.mean_value is not None: + object.__setattr__(self, "mean_value", float(self.mean_value)) + + @property + def content_key(self) -> str: + """Content identity used to deduplicate append-only evidence (order-free).""" + return content_sha(self.to_dict()) + + def to_dict(self) -> dict[str, Any]: + return { + "state_id": self.state_id, + "action_id": self.action_id, + "legal": self.legal, + "rollout_policy_sha": self.rollout_policy_sha, + "continuation_seeds": list(self.continuation_seeds), + "outcome_hashes": list(self.outcome_hashes), + "verifier_vectors": [[list(pair) for pair in vector] for vector in self.verifier_vectors], + "reward_vectors": [[list(pair) for pair in vector] for vector in self.reward_vectors], + "mean_value": self.mean_value, + "confidence_interval": list(self.confidence_interval) if self.confidence_interval else None, + "evidence_ids": list(self.evidence_ids), + "evidence_confidence": self.evidence_confidence, + "version": self.version, + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> ActionOutcomeV2: + fields = set(cls.__dataclass_fields__) + unknown = set(value) - fields + if unknown: + raise ValueError(f"unknown action outcome fields: {sorted(unknown)}") + version = int(value.get("version", SCHEMA_VERSION)) + if version != SCHEMA_VERSION: + raise ValueError(f"unsupported action outcome version: {version}") + legal = value["legal"] + if not isinstance(legal, bool): + raise ValueError("action outcome legal must be a boolean") + ci = value.get("confidence_interval") + return cls( + state_id=str(value["state_id"]), + action_id=int(value["action_id"]), + legal=legal, + rollout_policy_sha=str(value["rollout_policy_sha"]), + continuation_seeds=tuple(value.get("continuation_seeds") or ()), + outcome_hashes=tuple(value.get("outcome_hashes") or ()), + verifier_vectors=tuple( + tuple(tuple(pair) for pair in vector) for vector in value.get("verifier_vectors") or () + ), + reward_vectors=tuple( + tuple(tuple(pair) for pair in vector) for vector in value.get("reward_vectors") or () + ), + mean_value=value.get("mean_value"), + confidence_interval=tuple(ci) if ci else None, + evidence_ids=tuple(value.get("evidence_ids") or ()), + evidence_confidence=float(value.get("evidence_confidence", 0.0)), + version=version, + ) + + +def _normalize_gate_vector(vector: Sequence[Any]) -> tuple[tuple[str, str], ...]: + pairs = tuple((str(gate), str(status)) for gate, status in vector) + if tuple(gate for gate, _ in pairs) != GATE_ORDER: + raise ValueError("verifier vector must carry the complete ordered G0-G12 gate set") + for _, status in pairs: + if status not in _GATE_STATUSES: + raise ValueError(f"gate status must be one of {sorted(_GATE_STATUSES)}, got {status!r}") + return pairs + + +def _normalize_reward_vector(vector: Sequence[Any]) -> tuple[tuple[str, float], ...]: + # Validate at the boundary: a non-empty reward vector must carry exactly the + # expected metrics, so a downstream mean never fabricates a 0.0 for a missing one. + pairs = tuple((str(name), float(score)) for name, score in vector) + names = [name for name, _ in pairs] + if len(names) != len(set(names)): + raise ValueError("reward vector has duplicate metrics") + if set(names) != set(_METRICS): + raise ValueError(f"reward vector must carry exactly the metrics {sorted(_METRICS)}") + return pairs + + +def append_action_outcomes( + existing: Iterable[ActionOutcomeV2], new: Iterable[ActionOutcomeV2] +) -> tuple[ActionOutcomeV2, ...]: + """Append-only merge that deduplicates by content identity, not array order. + + Earlier evidence is never mutated; identical content appended twice collapses + to one row. The result is ordered deterministically by ``(action_id, key)``. + """ + merged: dict[str, ActionOutcomeV2] = {} + for outcome in (*existing, *new): + merged[outcome.content_key] = outcome + return tuple(sorted(merged.values(), key=lambda o: (o.action_id, o.content_key))) + + +# --------------------------------------------------------------------------- # +# Objective materialization. +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class ObjectiveView: + """A pure, versioned materialization of one state's action table.""" + + good_action_ids: tuple[int, ...] + bad_action_ids: tuple[int, ...] + ambiguous_action_ids: tuple[int, ...] + unobserved_action_ids: tuple[int, ...] + weights: tuple[tuple[int, float], ...] + materializer_id: str + materializer_config_hash: str + trainable: bool = True + version: int = SCHEMA_VERSION + + def __post_init__(self) -> None: + for name in ("good_action_ids", "bad_action_ids", "ambiguous_action_ids", "unobserved_action_ids"): + object.__setattr__(self, name, _ids(getattr(self, name))) + partitions = [ + set(self.good_action_ids), + set(self.bad_action_ids), + set(self.ambiguous_action_ids), + set(self.unobserved_action_ids), + ] + seen: set[int] = set() + for partition in partitions: + if partition & seen: + raise ValueError("objective-view action partitions must be disjoint") + seen |= partition + object.__setattr__( + self, "weights", tuple((int(a), float(w)) for a, w in self.weights) + ) + if any(action not in seen for action, _ in self.weights): + raise ValueError("objective-view weights reference an unpartitioned action") + + def to_dict(self) -> dict[str, Any]: + return { + "good_action_ids": list(self.good_action_ids), + "bad_action_ids": list(self.bad_action_ids), + "ambiguous_action_ids": list(self.ambiguous_action_ids), + "unobserved_action_ids": list(self.unobserved_action_ids), + "weights": [list(pair) for pair in self.weights], + "materializer_id": self.materializer_id, + "materializer_config_hash": self.materializer_config_hash, + "trainable": self.trainable, + "version": self.version, + } + + +def _observed_reward(outcome: ActionOutcomeV2) -> dict[str, float] | None: + """Mean reward vector across an outcome's rollouts, or None if unobserved.""" + if not outcome.reward_vectors: + return None + totals: dict[str, float] = {name: 0.0 for name in _METRICS} + for vector in outcome.reward_vectors: + row = dict(vector) + for name in _METRICS: + totals[name] += float(row.get(name, 0.0)) + count = len(outcome.reward_vectors) + return {name: totals[name] / count for name in _METRICS} + + +def _require_semantic_outcome(state: DecisionStateV2, outcome: ActionOutcomeV2) -> None: + """Reject an outcome that must never become a semantic good/bad label. + + A semantic materializer may only classify a *legal* action that belongs to this + state's legal set — matching ``state_id`` alone is not enough, or an arbitrary or + illegal ``action_id`` could be trained as a preference. + """ + if outcome.state_id != state.state_id: + raise ValueError("action outcome does not belong to this state") + if not outcome.legal: + raise ValueError("semantic materialization requires a legal action outcome") + if outcome.action_id not in state.legal_action_ids: + raise ValueError("action outcome action_id is outside the state's legal set") + + +def materialize_pareto( + state: DecisionStateV2, outcomes: Sequence[ActionOutcomeV2] +) -> ObjectiveView: + """Exact Pareto pass/fail over the ordered reward metrics. + + ``good`` = observed, reward-bearing, non-dominated actions; ``bad`` = observed + actions that are dominated or carry no reward evidence; ``ambiguous`` = + observed actions with no reward vector to score; ``unobserved`` = legal actions + with no evidence at all. + """ + by_action: dict[int, dict[str, float]] = {} + observed: set[int] = set() + for outcome in outcomes: + _require_semantic_outcome(state, outcome) + observed.add(outcome.action_id) + reward = _observed_reward(outcome) + if reward is not None: + by_action[outcome.action_id] = reward + + def dominates(left: dict[str, float], right: dict[str, float]) -> bool: + no_worse = all(left[name] >= right[name] for name in _METRICS) + better = any(left[name] > right[name] for name in _METRICS) + return no_worse and better + + scored = list(by_action.items()) + good = [ + action + for action, reward in scored + if not any(other != action and dominates(by_action[other], reward) for other, _ in scored) + ] + good_set = set(good) + bad = [action for action in observed if action not in good_set and action in by_action] + ambiguous = [action for action in observed if action not in by_action] + unobserved = [action for action in state.legal_action_ids if action not in observed] + return ObjectiveView( + good_action_ids=good, + bad_action_ids=bad, + ambiguous_action_ids=ambiguous, + unobserved_action_ids=unobserved, + weights=tuple((action, 1.0) for action in sorted(good_set)), + materializer_id="pareto_v2", + materializer_config_hash=content_sha({"metrics": list(_METRICS), "rule": "pareto_nondominated"}), + trainable=True, + ) + + +def materialize_constraint_shadow( + state: DecisionStateV2, outcomes: Sequence[ActionOutcomeV2] +) -> ObjectiveView: + """Legality-only diagnostic view. Explicitly non-semantic and NOT trainable. + + A semantic trainer must refuse ``trainable=False`` views: legality is not a + preference label (the E284 lesson). + """ + for outcome in outcomes: + if outcome.state_id != state.state_id: + raise ValueError("action outcome does not belong to this state") + observed = {outcome.action_id for outcome in outcomes} + return ObjectiveView( + good_action_ids=(), + bad_action_ids=(), + ambiguous_action_ids=tuple(sorted(observed)), + unobserved_action_ids=tuple(a for a in state.legal_action_ids if a not in observed), + weights=(), + materializer_id="constraint_shadow_diagnostic_v2", + materializer_config_hash=content_sha({"rule": "legality_only", "semantic": False}), + trainable=False, + ) + + +def _verified(outcome: ActionOutcomeV2) -> bool | None: + """True if every rollout G-vector has no failing gate; None if no evidence.""" + if not outcome.verifier_vectors: + return None + return all(status != "fail" for vector in outcome.verifier_vectors for _, status in vector) + + +def _action_summary( + state: DecisionStateV2, outcomes: Sequence[ActionOutcomeV2] +) -> dict[int, tuple[dict[str, float] | None, bool | None, float]]: + """Aggregate append-only evidence per action -> (mean reward, verified, confidence). + + Multiple rollout batches for one action are combined, so a materializer never + double-classifies an action that carries more than one outcome row. + """ + rewards: dict[int, list[dict[str, float]]] = {} + verified: dict[int, list[bool]] = {} + confidence: dict[int, float] = {} + for outcome in outcomes: + _require_semantic_outcome(state, outcome) + reward = _observed_reward(outcome) + if reward is not None: + rewards.setdefault(outcome.action_id, []).append(reward) + verdict = _verified(outcome) + if verdict is not None: + verified.setdefault(outcome.action_id, []).append(verdict) + confidence[outcome.action_id] = max( + confidence.get(outcome.action_id, 0.0), float(outcome.evidence_confidence) + ) + summary: dict[int, tuple[dict[str, float] | None, bool | None, float]] = {} + for action in set(confidence) | set(rewards) | set(verified): + rows = rewards.get(action) + mean = ( + {name: sum(row[name] for row in rows) / len(rows) for name in _METRICS} + if rows + else None + ) + verdicts = verified.get(action) + summary[action] = (mean, (all(verdicts) if verdicts else None), confidence.get(action, 0.0)) + return summary + + +def materialize_thresholded( + state: DecisionStateV2, + outcomes: Sequence[ActionOutcomeV2], + *, + metric: str = "reward", + threshold: float = 0.5, + min_confidence: float = 0.0, +) -> ObjectiveView: + """Threshold one scalar reward metric, gated by an evidence-confidence floor. + + Actions below the confidence floor, or with no reward evidence, stay + ``ambiguous`` rather than being labeled good or bad. + """ + if metric not in _METRICS: + raise ValueError(f"unknown metric {metric!r}; expected one of {list(_METRICS)}") + if not 0.0 <= min_confidence <= 1.0: + raise ValueError("min_confidence must be in [0, 1]") + summary = _action_summary(state, outcomes) + good: list[int] = [] + bad: list[int] = [] + ambiguous: list[int] = [] + for action, (mean, _verdict, confidence) in summary.items(): + if mean is None or confidence < min_confidence: + ambiguous.append(action) + elif mean[metric] >= threshold: + good.append(action) + else: + bad.append(action) + unobserved = [action for action in state.legal_action_ids if action not in summary] + return ObjectiveView( + good_action_ids=good, + bad_action_ids=bad, + ambiguous_action_ids=ambiguous, + unobserved_action_ids=unobserved, + weights=tuple((action, 1.0) for action in sorted(good)), + materializer_id="thresholded_v2", + materializer_config_hash=content_sha( + {"metric": metric, "threshold": threshold, "min_confidence": min_confidence} + ), + trainable=True, + ) + + +def materialize_single_best_worst( + state: DecisionStateV2, outcomes: Sequence[ActionOutcomeV2], *, metric: str = "reward" +) -> ObjectiveView: + """Single-best (good) and single-worst (bad) control over one metric.""" + if metric not in _METRICS: + raise ValueError(f"unknown metric {metric!r}; expected one of {list(_METRICS)}") + summary = _action_summary(state, outcomes) + scored = {action: mean[metric] for action, (mean, _v, _c) in summary.items() if mean is not None} + good: list[int] = [] + bad: list[int] = [] + if scored: + best = max(scored, key=lambda action: (scored[action], -action)) + worst = min(scored, key=lambda action: (scored[action], action)) + good = [best] + if worst != best: + bad = [worst] + chosen = set(good) | set(bad) + ambiguous = [action for action in summary if action not in chosen] + unobserved = [action for action in state.legal_action_ids if action not in summary] + return ObjectiveView( + good_action_ids=good, + bad_action_ids=bad, + ambiguous_action_ids=ambiguous, + unobserved_action_ids=unobserved, + weights=tuple((action, 1.0) for action in sorted(good)), + materializer_id="single_best_worst_v2", + materializer_config_hash=content_sha({"metric": metric, "rule": "single_best_worst"}), + trainable=True, + ) + + +def materialize_set_valued( + state: DecisionStateV2, outcomes: Sequence[ActionOutcomeV2] +) -> ObjectiveView: + """Set-valued good/bad partitions from verifier verdicts (all gates pass => good).""" + summary = _action_summary(state, outcomes) + good: list[int] = [] + bad: list[int] = [] + ambiguous: list[int] = [] + for action, (_mean, verdict, _confidence) in summary.items(): + if verdict is True: + good.append(action) + elif verdict is False: + bad.append(action) + else: + ambiguous.append(action) + unobserved = [action for action in state.legal_action_ids if action not in summary] + return ObjectiveView( + good_action_ids=good, + bad_action_ids=bad, + ambiguous_action_ids=ambiguous, + unobserved_action_ids=unobserved, + weights=tuple((action, 1.0) for action in sorted(good)), + materializer_id="set_valued_v2", + materializer_config_hash=content_sha({"rule": "verifier_all_pass"}), + trainable=True, + ) + + +# --------------------------------------------------------------------------- # +# V1 migration (one-way, lossless, evidence-honest). +# --------------------------------------------------------------------------- # +@dataclass(frozen=True) +class MigratedDecision: + """A V1 event migrated to V2 form, with an explicit completeness flag.""" + + state: DecisionStateV2 + outcomes: tuple[ActionOutcomeV2, ...] + complete: bool + note: str + + +def migrate_v1_event(event: DecisionEventV1) -> MigratedDecision: + """Migrate one :class:`DecisionEventV1` to V2 form without fabricating evidence. + + A V1 counterfactual becomes a state with *partial* action evidence (the good / + bad partition it recorded) and is marked incomplete because V1 kept no + replayable rollout/verifier vectors. A V1 constraint shadow becomes a legality + diagnostic. Neither path invents rollout, seed, or G-vector evidence. + """ + is_counterfactual = event.evidence_kind == "counterfactual" + bundle = ( + verifier_bundle_hash(SEMANTIC_VERIFIER_V1, _METRICS) + if is_counterfactual + else verifier_bundle_hash(_CONSTRAINT_SHADOW_VERIFIER, ()) + ) + state = DecisionStateV2( + group_id=event.group_id, + architecture="twotower", + context_text=event.context_text, + canvas_ids=event.canvas_ids, + decision_position=event.position, + legal_action_ids=event.legal_token_ids, + decision_kind=event.decision_kind, + abstract_state_role=event.decision_kind, + grammar_state_hash=content_sha( + {"decision_kind": event.decision_kind, "legal_token_ids": list(event.legal_token_ids)} + ), + policy_checkpoint_sha=event.policy_checkpoint_sha, + tokenizer_sha=event.tokenizer_sha, + decode_config_hash=event.decode_config_hash, + verifier_bundle_hash=bundle, + split=event.split, + ) + outcomes: list[ActionOutcomeV2] = [] + if is_counterfactual: + # Record the V1 good/bad partition as legal actions with NO replayable + # verifier/rollout evidence (empty vectors) -> incomplete by construction. + for action in event.good_token_ids: + outcomes.append( + ActionOutcomeV2( + state_id=state.state_id, + action_id=action, + legal=True, + rollout_policy_sha=event.policy_checkpoint_sha, + continuation_seeds=(event.seed,), + evidence_ids=(f"v1:{event.event_id}",), + evidence_confidence=event.evidence_confidence, + ) + ) + for action in event.bad_token_ids: + outcomes.append( + ActionOutcomeV2( + state_id=state.state_id, + action_id=action, + legal=action in set(event.legal_token_ids), + rollout_policy_sha=event.policy_checkpoint_sha, + continuation_seeds=(event.seed,), + evidence_ids=(f"v1:{event.event_id}",), + evidence_confidence=event.evidence_confidence, + ) + ) + note = "v1 counterfactual migrated with partial evidence (no replayable G-vectors)" + else: + note = "v1 constraint shadow migrated as a legality diagnostic (non-semantic)" + return MigratedDecision( + state=state, + outcomes=tuple(outcomes), + complete=False, + note=note, + ) + + +# --------------------------------------------------------------------------- # +# Atomic, duplicate-safe corpus IO (mirrors write_decision_events). +# --------------------------------------------------------------------------- # +def write_action_outcomes(path: Path | str, outcomes: Iterable[ActionOutcomeV2]) -> int: + """Atomically write action outcomes as sorted, deduplicated JSONL.""" + path = Path(path) + rows = append_action_outcomes((), outcomes) + handle = tempfile.NamedTemporaryFile( + "w", encoding="utf-8", dir=str(path.parent), delete=False, suffix=".tmp" + ) + try: + for outcome in rows: + handle.write(json.dumps(outcome.to_dict(), sort_keys=True) + "\n") + handle.flush() + os.fsync(handle.fileno()) + handle.close() + os.replace(handle.name, path) + except BaseException: + handle.close() + Path(handle.name).unlink(missing_ok=True) + raise + return len(rows) + + +def decision_state_manifest( + states: Sequence[DecisionStateV2], + outcomes: Sequence[ActionOutcomeV2], + views: Sequence[ObjectiveView] = (), +) -> dict[str, Any]: + """Separately fingerprint states, action evidence, and objective views. + + Rows are canonically ordered first, so a fingerprint changes only when the + underlying content changes, never when input row order changes. + """ + ordered_states = sorted(states, key=lambda state: state.state_id) + ordered_outcomes = append_action_outcomes((), outcomes) + ordered_views = sorted(views, key=lambda view: content_sha(view.to_dict())) + return { + "schema_version": SCHEMA_VERSION, + "state_count": len(ordered_states), + "outcome_count": len(ordered_outcomes), + "view_count": len(ordered_views), + "state_fingerprint": content_sha([state.to_dict() for state in ordered_states]), + "outcome_fingerprint": content_sha([outcome.to_dict() for outcome in ordered_outcomes]), + "view_fingerprint": content_sha([view.to_dict() for view in ordered_views]), + } + + +# --------------------------------------------------------------------------- # +# Objective-support admission (LDI0-03). +# --------------------------------------------------------------------------- # +def objective_view_signature(view: ObjectiveView) -> str: + """Objective-support signature: materializer identity + good/bad partition. + + Unlike a state-support signature (which keys on the state and judged positives + only), this INCLUDES the good/bad partition. A corpus can therefore pass + state-support coverage yet fail objective-support coverage when independently + judged legal alternatives yield different bad-action sets at one state — the + E284 blocker. + """ + return content_sha( + { + "materializer_id": view.materializer_id, + "materializer_config_hash": view.materializer_config_hash, + "good": list(view.good_action_ids), + "bad": list(view.bad_action_ids), + } + )[:16] + + +def objective_view_support( + items: Sequence[tuple[DecisionStateV2, ObjectiveView]], *, min_train_support: int = 1 +) -> dict[str, Any]: + """Held-out objective-view coverage — the V2 analogue of state-support coverage. + + ``items`` pairs each state with the objective view materialized for it. A + held-out objective signature is *covered* when at least ``min_train_support`` + train examples share it. + """ + if isinstance(min_train_support, bool) or int(min_train_support) < 1: + raise ValueError("min_train_support must be a positive integer") + counts: dict[str, dict[str, int]] = {"train": {}, "held_out": {}} + materializers: set[tuple[str, str]] = set() + non_trainable = 0 + for state, view in items: + if not view.trainable: + non_trainable += 1 + materializers.add((view.materializer_id, view.materializer_config_hash)) + bucket = counts.get(state.split) + if bucket is None: + raise ValueError(f"unexpected split {state.split!r}") + signature = objective_view_signature(view) + bucket[signature] = bucket.get(signature, 0) + 1 + held = set(counts["held_out"]) + covered = sorted(sig for sig in held if counts["train"].get(sig, 0) >= min_train_support) + uncovered = sorted(held - set(covered)) + return { + "train_signatures": len(counts["train"]), + "held_out_signatures": len(held), + "materializers": sorted(materializers), + "non_trainable_views": non_trainable, + "min_train_support": min_train_support, + "held_out_coverage": {"covered": covered, "uncovered": uncovered, "passed": not uncovered}, + } + + +def admit_semantic_corpus( + items: Sequence[tuple[DecisionStateV2, ObjectiveView]], + *, + materializer_id: str, + min_train_support: int = 1, +) -> dict[str, Any]: + """Fail closed before semantic training; return the objective-support report. + + Refuses when any view is non-trainable (e.g. a constraint-shadow diagnostic), + when a view's materializer does not match the requested objective, or when any + held-out objective signature lacks train support. Every refusal names what is + missing so objective support can be repaired (never by copying held-out + programs). + """ + report = objective_view_support(items, min_train_support=min_train_support) + non_trainable = [state.state_id for state, view in items if not view.trainable] + if non_trainable: + raise ValueError( + "semantic admission refused: " + f"{len(non_trainable)} non-trainable view(s) (e.g. constraint_shadow) " + "cannot supervise a semantic objective" + ) + mismatched = sorted( + {view.materializer_id for _, view in items if view.materializer_id != materializer_id} + ) + if mismatched: + raise ValueError( + "semantic admission refused: corpus materializer(s) " + f"{mismatched} do not match the requested objective {materializer_id!r}" + ) + if not report["held_out_coverage"]["passed"]: + missing = report["held_out_coverage"]["uncovered"] + raise ValueError( + "semantic admission refused: corpus lacks train support for " + f"{len(missing)} held-out objective signature(s); repair objective " + "support before training (E284 blocker)" + ) + return report diff --git a/src/slm_training/harnesses/preference/local_decisions.py b/src/slm_training/harnesses/preference/local_decisions.py index ac8805e8..94960290 100644 --- a/src/slm_training/harnesses/preference/local_decisions.py +++ b/src/slm_training/harnesses/preference/local_decisions.py @@ -162,6 +162,45 @@ def decision_signature_support( } +def objective_signature_support( + events: Iterable[DecisionEventV1], *, min_train_support: int = 1 +) -> dict[str, Any]: + """Held-out coverage of the OBJECTIVE signature (which includes sampled negatives). + + Unlike ``decision_signature_support`` (state support, good-only), this keys on + ``decision_signature`` (good + bad). A corpus can therefore pass state support + yet fail here when the sampled bad-action set at a held-out state has no train + support — the E284 blocker (stable state support does not imply objective + support). + """ + if min_train_support < 1: + raise ValueError("minimum train signature support must be positive") + rows = list(events) + counts = { + split: Counter( + decision_signature(event) for event in rows if event.split == split + ) + for split in ("train", "held_out") + } + held = set(counts["held_out"]) + covered = sorted( + signature for signature in held if counts["train"][signature] >= min_train_support + ) + uncovered = sorted(held - set(covered)) + return { + "minimum_train_support": min_train_support, + "counts": { + split: dict(sorted(split_counts.items())) + for split, split_counts in counts.items() + }, + "held_out_coverage": { + "covered": covered, + "uncovered": uncovered, + "passed": not uncovered, + }, + } + + def _event( trace: dict[str, Any], payload: dict[str, Any], diff --git a/src/slm_training/harnesses/preference/local_train.py b/src/slm_training/harnesses/preference/local_train.py index f108929b..15a836b0 100644 --- a/src/slm_training/harnesses/preference/local_train.py +++ b/src/slm_training/harnesses/preference/local_train.py @@ -19,6 +19,7 @@ decision_signature, decision_signature_metadata, load_decision_events, + objective_signature_support, ) from slm_training.models.twotower import TwoTowerModel @@ -1234,8 +1235,18 @@ def train_local_from_paths( block_by_decision_kind: bool = False, gradient_combination: GradientCombination = "proposal", optimizer_name: LocalOptimizer = "adamw", + require_objective_support: bool = False, ) -> dict: events = load_decision_events(events_path) + if require_objective_support: + support = objective_signature_support(events) + if not support["held_out_coverage"]["passed"]: + missing = support["held_out_coverage"]["uncovered"] + raise ValueError( + "decision event corpus lacks train support for " + f"{len(missing)} held-out objective signature(s); repair objective " + "support before training (E284 blocker)" + ) model = TwoTowerModel.from_checkpoint(checkpoint, device=device) _validate_identity(events, checkpoint, model) held_out_before = evaluate_local_decisions( diff --git a/src/slm_training/models/adapters/__init__.py b/src/slm_training/models/adapters/__init__.py new file mode 100644 index 00000000..d0689cf8 --- /dev/null +++ b/src/slm_training/models/adapters/__init__.py @@ -0,0 +1,11 @@ +"""Repository-owned removable model adapters (LDI2-01). + +``TwoTowerAdapterSpec`` is torch-free (config round-trip); import ``LowRankAdapter`` +from :mod:`slm_training.models.adapters.low_rank` where torch is available. +""" + +from __future__ import annotations + +from slm_training.models.adapters.spec import ADAPTER_SCHEMA_VERSION, TwoTowerAdapterSpec + +__all__ = ["ADAPTER_SCHEMA_VERSION", "TwoTowerAdapterSpec"] diff --git a/src/slm_training/models/adapters/low_rank.py b/src/slm_training/models/adapters/low_rank.py new file mode 100644 index 00000000..5cbc49a3 --- /dev/null +++ b/src/slm_training/models/adapters/low_rank.py @@ -0,0 +1,116 @@ +"""A removable LoRA-style low-rank delta wrapper for a single linear module (LDI2-01). + +For a frozen parent linear ``W`` (bias untouched): + + y = W x + scale * B(A(dropout(x))), scale = alpha / rank + +``B`` is zero-initialized, so a freshly attached adapter is output-identical to the +parent bit-for-bit. Only ``A``/``B`` receive gradients; the parent weight and bias are +frozen. The adapter can be disabled (restoring the exact parent map even after ``A``/``B`` +have changed), and merged one-way into a standalone ``nn.Linear`` on a copy — merging +never mutates the removable adapter. This is standard low-rank delta; it is deliberately +not called DoRA/PiSSA, which are not implemented here. +""" + +from __future__ import annotations + +import math + +import torch +from torch import nn + +__all__ = ["LowRankAdapter"] + + +class LowRankAdapter(nn.Module): + """Wrap a frozen ``nn.Linear`` with a removable low-rank delta.""" + + def __init__(self, base: nn.Linear, *, rank: int, alpha: float, dropout: float) -> None: + """Freeze ``base`` and attach zero-initialized ``A``/``B`` factors of the given rank. + + The parent weight and bias are frozen in place; ``lora_A`` is kaiming-initialized + and ``lora_B`` stays zero so the fresh adapter's delta is exactly zero. + """ + super().__init__() + if not isinstance(base, nn.Linear): + raise TypeError("LowRankAdapter wraps an nn.Linear module") + if int(rank) <= 0: + raise ValueError("adapter rank must be positive") + # Mirror the spec's trust-boundary checks so direct construction (which bypasses + # TwoTowerAdapterSpec) can never produce a silently-invalid adapter: a non-finite + # or non-positive alpha would poison `scaling`, and a negative dropout would + # silently degrade to Identity. + if not math.isfinite(float(alpha)) or float(alpha) <= 0.0: + raise ValueError("adapter alpha must be a positive finite number") + if not math.isfinite(float(dropout)) or not 0.0 <= float(dropout) < 1.0: + raise ValueError("adapter dropout must be a finite number in [0, 1)") + self.base = base + self.base.weight.requires_grad_(False) + if self.base.bias is not None: + self.base.bias.requires_grad_(False) + + self.rank = int(rank) + self.alpha = float(alpha) + self.scaling = self.alpha / self.rank + self.dropout: nn.Module = nn.Dropout(float(dropout)) if dropout > 0 else nn.Identity() + + weight = base.weight + self.lora_A = nn.Parameter( + torch.zeros(self.rank, base.in_features, dtype=weight.dtype, device=weight.device) + ) + self.lora_B = nn.Parameter( + torch.zeros(base.out_features, self.rank, dtype=weight.dtype, device=weight.device) + ) + # A is randomized; B stays zero so the initial delta is exactly zero. + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + self._enabled = True + + @property + def enabled(self) -> bool: + """Whether the low-rank delta is currently added to the parent's output.""" + return self._enabled + + def enable_adapter(self) -> None: + """Add the low-rank delta to the parent's output on subsequent forward passes.""" + self._enabled = True + + def disable_adapter(self) -> None: + """Bypass the delta so ``forward`` reproduces the frozen parent exactly.""" + self._enabled = False + + def adapter_parameters(self) -> list[nn.Parameter]: + """The two trainable factors ``[lora_A, lora_B]`` (the parent stays frozen).""" + return [self.lora_A, self.lora_B] + + def delta_weight(self) -> torch.Tensor: + """The dense delta ``scale * (B @ A)`` in the parent weight's dtype/device.""" + return self.scaling * (self.lora_B @ self.lora_A) + + def merged_weight(self) -> torch.Tensor: + """The parent weight with the delta folded in (does not mutate the parent).""" + return self.base.weight.data + self.delta_weight() + + def merged_linear(self) -> nn.Linear: + """A standalone ``nn.Linear`` equal to the adapter-enabled map, wrapper-free. + + The parent and this adapter are left untouched — merge is one-way on a copy. + """ + merged = nn.Linear( + self.base.in_features, + self.base.out_features, + bias=self.base.bias is not None, + ).to(self.base.weight.device, self.base.weight.dtype) + with torch.no_grad(): + merged.weight.copy_(self.merged_weight()) + if self.base.bias is not None: + merged.bias.copy_(self.base.bias.data) + return merged + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Return the parent output, plus the scaled low-rank delta when enabled.""" + out = self.base(x) + if not self._enabled: + return out + delta = nn.functional.linear(self.dropout(x), self.lora_A) + delta = nn.functional.linear(delta, self.lora_B) + return out + self.scaling * delta diff --git a/src/slm_training/models/adapters/spec.py b/src/slm_training/models/adapters/spec.py new file mode 100644 index 00000000..a8f1ea59 --- /dev/null +++ b/src/slm_training/models/adapters/spec.py @@ -0,0 +1,150 @@ +"""Versioned spec for the removable TwoTower low-rank adapter (LDI2-01). + +Torch-free so it round-trips through CLI / config / checkpoint metadata without +importing torch. The spec names *what* to adapt and the base identity it is bound to; +resolving the spec against a concrete model (and failing closed on a mismatch) is the +model's job, not the spec's. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Literal + +ADAPTER_SCHEMA_VERSION = 1 + +__all__ = ["ADAPTER_SCHEMA_VERSION", "TwoTowerAdapterSpec"] + + +@dataclass(frozen=True) +class TwoTowerAdapterSpec: + """A removable low-rank delta adapter configuration bound to a base checkpoint. + + ``target_modules`` names the linear-projection categories to adapt (matched + deterministically against the denoiser blocks); ``target_layer_indices`` optionally + restricts to selected layers (``None`` means every layer). ``include_output_head`` + opts the denoiser output head in. The context tower is never adapted by this spec. + """ + + method: Literal["low_rank"] + rank: int + alpha: float + dropout: float + target_modules: tuple[str, ...] + base_compatibility_fingerprint: str + base_checkpoint_sha: str + tokenizer_sha: str + target_layer_indices: tuple[int, ...] | None = None + include_output_head: bool = False + train_bias: Literal["none"] = "none" + schema_version: int = ADAPTER_SCHEMA_VERSION + + def __post_init__(self) -> None: + if self.schema_version != ADAPTER_SCHEMA_VERSION: + raise ValueError( + f"unsupported adapter schema version {self.schema_version}; " + f"expected {ADAPTER_SCHEMA_VERSION} — a newer schema must not be " + "reinterpreted as the current format" + ) + object.__setattr__(self, "rank", int(self.rank)) + object.__setattr__(self, "alpha", float(self.alpha)) + object.__setattr__(self, "dropout", float(self.dropout)) + object.__setattr__(self, "target_modules", tuple(self.target_modules)) + object.__setattr__( + self, + "target_layer_indices", + None + if self.target_layer_indices is None + else tuple(int(i) for i in self.target_layer_indices), + ) + if self.method != "low_rank": + raise ValueError("the first adapter backend supports only method='low_rank'") + if self.rank <= 0: + raise ValueError("adapter rank must be positive") + if self.alpha <= 0.0: + raise ValueError("adapter alpha must be positive") + if not 0.0 <= self.dropout < 1.0: + raise ValueError("adapter dropout must be in [0, 1)") + if self.train_bias != "none": + raise ValueError("the first adapter backend supports only train_bias='none'") + if not self.target_modules: + raise ValueError("adapter target_modules must not be empty") + if len(set(self.target_modules)) != len(self.target_modules): + raise ValueError("adapter target_modules must be unique") + if self.target_layer_indices is not None and any( + i < 0 for i in self.target_layer_indices + ): + raise ValueError("adapter target_layer_indices must be non-negative") + for name in ( + "base_compatibility_fingerprint", + "base_checkpoint_sha", + "tokenizer_sha", + ): + if not str(getattr(self, name)): + raise ValueError(f"adapter spec field {name!r} must be non-empty") + + @property + def scaling(self) -> float: + """The low-rank delta multiplier ``alpha / rank`` applied to ``B(A(x))``.""" + return self.alpha / self.rank + + def to_dict(self) -> dict[str, Any]: + """Return a plain-JSON dict (lists, not tuples) round-trippable via ``from_dict``.""" + return { + "schema_version": self.schema_version, + "method": self.method, + "rank": self.rank, + "alpha": self.alpha, + "dropout": self.dropout, + "target_modules": list(self.target_modules), + "target_layer_indices": ( + None + if self.target_layer_indices is None + else list(self.target_layer_indices) + ), + "include_output_head": self.include_output_head, + "train_bias": self.train_bias, + "base_compatibility_fingerprint": self.base_compatibility_fingerprint, + "base_checkpoint_sha": self.base_checkpoint_sha, + "tokenizer_sha": self.tokenizer_sha, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> TwoTowerAdapterSpec: + """Rebuild a spec from ``to_dict`` output, failing closed on unknown fields. + + An unrecognized key raises ``ValueError`` rather than being silently dropped, so a + newer or corrupted config never round-trips into a partially-applied adapter. + """ + known = { + "schema_version", + "method", + "rank", + "alpha", + "dropout", + "target_modules", + "target_layer_indices", + "include_output_head", + "train_bias", + "base_compatibility_fingerprint", + "base_checkpoint_sha", + "tokenizer_sha", + } + unknown = set(data) - known + if unknown: + raise ValueError(f"unknown adapter spec fields: {sorted(unknown)}") + layers = data.get("target_layer_indices") + return cls( + method=data["method"], + rank=data["rank"], + alpha=data["alpha"], + dropout=data["dropout"], + target_modules=tuple(data["target_modules"]), + base_compatibility_fingerprint=data["base_compatibility_fingerprint"], + base_checkpoint_sha=data["base_checkpoint_sha"], + tokenizer_sha=data["tokenizer_sha"], + target_layer_indices=None if layers is None else tuple(layers), + include_output_head=bool(data.get("include_output_head", False)), + train_bias=data.get("train_bias", "none"), + schema_version=int(data.get("schema_version", ADAPTER_SCHEMA_VERSION)), + ) diff --git a/src/slm_training/models/adapters/twotower_adapter.py b/src/slm_training/models/adapters/twotower_adapter.py new file mode 100644 index 00000000..0f6fddb2 --- /dev/null +++ b/src/slm_training/models/adapters/twotower_adapter.py @@ -0,0 +1,146 @@ +"""Deterministic resolution + attachment of low-rank adapters onto TwoTower (LDI2-01). + +Maps a :class:`TwoTowerAdapterSpec`'s target-module names onto the concrete linear +projections inside the denoiser's transformer blocks, wraps each with a +:class:`LowRankAdapter`, and reports which modules were resolved. Only the denoiser is +adapted here; the context tower is never touched. Missing, duplicate, or unsupported +targets fail closed with an actionable error. +""" + +from __future__ import annotations + +from typing import Any + +from torch import nn + +from slm_training.models.adapters.low_rank import LowRankAdapter +from slm_training.models.adapters.spec import TwoTowerAdapterSpec + +__all__ = [ + "TARGET_MODULE_PATHS", + "attach_low_rank_adapters", + "resolve_adapter_targets", +] + +# Canonical target-module name -> attribute path within one denoiser TransformerBlock. +TARGET_MODULE_PATHS: dict[str, tuple[Any, ...]] = { + "attn_q": ("self_attn", "q_proj"), + "attn_k": ("self_attn", "k_proj"), + "attn_v": ("self_attn", "v_proj"), + "attn_out": ("self_attn", "out_proj"), + "cross_attn_q": ("cross_attn", "q_proj"), + "cross_attn_k": ("cross_attn", "k_proj"), + "cross_attn_v": ("cross_attn", "v_proj"), + "cross_attn_out": ("cross_attn", "out_proj"), + "mlp_in": ("mlp", 0), + "mlp_out": ("mlp", 2), +} + + +def _child(module: Any, key: Any) -> Any: + """Index (int key) or attribute-access (str key) one step into ``module``, or ``None``.""" + if isinstance(key, int): + try: + return module[key] + except (IndexError, KeyError, TypeError): + return None + return getattr(module, key, None) + + +def _resolve_leaf(block: nn.Module, path: tuple[Any, ...]) -> tuple[Any, Any, Any]: + """Walk ``path`` to its leaf, returning ``(parent, leaf_key, leaf)`` or a ``None`` triple. + + Any missing intermediate short-circuits to ``(None, None, None)`` so the caller can + fail closed on an unresolved target. + """ + parent: Any = block + for key in path[:-1]: + parent = _child(parent, key) + if parent is None: + return None, None, None + leaf_key = path[-1] + return parent, leaf_key, _child(parent, leaf_key) + + +def resolve_adapter_targets( + denoiser: nn.Module, spec: TwoTowerAdapterSpec +) -> dict[str, tuple[nn.Module, Any, nn.Linear]]: + """Resolve a spec's targets to (parent, leaf-key, linear) triples, deterministically. + + Fails closed when a requested target names an unknown module, resolves to a + non-linear, or references a layer index outside the denoiser. + """ + layers = getattr(denoiser, "layers", None) + if layers is None: + raise ValueError("denoiser exposes no `layers` to adapt") + if spec.include_output_head: + # The output head is not among the resolvable transformer-block targets, so + # honoring the flag would require dedicated support. Fail closed rather than let + # `include_output_head=True` silently have no effect. + raise ValueError( + "include_output_head=True is not supported by the low_rank backend yet; " + "the denoiser output head is not an adaptable target" + ) + count = len(layers) + indices = ( + tuple(range(count)) + if spec.target_layer_indices is None + else spec.target_layer_indices + ) + resolved: dict[str, tuple[nn.Module, Any, nn.Linear]] = {} + for index in indices: + if index < 0 or index >= count: + raise ValueError( + f"adapter target layer index {index} is out of range [0, {count})" + ) + block = layers[index] + for name in spec.target_modules: + path = TARGET_MODULE_PATHS.get(name) + if path is None: + raise ValueError( + f"unsupported adapter target module {name!r}; " + f"expected one of {sorted(TARGET_MODULE_PATHS)}" + ) + parent, leaf_key, linear = _resolve_leaf(block, path) + if not isinstance(linear, nn.Linear): + raise ValueError( + f"adapter target {name!r} at denoiser layer {index} did not " + "resolve to an nn.Linear (unsupported module type or absent)" + ) + resolved[f"denoiser.layers.{index}.{name}"] = (parent, leaf_key, linear) + if not resolved: + raise ValueError("adapter spec matched no modules") + return resolved + + +def attach_low_rank_adapters( + denoiser: nn.Module, spec: TwoTowerAdapterSpec, *, seed: int +) -> dict[str, LowRankAdapter]: + """Wrap each resolved target with a LowRankAdapter, in place, and return the map. + + Adapter factors are initialized under a forked RNG so attaching an adapter never + shifts the model's training RNG stream (a fresh adapter is output-identical anyway, + since B is zero-initialized). + """ + import torch + + resolved = resolve_adapter_targets(denoiser, spec) + wrappers: dict[str, LowRankAdapter] = {} + for offset, (key, (parent, leaf_key, linear)) in enumerate(sorted(resolved.items())): + # Fork whichever generator the adapter factors are initialized on. LowRankAdapter + # builds `lora_A` on the base weight's device, so on CUDA an empty `devices` list + # would let attachment advance the external CUDA generator (breaking the + # "attachment never shifts the training RNG" guarantee); include the CUDA device. + device = linear.weight.device + fork_devices = [device] if device.type == "cuda" else [] + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(int(seed) + 4096 + offset) + wrapper = LowRankAdapter( + linear, rank=spec.rank, alpha=spec.alpha, dropout=spec.dropout + ) + if isinstance(leaf_key, int): + parent[leaf_key] = wrapper + else: + setattr(parent, leaf_key, wrapper) + wrappers[key] = wrapper + return wrappers diff --git a/src/slm_training/models/causal_lm_openui.py b/src/slm_training/models/causal_lm_openui.py index c9eef8d2..68412458 100644 --- a/src/slm_training/models/causal_lm_openui.py +++ b/src/slm_training/models/causal_lm_openui.py @@ -12,8 +12,19 @@ from slm_training.data.contract import GenerationRequest from slm_training.dsl.parser import stream_check, validate from slm_training.dsl.schema import ExampleRecord +from slm_training.harnesses.distill.trace_store import decode_config_hash from slm_training.lineage.records import content_sha from slm_training.lineage.tracks import CAUSAL_LORA_RECIPE +from slm_training.models.causal_trace import ( + AllowedIds, + CausalTracedGeneration, + CausalTraceIdentity, + CausalTraceWriter, + GeneratedOutcome, + TracePolicy, + capture_raw_steps, + fold_policy_identity, +) @dataclass(frozen=True) @@ -249,6 +260,189 @@ def _allowed_ids(self, prefix: tuple[int, ...]) -> tuple[int, ...]: self._grammar_mask_cache[prefix] = result return result + def _encode_prompt(self, prompt: str) -> Any: + """Encode the OpenUI chat prompt to input ids (shared by the traced paths).""" + system = "Return only one valid OpenUI program using placeholder content." + messages = [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ] + if hasattr(self.tokenizer, "apply_chat_template"): + return self.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt" + ).to(self.model.device) + return self.tokenizer(f"{system}\n{prompt}\n", return_tensors="pt")[ + "input_ids" + ].to(self.model.device) + + def active_adapter_identity(self) -> str: + """Content digest of the active adapter tensors, or "" when none are active. + + Hashing the tensor *values* (not just their names) means any mutation of an + adapter weight changes the folded ``policy_checkpoint_sha`` — so a trace + captured under one adapter cannot silently load as another. + """ + import torch + + adapters = [ + (name, param) + for name, param in self.model.named_parameters() + if "lora" in name.lower() + ] + if not adapters: + return "" + return content_sha( + { + name: param.detach().to(torch.float64).cpu().flatten().tolist() + for name, param in sorted(adapters, key=lambda item: item[0]) + } + ) + + def capture_identity(self, *, group_id: str, context_text: str) -> CausalTraceIdentity: + """Build the trajectory identity stamped onto every captured causal state. + + Base and adapter identity are folded into a single policy fingerprint because + the shared ``DecisionStateV2`` schema carries no dedicated adapter field, so an + adapter-enabled and adapter-disabled capture receive different state identities. + """ + adapter = self.active_adapter_identity() + return CausalTraceIdentity( + group_id=group_id, + context_text=context_text, + policy_checkpoint_sha=fold_policy_identity( + self.compatibility_fingerprint(), adapter + ), + tokenizer_sha=self.artifact_identity()["tokenizer_sha"], + decode_config_hash=decode_config_hash(getattr(self.model, "config", {})), + base_model_revision=self.config.base_model_revision, + adapter_identity=adapter, + ) + + def generate_constrained_traced( + self, + prompt: str, + *, + group_id: str, + policy: TracePolicy | None = None, + trace_writer: CausalTraceWriter | None = None, + max_new_tokens: int | None = None, + allowed_ids_fn: AllowedIds | None = None, + ) -> CausalTracedGeneration: + """Greedy constrained generation that captures exact per-step decision evidence. + + Unlike :meth:`generate_constrained` (unchanged and trace-free by default), this + drives a per-step loop so the raw pre-mask logits, raw argmax, legal set, and + constrained selection are all recoverable from integer prefix ids. The stored + ``context_ids`` are the full prefix (prompt + generated suffix), so a consumer + can replay ``model(context_ids).logits[:, -1, :]`` exactly. ``allowed_ids_fn`` + overrides the grammar legal-set seam (defaulting to :meth:`_allowed_ids`). + """ + import torch + + input_ids = self._encode_prompt(prompt) + prompt_row = tuple(int(value) for value in input_ids[0].tolist()) + prompt_len = len(prompt_row) + device = self.model.device + + def forward_logits(prefix: tuple[int, ...]) -> list[float]: + row = torch.tensor([list(prefix)], device=device) + with torch.inference_mode(): + logits = self.model(row).logits[0, -1, :] + return logits.to(torch.float32).cpu().tolist() + + def grammar_allowed(prefix: tuple[int, ...]) -> tuple[int, ...]: + return self._allowed_ids(tuple(int(token) for token in prefix[prompt_len:])) + + result = capture_raw_steps( + forward_logits=forward_logits, + allowed_ids=allowed_ids_fn or grammar_allowed, + eos_id=int(self.tokenizer.eos_token_id), + max_new_tokens=int( + max_new_tokens if max_new_tokens is not None else self.config.max_length + ), + initial_prefix=prompt_row, + policy=policy, + ) + text = self.tokenizer.decode( + result.generated_token_ids, skip_special_tokens=True + ).strip() + try: + program = validate(text) + final_text = (program.serialized or text).strip() + valid = True + except Exception: # noqa: BLE001 - honest: an unfinished decode is not valid + final_text = text + valid = False + if trace_writer is not None: + trace_writer.record_all(result) + return CausalTracedGeneration(text=final_text, result=result, valid=valid) + + def replay_causal_action( + self, + state: Any, + forced_action_id: int, + continuation_seed: int, + generation_config: dict[str, Any] | None = None, + *, + allowed_ids_fn: AllowedIds | None = None, + ) -> GeneratedOutcome: + """Replay a forced first action on the exact stored prefix, then continue. + + The forced action is applied to ``state.context_ids`` (the exact integer prefix + recovered without re-encoding text); continuation uses the deterministic + constrained policy, so the same seed/config reproduces the outcome. No judge + runs here — the pre-judge outcome is returned for the shared counterfactual + owner to score. + """ + import torch + + if state.context_ids is None: + raise ValueError("causal replay requires a stored integer prefix (context_ids)") + if int(forced_action_id) not in state.legal_action_ids: + raise ValueError("forced causal action is not legal for the stored state") + prompt_len = len(state.context_ids) - int(state.decision_position) + if prompt_len < 0: + raise ValueError("stored decision position exceeds the recorded prefix length") + base_prefix = tuple(int(token) for token in state.context_ids) + ( + int(forced_action_id), + ) + device = self.model.device + config = generation_config or {} + + def forward_logits(prefix: tuple[int, ...]) -> list[float]: + row = torch.tensor([list(prefix)], device=device) + with torch.inference_mode(): + logits = self.model(row).logits[0, -1, :] + return logits.to(torch.float32).cpu().tolist() + + def grammar_allowed(prefix: tuple[int, ...]) -> tuple[int, ...]: + return self._allowed_ids(tuple(int(token) for token in prefix[prompt_len:])) + + result = capture_raw_steps( + forward_logits=forward_logits, + allowed_ids=allowed_ids_fn or grammar_allowed, + eos_id=int(self.tokenizer.eos_token_id), + max_new_tokens=int(config.get("max_new_tokens", self.config.max_length)), + initial_prefix=base_prefix, + ) + # The stored prefix already holds the tokens generated before this decision; + # keep them so a decision_position > 0 replay materializes the full program. + generated_suffix = tuple(int(token) for token in state.context_ids[prompt_len:]) + generated = (*generated_suffix, int(forced_action_id), *result.generated_token_ids) + raw_text = self.tokenizer.decode(generated, skip_special_tokens=True).strip() + try: + program = validate(raw_text) + canonical: str | None = (program.serialized or raw_text).strip() + except Exception: # noqa: BLE001 - unfinished continuation has no canonical form + canonical = None + return GeneratedOutcome( + action_id=int(forced_action_id), + continuation_seed=int(continuation_seed), + finish_reason=result.stop_reason, + raw_program=raw_text, + canonical_program=canonical, + ) + def save(self, path: Path) -> None: path.mkdir(parents=True, exist_ok=False) self.model.save_pretrained(path) diff --git a/src/slm_training/models/causal_trace.py b/src/slm_training/models/causal_trace.py new file mode 100644 index 00000000..68d0be48 --- /dev/null +++ b/src/slm_training/models/causal_trace.py @@ -0,0 +1,613 @@ +"""Torch-free capture core for exact causal decision-state tracing (LDI1-01). + +The causal OpenUI plug-in decodes under a hard grammar constraint. To recover every +supervised decision from exact prefix token IDs and model logits — not from decoded +strings or later retokenization — the capture is split into a **torch-free core** +(this module) and a thin torch wiring on the plug-in. + +The core drives a greedy constrained decode through two injected callables: + +* ``forward_logits(prefix_ids) -> Sequence[float]`` — the model's next-token logits + for a prefix (abstracts ``model(input_ids).logits[:, -1, :]``); +* ``allowed_ids(prefix_ids) -> Sequence[int]`` — the grammar-legal token id set for a + prefix (abstracts the plug-in's ``_allowed_ids``). + +Because both are injected, the whole capture algorithm — raw argmax, legal masking, +constrained selection, constraint-shadow detection, per-step records, honest EOS/stop, +and bounded selection policies — is deterministic and unit-testable with fixture +logits, without torch, a real tokenizer, or the grammar. Integer token IDs are the +authority; decoded surfaces are diagnostic only. + +A ``constraint_shadow`` observation is emitted when the raw (pre-mask) argmax is +illegal while the constrained selection is legal. It is a non-semantic evidence class: +it measures where the hard constraint changes the policy and must never pass semantic +corpus admission or feed a semantic trainer. +""" + +from __future__ import annotations + +import json +import math +import os +import tempfile +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Any + +from slm_training.harnesses.distill.trace_store import TraceStore +from slm_training.harnesses.preference.decision_events_v2 import ( + ActionOutcomeV2, + DecisionStateV2, + ObjectiveView, + materialize_constraint_shadow, + verifier_bundle_hash, +) +from slm_training.harnesses.preference.local_decisions import split_for_group +from slm_training.lineage.records import content_sha + +ForwardLogits = Callable[[tuple[int, ...]], Sequence[float]] +AllowedIds = Callable[[tuple[int, ...]], Sequence[int]] +RoleOf = Callable[[tuple[int, ...]], str | None] + +CAUSAL_DECISION_KIND = "causal_decision" +CAUSAL_TRACE_ROW_KIND = "causal_decision" +_CONSTRAINT_SHADOW_VERIFIER = "constraint_shadow_legality_only" +_CAPTURE_VERIFIER = "causal_capture_no_semantic_verifier" + +__all__ = [ + "AllowedIds", + "CAUSAL_DECISION_KIND", + "CAUSAL_TRACE_ROW_KIND", + "CaptureResult", + "CausalTracedGeneration", + "CausalTraceError", + "CausalTraceIdentity", + "CausalTraceWriter", + "ForwardLogits", + "GeneratedOutcome", + "RawStepObservation", + "RoleOf", + "TracePolicy", + "TraceSelection", + "build_decision_state", + "capture_raw_steps", + "causal_trace_row", + "emit_causal_decision", + "fold_policy_identity", + "legal_set_reference", + "load_causal_decision_states", +] + + +class CausalTraceError(ValueError): + """Raised when a capture request is internally inconsistent.""" + + +def legal_set_reference(legal_token_ids: Sequence[int]) -> str: + """Content-addressed reference for a legal token set (order-independent).""" + return content_sha(sorted({int(token) for token in legal_token_ids})) + + +class TraceSelection(str, Enum): + """Bounded selection policies for which decisions are recorded.""" + + EVERY = "every" + CONSTRAINT_SHADOW_ONLY = "constraint_shadow_only" + MARGIN_THRESHOLD = "margin_threshold" + SAMPLED_POSITIONS = "sampled_positions" + NAMED_ROLES = "named_roles" + + +@dataclass(frozen=True) +class TracePolicy: + """A bounded selection policy plus the per-step telemetry width (``top_k``).""" + + selection: TraceSelection = TraceSelection.EVERY + margin_threshold: float = 0.0 + sampled_positions: tuple[int, ...] = () + named_roles: tuple[str, ...] = () + top_k: int = 5 + + def __post_init__(self) -> None: + if int(self.top_k) < 1: + raise ValueError("top_k must be >= 1") + object.__setattr__(self, "top_k", int(self.top_k)) + object.__setattr__(self, "margin_threshold", float(self.margin_threshold)) + object.__setattr__( + self, "sampled_positions", tuple(int(p) for p in self.sampled_positions) + ) + object.__setattr__( + self, "named_roles", tuple(str(r) for r in self.named_roles) + ) + + def records(self, obs: RawStepObservation) -> bool: + """Whether this observation is retained under the selection policy.""" + if self.selection is TraceSelection.EVERY: + return True + if self.selection is TraceSelection.CONSTRAINT_SHADOW_ONLY: + return obs.constraint_shadow + if self.selection is TraceSelection.MARGIN_THRESHOLD: + return obs.legal_margin <= self.margin_threshold + if self.selection is TraceSelection.SAMPLED_POSITIONS: + return obs.generated_ordinal in self.sampled_positions + if self.selection is TraceSelection.NAMED_ROLES: + return obs.grammar_role in self.named_roles + raise CausalTraceError(f"unknown selection policy: {self.selection}") + + +@dataclass(frozen=True) +class RawStepObservation: + """One committed next-token decision, recoverable from exact prefix IDs. + + ``decision_index`` counts only non-forced steps (a step with two or more legal + candidates is a real reversible decision); ``generated_ordinal`` counts every + emitted token. A ``forced`` step (a single legal candidate) is a deterministic + grammar deduction, not a decision, and carries the index of the decision it + precedes. + """ + + decision_index: int + generated_ordinal: int + prefix_token_ids: tuple[int, ...] + raw_argmax_id: int + selected_token_id: int + legal_token_ids: tuple[int, ...] + raw_topk: tuple[tuple[int, float, float], ...] + legal_topk: tuple[tuple[int, float], ...] + constraint_shadow: bool + forced: bool + grammar_role: str | None = None + + @property + def legal_set_reference(self) -> str: + """Content hash of this step's legal-candidate set (order-independent).""" + return legal_set_reference(self.legal_token_ids) + + @property + def legal_margin(self) -> float: + """Post-mask top1 − top2 legal probability (top1 when a single candidate).""" + if not self.legal_topk: + return 0.0 + if len(self.legal_topk) == 1: + return self.legal_topk[0][1] + return self.legal_topk[0][1] - self.legal_topk[1][1] + + def to_dict(self) -> dict[str, Any]: + """Serialize this observation to plain JSON (tuples flattened to lists).""" + return { + "decision_index": self.decision_index, + "generated_ordinal": self.generated_ordinal, + "prefix_token_ids": list(self.prefix_token_ids), + "raw_argmax_id": self.raw_argmax_id, + "selected_token_id": self.selected_token_id, + "legal_token_ids": list(self.legal_token_ids), + "legal_set_reference": self.legal_set_reference, + "raw_topk": [list(row) for row in self.raw_topk], + "legal_topk": [list(row) for row in self.legal_topk], + "constraint_shadow": self.constraint_shadow, + "forced": self.forced, + "grammar_role": self.grammar_role, + } + + +@dataclass(frozen=True) +class CaptureResult: + """The outcome of a traced constrained decode. + + ``observations`` holds only the policy-retained decisions; ``generated_token_ids`` + is the complete emitted continuation (every selected token, unaffected by the + selection policy) so the terminal program can be materialized exactly; + ``stop_reason`` records why decoding stopped. + """ + + observations: tuple[RawStepObservation, ...] + generated_token_ids: tuple[int, ...] + stop_reason: str + + @property + def constraint_shadow_count(self) -> int: + """How many retained observations are constraint shadows (raw pick was illegal).""" + return sum(1 for obs in self.observations if obs.constraint_shadow) + + def to_dict(self) -> dict[str, Any]: + """Serialize the retained observations and full continuation to plain JSON.""" + return { + "observations": [obs.to_dict() for obs in self.observations], + "generated_token_ids": list(self.generated_token_ids), + "stop_reason": self.stop_reason, + } + + +def _log_softmax(logits: Sequence[float]) -> list[float]: + """Numerically stable log-softmax over the full vocabulary logit row.""" + peak = max(logits) + shifted = [float(value) - peak for value in logits] + log_denom = math.log(sum(math.exp(value) for value in shifted)) + return [value - log_denom for value in shifted] + + +def _raw_topk( + logits: Sequence[float], log_probs: Sequence[float], top_k: int +) -> tuple[tuple[int, float, float], ...]: + """Top-``k`` ``(token_id, logit, log_prob)`` over the unmasked row, ties broken by id.""" + order = sorted(range(len(logits)), key=lambda i: (logits[i], -i), reverse=True) + return tuple( + (int(i), float(logits[i]), float(log_probs[i])) for i in order[:top_k] + ) + + +def _legal_topk( + logits: Sequence[float], legal: Sequence[int], top_k: int +) -> tuple[tuple[int, float], ...]: + """Top-``k`` ``(token_id, prob)`` re-normalized over just the legal candidate set.""" + peak = max(logits[token] for token in legal) + weights = {int(token): math.exp(float(logits[token]) - peak) for token in legal} + total = sum(weights.values()) + order = sorted(weights, key=lambda token: (weights[token], -token), reverse=True) + return tuple((token, weights[token] / total) for token in order[:top_k]) + + +def capture_raw_steps( + *, + forward_logits: ForwardLogits, + allowed_ids: AllowedIds, + eos_id: int, + max_new_tokens: int, + initial_prefix: Sequence[int] = (), + policy: TracePolicy | None = None, + role_of: RoleOf | None = None, +) -> CaptureResult: + """Greedy constrained decode that records exact per-step decision evidence. + + At each step the raw (pre-mask) argmax and the constrained selection (greedy over + the legal set) are both computed from the same logits, so a ``constraint_shadow`` + — a raw winner outside the legal set overridden by a legal selection — is exact. + Decoding stops honestly when EOS is selected (``eos``), when the legal set is empty + (``no_legal_continuation``), or when ``max_new_tokens`` is reached + (``max_new_tokens``). The returned ``observations`` are only those retained by + ``policy``, but selection never changes which token is emitted — the full emitted + continuation is returned in ``generated_token_ids``. + """ + policy = policy or TracePolicy() + if int(max_new_tokens) < 0: + raise CausalTraceError("max_new_tokens must be non-negative") + prefix = tuple(int(token) for token in initial_prefix) + eos = int(eos_id) + observations: list[RawStepObservation] = [] + generated: list[int] = [] + decisions_seen = 0 + stop_reason = "max_new_tokens" + for ordinal in range(int(max_new_tokens)): + logits = [float(value) for value in forward_logits(prefix)] + if not logits: + raise CausalTraceError("forward_logits returned an empty logit vector") + if not all(math.isfinite(value) for value in logits): + raise CausalTraceError("forward_logits returned non-finite logits") + legal = tuple(int(token) for token in allowed_ids(prefix)) + if not legal: + stop_reason = "no_legal_continuation" + break + legal_set = set(legal) + if any(token < 0 or token >= len(logits) for token in legal_set): + raise CausalTraceError("legal token id out of logit range") + log_probs = _log_softmax(logits) + raw_argmax = max(range(len(logits)), key=lambda i: (logits[i], -i)) + selected = max(legal, key=lambda token: (logits[token], -token)) + forced = len(legal_set) == 1 + obs = RawStepObservation( + decision_index=decisions_seen, + generated_ordinal=ordinal, + prefix_token_ids=prefix, + raw_argmax_id=int(raw_argmax), + selected_token_id=int(selected), + legal_token_ids=legal, + raw_topk=_raw_topk(logits, log_probs, policy.top_k), + legal_topk=_legal_topk(logits, legal, policy.top_k), + constraint_shadow=(raw_argmax not in legal_set) and (selected in legal_set), + forced=forced, + grammar_role=role_of(prefix) if role_of is not None else None, + ) + if policy.records(obs): + observations.append(obs) + if not forced: + decisions_seen += 1 + generated.append(int(selected)) + prefix = (*prefix, int(selected)) + if selected == eos: + stop_reason = "eos" + break + return CaptureResult( + observations=tuple(observations), + generated_token_ids=tuple(generated), + stop_reason=stop_reason, + ) + + +def fold_policy_identity(base_checkpoint_sha: str, adapter_identity: str) -> str: + """Fold base-checkpoint and active-adapter identity into one policy fingerprint. + + ``DecisionStateV2`` carries a single ``policy_checkpoint_sha`` with no dedicated + adapter field, so both the base checkpoint and the active adapter are folded here. + Because the state id hashes ``policy_checkpoint_sha``, an adapter-enabled and an + adapter-disabled capture over the same prefix receive **different** state + identities — the acceptance requirement that adapter fingerprints be part of state + identity. + """ + return content_sha( + {"base_checkpoint_sha": str(base_checkpoint_sha), "adapter": str(adapter_identity)} + ) + + +@dataclass(frozen=True) +class CausalTraceIdentity: + """The trajectory-level identity stamped onto every captured causal state.""" + + group_id: str + context_text: str + policy_checkpoint_sha: str + tokenizer_sha: str + decode_config_hash: str + base_model_revision: str = "" + adapter_identity: str = "" + + def __post_init__(self) -> None: + for name in ( + "group_id", + "policy_checkpoint_sha", + "tokenizer_sha", + "decode_config_hash", + ): + if not str(getattr(self, name)): + raise CausalTraceError(f"causal trace identity field {name!r} must be non-empty") + + +def build_decision_state( + obs: RawStepObservation, identity: CausalTraceIdentity +) -> DecisionStateV2: + """Materialize a causal ``DecisionStateV2`` from one raw observation. + + Integer prefix ids are the state authority (``context_ids``); the grammar state is + the content-addressed legal set. A ``constraint_shadow`` step is named as such and + tagged with the legality-only verifier so downstream code cannot mistake it for a + semantic verdict. + """ + shadow = obs.constraint_shadow + if shadow: + decision_kind = "constraint_shadow" + verifier = _CONSTRAINT_SHADOW_VERIFIER + elif obs.forced: + decision_kind = "forced_deduction" + verifier = _CAPTURE_VERIFIER + else: + decision_kind = CAUSAL_DECISION_KIND + verifier = _CAPTURE_VERIFIER + role = obs.grammar_role or decision_kind + return DecisionStateV2( + group_id=identity.group_id, + architecture="causal", + context_text=identity.context_text, + decision_position=obs.generated_ordinal, + legal_action_ids=obs.legal_token_ids, + decision_kind=decision_kind, + abstract_state_role=role, + grammar_state_hash=obs.legal_set_reference, + policy_checkpoint_sha=identity.policy_checkpoint_sha, + tokenizer_sha=identity.tokenizer_sha, + decode_config_hash=identity.decode_config_hash, + verifier_bundle_hash=verifier_bundle_hash(verifier, ()), + split=split_for_group(identity.group_id), + context_ids=obs.prefix_token_ids, + generation_step=obs.generated_ordinal, + ) + + +def _constraint_shadow_outcome( + state: DecisionStateV2, obs: RawStepObservation, identity: CausalTraceIdentity +) -> ActionOutcomeV2: + """A legality-only outcome for the selected (legal) action of a shadow step. + + It carries **no** reward or verifier vectors — legality is the only evidence — so + the view it feeds is non-trainable and cannot supervise a semantic objective. + """ + return ActionOutcomeV2( + state_id=state.state_id, + action_id=obs.selected_token_id, + legal=True, + rollout_policy_sha=identity.policy_checkpoint_sha, + ) + + +def emit_causal_decision( + obs: RawStepObservation, identity: CausalTraceIdentity +) -> tuple[DecisionStateV2, tuple[ActionOutcomeV2, ...], ObjectiveView | None]: + """Emit the schema-compatible evidence for one raw observation. + + A ``constraint_shadow`` step yields a legality-only outcome and a **non-trainable** + ``constraint_shadow`` view (via ``materialize_constraint_shadow``); every other step + yields the replayable state alone (no rollout evidence is invented). The returned + view is ``None`` unless a shadow was observed. + """ + state = build_decision_state(obs, identity) + if not obs.constraint_shadow: + return state, (), None + outcome = _constraint_shadow_outcome(state, obs, identity) + view = materialize_constraint_shadow(state, (outcome,)) + return state, (outcome,), view + + +def causal_trace_row( + state: DecisionStateV2, + outcomes: Sequence[ActionOutcomeV2], + obs: RawStepObservation, + identity: CausalTraceIdentity, +) -> dict[str, Any]: + """A ``TraceStore`` row for one captured causal decision. + + Identity hashes are lifted to the row top level so a consumer can fail closed on a + mismatched checkpoint/tokenizer without parsing the full state. + """ + return { + "kind": CAUSAL_TRACE_ROW_KIND, + "state": state.to_dict(), + "outcomes": [outcome.to_dict() for outcome in outcomes], + "raw_observation": obs.to_dict(), + "constraint_shadow": obs.constraint_shadow, + "policy_checkpoint_sha": identity.policy_checkpoint_sha, + "tokenizer_sha": identity.tokenizer_sha, + "decode_config_hash": identity.decode_config_hash, + } + + +class CausalTraceWriter: + """Append captured causal decisions to a ``TraceStore`` and track a manifest.""" + + def __init__(self, store: TraceStore, identity: CausalTraceIdentity) -> None: + """Bind the writer to an append-only ``store`` and the shared trajectory identity.""" + self._store = store + self._identity = identity + self._state_count = 0 + self._shadow_count = 0 + self._legal_set_refs: dict[str, int] = {} + self.trajectory_ids: list[str] = [] + + def record(self, obs: RawStepObservation) -> str: + """Append one observation's causal state to the store and return its trajectory id.""" + state, outcomes, _view = emit_causal_decision(obs, self._identity) + trajectory_id = self._store.append( + causal_trace_row(state, outcomes, obs, self._identity) + ) + self.trajectory_ids.append(trajectory_id) + self._state_count += 1 + if obs.constraint_shadow: + self._shadow_count += 1 + ref = obs.legal_set_reference + self._legal_set_refs[ref] = self._legal_set_refs.get(ref, 0) + 1 + return trajectory_id + + def record_all(self, result: CaptureResult) -> list[str]: + """Append every retained observation from a capture, returning their trajectory ids.""" + return [self.record(obs) for obs in result.observations] + + def manifest(self) -> dict[str, Any]: + """Summarize what was written: identity, counts, legal-set reuse, and byte totals.""" + traces_path = Path(self._store.traces_path) + total_bytes = traces_path.stat().st_size if traces_path.exists() else 0 + duplicate_set_reuse = sum(count - 1 for count in self._legal_set_refs.values()) + return { + "kind": "causal_trace_manifest", + "group_id": self._identity.group_id, + "policy_checkpoint_sha": self._identity.policy_checkpoint_sha, + "tokenizer_sha": self._identity.tokenizer_sha, + "decode_config_hash": self._identity.decode_config_hash, + "base_model_revision": self._identity.base_model_revision, + "adapter_identity": self._identity.adapter_identity, + "state_count": self._state_count, + "constraint_shadow_count": self._shadow_count, + "unique_legal_sets": len(self._legal_set_refs), + "duplicate_set_reuse": duplicate_set_reuse, + "total_bytes": total_bytes, + "bytes_per_state": ( + total_bytes / self._state_count if self._state_count else 0.0 + ), + } + + def write_manifest(self, path: Path | str) -> None: + """Atomically write the manifest JSON to ``path`` (tmp file + fsync + replace).""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + fd, raw = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(path.parent)) + tmp = Path(raw) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(self.manifest(), handle, indent=2, sort_keys=True, allow_nan=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + finally: + tmp.unlink(missing_ok=True) + + +def load_causal_decision_states( + store: TraceStore, + *, + expected_checkpoint_sha: str, + expected_tokenizer_sha: str, +) -> list[DecisionStateV2]: + """Load captured causal states, failing closed on a mismatched identity. + + Loading against a checkpoint or tokenizer other than the one that produced the + trace raises before any state is returned — rollouts from different checkpoints + must never be mixed. ``DecisionStateV2.from_dict`` additionally re-verifies each + state id (tamper check). + """ + states: list[DecisionStateV2] = [] + for row in store.iter_kind(CAUSAL_TRACE_ROW_KIND): + if row.get("policy_checkpoint_sha") != expected_checkpoint_sha: + raise ValueError("causal decision states do not match the policy checkpoint") + if row.get("tokenizer_sha") != expected_tokenizer_sha: + raise ValueError("causal decision states do not match the checkpoint tokenizer") + state = DecisionStateV2.from_dict(row["state"]) + # The nested state must agree with the row envelope and the requested identity, + # not merely be internally consistent. + if state.policy_checkpoint_sha != expected_checkpoint_sha: + raise ValueError("causal decision state does not match the policy checkpoint") + if state.tokenizer_sha != expected_tokenizer_sha: + raise ValueError("causal decision state does not match the checkpoint tokenizer") + states.append(state) + return states + + +@dataclass(frozen=True) +class GeneratedOutcome: + """A forced-action replay outcome, handed to the shared counterfactual owner. + + The plug-in never runs a judge: it returns the replay-side fields only (action, + seed, finish reason, raw program, canonical program when valid). The counterfactual + owner scores it with ``semantic_outcome`` and partitions with + ``label_pareto_candidates``. ``to_candidate`` yields the pre-judge candidate dict + that owner consumes. + """ + + action_id: int + continuation_seed: int + finish_reason: str + raw_program: str + canonical_program: str | None + selected: bool = True + completion_source: str = "policy" + + def to_dict(self) -> dict[str, Any]: + """Serialize the full replay outcome (including the raw/canonical program) to JSON.""" + return { + "action_id": self.action_id, + "continuation_seed": self.continuation_seed, + "finish_reason": self.finish_reason, + "raw_program": self.raw_program, + "canonical_program": self.canonical_program, + "selected": self.selected, + "completion_source": self.completion_source, + } + + def to_candidate(self) -> dict[str, Any]: + """Project to the pre-judge candidate dict the counterfactual owner scores.""" + canonical = self.canonical_program + return { + "token_id": self.action_id, + "selected": self.selected, + "completion_source": self.completion_source, + "raw_text": self.raw_program, + "text": canonical if canonical is not None else self.raw_program, + "finalization_changed": canonical is not None and canonical != self.raw_program, + } + + +@dataclass(frozen=True) +class CausalTracedGeneration: + """The result of a traced constrained generation on the plug-in.""" + + text: str + result: CaptureResult + valid: bool diff --git a/src/slm_training/models/twotower.py b/src/slm_training/models/twotower.py index 982b94dd..d419b12b 100644 --- a/src/slm_training/models/twotower.py +++ b/src/slm_training/models/twotower.py @@ -856,6 +856,199 @@ def optimizer_parameter_groups(self) -> list[dict[str, list[nn.Parameter]]]: if grouped.get(owner) ] + def attach_adapter(self, spec: Any) -> None: + """Attach a removable low-rank adapter to the denoiser (adapter-only training). + + Fails closed if an adapter is already attached or the spec's base compatibility + fingerprint does not match this model. Every non-adapter parameter is frozen, so + ``trainable_parameters()`` yields only the adapter tensors while the parent + weights stay untouched — ``disable_adapter()`` therefore restores the exact + parent map. The context tower is never adapted. + """ + from slm_training.models.adapters.twotower_adapter import attach_low_rank_adapters + + if getattr(self, "_adapter_spec", None) is not None: + raise ValueError("an adapter is already attached to this model") + expected = getattr(spec, "base_compatibility_fingerprint", "") + if expected and expected != self.compatibility_fingerprint(): + raise ValueError( + "adapter base compatibility fingerprint does not match this model" + ) + # NB: spec.base_checkpoint_sha is provenance only. The architecture fingerprint + # above does not distinguish two checkpoints of the same shape; enforcing the + # exact base checkpoint needs the model to carry its loaded checkpoint identity, + # which is deferred with the checkpoint-interplay work (see the LDI2-01 memo). + self._adapter_modules = attach_low_rank_adapters( + self.denoiser, spec, seed=int(self.config.seed) + ) + for name, parameter in self.named_parameters(): + parameter.requires_grad_("lora_" in name.lower()) + self._adapter_spec = spec + + def has_adapter(self) -> bool: + return getattr(self, "_adapter_spec", None) is not None + + def enable_adapter(self) -> None: + for wrapper in getattr(self, "_adapter_modules", {}).values(): + wrapper.enable_adapter() + + def disable_adapter(self) -> None: + for wrapper in getattr(self, "_adapter_modules", {}).values(): + wrapper.disable_adapter() + + def adapter_parameters(self): + for wrapper in getattr(self, "_adapter_modules", {}).values(): + yield from wrapper.adapter_parameters() + + def active_adapter_identity(self) -> str: + """Content digest of the active adapter tensors, or "" when none are attached.""" + from slm_training.lineage.records import content_sha + + modules = getattr(self, "_adapter_modules", {}) + if not modules: + return "" + return content_sha( + { + key: [ + parameter.detach().to(torch.float64).cpu().flatten().tolist() + for parameter in wrapper.adapter_parameters() + ] + for key, wrapper in sorted(modules.items()) + } + ) + + def merge_adapter_copy(self) -> "TwoTowerModel": + """Return a wrapper-free copy with the adapter delta folded into the weights. + + Merge is one-way and on a **copy**: this model and its removable adapter are + left untouched. Every ``LowRankAdapter`` in the copy's denoiser is replaced by a + plain ``nn.Linear`` equal to the adapter-enabled map, so the merged model carries + no active wrappers and trains as an ordinary full model. + """ + import copy + + from slm_training.models.adapters.low_rank import LowRankAdapter + + if not self.has_adapter(): + raise ValueError("no adapter is attached to merge") + merged = copy.deepcopy(self) + + def _fold(module: nn.Module) -> None: + for name, child in list(module.named_children()): + if isinstance(child, LowRankAdapter): + setattr(module, name, child.merged_linear()) + else: + _fold(child) + + _fold(merged.denoiser) + merged._adapter_modules = {} + merged._adapter_spec = None + for parameter in merged.parameters(): + parameter.requires_grad_(True) + return merged + + def save_adapter(self, path: Path | str, *, provenance: dict[str, Any] | None = None) -> None: + """Write the removable adapter (config + tensors + manifest) to its own directory. + + The base checkpoint is not duplicated — only the adapter tensors and the identity + needed to fail closed on load. Requires an attached adapter. + """ + import json + + if not self.has_adapter(): + raise ValueError("no adapter is attached to save") + path = Path(path) + path.mkdir(parents=True, exist_ok=True) + spec = self._adapter_spec + (path / "adapter_config.json").write_text( + json.dumps(spec.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + tensors = { + name: parameter.detach().cpu() + for name, parameter in self.named_parameters() + if "lora_" in name.lower() + } + torch.save(tensors, path / "adapter_model.pt") + manifest = { + "kind": "twotower_low_rank_adapter", + "schema_version": spec.schema_version, + "module_map": sorted(self._adapter_modules), + "parameter_names": sorted(tensors), + "parameter_shapes": {name: list(t.shape) for name, t in tensors.items()}, + "trainable_parameter_count": int(sum(t.numel() for t in tensors.values())), + "adapter_bytes": int((path / "adapter_model.pt").stat().st_size), + "base_compatibility_fingerprint": spec.base_compatibility_fingerprint, + "base_checkpoint_sha": spec.base_checkpoint_sha, + "tokenizer_sha": spec.tokenizer_sha, + "provenance": provenance or {}, + } + (path / "adapter_manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + def load_adapter(self, path: Path | str, *, trainable: bool = False) -> Any: + """Attach and load a removable adapter, failing closed on an identity mismatch. + + The adapter's base fingerprint, tokenizer identity, and resolved module map must + match this model, or loading raises before any tensor is copied. + """ + import json + + from slm_training.models.adapters.spec import TwoTowerAdapterSpec + + path = Path(path) + spec = TwoTowerAdapterSpec.from_dict( + json.loads((path / "adapter_config.json").read_text(encoding="utf-8")) + ) + manifest_path = path / "adapter_manifest.json" + if not manifest_path.exists(): + raise ValueError("adapter directory is missing adapter_manifest.json") + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + + # --- Validate the whole artifact BEFORE mutating the model --------------------- + # Tokenizer + config/manifest integrity are all checked up front so a mismatched + # or truncated adapter never leaves the model half-attached. + expected_tokenizer = self.artifact_identity()["tokenizer_sha"] + if spec.tokenizer_sha and spec.tokenizer_sha != expected_tokenizer: + raise ValueError("adapter tokenizer identity does not match this model") + if manifest.get("base_compatibility_fingerprint") != spec.base_compatibility_fingerprint: + raise ValueError("adapter manifest disagrees with its config on base identity") + tensors = torch.load(path / "adapter_model.pt", map_location="cpu", weights_only=True) + loaded_names = set(tensors) + expected_names = set(manifest.get("parameter_names", [])) + if loaded_names != expected_names: + raise ValueError( + "adapter tensors do not match the manifest parameter set " + f"(missing={sorted(expected_names - loaded_names)}, " + f"unexpected={sorted(loaded_names - expected_names)})" + ) + expected_shapes = manifest.get("parameter_shapes", {}) + for name, tensor in tensors.items(): + if list(tensor.shape) != list(expected_shapes.get(name, [])): + raise ValueError( + f"adapter tensor {name!r} shape {list(tensor.shape)} does not match " + f"manifest shape {expected_shapes.get(name)}" + ) + + # --- Mutate only after every check above has passed --------------------------- + # attach_adapter fails closed on a base-fingerprint mismatch before wrapping. + self.attach_adapter(spec) + parameters = dict(self.named_parameters()) + model_adapter_names = {name for name in parameters if "lora_" in name.lower()} + if loaded_names != model_adapter_names: + raise ValueError( + "adapter tensors do not match the attached module map " + f"(missing={sorted(model_adapter_names - loaded_names)}, " + f"unexpected={sorted(loaded_names - model_adapter_names)})" + ) + with torch.no_grad(): + for name, tensor in tensors.items(): + target = parameters[name] + target.copy_(tensor.to(target.device, target.dtype)) + for name, parameter in self.named_parameters(): + parameter.requires_grad_(bool(trainable) and "lora_" in name.lower()) + return spec + def _count_context_tokens(self, text: str) -> int: """Context token count under the active backend (capped at max_prompt_len).""" if is_hf_context(self.context): diff --git a/tests/test_harnesses/preference/test_decision_diagnostics.py b/tests/test_harnesses/preference/test_decision_diagnostics.py new file mode 100644 index 00000000..cbc44f2a --- /dev/null +++ b/tests/test_harnesses/preference/test_decision_diagnostics.py @@ -0,0 +1,129 @@ +"""Tests for bounded local-decision diagnostics (LDI0-03). + +Exercises the cumulative wall-time deadline with a deterministic fake clock: an +expired run produces no result artifact, and a full-parameter request is refused +as not_authorized. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +from slm_training.harnesses.preference import decision_diagnostics as dd +from slm_training.harnesses.preference.decision_diagnostics import ( + DiagnosticBudget, + not_authorized_report, + run_bounded_stages, + tier1_objective_geometry, + tier2_subspace_gradients, + write_diagnostic_report, +) + + +def _view(good: tuple[int, ...], bad: tuple[int, ...]) -> SimpleNamespace: + """A getattr-compatible stand-in for ObjectiveView (good/bad action ids).""" + return SimpleNamespace(good_action_ids=tuple(good), bad_action_ids=tuple(bad)) + + +class _FakeClock: + """Deterministic monotonic clock: pops queued readings, then holds the last.""" + + def __init__(self, values: list[float]) -> None: + self._values = list(values) + self._last = self._values[-1] if self._values else 0.0 + + def __call__(self) -> float: + if self._values: + self._last = self._values.pop(0) + return self._last + + +def test_budget_capped_at_five_minutes() -> None: + assert DiagnosticBudget().max_wall_minutes == 5.0 + with pytest.raises(ValueError, match=r"\(0, 5\]"): + DiagnosticBudget(max_wall_minutes=6.0) + with pytest.raises(ValueError): + DiagnosticBudget(max_wall_minutes=0.0) + + +def test_completed_run_reports_all_stages(monkeypatch) -> None: + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0])) + report = run_bounded_stages([("a", lambda: {"x": 1}), ("b", lambda: {"y": 2})]) + assert report["status"] == "completed" + assert report["result"] == {"a": {"x": 1}, "b": {"y": 2}} + assert [stage["stage"] for stage in report["stage_telemetry"]] == ["a", "b"] + + +def test_deadline_expiry_produces_no_result(monkeypatch) -> None: + # start=0 -> deadline=300s; the next monotonic read is 500 -> expired before any stage. + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0, 500.0])) + ran: list[str] = [] + report = run_bounded_stages( + [("a", lambda: ran.append("a")), ("b", lambda: ran.append("b"))] + ) + assert report["status"] == "expired" + assert report["result"] is None + assert ran == [] + + +def test_not_authorized_report_has_no_result() -> None: + report = not_authorized_report("full-parameter Tier-2 is not authorized") + assert report["status"] == "not_authorized" + assert report["result"] is None + assert "not authorized" in report["reason"] + + +def test_report_write_is_atomic(tmp_path) -> None: + path = tmp_path / "diag.json" + write_diagnostic_report(path, {"status": "expired", "result": None}) + assert json.loads(path.read_text())["result"] is None + + +def test_tier1_objective_geometry_flags_contradiction(monkeypatch) -> None: + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0])) + # One state, two views: action 4 is good in the first view but bad in the second. + per_state = [[_view((4,), (9,)), _view((5,), (4,))]] + report = tier1_objective_geometry(per_state) + assert report["status"] == "completed" + geometry = report["result"]["objective_geometry"] + assert geometry["states"] == 1 + assert geometry["objective_contradictions"] == 1 + assert geometry["mean_good_set_overlap"] == 0.0 + + +def test_tier1_objective_geometry_agrees_when_views_match(monkeypatch) -> None: + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0])) + per_state = [[_view((4,), (9,)), _view((4,), (9,))]] + geometry = tier1_objective_geometry(per_state)["result"]["objective_geometry"] + assert geometry["objective_contradictions"] == 0 + assert geometry["mean_good_set_overlap"] == 1.0 + + +def test_tier1_objective_geometry_respects_deadline(monkeypatch) -> None: + # start=0 -> deadline=300s; the next read is 500 -> expired before the stage runs. + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0, 500.0])) + report = tier1_objective_geometry([[_view((4,), (9,))]]) + assert report["status"] == "expired" + assert report["result"] is None + + +def test_tier2_refuses_full_parameter_request() -> None: + for subset in (None, []): + report = tier2_subspace_gradients(trainable_parameter_subset=subset) + assert report["status"] == "not_authorized" + assert report["result"] is None + assert "adapter subset" in report["reason"] + + +def test_tier2_plans_over_explicit_subset(monkeypatch) -> None: + monkeypatch.setattr(dd.time, "monotonic", _FakeClock([0.0])) + report = tier2_subspace_gradients( + trainable_parameter_subset=["adapter.0.weight", "adapter.0.bias"] + ) + assert report["status"] == "completed" + plan = report["result"]["subspace_plan"] + assert plan["parameter_subset_size"] == 2 + assert plan["gradients"] == "deferred_to_model_stage" diff --git a/tests/test_harnesses/preference/test_decision_events_v2.py b/tests/test_harnesses/preference/test_decision_events_v2.py new file mode 100644 index 00000000..ca753f11 --- /dev/null +++ b/tests/test_harnesses/preference/test_decision_events_v2.py @@ -0,0 +1,393 @@ +"""Tests for DecisionEventV2 (LDI0-02). + +Contract: ``docs/design/local-decision-interventions.md``. These exercise the +E284-blocker fix — a decision state's identity is separated from its per-action +verifier evidence and from any materialized objective view — plus the fail-closed +validation, append-only evidence merge, honest V1 migration, and the +non-trainable constraint-shadow materializer. +""" + +from __future__ import annotations + +import pytest + +from slm_training.harnesses.preference.counterfactuals import _METRICS +from slm_training.harnesses.preference.decision_events_v2 import ( + GATE_ORDER, + ActionOutcomeV2, + DecisionStateV2, + ObjectiveView, + admit_semantic_corpus, + append_action_outcomes, + decision_state_manifest, + materialize_constraint_shadow, + materialize_pareto, + materialize_set_valued, + materialize_single_best_worst, + materialize_thresholded, + migrate_v1_event, + objective_view_support, + write_action_outcomes, +) +from slm_training.harnesses.preference.local_decisions import ( + DecisionEventV1, + split_for_group, +) + +GROUP = "ldi0-02-group" +GATE_VECTOR = tuple((gate, "skip") for gate in GATE_ORDER) + + +def _state(**overrides: object) -> DecisionStateV2: + kwargs: dict[str, object] = { + "group_id": GROUP, + "architecture": "twotower", + "context_text": "root=Stack([", + "canvas_ids": (1, 2, 3), + "decision_position": 1, + "legal_action_ids": (4, 9, 10), + "decision_kind": "component", + "abstract_state_role": "component_slot", + "grammar_state_hash": "gsh", + "policy_checkpoint_sha": "pcs", + "tokenizer_sha": "tsha", + "decode_config_hash": "dch", + "verifier_bundle_hash": "vbh", + "split": split_for_group(GROUP), + } + kwargs.update(overrides) + return DecisionStateV2(**kwargs) # type: ignore[arg-type] + + +def _outcome(action_id: int, *, reward: float = 0.5, **overrides: object) -> ActionOutcomeV2: + kwargs: dict[str, object] = { + "state_id": _state().state_id, + "action_id": action_id, + "legal": True, + "rollout_policy_sha": "pcs", + "continuation_seeds": (0,), + "verifier_vectors": (GATE_VECTOR,), + "reward_vectors": (tuple((metric, reward) for metric in _METRICS),), + "evidence_ids": (f"probe:{action_id}",), + "evidence_confidence": 0.9, + } + kwargs.update(overrides) + return ActionOutcomeV2(**kwargs) # type: ignore[arg-type] + + +def _v1_event(**overrides: object) -> DecisionEventV1: + kwargs: dict[str, object] = { + "event_id": "e1", + "group_id": GROUP, + "context_text": "root=Stack([", + "canvas_ids": (1, 2, 3), + "position": 1, + "good_token_ids": (4,), + "bad_token_ids": (9,), + "legal_token_ids": (4, 9, 10), + "evidence_kind": "counterfactual", + "evidence_confidence": 0.9, + "decision_kind": "component", + "split": split_for_group(GROUP), + "policy_checkpoint_sha": "pcs", + "tokenizer_sha": "tsha", + "decode_config_hash": "dch", + "seed": 0, + "trajectory_id": "traj-1", + } + kwargs.update(overrides) + return DecisionEventV1(**kwargs) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- # +# State identity is separate from action evidence (the E284 fix). +# --------------------------------------------------------------------------- # +def test_state_id_is_independent_of_action_evidence_order_and_augmentation() -> None: + state = _state() + a = [_outcome(4), _outcome(9)] + b = [_outcome(9), _outcome(4), _outcome(4)] # reordered + duplicated + # The state id depends only on the state; evidence never touches it. + assert _state().state_id == state.state_id + merged_a = append_action_outcomes((), a) + merged_b = append_action_outcomes((), b) + assert {o.action_id for o in merged_a} == {o.action_id for o in merged_b} + assert state.state_id == _state().state_id # unchanged by any evidence handling + + +def test_two_label_samples_of_one_state_share_one_state_id() -> None: + # Same exact state, different sampled good/bad labels -> one state row. + assert _state().state_id == _state().state_id + # A different legal set is a different state. + assert _state(legal_action_ids=(4, 9)).state_id != _state().state_id + + +def test_state_id_changes_for_every_hard_state_field() -> None: + base = _state().state_id + assert _state(context_text="other").state_id != base + assert _state(decision_position=2).state_id != base + assert _state(decision_kind="bind").state_id != base + assert _state(policy_checkpoint_sha="other").state_id != base + assert _state(verifier_bundle_hash="other").state_id != base + assert _state(grammar_state_hash="other").state_id != base + + +def test_state_round_trips_and_rejects_tampered_id() -> None: + state = _state() + assert DecisionStateV2.from_dict(state.to_dict()) == state + data = state.to_dict() + data["state_id"] = "deadbeef" + with pytest.raises(ValueError, match="tampered"): + DecisionStateV2.from_dict(data) + + +# --------------------------------------------------------------------------- # +# Fail-closed validation. +# --------------------------------------------------------------------------- # +def test_state_rejects_split_not_derived_from_group() -> None: + other = "train" if split_for_group(GROUP) == "held_out" else "held_out" + with pytest.raises(ValueError, match="split must be derived"): + _state(split=other) + + +def test_state_rejects_unknown_fields() -> None: + data = _state().to_dict() + data["surprise"] = 1 + with pytest.raises(ValueError, match="unknown decision state fields"): + DecisionStateV2.from_dict(data) + + +def test_causal_state_requires_prefix_and_twotower_requires_canvas() -> None: + with pytest.raises(ValueError, match="causal state requires context_ids"): + _state(architecture="causal", canvas_ids=None, context_ids=None, decision_position=3) + with pytest.raises(ValueError, match="twotower state requires canvas_ids"): + _state(canvas_ids=None) + # A valid causal state keys off the prefix. + causal = _state(architecture="causal", canvas_ids=None, context_ids=(1, 2, 3), decision_position=3) + assert causal.state_id != _state().state_id + + +def test_action_outcome_requires_complete_ordered_gate_vector() -> None: + partial = GATE_VECTOR[:5] + with pytest.raises(ValueError, match="complete ordered G0-G12"): + _outcome(4, verifier_vectors=(partial,)) + with pytest.raises(ValueError, match="gate status"): + _outcome(4, verifier_vectors=(tuple((gate, "maybe") for gate in GATE_ORDER),)) + + +def test_action_outcome_round_trips() -> None: + outcome = _outcome(4) + assert ActionOutcomeV2.from_dict(outcome.to_dict()) == outcome + + +# --------------------------------------------------------------------------- # +# Append-only evidence merge. +# --------------------------------------------------------------------------- # +def test_append_only_merge_dedups_by_content_not_order(tmp_path) -> None: + first = [_outcome(4)] + # Re-observing the identical evidence must collapse; new evidence appends. + second = [_outcome(4), _outcome(9, reward=0.2)] + merged = append_action_outcomes(first, second) + assert [o.action_id for o in merged] == [4, 9] + path = tmp_path / "outcomes.jsonl" + assert write_action_outcomes(path, merged) == 2 + + +# --------------------------------------------------------------------------- # +# Objective materialization. +# --------------------------------------------------------------------------- # +def test_pareto_view_is_trainable_and_partitions_actions() -> None: + state = _state() + good = _outcome(4, reward=0.9) + bad = _outcome(9, reward=0.1) + view = materialize_pareto(state, [good, bad]) + assert view.trainable is True + assert view.good_action_ids == (4,) + assert view.bad_action_ids == (9,) + # Legal action 10 was never observed. + assert 10 in view.unobserved_action_ids + + +def test_constraint_shadow_view_is_not_trainable() -> None: + state = _state() + view = materialize_constraint_shadow(state, [_outcome(4), _outcome(9)]) + assert view.trainable is False + assert view.good_action_ids == () + assert view.bad_action_ids == () + assert set(view.ambiguous_action_ids) == {4, 9} + + +def test_materializer_rejects_foreign_state_outcome() -> None: + other_state = _state(context_text="different") + foreign = _outcome(4, state_id=other_state.state_id) + with pytest.raises(ValueError, match="does not belong to this state"): + materialize_pareto(_state(), [foreign]) + + +def test_semantic_materializer_rejects_illegal_or_out_of_domain_outcome() -> None: + state = _state() + # An illegal outcome must never become a semantic good/bad label. + with pytest.raises(ValueError, match="requires a legal action outcome"): + materialize_pareto(state, [_outcome(4, legal=False)]) + # An action outside the state's legal set is rejected before materialization. + with pytest.raises(ValueError, match="outside the state's legal set"): + materialize_thresholded(state, [_outcome(99)]) + + +def test_action_outcome_rejects_incomplete_reward_vector() -> None: + # A partial reward vector must fail at construction, not fabricate a 0.0 later. + with pytest.raises(ValueError, match="exactly the metrics"): + _outcome(4, reward_vectors=((("reward", 0.5),),)) + + +# --------------------------------------------------------------------------- # +# V1 migration is honest about incompleteness. +# --------------------------------------------------------------------------- # +def test_v1_counterfactual_migrates_with_partial_incomplete_evidence() -> None: + migrated = migrate_v1_event(_v1_event()) + assert migrated.complete is False + assert migrated.state.architecture == "twotower" + assert migrated.state.split == split_for_group(GROUP) + actions = {o.action_id for o in migrated.outcomes} + assert actions == {4, 9} + # No replayable verifier/reward evidence may be fabricated from V1. + assert all(o.verifier_vectors == () for o in migrated.outcomes) + assert all(o.reward_vectors == () for o in migrated.outcomes) + + +def test_v1_constraint_shadow_migrates_as_legality_diagnostic() -> None: + migrated = migrate_v1_event(_v1_event(evidence_kind="constraint_shadow")) + assert migrated.complete is False + assert "legality diagnostic" in migrated.note + + +def test_v1_migration_is_idempotent() -> None: + event = _v1_event() + assert migrate_v1_event(event) == migrate_v1_event(event) + + +# --------------------------------------------------------------------------- # +# Remaining materializer views. +# --------------------------------------------------------------------------- # +def test_thresholded_view_respects_threshold_and_confidence() -> None: + state = _state() + high = _outcome(4, reward=0.8) + low = _outcome(9, reward=0.2) + unsure = _outcome(10, reward=0.9, evidence_confidence=0.1) + view = materialize_thresholded( + state, [high, low, unsure], metric="reward", threshold=0.5, min_confidence=0.5 + ) + assert view.good_action_ids == (4,) + assert view.bad_action_ids == (9,) + assert view.ambiguous_action_ids == (10,) # below the confidence floor + assert view.trainable is True + + +def test_single_best_worst_view() -> None: + state = _state() + view = materialize_single_best_worst( + state, [_outcome(4, reward=0.9), _outcome(9, reward=0.1)], metric="reward" + ) + assert view.good_action_ids == (4,) + assert view.bad_action_ids == (9,) + + +def test_set_valued_view_uses_verifier_verdicts() -> None: + state = _state() + passing = _outcome(4) # all-skip gate vector -> no failing gate -> verified + failing = _outcome( + 9, verifier_vectors=(tuple((gate, "fail" if gate == "G2" else "skip") for gate in GATE_ORDER),) + ) + view = materialize_set_valued(state, [passing, failing]) + assert view.good_action_ids == (4,) + assert view.bad_action_ids == (9,) + assert view.trainable is True + + +# --------------------------------------------------------------------------- # +# Manifest separately fingerprints states / evidence / views. +# --------------------------------------------------------------------------- # +def test_manifest_fingerprints_are_separate_and_order_independent() -> None: + state = _state() + outcomes = [_outcome(4), _outcome(9)] + view = materialize_pareto(state, outcomes) + manifest = decision_state_manifest([state], outcomes, [view]) + reordered = decision_state_manifest([state], list(reversed(outcomes)), [view]) + # Row order does not change the fingerprint (evidence is content-sorted). + assert manifest["outcome_fingerprint"] == reordered["outcome_fingerprint"] + # The three concerns fingerprint independently. + assert manifest["state_fingerprint"] != manifest["outcome_fingerprint"] + assert manifest["view_fingerprint"] != manifest["state_fingerprint"] + + +# --------------------------------------------------------------------------- # +# Objective-support admission (the E284 fix). +# --------------------------------------------------------------------------- # +def _split_group(split: str) -> str: + index = 0 + while split_for_group(f"obj{index}") != split: + index += 1 + return f"obj{index}" + + +def _state_for(group: str) -> DecisionStateV2: + return _state(group_id=group, split=split_for_group(group)) + + +def _view( + good: tuple[int, ...], + bad: tuple[int, ...], + *, + materializer_id: str = "pareto_v2", + trainable: bool = True, +) -> ObjectiveView: + return ObjectiveView( + good_action_ids=good, + bad_action_ids=bad, + ambiguous_action_ids=(), + unobserved_action_ids=(), + weights=(), + materializer_id=materializer_id, + materializer_config_hash="cfg", + trainable=trainable, + ) + + +def test_objective_support_gap_matches_e284_pattern() -> None: + train_state = _state_for(_split_group("train")) + held_state = _state_for(_split_group("held_out")) + # State support (good-only) matches, but objective support (good+bad) differs, + # so a held-out objective signature is uncovered — exactly the E284 mechanism. + train_view = _view((4,), (9,)) + held_view = _view((4,), (10,)) + report = objective_view_support([(train_state, train_view), (held_state, held_view)]) + assert report["held_out_coverage"]["passed"] is False + assert len(report["held_out_coverage"]["uncovered"]) == 1 + with pytest.raises(ValueError, match="lacks train support"): + admit_semantic_corpus( + [(train_state, train_view), (held_state, held_view)], materializer_id="pareto_v2" + ) + + +def test_objective_support_passes_when_held_out_signature_is_covered() -> None: + view = _view((4,), (9,)) + report = admit_semantic_corpus( + [(_state_for(_split_group("train")), view), (_state_for(_split_group("held_out")), view)], + materializer_id="pareto_v2", + ) + assert report["held_out_coverage"]["passed"] is True + + +def test_admission_refuses_non_trainable_constraint_shadow() -> None: + shadow = _view((), (), materializer_id="constraint_shadow_diagnostic_v2", trainable=False) + with pytest.raises(ValueError, match="non-trainable"): + admit_semantic_corpus( + [(_state_for(_split_group("train")), shadow)], materializer_id="pareto_v2" + ) + + +def test_admission_refuses_materializer_mismatch() -> None: + view = _view((4,), (9,), materializer_id="thresholded_v2") + with pytest.raises(ValueError, match="do not match the requested"): + admit_semantic_corpus( + [(_state_for(_split_group("train")), view)], materializer_id="pareto_v2" + ) diff --git a/tests/test_harnesses/preference/test_local_decisions.py b/tests/test_harnesses/preference/test_local_decisions.py index 129e2281..5ca88242 100644 --- a/tests/test_harnesses/preference/test_local_decisions.py +++ b/tests/test_harnesses/preference/test_local_decisions.py @@ -14,6 +14,7 @@ events_from_trace, load_decision_events, load_trace_rows, + objective_signature_support, split_for_group, write_decision_events, ) @@ -284,3 +285,68 @@ def test_support_signature_ignores_sampled_negative_variation() -> None: assert decision_support_signature(changed_bad) == decision_support_signature(event) assert decision_support_signature(changed_legal) != decision_support_signature(event) + + +def _objective_event(group: str, split: str, bad: int) -> DecisionEventV1: + return DecisionEventV1( + event_id=f"e-{group}-{bad}", + group_id=group, + context_text="root=Stack([", + canvas_ids=(1, 2, 3), + position=1, + good_token_ids=(4,), + bad_token_ids=(bad,), + legal_token_ids=(4, 9, 10), + evidence_kind="counterfactual", + evidence_confidence=0.9, + decision_kind="component", + split=split, + policy_checkpoint_sha="pcs", + tokenizer_sha="tsha", + decode_config_hash="dch", + seed=0, + trajectory_id="traj", + ) + + +def _grow_group(prefix: str, split: str) -> str: + group = prefix + while split_for_group(group) != split: + group += "x" + return group + + +def test_objective_signature_support_catches_e284_pattern() -> None: + train = _objective_event(_grow_group("obj-train", "train"), "train", 9) + held = _objective_event(_grow_group("obj-held", "held_out"), "held_out", 10) + # The state-support signature (good-only) matches, so state support passes... + assert decision_signature_support([train, held])["held_out_coverage"]["passed"] is True + # ...but the objective signature (good + bad) differs, so objective support fails. + objective = objective_signature_support([train, held]) + assert objective["held_out_coverage"]["passed"] is False + assert len(objective["held_out_coverage"]["uncovered"]) == 1 + + +def test_train_local_from_paths_refuses_objective_unsupported_corpus(tmp_path) -> None: + from pathlib import Path + + from slm_training.harnesses.preference.local_train import train_local_from_paths + + events_path = tmp_path / "events.jsonl" + write_decision_events( + events_path, + [ + _objective_event(_grow_group("obj-train", "train"), "train", 9), + _objective_event(_grow_group("obj-held", "held_out"), "held_out", 10), + ], + ) + # The admission refusal fires before the checkpoint is loaded, so the missing + # checkpoint is never reached. + with pytest.raises(ValueError, match="objective signature"): + train_local_from_paths( + Path("nonexistent.pt"), + events_path, + out_dir=tmp_path / "out", + objective="ce_margin", + require_objective_support=True, + ) diff --git a/tests/test_models/fixtures/ldi1/README.md b/tests/test_models/fixtures/ldi1/README.md new file mode 100644 index 00000000..555fdbd1 --- /dev/null +++ b/tests/test_models/fixtures/ldi1/README.md @@ -0,0 +1,20 @@ +# LDI1-01 fixture-grade causal decision trace + +Deterministic, torch-free evidence for exact causal decision-state capture +(SLM-119 / LDI1-01). Regenerate with the recipe in +`docs/design/iter-ldi1-01-causal-capture-20260718.md`; `test_causal_trace_fixture.py` +loads and checks it. + +- `traces.jsonl` — four captured causal decisions (a constraint shadow, a forced + deduction, an ordinary decision, and the EOS stop) through the shared `TraceStore`. + Integer prefix ids are the state authority; per-step raw/legal distribution telemetry + is recorded. +- `causal_trace_manifest.json` — model/tokenizer/adapter/decode identities and state + counts (bytes/state and duplicate-set reuse). No timestamps, so it is reproducible. +- `forced_counterfactual_outcome.json` — one forced legal action replay outcome handed to + the counterfactual owner (no judge, no label). `canonical_program` is `null`: + canonicalization is the strict validator's job and the component catalog is + environment/pack-dependent, so the fixture does not assert the raw program is valid. + +This artifact carries **no semantic label and makes no model-quality claim**; the +constraint-shadow evidence is legality-only and non-trainable by construction. diff --git a/tests/test_models/fixtures/ldi1/causal_trace_manifest.json b/tests/test_models/fixtures/ldi1/causal_trace_manifest.json new file mode 100644 index 00000000..14f87727 --- /dev/null +++ b/tests/test_models/fixtures/ldi1/causal_trace_manifest.json @@ -0,0 +1,15 @@ +{ + "adapter_identity": "adapter-A", + "base_model_revision": "rev-fixture", + "bytes_per_state": 1788.75, + "constraint_shadow_count": 1, + "decode_config_hash": "decode-config-hash-fixture", + "duplicate_set_reuse": 0, + "group_id": "ldi1-fixture-group", + "kind": "causal_trace_manifest", + "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", + "state_count": 4, + "tokenizer_sha": "tokenizer-sha-fixture", + "total_bytes": 7155, + "unique_legal_sets": 4 +} diff --git a/tests/test_models/fixtures/ldi1/forced_counterfactual_outcome.json b/tests/test_models/fixtures/ldi1/forced_counterfactual_outcome.json new file mode 100644 index 00000000..fb34dfb2 --- /dev/null +++ b/tests/test_models/fixtures/ldi1/forced_counterfactual_outcome.json @@ -0,0 +1,9 @@ +{ + "action_id": 3, + "canonical_program": null, + "completion_source": "policy", + "continuation_seed": 7, + "finish_reason": "eos", + "raw_program": "root=Stack([Text()])", + "selected": true +} diff --git a/tests/test_models/fixtures/ldi1/traces.jsonl b/tests/test_models/fixtures/ldi1/traces.jsonl new file mode 100644 index 00000000..13b4a13b --- /dev/null +++ b/tests/test_models/fixtures/ldi1/traces.jsonl @@ -0,0 +1,4 @@ +{"trajectory_id": "00000000-c8f7847e86a4", "trace_id": "ldi1-trace", "span_id": "ldi1-span", "run_id": "ldi1-fixture", "kind": "causal_decision", "state": {"group_id": "ldi1-fixture-group", "architecture": "causal", "context_text": "root=Stack([", "decision_position": 0, "legal_action_ids": [2, 3], "decision_kind": "constraint_shadow", "abstract_state_role": "constraint_shadow", "grammar_state_hash": "18b058689d010062a25cd7fb949b2abfea1757d2de45ca3abdda17db956e506b", "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture", "verifier_bundle_hash": "51669ede9e42f7ebb995ef7ab1cd410ef453e2838d3d84cd006381e05141238f", "split": "train", "context_ids": [], "canvas_ids": null, "generation_step": 0, "version": 2, "state_id": "ac40cdff11fa60df878f5f40cbba889d292f5ab2b71d4c02023cb00ea3616b26"}, "outcomes": [{"state_id": "ac40cdff11fa60df878f5f40cbba889d292f5ab2b71d4c02023cb00ea3616b26", "action_id": 3, "legal": true, "rollout_policy_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "continuation_seeds": [], "outcome_hashes": [], "verifier_vectors": [], "reward_vectors": [], "mean_value": null, "confidence_interval": null, "evidence_ids": [], "evidence_confidence": 0.0, "version": 2}], "raw_observation": {"decision_index": 0, "generated_ordinal": 0, "prefix_token_ids": [], "raw_argmax_id": 5, "selected_token_id": 3, "legal_token_ids": [2, 3], "legal_set_reference": "18b058689d010062a25cd7fb949b2abfea1757d2de45ca3abdda17db956e506b", "raw_topk": [[5, 3.0, -0.3149288189460638], [3, 0.9, -2.414928818946064], [2, 0.5, -2.814928818946064], [4, 0.2, -3.1149288189460638], [1, 0.1, -3.214928818946064]], "legal_topk": [[3, 0.598687660112452], [2, 0.401312339887548]], "constraint_shadow": true, "forced": false, "grammar_role": null}, "constraint_shadow": true, "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture"} +{"trajectory_id": "00000001-3c3543975db5", "trace_id": "ldi1-trace", "span_id": "ldi1-span", "run_id": "ldi1-fixture", "kind": "causal_decision", "state": {"group_id": "ldi1-fixture-group", "architecture": "causal", "context_text": "root=Stack([", "decision_position": 1, "legal_action_ids": [4], "decision_kind": "forced_deduction", "abstract_state_role": "forced_deduction", "grammar_state_hash": "46b1884167c4edd308bcf0c04163dd02d05c9742b35e86b57b5f7ed1b82f3850", "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture", "verifier_bundle_hash": "62adf959ad468c3976f29b5b0655e4871ef918d3cb0fa191f76059314c72eafa", "split": "train", "context_ids": [3], "canvas_ids": null, "generation_step": 1, "version": 2, "state_id": "c4b286d51fb2f04dd95f7ba9f0fa9c44970abb8926d61dfc4f35e1c10e6c405b"}, "outcomes": [], "raw_observation": {"decision_index": 1, "generated_ordinal": 1, "prefix_token_ids": [3], "raw_argmax_id": 4, "selected_token_id": 4, "legal_token_ids": [4], "legal_set_reference": "46b1884167c4edd308bcf0c04163dd02d05c9742b35e86b57b5f7ed1b82f3850", "raw_topk": [[4, 5.0, -0.04072092377952481], [5, 0.4, -4.640720923779525], [3, 0.3, -4.740720923779525], [2, 0.2, -4.840720923779525], [1, 0.1, -4.9407209237795255]], "legal_topk": [[4, 1.0]], "constraint_shadow": false, "forced": true, "grammar_role": null}, "constraint_shadow": false, "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture"} +{"trajectory_id": "00000002-40cc3d30f3a3", "trace_id": "ldi1-trace", "span_id": "ldi1-span", "run_id": "ldi1-fixture", "kind": "causal_decision", "state": {"group_id": "ldi1-fixture-group", "architecture": "causal", "context_text": "root=Stack([", "decision_position": 2, "legal_action_ids": [2, 4], "decision_kind": "causal_decision", "abstract_state_role": "causal_decision", "grammar_state_hash": "f25ed8cd5464710aff6e5d655a575ac273cb9c19d696e53fa13de05db97454ca", "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture", "verifier_bundle_hash": "62adf959ad468c3976f29b5b0655e4871ef918d3cb0fa191f76059314c72eafa", "split": "train", "context_ids": [3, 4], "canvas_ids": null, "generation_step": 2, "version": 2, "state_id": "764f01b21081088602777196ea7e0026dceb0166aaa45c607f2b66f24f33c1d2"}, "outcomes": [], "raw_observation": {"decision_index": 1, "generated_ordinal": 2, "prefix_token_ids": [3, 4], "raw_argmax_id": 2, "selected_token_id": 2, "legal_token_ids": [2, 4], "legal_set_reference": "f25ed8cd5464710aff6e5d655a575ac273cb9c19d696e53fa13de05db97454ca", "raw_topk": [[2, 2.0, -0.6935299852333918], [4, 1.0, -1.6935299852333918], [5, 0.3, -2.3935299852333918], [3, 0.2, -2.493529985233392], [1, 0.1, -2.593529985233392]], "legal_topk": [[2, 0.7310585786300049], [4, 0.2689414213699951]], "constraint_shadow": false, "forced": false, "grammar_role": null}, "constraint_shadow": false, "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture"} +{"trajectory_id": "00000003-66bff48e87d0", "trace_id": "ldi1-trace", "span_id": "ldi1-span", "run_id": "ldi1-fixture", "kind": "causal_decision", "state": {"group_id": "ldi1-fixture-group", "architecture": "causal", "context_text": "root=Stack([", "decision_position": 3, "legal_action_ids": [0], "decision_kind": "forced_deduction", "abstract_state_role": "forced_deduction", "grammar_state_hash": "d0bca111f8628137adc4c16f123496dcdd1d590d06cb5d9acd68b39fe656fb97", "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture", "verifier_bundle_hash": "62adf959ad468c3976f29b5b0655e4871ef918d3cb0fa191f76059314c72eafa", "split": "train", "context_ids": [3, 4, 2], "canvas_ids": null, "generation_step": 3, "version": 2, "state_id": "a561501189d5960168cac1cca543b2fb56622d709dbeda503e2b50a598fb08d4"}, "outcomes": [], "raw_observation": {"decision_index": 2, "generated_ordinal": 3, "prefix_token_ids": [3, 4, 2], "raw_argmax_id": 0, "selected_token_id": 0, "legal_token_ids": [0], "legal_set_reference": "d0bca111f8628137adc4c16f123496dcdd1d590d06cb5d9acd68b39fe656fb97", "raw_topk": [[0, 4.0, -0.1176561970721951], [5, 0.5, -3.617656197072195], [4, 0.4, -3.7176561970721953], [3, 0.3, -3.8176561970721954], [2, 0.2, -3.917656197072195]], "legal_topk": [[0, 1.0]], "constraint_shadow": false, "forced": true, "grammar_role": null}, "constraint_shadow": false, "policy_checkpoint_sha": "f112adcec23d6510f03c0499326c3a58e2ac963d4f3401f033422651310e95c2", "tokenizer_sha": "tokenizer-sha-fixture", "decode_config_hash": "decode-config-hash-fixture"} diff --git a/tests/test_models/test_adapter_spec.py b/tests/test_models/test_adapter_spec.py new file mode 100644 index 00000000..41d3fb50 --- /dev/null +++ b/tests/test_models/test_adapter_spec.py @@ -0,0 +1,55 @@ +"""Torch-free tests for the TwoTower adapter spec (LDI2-01 / SLM-123).""" + +from __future__ import annotations + +import pytest + +from slm_training.models.adapters import TwoTowerAdapterSpec + + +def _spec(**overrides) -> TwoTowerAdapterSpec: + base = dict( + method="low_rank", + rank=4, + alpha=8.0, + dropout=0.0, + target_modules=("attn_q", "attn_v"), + base_compatibility_fingerprint="base-fp", + base_checkpoint_sha="ckpt-sha", + tokenizer_sha="tok-sha", + ) + base.update(overrides) + return TwoTowerAdapterSpec(**base) + + +def test_spec_round_trips_through_dict() -> None: + spec = _spec(target_layer_indices=(0, 2), include_output_head=True) + restored = TwoTowerAdapterSpec.from_dict(spec.to_dict()) + assert restored == spec + assert restored.scaling == 8.0 / 4 + + +@pytest.mark.parametrize( + "overrides,match", + [ + ({"rank": 0}, "rank must be positive"), + ({"alpha": 0.0}, "alpha must be positive"), + ({"dropout": 1.0}, r"dropout must be in \[0, 1\)"), + ({"method": "dora"}, "only method='low_rank'"), + ({"train_bias": "all"}, "only train_bias='none'"), + ({"target_modules": ()}, "must not be empty"), + ({"target_modules": ("attn_q", "attn_q")}, "must be unique"), + ({"base_checkpoint_sha": ""}, "must be non-empty"), + ({"schema_version": 2}, "unsupported adapter schema version"), + ], +) +def test_spec_validation_rejects_bad_configs(overrides, match) -> None: + with pytest.raises(ValueError, match=match): + _spec(**overrides) + + +def test_from_dict_rejects_unknown_fields() -> None: + data = _spec().to_dict() + data["mystery"] = 1 + with pytest.raises(ValueError, match="unknown adapter spec fields"): + TwoTowerAdapterSpec.from_dict(data) diff --git a/tests/test_models/test_causal_trace.py b/tests/test_models/test_causal_trace.py new file mode 100644 index 00000000..be237ecd --- /dev/null +++ b/tests/test_models/test_causal_trace.py @@ -0,0 +1,251 @@ +"""Torch-free tests for exact causal decision-state capture (LDI1-01 / SLM-119). + +A tiny deterministic decode over fixture logits/legal-sets exercises the whole +capture→emit→persist path without torch, a real tokenizer, or the grammar: + +* constraint shadows fire only when the raw winner is illegal and the selection legal; +* shadow evidence is non-trainable and refused by semantic admission; +* forced (single-legal) steps are deductions, not decisions; +* bounded selection policies retain the right subset without changing the emission; +* integer prefix ids — never decoded text — are the state authority; +* traces round-trip through the shared TraceStore and fail closed on identity mismatch. +""" + +from __future__ import annotations + +import json + +import pytest + +from slm_training.harnesses.distill.trace_store import TraceStore +from slm_training.harnesses.preference.decision_events_v2 import admit_semantic_corpus +from slm_training.models.causal_trace import ( + CausalTraceIdentity, + CausalTraceWriter, + GeneratedOutcome, + TracePolicy, + TraceSelection, + capture_raw_steps, + emit_causal_decision, + fold_policy_identity, + legal_set_reference, + load_causal_decision_states, +) + +# vocab 6, eos=0. A four-step decode over exact prefixes: +# () raw argmax 5 (illegal), legal {2,3}, select 3 -> constraint shadow +# (3,) legal {4} only, select 4 -> forced deduction +# (3,4) raw argmax 2 (legal), legal {2,4}, select 2 -> ordinary decision +# (3,4,2) legal {0} (eos now valid), select 0 -> stop +_LOGITS = { + (): [0.0, 0.1, 0.5, 0.9, 0.2, 3.0], + (3,): [0.0, 0.1, 0.2, 0.3, 5.0, 0.4], + (3, 4): [0.0, 0.1, 2.0, 0.2, 1.0, 0.3], + (3, 4, 2): [4.0, 0.1, 0.2, 0.3, 0.4, 0.5], +} +_LEGAL = { + (): (2, 3), + (3,): (4,), + (3, 4): (2, 4), + (3, 4, 2): (0,), +} + + +def _forward(prefix: tuple[int, ...]) -> list[float]: + return _LOGITS[prefix] + + +def _allowed(prefix: tuple[int, ...]) -> tuple[int, ...]: + return _LEGAL[prefix] + + +def _capture(policy: TracePolicy | None = None, role_of=None): + return capture_raw_steps( + forward_logits=_forward, + allowed_ids=_allowed, + eos_id=0, + max_new_tokens=16, + policy=policy, + role_of=role_of, + ) + + +def _identity(adapter: str = "adapterA") -> CausalTraceIdentity: + return CausalTraceIdentity( + group_id="grp", + context_text="root=Stack([", + policy_checkpoint_sha=fold_policy_identity("base-ckpt", adapter), + tokenizer_sha="tok-sha", + decode_config_hash="dch", + base_model_revision="rev", + adapter_identity=adapter, + ) + + +def test_decode_emits_full_generation_and_stop_reason() -> None: + result = _capture() + assert result.generated_token_ids == (3, 4, 2, 0) + assert result.stop_reason == "eos" + assert result.constraint_shadow_count == 1 + assert len(result.observations) == 4 + + +def test_constraint_shadow_only_when_raw_illegal_and_selection_legal() -> None: + obs = _capture().observations + shadow, forced_step, ordinary, eos_step = obs + assert shadow.constraint_shadow is True + assert shadow.raw_argmax_id == 5 and shadow.selected_token_id == 3 + assert shadow.forced is False + # The ordinary step's raw winner is legal, so no shadow is recorded. + assert ordinary.constraint_shadow is False + assert ordinary.raw_argmax_id == 2 == ordinary.selected_token_id + # Forced (single-legal) steps are deductions, never shadows. + assert forced_step.forced is True and forced_step.constraint_shadow is False + assert eos_step.selected_token_id == 0 + + +def test_decision_index_counts_only_non_forced_steps() -> None: + obs = _capture().observations + # shadow(non-forced)=0, forced keeps 1, ordinary(non-forced)=1, eos forced keeps 2. + assert [o.decision_index for o in obs] == [0, 1, 1, 2] + assert [o.generated_ordinal for o in obs] == [0, 1, 2, 3] + + +def test_eos_selected_only_after_the_prefix_validates() -> None: + # EOS (0) is absent from every legal set until the final validating prefix. + for prefix, legal in _LEGAL.items(): + if prefix != (3, 4, 2): + assert 0 not in legal + assert _capture().generated_token_ids[-1] == 0 + + +def test_integer_prefix_ids_are_the_state_authority() -> None: + shadow = _capture().observations[0] + state, _outcomes, _view = emit_causal_decision(shadow, _identity()) + # The exact integer prefix — not decoded text — identifies the causal state. + assert state.architecture == "causal" + assert state.context_ids == shadow.prefix_token_ids == () + assert state.legal_action_ids == (2, 3) + assert state.decision_position == 0 + + +def test_shadow_view_is_non_trainable_and_refused_by_admission() -> None: + shadow = _capture().observations[0] + state, outcomes, view = emit_causal_decision(shadow, _identity()) + assert state.decision_kind == "constraint_shadow" + assert view is not None and view.trainable is False + assert len(outcomes) == 1 and outcomes[0].legal is True + # A legality-only shadow view cannot supervise a semantic objective. + with pytest.raises(ValueError, match="non-trainable"): + admit_semantic_corpus([(state, view)], materializer_id=view.materializer_id) + + +def test_ordinary_step_emits_replayable_state_without_invented_outcomes() -> None: + ordinary = _capture().observations[2] + state, outcomes, view = emit_causal_decision(ordinary, _identity()) + assert state.decision_kind == "causal_decision" + assert outcomes == () and view is None + + +def test_adapter_identity_is_part_of_state_identity() -> None: + shadow = _capture().observations[0] + on, _o, _v = emit_causal_decision(shadow, _identity("adapterA")) + off, _o2, _v2 = emit_causal_decision(shadow, _identity("")) + assert on.state_id != off.state_id + + +@pytest.mark.parametrize( + "policy,expected_ordinals", + [ + (TracePolicy(selection=TraceSelection.EVERY), [0, 1, 2, 3]), + (TracePolicy(selection=TraceSelection.CONSTRAINT_SHADOW_ONLY), [0]), + (TracePolicy(selection=TraceSelection.MARGIN_THRESHOLD, margin_threshold=0.3), [0]), + (TracePolicy(selection=TraceSelection.SAMPLED_POSITIONS, sampled_positions=(2,)), [2]), + ], +) +def test_selection_policies_retain_the_right_subset(policy, expected_ordinals) -> None: + result = _capture(policy=policy) + assert [o.generated_ordinal for o in result.observations] == expected_ordinals + # Selection never changes what is emitted. + assert result.generated_token_ids == (3, 4, 2, 0) + + +def test_named_role_policy_uses_the_role_callback() -> None: + role_of = {(): "root"}.get + result = _capture( + policy=TracePolicy(selection=TraceSelection.NAMED_ROLES, named_roles=("root",)), + role_of=role_of, + ) + assert [o.generated_ordinal for o in result.observations] == [0] + assert result.observations[0].grammar_role == "root" + + +def test_legal_set_reference_is_order_independent_content_address() -> None: + assert legal_set_reference((2, 3)) == legal_set_reference((3, 2, 3)) + assert legal_set_reference((2, 3)) != legal_set_reference((2, 4)) + + +def test_trace_round_trips_through_shared_store(tmp_path) -> None: + identity = _identity() + store = TraceStore(tmp_path / "traces", run_id="ldi1-fixture") + writer = CausalTraceWriter(store, identity) + writer.record_all(_capture()) + + assert len(store) == 4 + rows = list(store.iter_kind("causal_decision")) + assert len(rows) == 4 + assert all(row["run_id"] == "ldi1-fixture" for row in rows) + + states = load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha=identity.tokenizer_sha, + ) + assert [s.architecture for s in states] == ["causal"] * 4 + assert states[0].context_ids == () + + manifest = writer.manifest() + assert manifest["state_count"] == 4 + assert manifest["constraint_shadow_count"] == 1 + assert manifest["duplicate_set_reuse"] == ( + manifest["state_count"] - manifest["unique_legal_sets"] + ) + assert manifest["bytes_per_state"] > 0 + + manifest_path = tmp_path / "causal_manifest.json" + writer.write_manifest(manifest_path) + assert json.loads(manifest_path.read_text())["state_count"] == 4 + + +def test_load_fails_closed_on_identity_mismatch(tmp_path) -> None: + identity = _identity() + store = TraceStore(tmp_path / "traces", run_id="ldi1-fixture") + CausalTraceWriter(store, identity).record_all(_capture()) + + with pytest.raises(ValueError, match="policy checkpoint"): + load_causal_decision_states( + store, + expected_checkpoint_sha="wrong", + expected_tokenizer_sha=identity.tokenizer_sha, + ) + with pytest.raises(ValueError, match="tokenizer"): + load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha="wrong", + ) + + +def test_generated_outcome_returns_pre_judge_candidate() -> None: + outcome = GeneratedOutcome( + action_id=3, + continuation_seed=7, + finish_reason="eos", + raw_program="root=Stack([Text()])", + canonical_program="root=Stack([Text()])", + ) + candidate = outcome.to_candidate() + # The plug-in never judges: no verified/metrics keys leak into the candidate. + assert candidate["token_id"] == 3 + assert candidate["finalization_changed"] is False + assert "verified" not in candidate and "metrics" not in candidate diff --git a/tests/test_models/test_causal_trace_fixture.py b/tests/test_models/test_causal_trace_fixture.py new file mode 100644 index 00000000..c7d97a75 --- /dev/null +++ b/tests/test_models/test_causal_trace_fixture.py @@ -0,0 +1,79 @@ +"""The committed LDI1-01 fixture trace is load-bearing evidence (SLM-119). + +Asserts the fixture-grade artifact shows exact prefix replay, raw/legal distribution +telemetry, at least one non-admittable constraint shadow, and a forced legal +counterfactual outcome — with no semantic label or model-quality claim. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from slm_training.harnesses.distill.trace_store import TraceStore +from slm_training.harnesses.preference.decision_events_v2 import ( + admit_semantic_corpus, + materialize_constraint_shadow, +) +from slm_training.models.causal_trace import fold_policy_identity, load_causal_decision_states + +_FIXTURE = Path(__file__).parent / "fixtures" / "ldi1" + + +def test_fixture_states_replay_from_exact_integer_prefixes() -> None: + store = TraceStore(_FIXTURE) # read-only: construction does not write + states = load_causal_decision_states( + store, + expected_checkpoint_sha=fold_policy_identity("base-ckpt-sha", "adapter-A"), + expected_tokenizer_sha="tokenizer-sha-fixture", + ) + assert [s.architecture for s in states] == ["causal"] * 4 + # Integer prefixes are the authority; the first decision is at the empty suffix. + assert states[0].context_ids == () + assert states[0].decision_position == 0 + + +def test_fixture_rows_carry_raw_and_legal_telemetry() -> None: + rows = list(TraceStore(_FIXTURE).iter_kind("causal_decision")) + assert len(rows) == 4 + for row in rows: + raw = row["raw_observation"] + assert raw["raw_topk"] and len(raw["raw_topk"][0]) == 3 # (id, logit, logprob) + assert raw["legal_topk"] + assert raw["legal_set_reference"] + + +def test_fixture_has_a_non_admittable_constraint_shadow() -> None: + store = TraceStore(_FIXTURE) + shadow_rows = [ + row for row in store.iter_kind("causal_decision") if row["constraint_shadow"] + ] + assert len(shadow_rows) >= 1 + states = load_causal_decision_states( + store, + expected_checkpoint_sha=fold_policy_identity("base-ckpt-sha", "adapter-A"), + expected_tokenizer_sha="tokenizer-sha-fixture", + ) + shadow_state = next(s for s in states if s.decision_kind == "constraint_shadow") + view = materialize_constraint_shadow(shadow_state, ()) + assert view.trainable is False + with pytest.raises(ValueError, match="non-trainable"): + admit_semantic_corpus([(shadow_state, view)], materializer_id=view.materializer_id) + + +def test_fixture_manifest_and_forced_outcome_are_honest() -> None: + manifest = json.loads((_FIXTURE / "causal_trace_manifest.json").read_text()) + assert manifest["state_count"] == 4 + assert manifest["constraint_shadow_count"] == 1 + assert manifest["bytes_per_state"] > 0 + + outcome = json.loads((_FIXTURE / "forced_counterfactual_outcome.json").read_text()) + # A forced legal action replay outcome handed to the counterfactual owner: no judge + # label, and canonicalization is deferred to the strict validator (the raw program is + # not asserted valid here, since the component catalog is environment/pack-dependent). + assert outcome["action_id"] == 3 + assert outcome["finish_reason"] == "eos" + assert outcome["canonical_program"] is None + assert "verified" not in outcome and "metrics" not in outcome diff --git a/tests/test_models/test_causal_trace_plugin.py b/tests/test_models/test_causal_trace_plugin.py new file mode 100644 index 00000000..fe093928 --- /dev/null +++ b/tests/test_models/test_causal_trace_plugin.py @@ -0,0 +1,222 @@ +"""Torch-gated tests for the causal plug-in traced decode (LDI1-01 / SLM-119). + +A tiny deterministic causal model drives the real per-step torch forward through the +capture loop. The grammar legal-set seam is injected so the loop is exercised without +the OpenUI parser. These tests assert the torch wiring, exact logit replay, reproducible +forced-action replay, and that tracing never changes which token is emitted. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from slm_training.harnesses.distill.trace_store import TraceStore # noqa: E402 +from slm_training.models.causal_lm_openui import ( # noqa: E402 + CausalLMOpenUIConfig, + CausalLMOpenUIPlugin, +) +from slm_training.models.causal_trace import ( # noqa: E402 + CausalTraceWriter, + load_causal_decision_states, +) + +# Legality by generated-suffix length (prompt is a single token, so prompt_len == 1). +_LEGAL_BY_SUFFIX = {0: (2, 3), 1: (4,), 2: (0,)} + + +class _Output: + def __init__(self, logits: "torch.Tensor") -> None: + self.logits = logits + + +class _TinyCausalModel(torch.nn.Module): + def __init__(self, vocab: int = 6, hidden: int = 4) -> None: + super().__init__() + self.emb = torch.nn.Embedding(vocab, hidden) + self.head = torch.nn.Linear(hidden, vocab) + self.device = torch.device("cpu") + self.config = torch.nn.Module() # a plain attr-bag for decode_config_hash + + def forward(self, input_ids: "torch.Tensor") -> _Output: + return _Output(self.head(self.emb(input_ids))) + + def state_dict(self, *args, **kwargs): # keep compatibility_fingerprint cheap + return {} + + +class _TinyTokenizer: + eos_token_id = 0 + pad_token_id = 0 + init_kwargs: dict = {} + + def __len__(self) -> int: + return 6 + + def __call__(self, text: str, return_tensors: str = "pt"): + return {"input_ids": torch.tensor([[1]])} + + def decode(self, ids, skip_special_tokens: bool = True) -> str: + return "".join(str(int(i)) for i in ids) + + +def _plugin() -> CausalLMOpenUIPlugin: + torch.manual_seed(0) + config = CausalLMOpenUIConfig(base_model_id="tiny", base_model_revision="rev") + return CausalLMOpenUIPlugin(_TinyCausalModel(), _TinyTokenizer(), config) + + +def _allowed(prefix: tuple[int, ...]) -> tuple[int, ...]: + return _LEGAL_BY_SUFFIX.get(len(prefix) - 1, (0,)) # prompt_len == 1; default EOS + + +def test_traced_decode_records_states_and_stops_on_eos(tmp_path) -> None: + plugin = _plugin() + identity = plugin.capture_identity(group_id="grp", context_text="root=Stack([") + store = TraceStore(tmp_path / "traces", run_id="ldi1-plugin") + writer = CausalTraceWriter(store, identity) + + out = plugin.generate_constrained_traced( + "Make a card", group_id="grp", trace_writer=writer, allowed_ids_fn=_allowed + ) + assert out.result.generated_token_ids[-1] == 0 # EOS + assert out.result.stop_reason == "eos" + # Every emitted token was legal under the injected grammar seam. + for obs in out.result.observations: + assert obs.selected_token_id in _allowed(obs.prefix_token_ids) + + states = load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha=identity.tokenizer_sha, + ) + assert len(states) == len(out.result.observations) + # context_ids are the FULL prefix (prompt + suffix), starting at the prompt token. + assert states[0].context_ids == (1,) + + +def test_stored_logits_replay_within_tolerance(tmp_path) -> None: + plugin = _plugin() + out = plugin.generate_constrained_traced( + "Make a card", group_id="grp", allowed_ids_fn=_allowed + ) + for obs in out.result.observations: + row = torch.tensor([list(obs.prefix_token_ids)]) + with torch.inference_mode(): + recomputed = plugin.model(row).logits[0, -1, :] + stored_id, stored_logit, _lp = obs.raw_topk[0] + assert obs.raw_argmax_id == int(recomputed.argmax(-1)) + assert abs(float(recomputed[stored_id]) - stored_logit) < 1e-4 + + +def test_tracing_does_not_change_emitted_tokens() -> None: + plugin = _plugin() + out = plugin.generate_constrained_traced( + "Make a card", group_id="grp", allowed_ids_fn=_allowed + ) + # Independently reproduce greedy-over-legal selection from the same model. + prefix: tuple[int, ...] = (1,) + expected: list[int] = [] + for _ in range(len(out.result.generated_token_ids)): + with torch.inference_mode(): + logits = plugin.model(torch.tensor([list(prefix)])).logits[0, -1, :] + legal = _allowed(prefix) + chosen = max(legal, key=lambda t: (float(logits[t]), -t)) + expected.append(chosen) + prefix = (*prefix, chosen) + if chosen == 0: + break + assert list(out.result.generated_token_ids) == expected + + +def test_forced_action_replay_is_reproducible(tmp_path) -> None: + plugin = _plugin() + identity = plugin.capture_identity(group_id="grp", context_text="root=Stack([") + store = TraceStore(tmp_path / "traces", run_id="ldi1-plugin") + writer = CausalTraceWriter(store, identity) + plugin.generate_constrained_traced( + "Make a card", group_id="grp", trace_writer=writer, allowed_ids_fn=_allowed + ) + state = load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha=identity.tokenizer_sha, + )[0] + + first = plugin.replay_causal_action(state, 3, 7, allowed_ids_fn=_allowed) + second = plugin.replay_causal_action(state, 3, 7, allowed_ids_fn=_allowed) + assert first == second # deterministic continuation for the same seed + assert first.action_id == 3 # forced action applied to the exact prefix + assert first.continuation_seed == 7 + + +def test_replay_keeps_tokens_generated_before_the_decision(tmp_path) -> None: + plugin = _plugin() + identity = plugin.capture_identity(group_id="grp", context_text="root=Stack([") + store = TraceStore(tmp_path / "traces", run_id="ldi1-plugin") + writer = CausalTraceWriter(store, identity) + plugin.generate_constrained_traced( + "Make a card", group_id="grp", trace_writer=writer, allowed_ids_fn=_allowed + ) + states = load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha=identity.tokenizer_sha, + ) + # A decision after ordinal zero: its stored prefix already holds an earlier token. + state = next(s for s in states if s.decision_position >= 1) + prompt_len = len(state.context_ids) - state.decision_position + earlier = state.context_ids[prompt_len:] + assert earlier # non-empty: tokens were generated before this decision + forced = state.legal_action_ids[0] + outcome = plugin.replay_causal_action(state, forced, 0, allowed_ids_fn=_allowed) + # The earlier generated tokens survive into the materialized program. + assert outcome.raw_program.startswith(plugin.tokenizer.decode(earlier)) + + +def test_replay_rejects_illegal_forced_action(tmp_path) -> None: + plugin = _plugin() + identity = plugin.capture_identity(group_id="grp", context_text="root=Stack([") + store = TraceStore(tmp_path / "traces", run_id="ldi1-plugin") + writer = CausalTraceWriter(store, identity) + plugin.generate_constrained_traced( + "Make a card", group_id="grp", trace_writer=writer, allowed_ids_fn=_allowed + ) + state = load_causal_decision_states( + store, + expected_checkpoint_sha=identity.policy_checkpoint_sha, + expected_tokenizer_sha=identity.tokenizer_sha, + )[0] + illegal = max(state.legal_action_ids) + 100 + with pytest.raises(ValueError, match="not legal"): + plugin.replay_causal_action(state, illegal, 0, allowed_ids_fn=_allowed) + + +def test_adapter_identity_changes_when_a_lora_weight_mutates() -> None: + plugin = _plugin() + lora = torch.zeros(2, 2) + plugin.model.named_parameters = lambda: [ # type: ignore[assignment] + ("base_model.layers.0.lora_A.weight", lora) + ] + before = plugin.active_adapter_identity() + assert before != "" + with torch.no_grad(): + lora.add_(1.0) # mutate the tensor value without renaming it + assert plugin.active_adapter_identity() != before + + +def test_active_adapter_identity_reflects_lora_params() -> None: + plugin = _plugin() + assert plugin.active_adapter_identity() == "" # no adapter params on the base model + base_identity = plugin.capture_identity(group_id="grp", context_text="ctx") + + real = plugin.model.named_parameters + plugin.model.named_parameters = lambda: [ # type: ignore[assignment] + ("base_model.layers.0.lora_A.weight", torch.zeros(1)), + *list(real()), + ] + assert plugin.active_adapter_identity() != "" + adapter_identity = plugin.capture_identity(group_id="grp", context_text="ctx") + # Adapter presence changes the folded policy fingerprint (hence state identity). + assert adapter_identity.policy_checkpoint_sha != base_identity.policy_checkpoint_sha diff --git a/tests/test_models/test_low_rank_adapter.py b/tests/test_models/test_low_rank_adapter.py new file mode 100644 index 00000000..f99c67f8 --- /dev/null +++ b/tests/test_models/test_low_rank_adapter.py @@ -0,0 +1,92 @@ +"""Torch-gated tests for the removable low-rank adapter wrapper (LDI2-01 / SLM-123).""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +from torch import nn # noqa: E402 + +from slm_training.models.adapters.low_rank import LowRankAdapter # noqa: E402 + + +def _base() -> nn.Linear: + torch.manual_seed(0) + return nn.Linear(6, 4) + + +def _adapter(rank: int = 2, alpha: float = 4.0, dropout: float = 0.0) -> LowRankAdapter: + return LowRankAdapter(_base(), rank=rank, alpha=alpha, dropout=dropout) + + +def test_fresh_enabled_adapter_is_bit_identical_to_parent() -> None: + adapter = _adapter() + x = torch.randn(3, 6) + assert torch.equal(adapter(x), adapter.base(x)) # B is zero-initialized + + +def test_nonzero_adapter_effect_and_disable_restore() -> None: + adapter = _adapter() + x = torch.randn(3, 6) + baseline = adapter.base(x).clone() + with torch.no_grad(): + adapter.lora_B.fill_(0.1) # simulate a trained delta + assert not torch.allclose(adapter(x), baseline) + adapter.disable_adapter() + assert torch.equal(adapter(x), baseline) # disable restores the exact parent map + adapter.enable_adapter() + assert not torch.allclose(adapter(x), baseline) + + +def test_only_adapter_parameters_receive_gradients() -> None: + adapter = _adapter() + with torch.no_grad(): + adapter.lora_B.fill_(0.1) + adapter(torch.randn(2, 6)).sum().backward() + assert adapter.base.weight.grad is None + assert adapter.base.bias.grad is None + assert adapter.lora_A.grad is not None + assert adapter.lora_B.grad is not None + + +def test_adapter_params_follow_base_dtype_and_batched_dims() -> None: + adapter = _adapter() + assert adapter.lora_A.dtype == adapter.base.weight.dtype + assert adapter.lora_B.device == adapter.base.weight.device + out = adapter(torch.randn(2, 3, 6)) # batched leading dims + assert out.shape == (2, 3, 4) + + +def test_merged_linear_matches_enabled_output_without_mutating_base() -> None: + adapter = _adapter() + with torch.no_grad(): + adapter.lora_B.fill_(0.1) + base_weight_before = adapter.base.weight.detach().clone() + x = torch.randn(5, 6) + + merged = adapter.merged_linear() + assert isinstance(merged, nn.Linear) and not isinstance(merged, LowRankAdapter) + assert torch.allclose(merged(x), adapter(x), atol=1e-6) + # Merge is one-way on a copy: the parent weight and the removable adapter are intact. + assert torch.equal(adapter.base.weight, base_weight_before) + assert torch.allclose(adapter(x), adapter.base(x) + adapter.scaling * nn.functional.linear( + nn.functional.linear(x, adapter.lora_A), adapter.lora_B + )) + + +def test_rank_must_be_positive_and_wraps_only_linear() -> None: + with pytest.raises(ValueError, match="rank must be positive"): + LowRankAdapter(_base(), rank=0, alpha=1.0, dropout=0.0) + with pytest.raises(TypeError, match="wraps an nn.Linear"): + LowRankAdapter(nn.ReLU(), rank=2, alpha=1.0, dropout=0.0) # type: ignore[arg-type] + + +def test_alpha_and_dropout_are_validated() -> None: + # Direct construction bypasses TwoTowerAdapterSpec, so the wrapper mirrors its checks. + for bad_alpha in (0.0, -1.0, float("nan"), float("inf")): + with pytest.raises(ValueError, match="alpha must be a positive finite"): + LowRankAdapter(_base(), rank=2, alpha=bad_alpha, dropout=0.0) + for bad_dropout in (-0.1, 1.0, float("nan")): + with pytest.raises(ValueError, match=r"dropout must be a finite number in \[0, 1\)"): + LowRankAdapter(_base(), rank=2, alpha=1.0, dropout=bad_dropout) diff --git a/tests/test_models/test_twotower_adapter.py b/tests/test_models/test_twotower_adapter.py new file mode 100644 index 00000000..df387a65 --- /dev/null +++ b/tests/test_models/test_twotower_adapter.py @@ -0,0 +1,221 @@ +"""Torch-gated tests for TwoTower low-rank adapter attachment (LDI2-01 / SLM-123).""" + +from __future__ import annotations + +import json + +import pytest + +torch = pytest.importorskip("torch") + +from slm_training.dsl.schema import ExampleRecord # noqa: E402 +from slm_training.models.adapters import TwoTowerAdapterSpec # noqa: E402 +from slm_training.models.adapters.low_rank import LowRankAdapter # noqa: E402 +from slm_training.models.twotower import TwoTowerConfig, TwoTowerModel # noqa: E402 + +_HERO = ( + 'root = Stack([hero], "column")\n' + 'hero_title = TextContent(":hero.title")\n' + 'hero_body = TextContent(":hero.body")\n' + "hero = Card([hero_title, hero_body])" +) + + +def _records() -> list[ExampleRecord]: + return [ExampleRecord(id="a", prompt="Hero", openui=_HERO, split="train")] + + +def _model() -> TwoTowerModel: + return TwoTowerModel.from_records( + _records(), + config=TwoTowerConfig(d_model=32, n_heads=4, context_layers=1, denoiser_layers=1), + device="cpu", + ) + + +def _spec(model: TwoTowerModel, **overrides) -> TwoTowerAdapterSpec: + base = dict( + method="low_rank", + rank=2, + alpha=4.0, + dropout=0.0, + target_modules=("attn_q", "attn_v"), + base_compatibility_fingerprint=model.compatibility_fingerprint(), + base_checkpoint_sha="ckpt", + tokenizer_sha=model.artifact_identity()["tokenizer_sha"], + ) + base.update(overrides) + return TwoTowerAdapterSpec(**base) + + +def test_attach_resolves_targets_and_freezes_base() -> None: + model = _model() + model.attach_adapter(_spec(model)) + assert model.has_adapter() + # Two targets on the single denoiser layer -> two wrapped linears. + assert len(model._adapter_modules) == 2 + assert isinstance(model.denoiser.layers[0].self_attn.q_proj, LowRankAdapter) + # Adapter-only mode: every trainable parameter is an adapter tensor. + trainable = list(model.trainable_parameters()) + adapter = list(model.adapter_parameters()) + assert trainable and len(trainable) == len(adapter) + + +def test_fresh_adapter_is_output_identical_at_every_site() -> None: + model = _model() + model.attach_adapter(_spec(model)) + model.enable_adapter() + x = torch.randn(2, 5, 32) # d_model == 32 + for wrapper in model._adapter_modules.values(): + assert torch.equal(wrapper(x), wrapper.base(x)) # B is zero-initialized + + +def test_disable_restores_parent_after_weights_change() -> None: + model = _model() + model.attach_adapter(_spec(model)) + x = torch.randn(3, 32) + wrapper = next(iter(model._adapter_modules.values())) + baseline = wrapper.base(x).clone() + with torch.no_grad(): + wrapper.lora_B.fill_(0.05) + model.enable_adapter() + assert not torch.allclose(wrapper(x), baseline) + model.disable_adapter() + assert torch.equal(wrapper(x), baseline) + + +def test_only_adapter_parameters_receive_gradients() -> None: + model = _model() + model.attach_adapter(_spec(model)) + with torch.no_grad(): + for wrapper in model._adapter_modules.values(): + wrapper.lora_B.add_(0.05) + model.training_loss(_records()).backward() + base = model.denoiser.layers[0].self_attn.q_proj.base.weight + assert base.grad is None + assert any(p.grad is not None for p in model.adapter_parameters()) + + +def test_attach_does_not_shift_training_rng() -> None: + model = _model() + spec = _spec(model) + before = torch.random.get_rng_state() + model.attach_adapter(spec) + assert torch.equal(before, torch.random.get_rng_state()) + + +def test_attach_reports_actionable_target_errors() -> None: + model = _model() + with pytest.raises(ValueError, match="unsupported adapter target"): + model.attach_adapter(_spec(model, target_modules=("bogus",))) + other = _model() + with pytest.raises(ValueError, match="out of range"): + other.attach_adapter(_spec(other, target_layer_indices=(9,))) + + +def test_attach_rejects_mismatched_base_fingerprint() -> None: + model = _model() + with pytest.raises(ValueError, match="base compatibility fingerprint"): + model.attach_adapter(_spec(model, base_compatibility_fingerprint="not-this-model")) + + +def test_attach_rejects_include_output_head() -> None: + # The output head is not a resolvable target, so the flag must fail closed rather + # than silently have no effect. + model = _model() + with pytest.raises(ValueError, match="include_output_head"): + model.attach_adapter(_spec(model, include_output_head=True)) + assert not model.has_adapter() + + +def test_merge_adapter_copy_matches_enabled_and_leaves_original() -> None: + model = _model() + model.attach_adapter(_spec(model)) + with torch.no_grad(): + for wrapper in model._adapter_modules.values(): + wrapper.lora_B.add_(0.03) + model.enable_adapter() + + merged = model.merge_adapter_copy() + assert not merged.has_adapter() + merged_linear = merged.denoiser.layers[0].self_attn.q_proj + assert not isinstance(merged_linear, LowRankAdapter) + + x = torch.randn(2, 5, 32) + wrapper = model.denoiser.layers[0].self_attn.q_proj + assert torch.allclose(merged_linear(x), wrapper(x), atol=1e-6) + # Merge is one-way on a copy: the original still carries its removable adapter. + assert model.has_adapter() + assert isinstance(model.denoiser.layers[0].self_attn.q_proj, LowRankAdapter) + + +def test_merge_requires_an_attached_adapter() -> None: + with pytest.raises(ValueError, match="no adapter is attached"): + _model().merge_adapter_copy() + + +def test_save_and_load_adapter_round_trip(tmp_path) -> None: + model = _model() + model.attach_adapter(_spec(model)) + with torch.no_grad(): + for wrapper in model._adapter_modules.values(): + wrapper.lora_B.add_(0.02) + model.enable_adapter() + x = torch.randn(3, 32) + trained = model.denoiser.layers[0].self_attn.q_proj(x).clone() + + model.save_adapter(tmp_path / "adapter") + assert (tmp_path / "adapter" / "adapter_config.json").exists() + manifest = json.loads((tmp_path / "adapter" / "adapter_manifest.json").read_text()) + assert manifest["trainable_parameter_count"] > 0 + assert manifest["module_map"] + + # A fresh model with the same seed reproduces the adapted logits after loading. + fresh = _model() + fresh.load_adapter(tmp_path / "adapter") + fresh.enable_adapter() + assert torch.allclose( + fresh.denoiser.layers[0].self_attn.q_proj(x), trained, atol=1e-6 + ) + + +def test_load_adapter_fails_closed_on_base_mismatch(tmp_path) -> None: + model = _model() + model.attach_adapter(_spec(model)) + model.save_adapter(tmp_path / "adapter") + # A different width yields a different compatibility fingerprint. + other = TwoTowerModel.from_records( + _records(), + config=TwoTowerConfig(d_model=64, n_heads=4, context_layers=1, denoiser_layers=1), + device="cpu", + ) + with pytest.raises(ValueError, match="fingerprint"): + other.load_adapter(tmp_path / "adapter") + + +def test_load_adapter_rejects_truncated_artifact_without_mutating(tmp_path) -> None: + model = _model() + model.attach_adapter(_spec(model)) + model.save_adapter(tmp_path / "adapter") + # Drop a required lora factor from the saved tensors (a truncated/corrupt artifact). + tensors = torch.load(tmp_path / "adapter" / "adapter_model.pt", weights_only=True) + del tensors[sorted(tensors)[0]] + torch.save(tensors, tmp_path / "adapter" / "adapter_model.pt") + + fresh = _model() + with pytest.raises(ValueError, match="manifest parameter set"): + fresh.load_adapter(tmp_path / "adapter") + # The whole artifact is validated before any mutation, so the failed load leaves the + # model untouched (no adapter attached, base parameters still trainable). + assert not fresh.has_adapter() + + +def test_active_adapter_identity_tracks_adapter_weights() -> None: + model = _model() + assert model.active_adapter_identity() == "" + model.attach_adapter(_spec(model)) + before = model.active_adapter_identity() + assert before != "" + with torch.no_grad(): + next(iter(model._adapter_modules.values())).lora_B.add_(1.0) + assert model.active_adapter_identity() != before