diff --git a/CHANGELOG.md b/CHANGELOG.md index cec0f74..dd1aa84 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,45 @@ reproducing 70+ versions of notes. ## [Unreleased] +## [0.71.16] - 2026-06-07 + +### Added +- **Covariance-preconditioned ROME via `--cov-corpus`** (closes #250). `soup edit + set --method rome --cov-corpus ` now estimates the key covariance + `C = E[k kᵀ] + λI` over a stats corpus and uses the preconditioned update + `u = C⁻¹ k*` instead of the covariance-free `C = I` path — the genuine ROME + closed form, which spreads the rank-1 update mass to reduce collateral + interference with other facts. Falls back to `C = I` when no corpus is given. + The exact post-condition `down(k*) += delta` is preserved either way. The + corpus loader is cwd-contained, symlink-rejected (O_NOFOLLOW + raw-path + lstat), and size/line-capped; `--cov-corpus` is rejected (fail-loud) for any + method other than `rome`. Verified on real `gpt2` (prob 0.005 → 0.9997) and + SmolLM2-135M. +- **GPT-2 (`transformer.h` / `mlp.c_proj`) support in the edit kernels** (closes + #251). ROME / MEMIT / AlphaEdit now edit GPT-2-family models, not just + Llama-family. The `Conv1D` weight layout (`[in, out]`, transposed relative to + `nn.Linear`'s `[out, in]`) gets a transpose-aware rank-1 update, AlphaEdit + null-space projection, and MEMIT band dim-check. PEFT-wrapped GPT-2 / Llama + models are unwrapped via `get_base_model`. Verified end-to-end on real `gpt2`. +- **Mixtral joins the LongLoRA architecture allowlist** (closes #147). A bare + `mistral` token does not appear in `mixtral` (m-i-x vs m-i-s), so the existing + `is_mistral_model` detector excluded the MoE variant. A dedicated + `is_mixtral_model` helper + `MixtralAttention` entry in the S² forward-override + regex + `_SEPARATE_QKV_FAMILIES` now cover Mixtral-8x7B / 8x22B (the attention + is the standard separate-QKV shell; the MoE lives in the MLP). + +### Fixed +- **Atomic `EditGovernor` edit-count increment** (closes #252). Two concurrent + `soup edit set` runs on the same base model could lose an increment: each read + the persisted count, added locally, and the last writer clobbered the first. + `save_state` now re-reads the persisted count INSIDE the cross-process lock and + merges this run's delta (`edit_count − persisted_baseline`), mirroring the + v0.60.0 `namespace_pin` pattern. Verified: two governors recording 3 + 2 edits + from the same baseline persist a merged 5 (not a clobbered 2 or a naive +1). + +### Notes +- Test count: 13511 → 13595 (+84 net; +81 in `tests/test_v07116.py`). + ## [0.71.15] - 2026-06-07 ### Fixed diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3e8e935..50bbce4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -120,7 +120,7 @@ src/soup_cli/ templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (286 files, 13511 tests) +tests/ - Test suite (287 files, 13595 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index b21ad58..7a31a12 100644 --- a/README.md +++ b/README.md @@ -49,19 +49,19 @@ infrastructure instead of improving models. Soup fixes that. ## What's New -**v0.71.15 — Loop & lifecycle polish.** Five fixes across the training-loop, merge, and governance -surfaces (all validated on the real RTX 3050 + SmolLM2-135M): +**v0.71.16 — Knowledge edit depth.** Deeper, safer surgical knowledge editing + a wider LongLoRA +allowlist (validated on real `gpt2` + SmolLM2-135M): -- **Iterative-DPO config bug fixed** — `soup iterative-dpo`'s per-round trainer rendered an invalid - `output:` shape that the spawned `soup train` rejected; now it round-trips cleanly. -- **CMA-ES merge loads the base once** — `soup adapters merge --strategy cmaes` now reuses the base - model across the whole candidate population instead of reloading it per candidate. -- **`soup loop` budget gate estimates real cost** — the pre-wired loop's per-iteration dollar estimate - was a `0.0` placeholder; it now derives a forward cost from the last run's GPU + duration. -- **`--diagnose-gate` is multi-node aware** — the post-training gate (and `--annex-xi` / - `--repro-receipt`) now fire once per *cluster* (`RANK==0`), not once per node. -- **`soup train --track-energy --energy-out `** — persists the measured energy/CO2 so - `soup bom emit --energy ` can attach it to an ML-BOM. +- **GPT-2 models are now editable** — `soup edit set` (ROME / MEMIT / AlphaEdit) handles the GPT-2 + `Conv1D` (`mlp.c_proj`) weight layout alongside Llama, not just Llama-family. Real `gpt2` smoke: + target probability 0.005 → 0.9996. +- **Covariance-preconditioned ROME** — `soup edit set --method rome --cov-corpus ` estimates + the key covariance and uses the genuine `C⁻¹ k*` update (spreads the edit mass, less collateral + interference). Falls back to `C = I` with no corpus. +- **Atomic edit-governor count** — two concurrent `soup edit set` runs on the same base can no longer + lose an increment; the count is merged under the cross-process lock. +- **Mixtral joins the LongLoRA allowlist** — `use_longlora: true` now accepts Mixtral-8x7B / 8x22B + (the bare `mistral` token never matched the MoE variant before). Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). diff --git a/docs/adapters-and-governance.md b/docs/adapters-and-governance.md index 2043e18..8b34051 100644 --- a/docs/adapters-and-governance.md +++ b/docs/adapters-and-governance.md @@ -147,6 +147,11 @@ soup edit set \ --target "Lyon" \ --output ./edited --device cuda +# Covariance-preconditioned ROME (genuine C^-1 k* update) from a stats corpus. +soup edit set --base ./model --method rome \ + --subject "..." --target "..." \ + --cov-corpus ./stats.jsonl --output ./edited + # Plan-only mode validates the request + prints the resolved EditPlan + exits 0. soup edit set --base ./model --method rome --subject "..." --target "..." --plan-only @@ -156,9 +161,11 @@ soup edit diff before after \ --probes probes.jsonl --output diff.json ``` -The kernels are a covariance-free (`C = I`) variant of the ROME family — well-defined, tractable on a 4 GB box, and validated on SmolLM2-135M (a ROME edit moved `P("Lyon" | "The capital of France is")` from 0.0016 → 0.96). **ROME** edits a single layer; **MEMIT** distributes the residual across a layer band; **AlphaEdit** projects the update orthogonal to the down-proj's top singular direction. +The kernels run a real rank-1 weight update at an MLP down-projection. They support both the **Llama-family** (`model.model.layers[L].mlp.down_proj`, an `nn.Linear` with `[out, in]` weights) and the **GPT-2-family** (`model.transformer.h[L].mlp.c_proj`, a `Conv1D` with the transposed `[in, out]` layout) — the update is transpose-aware so the post-condition `down(k*) += delta` holds exactly for either layout (validated on real `gpt2`: `P(target)` 0.005 → 0.9996, and on SmolLM2-135M). PEFT-wrapped models are unwrapped automatically. **ROME** edits a single layer; **MEMIT** distributes the residual across a layer band; **AlphaEdit** projects the update orthogonal to the down-proj's top singular direction. -The sequential edit governor is persisted (SQLite, cross-process-locked, `SOUP_EDIT_GOVERNOR_DB` override) so the per-base-model edit count + norm-blowup verdict survive across separate `soup edit set` runs. `soup edit set` consults it automatically: it refuses BEFORE the model load past the per-base cap or after a BLOWUP verdict, and records the measured `||ΔW||_F` after each edit. Pass `--no-governor` to opt out. `--registry-id ` attaches the edited model (or GRACE codebook) into the Registry lineage. +By default ROME uses the covariance-free `C = I` form. Pass `--cov-corpus ` to estimate the key covariance `C = E[k kᵀ] + λI` over a stats corpus and apply the genuine ROME closed form `u = C⁻¹ k*` — this spreads the rank-1 update mass per the closed form and reduces collateral interference with other facts (the exact post-condition is preserved either way). The corpus loader is cwd-contained, symlink-rejected, and size/line-capped; `--cov-corpus` is rejected for any method other than `rome`. + +The sequential edit governor is persisted (SQLite, cross-process-locked, `SOUP_EDIT_GOVERNOR_DB` override) so the per-base-model edit count + norm-blowup verdict survive across separate `soup edit set` runs. The count increment is atomic — two concurrent `soup edit set` runs on the same base are merged under the cross-process lock (no lost increment). `soup edit set` consults the governor automatically: it refuses BEFORE the model load past the per-base cap or after a BLOWUP verdict, and records the measured `||ΔW||_F` after each edit. Pass `--no-governor` to opt out. `--registry-id ` attaches the edited model (or GRACE codebook) into the Registry lineage. ## Activation Steering (`soup steer`) diff --git a/docs/commands.md b/docs/commands.md index 509df1e..1f30224 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -171,7 +171,7 @@ soup can pack --entry-id --out r.can --attest Embed in-to soup audit-log tail / rotate Tail / rotate the per-command HIPAA/SOC2 audit log (~/.soup/audit.jsonl) soup --no-audit-log / SOUP_NO_AUDIT_LOG=1 Opt out of the per-command audit line soup eval unlearning --benchmark tofu|muse|wmdp Forget Quality + Model Utility + PrivLeak verdict -soup edit set --base --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output ] [--device cpu] [--governor/--no-governor] [--registry-id ] Live surgical knowledge edit (--plan-only available) +soup edit set --base --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output ] [--device cpu] [--governor/--no-governor] [--registry-id ] [--cov-corpus ] Live surgical knowledge edit (GPT-2 Conv1D + Llama; --cov-corpus = covariance-preconditioned ROME, rome-only; --plan-only available) soup edit diff --probes p.jsonl [--before-model --after-model ] Knowledge-injection diff (live before/after generation when both models given) soup train --task unlearn NPO/SimNPO/RMU unlearning from data.forget_set (+ optional data.retain_set) soup train # data.format='raft' Answer-only span-mask RAFT training (golden+distractor docs, [doc-N] citations); generator-stage configs auto-link the latest RA-DIT retriever diff --git a/docs/peft-and-efficiency.md b/docs/peft-and-efficiency.md index aad0068..e7839bd 100644 --- a/docs/peft-and-efficiency.md +++ b/docs/peft-and-efficiency.md @@ -30,12 +30,15 @@ ## LongLoRA Forward Override When `use_longlora: true` is set on an SFT config with a Llama / CodeLlama / -Mistral / Qwen / Phi base, the trainer wraps the model in a +Mistral / Mixtral / Qwen / Phi base, the trainer wraps the model in a `LongLoRAForwardOverride` context that monkey-patches every attention forward to apply the S² shifted-sparse shift (paper §3.2) — half the heads are rolled -by `group_size // 2` along the sequence dim. Restoration on context exit is -idempotent and best-effort safe; FlashAttention v3 builds are rejected at the -schema gate (the custom-mask kernels conflict). +by `group_size // 2` along the sequence dim. Mixtral joined the allowlist in +v0.71.16 (a bare `mistral` token never matched the MoE variant); its attention +is the standard separate-QKV shell — the MoE lives in the MLP — so the same +Q/K projection-shift path is reused. Restoration on context exit is idempotent +and best-effort safe; FlashAttention v3 builds are rejected at the schema gate +(the custom-mask kernels conflict). ## Multipack — FFD Bin-Packing Sampler diff --git a/pyproject.toml b/pyproject.toml index 6c567ae..311dcf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.71.15" +version = "0.71.16" description = "Fine-tune and post-train LLMs in one command. No SSH, no config hell." readme = "README.md" license = "Apache-2.0" diff --git a/src/soup_cli/__init__.py b/src/soup_cli/__init__.py index f10f2ec..6ac7390 100644 --- a/src/soup_cli/__init__.py +++ b/src/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune and post-train LLMs in one command.""" -__version__ = "0.71.15" +__version__ = "0.71.16" diff --git a/src/soup_cli/commands/edit.py b/src/soup_cli/commands/edit.py index 3e51879..90227ad 100644 --- a/src/soup_cli/commands/edit.py +++ b/src/soup_cli/commands/edit.py @@ -26,6 +26,89 @@ app = typer.Typer( rich_markup_mode="rich", ) +# v0.71.16 #250 — covariance-corpus loader caps. +_MAX_COV_CORPUS_BYTES = 64 * 1024 * 1024 # 64 MiB +_MAX_COV_CORPUS_LINES = 100_000 + + +def _load_cov_corpus(path: str) -> list[str]: + """Load a covariance corpus (JSONL or plain text) for ``--cov-corpus`` (#250). + + Each JSONL object contributes its ``text`` / ``prompt`` / ``content`` field; + every other line (non-dict JSON or non-JSON) contributes its raw stripped + text. cwd-contained, symlink-rejected (TOCTOU on the raw path), size + line + capped. Raises ``FileNotFoundError`` for a missing file and ``ValueError`` + for a containment / symlink / oversize / empty-corpus failure. + """ + import errno + import json as _json + import os + import stat + + from soup_cli.utils.paths import is_under_cwd + + if not isinstance(path, str) or not path: + raise ValueError("--cov-corpus path must be a non-empty string") + if "\x00" in path: + raise ValueError("--cov-corpus path must not contain null bytes") + if not is_under_cwd(path): + raise ValueError( + f"--cov-corpus path must stay under cwd: {os.path.basename(path)!r}" + ) + # Raw-path symlink guard (the Windows guard, where O_NOFOLLOW is a no-op). + try: + if os.path.lexists(path) and stat.S_ISLNK(os.lstat(path).st_mode): + raise ValueError("--cov-corpus path must not be a symlink") + except OSError: + pass # lexists/lstat race — os.open below surfaces the real error + # O_NOFOLLOW closes the TOCTOU window between the check and the read: os.open + # fails with ELOOP if the final component is a symlink (parity with the + # diagnose.live / tunability readers). + try: + fd = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + except FileNotFoundError as exc: + raise FileNotFoundError( + f"--cov-corpus file not found: {os.path.basename(path)!r}" + ) from exc + except OSError as exc: + if getattr(exc, "errno", None) == errno.ELOOP: + raise ValueError("--cov-corpus path must not be a symlink") from exc + raise ValueError( + f"--cov-corpus path is not readable ({type(exc).__name__})" + ) from exc + rows: list[str] = [] + with os.fdopen(fd, encoding="utf-8") as fh: + st = os.fstat(fh.fileno()) + if not stat.S_ISREG(st.st_mode): + raise ValueError("--cov-corpus path must be a regular file") + if st.st_size > _MAX_COV_CORPUS_BYTES: + raise ValueError( + f"--cov-corpus file too large (> {_MAX_COV_CORPUS_BYTES} bytes)" + ) + for i, line in enumerate(fh): + if i >= _MAX_COV_CORPUS_LINES: + break + stripped = line.strip() + if not stripped: + continue + parsed = None + is_json = True + try: + parsed = _json.loads(stripped) + except _json.JSONDecodeError: + is_json = False + if is_json and isinstance(parsed, dict): + for fld in ("text", "prompt", "content"): + val = parsed.get(fld) + if isinstance(val, str) and val.strip(): + rows.append(val) + break + else: + rows.append(stripped) + if not rows: + raise ValueError("--cov-corpus has no usable text rows") + return rows + @app.command(name="set") def set_edit( @@ -61,6 +144,13 @@ def set_edit( True, "--governor/--no-governor", help="Consult the sequential-edit governor (refuse on norm blowup).", ), + cov_corpus: Optional[str] = typer.Option( + None, "--cov-corpus", + help=( + "ROME only: JSONL/text corpus to estimate the key covariance C " + "for a C^{-1}-preconditioned update (reduces collateral damage)." + ), + ), plan_only: bool = typer.Option( False, "--plan-only", help="Print the resolved EditPlan and exit without applying.", @@ -113,6 +203,16 @@ def set_edit( console.print("[green]Plan-only mode — exiting without applying.[/]") return + # v0.71.16 #250 — load the optional covariance corpus (path security lives + # here at the CLI boundary; apply_edit rejects it for non-ROME methods). + cov_corpus_rows = None + if cov_corpus is not None: + try: + cov_corpus_rows = _load_cov_corpus(cov_corpus) + except (ValueError, FileNotFoundError, OSError) as exc: + console.print(f"[red]Cannot load --cov-corpus:[/] {escape(str(exc))}") + raise typer.Exit(2) from exc + # Load a persisted governor so sequential edits across separate `soup edit # set` runs accumulate (#196 / #197). Best-effort — a missing / unreadable # governor DB never blocks an edit. @@ -142,6 +242,7 @@ def set_edit( output_dir=output, governor=governor, device=device, + cov_corpus=cov_corpus_rows, ) except GovernedEditError as exc: console.print( diff --git a/src/soup_cli/utils/edit_governor.py b/src/soup_cli/utils/edit_governor.py index fe86ac7..b6c4624 100644 --- a/src/soup_cli/utils/edit_governor.py +++ b/src/soup_cli/utils/edit_governor.py @@ -240,6 +240,14 @@ class EditGovernor: last_method: str = "" last_verdict: str = "OK" last_norm_delta: float = 0.0 + # v0.71.16 #252 — the edit_count this governor was loaded / constructed + # with. The atomic store save merges THIS run's increments + # (``edit_count - _persisted_edit_count``) onto the freshly-read persisted + # count so two concurrent ``soup edit set`` runs cannot lose an increment. + # Defaults to ``-1`` (sentinel) and is initialised to ``edit_count`` in + # ``__post_init__``. compare/repr-excluded so it never leaks into equality + # or snapshots. + _persisted_edit_count: int = field(default=-1, compare=False, repr=False) def __post_init__(self) -> None: if not isinstance(self.base_model, str): @@ -267,6 +275,10 @@ class EditGovernor: raise ValueError( f"max_sequential_edits must be <= {_MAX_SEQ_EDITS}" ) + # v0.71.16 #252 — seed the baseline at construction time. ``-1`` is the + # "not supplied" sentinel; otherwise honour an explicit baseline. + if self._persisted_edit_count < 0: + self._persisted_edit_count = self.edit_count def record_edit(self, *, method: str, norm_delta: float) -> None: """Append a completed edit to the governor's history.""" @@ -508,10 +520,27 @@ class EditGovernorStore: } def save_state(self, governor: "EditGovernor") -> None: - """Upsert a governor's state under the cross-process lock.""" + """Upsert a governor's state with an atomic read-modify-write (#252). + + The persisted ``edit_count`` is re-read INSIDE the cross-process lock + and merged with THIS run's increments + (``governor.edit_count - governor._persisted_edit_count``) so two + concurrent ``soup edit set`` runs on the same base cannot lose a count + (the pre-#252 absolute write let the last writer clobber the first). + The governor's in-memory ``edit_count`` and baseline are then advanced + to the merged value so a subsequent save does not double-count. + """ if not isinstance(governor, EditGovernor): raise TypeError("governor must be an EditGovernor") with self._cross_process_lock(): + # Read INSIDE the lock so a racing writer's committed row is + # observed (closes the get→insert TOCTOU; mirrors namespace_pin). + existing = self.get_state(governor.base_model) + persisted = existing["edit_count"] if existing is not None else 0 + pending = governor.edit_count - governor._persisted_edit_count + if pending < 0: + pending = 0 + merged = persisted + pending self._conn.execute( "INSERT OR REPLACE INTO edit_governors " "(base_model, edit_count, last_method, last_verdict, " @@ -519,7 +548,7 @@ class EditGovernorStore: "VALUES (?, ?, ?, ?, ?, ?, ?)", ( governor.base_model, - governor.edit_count, + merged, governor.last_method, governor.last_verdict, governor.last_norm_delta, @@ -528,6 +557,8 @@ class EditGovernorStore: ), ) self._conn.commit() + governor.edit_count = merged + governor._persisted_edit_count = merged def save_governor(store: EditGovernorStore, governor: "EditGovernor") -> None: diff --git a/src/soup_cli/utils/edit_kernels.py b/src/soup_cli/utils/edit_kernels.py index b76ebfe..1056734 100644 --- a/src/soup_cli/utils/edit_kernels.py +++ b/src/soup_cli/utils/edit_kernels.py @@ -22,8 +22,13 @@ Every heavy import (``torch``) is local so importing this module stays cheap from __future__ import annotations +import logging +import math +from collections.abc import Iterator from dataclasses import dataclass -from typing import List +from typing import List, Optional, Sequence + +_LOG = logging.getLogger("soup.edit_kernels") # Optimisation hyper-parameters (tuned for tiny models; deliberately modest so # a single edit runs in a few seconds on CPU/tiny-GPU). @@ -32,6 +37,11 @@ _DEFAULT_LR = 0.5 _MEMIT_BAND = 3 # number of layers MEMIT distributes the residual across _MAX_PROMPT_TOKENS = 256 +# v0.71.16 #250 — covariance-preconditioned ROME defaults. +_DEFAULT_COV_RIDGE = 0.01 # λ in C = E[k k^T] + λI (keeps C invertible) +_DEFAULT_COV_MAX_PROMPTS = 256 +_DEFAULT_COV_MAX_TOKENS = 64 + @dataclass(frozen=True) class EditKernelResult: @@ -43,42 +53,95 @@ class EditKernelResult: layers_edited: tuple[int, ...] -def _locate_decoder_layers(model: object) -> object: - """Return the decoder-layer ``ModuleList`` for a Llama-family model. +def _is_transposed_proj(module: object) -> bool: + """Return True for a GPT-2 ``Conv1D`` down-projection. - Raises ``ValueError`` for architectures we cannot locate (e.g. GPT-2 style - ``transformer.h`` — out of scope for v0.71.9; the ROME family is wired for - the Llama-shape ``mlp.down_proj`` linear). + transformers' ``Conv1D`` stores a transposed weight (``[in, out]`` rather + than nn.Linear's ``[out, in]``) and an ``nf`` int (the output feature + count); nn.Linear has neither. Detecting via the ``nf`` attribute is robust + to the transformers internal class path (v0.71.16 #251). + """ + nf = getattr(module, "nf", None) + return isinstance(nf, int) and not isinstance(nf, bool) + + +def _proj_out_dim(module: object) -> int: + """Output (hidden) dimension of a down-projection module. + + * nn.Linear: weight is ``[out, in]`` → out = ``weight.shape[0]``. + * GPT-2 Conv1D: weight is ``[in, out]`` → out = ``nf`` (= ``weight.shape[1]``). + """ + if _is_transposed_proj(module): + return int(module.nf) # type: ignore[attr-defined] + return int(module.weight.shape[0]) # type: ignore[attr-defined] + + +def _candidate_models(model: object) -> Iterator[object]: + """Yield ``model`` then its PEFT base (if any) for layer lookup.""" + yield model + get_base = getattr(model, "get_base_model", None) + if callable(get_base): + try: + yield get_base() + except Exception as exc: # noqa: BLE001 — best-effort PEFT unwrap + # Don't mask a real get_base_model() crash silently — surface it at + # DEBUG so it's inspectable; _locate_decoder_layers still raises a + # clear ValueError downstream when no layers are found. + _LOG.debug("get_base_model() failed during layer lookup: %s", exc) + + +def _layers_from(model: object) -> object: + """Return the decoder-layer container on ``model`` or ``None``. + + Llama-family: ``model.model.layers``. GPT-2-family: ``model.transformer.h``. """ inner = getattr(model, "model", None) layers = getattr(inner, "layers", None) if inner is not None else None - if layers is None: - # PEFT-wrapped or unusual nesting — try get_base_model. - get_base = getattr(model, "get_base_model", None) - if callable(get_base): - base = get_base() - inner = getattr(base, "model", None) - layers = getattr(inner, "layers", None) if inner is not None else None - if layers is None: - raise ValueError( - "could not locate decoder layers (expected a Llama-family " - "model.model.layers); ROME/MEMIT/AlphaEdit support the " - "mlp.down_proj architecture in v0.71.9" - ) - return layers + if layers is not None: + return layers + transformer = getattr(model, "transformer", None) + h = getattr(transformer, "h", None) if transformer is not None else None + if h is not None: + return h + return None + + +def _locate_decoder_layers(model: object) -> object: + """Return the decoder-layer ``ModuleList`` for a supported model. + + Supports the Llama-family ``model.model.layers`` and the GPT-2-family + ``model.transformer.h`` (v0.71.16 #251). PEFT-wrapped models are unwrapped + via ``get_base_model``. Raises ``ValueError`` for unsupported architectures. + """ + for candidate in _candidate_models(model): + layers = _layers_from(candidate) + if layers is not None: + return layers + raise ValueError( + "could not locate decoder layers (expected a Llama-family " + "model.model.layers or a GPT-2-family model.transformer.h); " + "ROME/MEMIT/AlphaEdit support the mlp.down_proj (Llama) and " + "mlp.c_proj (GPT-2) architectures" + ) def _down_proj(layers: object, layer: int) -> object: - """Return the ``mlp.down_proj`` linear for decoder ``layer``.""" + """Return the MLP down-projection for decoder ``layer``. + + Llama: ``mlp.down_proj`` (nn.Linear, weight ``[out, in]``). GPT-2: + ``mlp.c_proj`` (Conv1D, weight ``[in, out]``) — v0.71.16 #251. + """ try: block = layers[layer] # type: ignore[index] except (IndexError, TypeError) as exc: raise ValueError(f"layer index {layer} out of range") from exc mlp = getattr(block, "mlp", None) down = getattr(mlp, "down_proj", None) if mlp is not None else None + if down is None and mlp is not None: + down = getattr(mlp, "c_proj", None) # GPT-2 Conv1D if down is None or not hasattr(down, "weight"): raise ValueError( - f"decoder layer {layer} has no mlp.down_proj weight " + f"decoder layer {layer} has no mlp.down_proj / mlp.c_proj weight " "(unsupported architecture for ROME-family edits)" ) return down @@ -142,7 +205,9 @@ def _optimise_residual( # residual feeds the prediction of the first target token). inject_pos = len(subj_ids) - 1 - hidden = down.weight.shape[0] + # Delta lives in the OUTPUT (hidden) space — for a GPT-2 Conv1D that is + # ``nf``, NOT ``weight.shape[0]`` (which is the input dim there). + hidden = _proj_out_dim(down) delta = torch.zeros( hidden, device=device, dtype=down.weight.dtype, requires_grad=True ) @@ -172,34 +237,79 @@ def _optimise_residual( return delta.detach() -def _rank1_update(down, key, delta) -> float: - """Apply ``W += delta @ key^T / ||key||^2`` in place. Returns Frobenius norm. +def _rank1_update(down, key, delta, *, cov=None) -> float: + """Apply a (covariance-preconditioned) rank-1 update in place. - Makes ``W @ key == (old W @ key) + delta`` exactly, so the down-proj now - produces the optimised residual for this key. + Makes ``down(key) == old + delta`` exactly, so the down-proj now produces + the optimised residual for this key. Returns the update's Frobenius norm. + + ``cov`` (optional, v0.71.16 #250): the key covariance matrix ``C``. When + given, the update uses the preconditioned key ``u = C^{-1} k*`` (computed + via a stable linear solve) so the update mass is spread per the ROME + closed form, reducing collateral interference with other keys. The exact + post-condition ``down(k*) += delta`` is preserved either way because + ``denom = u·k*`` normalises it. + + Handles both weight layouts (v0.71.16 #251): + + * nn.Linear ``[out, in]`` — ``W += outer(delta, u) / denom``. + * GPT-2 Conv1D ``[in, out]`` — ``W += outer(u, delta) / denom`` (transposed). """ import torch - key = key.to(down.weight.dtype) - denom = float(torch.dot(key, key).item()) - if denom <= 0.0: - raise ValueError("key vector has zero norm; cannot apply rank-1 update") - update = torch.outer(delta.to(down.weight.dtype), key) / denom + key_w = key.to(down.weight.dtype) + if cov is not None: + # u = C^{-1} k* via a linear solve (cheaper + more stable than a full + # inverse — we only need one column). Solve in fp32 then cast back. A + # singular / non-finite covariance makes solve raise a torch LinAlgError + # (a RuntimeError subclass) — surface it as a clean ValueError. + try: + u = torch.linalg.solve( + cov.to(torch.float32), key.to(torch.float32) + ).to(down.weight.dtype) + except RuntimeError as exc: + raise ValueError( + f"covariance solve failed (singular / non-finite C): {exc}" + ) from exc + else: + u = key_w + denom = float(torch.dot(u, key_w).item()) + # Reject zero norm AND non-finite denom — with the #250 covariance path a + # pathological solve can yield NaN, and ``NaN <= 0.0`` is False, which would + # otherwise let a NaN update silently corrupt the weights. + if not math.isfinite(denom) or denom <= 0.0: + raise ValueError( + "key vector has zero norm (or degenerate covariance); " + "cannot apply rank-1 update" + ) + delta_w = delta.to(down.weight.dtype) + if _is_transposed_proj(down): + update = torch.outer(u, delta_w) / denom + else: + update = torch.outer(delta_w, u) / denom with torch.no_grad(): down.weight.add_(update) return float(torch.linalg.norm(update).item()) def _alphaedit_project(down, update): - """Project ``update`` orthogonal to ``W``'s top left-singular direction. + """Project a logical ``[out, in]`` update orthogonal to ``W``'s top + left-singular direction. A light null-space projection (covariance-free AlphaEdit) that reduces - interference with the down-proj's dominant feature direction. + interference with the down-proj's dominant feature direction. Operates in + the logical ``[out, in]`` orientation regardless of weight layout: for a + GPT-2 Conv1D (weight ``[in, out]``) the weight is transposed to ``[out, in]`` + before the power iteration so the projection removes the OUTPUT-direction + component consistently. The returned update is ``[out, in]`` — the caller + transposes it back when applying to a Conv1D (v0.71.16 #251). """ import torch with torch.no_grad(): w = down.weight.detach().to(torch.float32) + if _is_transposed_proj(down): + w = w.t() # Conv1D [in, out] -> logical [out, in] # Top left-singular vector via a couple of power iterations (cheap). # Seed the generator on CPU so the projection is deterministic / # reproducible across runs (review MEDIUM M3). @@ -218,6 +328,81 @@ def _alphaedit_project(down, update): return proj.to(down.weight.dtype) +def estimate_key_covariance( + model, + tokenizer, + down, + corpus: Sequence[str], + *, + device: str, + ridge: float = _DEFAULT_COV_RIDGE, + max_prompts: int = _DEFAULT_COV_MAX_PROMPTS, + max_tokens: int = _DEFAULT_COV_MAX_TOKENS, +): + """Estimate the key covariance ``C = E[k k^T] + ridge*I`` over ``corpus``. + + v0.71.16 #250 — captures the down-projection INPUT (the key) at every token + position for each corpus prompt and accumulates the second moment. The + ridge term keeps ``C`` strictly positive-definite (invertible) for the ROME + ``C^{-1} k*`` solve. Returns a ``[in, in]`` float32 matrix where ``in`` is + the down-proj input (intermediate) dim. + """ + import torch + + # Range-guard the caps (defence-in-depth for direct callers — the CLI only + # passes the corpus, so these stay at their defaults in the normal flow). + for _name, _val in (("max_prompts", max_prompts), ("max_tokens", max_tokens)): + if isinstance(_val, bool) or not isinstance(_val, int) or _val < 1: + raise ValueError(f"{_name} must be a positive int") + if ( + isinstance(ridge, bool) + or not isinstance(ridge, (int, float)) + or not math.isfinite(float(ridge)) + or ridge < 0.0 + ): + raise ValueError("ridge must be a finite non-negative number") + + captured: List[object] = [] + + def _pre_hook(_mod, args): + captured.append(args[0][0].detach().to(torch.float32)) # [seq, in] + + handle = down.register_forward_pre_hook(_pre_hook) + cov = None + count = 0 + try: + for prompt in list(corpus)[:max_prompts]: + if not isinstance(prompt, str) or not prompt.strip(): + continue + captured.clear() + inputs = tokenizer( + prompt, + return_tensors="pt", + truncation=True, + max_length=max_tokens, + ).to(device) + with torch.no_grad(): + model(**inputs) + if not captured: + continue + ks = captured[-1] # [seq, in] + if cov is None: + cov = torch.zeros( + ks.shape[1], ks.shape[1], dtype=torch.float32, device=ks.device + ) + cov += ks.t() @ ks + count += ks.shape[0] + finally: + handle.remove() + if cov is None or count == 0: + raise ValueError("covariance corpus produced no key vectors") + cov = cov / float(count) + cov = cov + ridge * torch.eye( + cov.shape[0], dtype=torch.float32, device=cov.device + ) + return cov + + def apply_rome_edit( model, tokenizer, @@ -228,8 +413,16 @@ def apply_rome_edit( device: str, grad_steps: int = _DEFAULT_GRAD_STEPS, lr: float = _DEFAULT_LR, + cov_corpus: Optional[Sequence[str]] = None, + cov_ridge: float = _DEFAULT_COV_RIDGE, ) -> EditKernelResult: - """Single-layer rank-1 ROME edit. Mutates ``model`` in place.""" + """Single-layer rank-1 ROME edit. Mutates ``model`` in place. + + When ``cov_corpus`` is supplied (v0.71.16 #250) the rank-1 update is + preconditioned with the key covariance ``C`` estimated over the corpus + (``C^{-1} k*``), reducing collateral interference with other facts. Falls + back to ``C = I`` (the v0.71.9 covariance-free path) otherwise. + """ layers = _locate_decoder_layers(model) down = _down_proj(layers, layer) key = _capture_key(model, tokenizer, down, subject, device) @@ -238,7 +431,12 @@ def apply_rome_edit( subject=subject, target=target, device=device, grad_steps=grad_steps, lr=lr, ) - norm = _rank1_update(down, key, delta) + cov = None + if cov_corpus: + cov = estimate_key_covariance( + model, tokenizer, down, cov_corpus, device=device, ridge=cov_ridge, + ) + norm = _rank1_update(down, key, delta, cov=cov) return EditKernelResult( method="rome", layer=layer, norm_delta=norm, layers_edited=(layer,), ) @@ -278,9 +476,10 @@ def apply_memit_edit( down = _down_proj(layers, idx) # Re-capture the key at THIS layer (its intermediate dim matches its W). key = _capture_key(model, tokenizer, down, subject, device) - if share.shape[0] != down.weight.shape[0]: - # Hidden dims must match across layers for a residual share; skip - # any layer whose output width differs (defensive). + if share.shape[0] != _proj_out_dim(down): + # Hidden (output) dims must match across layers for a residual + # share; skip any layer whose output width differs (defensive). + # ``_proj_out_dim`` handles the Conv1D transposed layout (#251). continue total_norm += _rank1_update(down, key, share) edited.append(idx) @@ -323,10 +522,14 @@ def apply_alphaedit_edit( denom = float(torch.dot(key_t, key_t).item()) if denom <= 0.0: raise ValueError("key vector has zero norm; cannot apply AlphaEdit update") + # Logical [out, in] ROME update (delta is OUTPUT-space, key is INPUT-space). rome_update = torch.outer(delta.to(down.weight.dtype), key_t) / denom - projected = _alphaedit_project(down, rome_update) + projected = _alphaedit_project(down, rome_update) # logical [out, in] with torch.no_grad(): - down.weight.add_(projected) + if _is_transposed_proj(down): + down.weight.add_(projected.t()) # Conv1D: back to [in, out] (#251) + else: + down.weight.add_(projected) norm = float(torch.linalg.norm(projected).item()) return EditKernelResult( method="alphaedit", layer=layer, norm_delta=norm, layers_edited=(layer,), @@ -344,12 +547,17 @@ def run_edit_kernel( device: str, grad_steps: int = _DEFAULT_GRAD_STEPS, lr: float = _DEFAULT_LR, + cov_corpus: Optional[Sequence[str]] = None, ) -> EditKernelResult: - """Dispatch to the per-method kernel. ``method`` must already be canonical.""" + """Dispatch to the per-method kernel. ``method`` must already be canonical. + + ``cov_corpus`` (v0.71.16 #250) is only consumed by the ROME kernel — the + caller (``apply_edit``) rejects it for other methods before reaching here. + """ if method == "rome": return apply_rome_edit( model, tokenizer, subject=subject, target=target, layer=layer, - device=device, grad_steps=grad_steps, lr=lr, + device=device, grad_steps=grad_steps, lr=lr, cov_corpus=cov_corpus, ) if method == "memit": return apply_memit_edit( diff --git a/src/soup_cli/utils/knowledge_edit.py b/src/soup_cli/utils/knowledge_edit.py index b089398..f7ba5e2 100644 --- a/src/soup_cli/utils/knowledge_edit.py +++ b/src/soup_cli/utils/knowledge_edit.py @@ -20,7 +20,7 @@ from __future__ import annotations from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Mapping, Optional, Tuple +from typing import TYPE_CHECKING, Mapping, Optional, Sequence, Tuple if TYPE_CHECKING: # pragma: no cover — type-only import, no runtime cost from soup_cli.utils.edit_governor import EditGovernor @@ -274,8 +274,9 @@ def apply_edit( trust_remote_code: bool = False, grad_steps: int = 25, lr: float = 0.5, + cov_corpus: Optional[Sequence[str]] = None, ) -> EditResult: - """Apply a knowledge edit live (v0.71.9 #194 / #197 / #203). + """Apply a knowledge edit live (v0.71.9 #194 / #197 / #203; v0.71.16 #250). ROME / MEMIT / AlphaEdit run the rank-1 weight-edit kernel; ``grace`` builds a discrete codebook sidecar. When a ``governor`` is supplied it is @@ -283,12 +284,26 @@ def apply_edit( BEFORE the edit (raising :class:`GovernedEditError` on refusal) and :meth:`EditGovernor.record_edit` runs AFTER with the measured norm delta. + ``cov_corpus`` (v0.71.16 #250) preconditions the ROME update with the key + covariance estimated over the supplied prompts (``C^{-1} k*``). It is only + valid for ``method='rome'`` — covariance preconditioning is the ROME + closed-form update — and is rejected (fail-loud) for any other method. + Re-validates the method so callers passing a duck-typed plan still hit a meaningful error before any model load. Returns an :class:`EditResult`. """ method_attr = getattr(plan, "method", None) canonical = validate_edit_method(method_attr) + # v0.71.16 #250 — reject --cov-corpus for non-ROME methods (fail-loud, no + # silent no-op). Runs BEFORE any model load. + if cov_corpus is not None and canonical != "rome": + raise ValueError( + "--cov-corpus is only supported for method='rome' (got " + f"method={canonical!r}); covariance preconditioning is the ROME " + "closed-form update." + ) + # #197: consult the governor BEFORE doing any work. A refusal must abort # before we burn a model load. if governor is not None: @@ -320,6 +335,7 @@ def apply_edit( model, tokenizer, method=canonical, subject=plan.subject, target=plan.target, layer=plan.layer, device=dev, grad_steps=grad_steps, lr=lr, + cov_corpus=cov_corpus, ) prob_after = measure_target_prob( model, tokenizer, subject=plan.subject, target=plan.target, device=dev, diff --git a/src/soup_cli/utils/longlora.py b/src/soup_cli/utils/longlora.py index e0e6d69..548e378 100644 --- a/src/soup_cli/utils/longlora.py +++ b/src/soup_cli/utils/longlora.py @@ -36,6 +36,14 @@ _LLAMA_REGEX = re.compile(r"(?:^|[^a-z0-9])(?:code)?-?llama(?:-?\d+(?:\.\d+)?)?( _MISTRAL_REGEX = re.compile( r"(?:^|[^a-z0-9])mistral(?:-?\d+(?:\.\d+)?)?(?:[^a-z0-9]|$)" ) +# v0.71.16 #147 — Mixtral needs its OWN regex: the bare ``mistral`` token does +# not appear in ``mixtral`` (m-i-x vs m-i-s), so ``is_mistral_model`` excludes +# the MoE variant. The Mixtral attention is structurally identical to Mistral's +# (separate q/k/v projections, GQA) — only the MLP is a sparse MoE — so the S² +# forward override reuses the separate-QKV projection-shift path. +_MIXTRAL_REGEX = re.compile( + r"(?:^|[^a-z0-9])mixtral(?:-?\d+(?:\.\d+)?)?(?:[^a-z0-9]|$)" +) _QWEN_REGEX = re.compile(r"(?:^|[^a-z0-9])qwen(?:-?\d+(?:\.\d+)?)?(?:[^a-z0-9]|$)") _PHI_REGEX = re.compile(r"(?:^|[^a-z0-9])phi(?:-?\d+(?:\.\d+)?)?(?:[^a-z0-9]|$)") @@ -70,9 +78,12 @@ def _check_model_name(model_name: str) -> str | None: def is_mistral_model(model_name: str) -> bool: - """Return True if ``model_name`` belongs to the Mistral / Mixtral family. + """Return True if ``model_name`` belongs to the dense Mistral family. - v0.53.4 #120 — LongLoRA architecture allowlist expansion. + v0.53.4 #120 — LongLoRA architecture allowlist expansion. NOTE: this + deliberately EXCLUDES the Mixtral MoE variant (``mixtral`` does not contain + the ``mistral`` token) — use :func:`is_mixtral_model` for Mixtral + (v0.71.16 #147). """ lowered = _check_model_name(model_name) if lowered is None: @@ -80,6 +91,22 @@ def is_mistral_model(model_name: str) -> bool: return _MISTRAL_REGEX.search(lowered) is not None +def is_mixtral_model(model_name: str) -> bool: + """Return True if ``model_name`` belongs to the Mixtral MoE family. + + v0.71.16 #147 — dedicated detector. ``is_mistral_model`` deliberately + excludes Mixtral (different token), so the LongLoRA allowlist needs this + helper to cover Mixtral-8x7B / Mixtral-8x22B. The Mixtral attention is the + standard separate-QKV (GQA) shell — the MoE lives in the MLP — so the S² + forward override needs no special attention handling, only the family + + attention-class-name allowlist entries. + """ + lowered = _check_model_name(model_name) + if lowered is None: + return False + return _MIXTRAL_REGEX.search(lowered) is not None + + def is_qwen_model(model_name: str) -> bool: """Return True if ``model_name`` belongs to the Qwen family. @@ -105,15 +132,12 @@ def is_phi_model(model_name: str) -> bool: def is_supported_longlora_arch(model_name: object) -> bool: """Return True if ``model_name`` is in the LongLoRA allowlist. - The v0.53.4 #120 allowlist covers Llama / CodeLlama (Llama 3.x heritage), - Mistral, Qwen, and Phi. The S² forward override (deferred to v0.49.1) - attaches per-arch; the schema gate uses this helper. + The allowlist covers Llama / CodeLlama (Llama 3.x heritage), Mistral, + Mixtral (v0.71.16 #147), Qwen, and Phi. The S² forward override attaches + per-arch; the schema gate uses this helper. Returns ``False`` (never raises) on non-string input — matches ``is_known_vlm_base`` / ``is_bitnet_model`` defensive-surface policy. - Mixtral is intentionally NOT covered (regex matches the bare token - ``mistral``, not the Mixtral MoE variant); add a dedicated helper when - Mixtral attention upstream lands a stable forward signature. """ if not isinstance(model_name, str): return False @@ -121,6 +145,7 @@ def is_supported_longlora_arch(model_name: object) -> bool: return ( is_llama_model(model_name) or is_mistral_model(model_name) + or is_mixtral_model(model_name) or is_qwen_model(model_name) or is_phi_model(model_name) ) @@ -173,7 +198,7 @@ def validate_longlora_compat( LongLoRA requires: * ``task='sft'`` (preference / RL trainers have a different forward path) * ``backend='transformers'`` (Unsloth has its own attention; MLX has none) - * Llama / Mistral / Qwen / Phi base model (v0.53.4 #120 allowlist) + * Llama / Mistral / Mixtral / Qwen / Phi base model (v0.71.16 #147 allowlist) * ``use_ring_attention=False`` (mutually exclusive custom attention kernel) * No FlashAttention v3 build present (v0.53.4 #122 — incompatible custom-mask) """ @@ -209,7 +234,7 @@ def validate_longlora_compat( safe_name = _truncate_for_message(model_name) raise ValueError( "LongLoRA architecture allowlist currently covers Llama / " - "CodeLlama, Mistral, Qwen, and Phi families (got " + "CodeLlama, Mistral, Mixtral, Qwen, and Phi families (got " f"base={safe_name!r}). Open a feature request to add another " "architecture." ) @@ -287,7 +312,7 @@ def shift_heads_for_s2( # Mistral, Qwen2 — GQA handled by deriving the head count per-projection from # its output dim) patch q_proj + k_proj independently. Phi-3 fuses Q/K/V into a # single ``qkv_proj`` and needs the slice split before shifting. -_SEPARATE_QKV_FAMILIES: tuple[str, ...] = ("Llama", "Mistral", "Qwen") +_SEPARATE_QKV_FAMILIES: tuple[str, ...] = ("Llama", "Mistral", "Mixtral", "Qwen") _FUSED_QKV_FAMILIES: tuple[str, ...] = ("Phi",) @@ -426,8 +451,11 @@ class LongLoRAForwardOverride: # crafted model class with an arbitrarily long name does not feed # an unbounded string into the regex. max_class_name_len = 256 + # v0.71.16 #147 — ``Mixtral`` added: ``Mistral\w*Attention`` does NOT + # match ``MixtralAttention`` (different token), so the override would + # silently skip Mixtral without this alternative. attention_class_re = re.compile( - r"(?:Llama|Mistral|Qwen|Phi)\w*Attention$" + r"(?:Llama|Mistral|Mixtral|Qwen|Phi)\w*Attention$" ) for module in _walk_modules(self.model): diff --git a/tests/test_v07116.py b/tests/test_v07116.py new file mode 100644 index 0000000..608be39 --- /dev/null +++ b/tests/test_v07116.py @@ -0,0 +1,1247 @@ +"""v0.71.16 — Knowledge edit depth. + +Closes: + * #251 — edit kernels gain GPT-2 ``transformer.h`` / ``mlp.c_proj`` support + (transpose-aware rank-1 update for the Conv1D weight layout). + * #252 — EditGovernor edit-count increment is now atomic (baseline-delta + merge under the cross-process lock) so concurrent ``soup edit set`` + runs cannot lose an increment. + * #250 — covariance-preconditioned ROME via ``--cov-corpus`` (estimate the + key covariance C from a stats corpus and use C^{-1} k* instead of + k*; falls back to C=I when no corpus). + * #147 — Mixtral joins the LongLoRA architecture allowlist (dedicated + ``is_mixtral_model`` helper + ``MixtralAttention`` forward override). + +Kernel maths are exercised with real torch (the [dev] extra) on tiny CPU +fakes. Full real-model apply paths are covered by the release step-6 smoke +(tiny-gpt2 + SmolLM2-135M). +""" + +from __future__ import annotations + +import json +import os +import sys +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +nn = torch.nn # assignment (not an import) so it works after importorskip + + +# =========================================================================== +# Shared fakes — a faithful tiny GPT-2 / Llama LM (forward + loss + tokenizer) +# =========================================================================== + + +class _Conv1D(nn.Module): + """Faithful mini ``transformers.pytorch_utils.Conv1D``. + + Weight is ``[nx, nf]`` (transposed relative to nn.Linear's ``[out, in]``), + ``nf`` is the output feature count, and ``forward`` is ``x @ weight + bias``. + """ + + def __init__(self, nx: int, nf: int): + super().__init__() + self.nf = nf + self.weight = nn.Parameter(torch.randn(nx, nf) * 0.05) + self.bias = nn.Parameter(torch.zeros(nf)) + + def forward(self, x): # noqa: D401 + return x @ self.weight + self.bias + + +class _FakeEnc(dict): + """BatchEncoding-like: dict for ``**`` unpacking + a no-op ``.to``.""" + + def to(self, _device): + return self + + +class _FakeTok: + vocab = 24 + + def _ids(self, text: str) -> list[int]: + out = [(sum(ord(c) for c in t) % (self.vocab - 4)) + 1 for t in text.split()] + return out or [1] + + def __call__( + self, + text, + return_tensors=None, + truncation=False, + max_length=None, + add_special_tokens=True, + ): + ids = self._ids(text) + if add_special_tokens: + ids = [0] + ids + if max_length is not None: + ids = ids[:max_length] + if not ids: + ids = [1] + if return_tensors == "pt": + return _FakeEnc(input_ids=torch.tensor([ids], dtype=torch.long)) + return {"input_ids": ids} + + +def _lm_forward(hidden, lm_head, input_ids, labels): + logits = lm_head(hidden) + out = SimpleNamespace(logits=logits) + if labels is not None: + shift_logits = logits[:, :-1, :].reshape(-1, logits.shape[-1]) + shift_labels = labels[:, 1:].reshape(-1) + out.loss = nn.functional.cross_entropy( + shift_logits, shift_labels, ignore_index=-100 + ) + return out + + +class _LlamaMLP(nn.Module): + def __init__(self, hidden, inter): + super().__init__() + self.gate_proj = nn.Linear(hidden, inter, bias=False) + self.down_proj = nn.Linear(inter, hidden, bias=False) + + def forward(self, x): + return self.down_proj(torch.relu(self.gate_proj(x))) + + +class _LlamaBlock(nn.Module): + def __init__(self, hidden, inter): + super().__init__() + self.mlp = _LlamaMLP(hidden, inter) + + def forward(self, x): + return x + self.mlp(x) + + +class _LlamaInner(nn.Module): + def __init__(self, vocab, hidden, inter, layers): + super().__init__() + self.embed_tokens = nn.Embedding(vocab, hidden) + self.layers = nn.ModuleList( + [_LlamaBlock(hidden, inter) for _ in range(layers)] + ) + + def forward(self, input_ids): + x = self.embed_tokens(input_ids) + for blk in self.layers: + x = blk(x) + return x + + +class _FakeLlamaLM(nn.Module): + def __init__(self, vocab=24, hidden=8, inter=16, layers=3): + super().__init__() + self.model = _LlamaInner(vocab, hidden, inter, layers) + self.lm_head = nn.Linear(hidden, vocab, bias=False) + + def forward(self, input_ids=None, labels=None, **_kw): + return _lm_forward(self.model(input_ids), self.lm_head, input_ids, labels) + + +class _GPT2MLP(nn.Module): + def __init__(self, hidden, inter): + super().__init__() + self.c_fc = _Conv1D(hidden, inter) + self.c_proj = _Conv1D(inter, hidden) + + def forward(self, x): + return self.c_proj(torch.relu(self.c_fc(x))) + + +class _GPT2Block(nn.Module): + def __init__(self, hidden, inter): + super().__init__() + self.mlp = _GPT2MLP(hidden, inter) + + def forward(self, x): + return x + self.mlp(x) + + +class _GPT2Inner(nn.Module): + def __init__(self, vocab, hidden, inter, layers): + super().__init__() + self.wte = nn.Embedding(vocab, hidden) + self.h = nn.ModuleList([_GPT2Block(hidden, inter) for _ in range(layers)]) + + def forward(self, input_ids): + x = self.wte(input_ids) + for blk in self.h: + x = blk(x) + return x + + +class _FakeGPT2LM(nn.Module): + def __init__(self, vocab=24, hidden=8, inter=16, layers=3): + super().__init__() + self.transformer = _GPT2Inner(vocab, hidden, inter, layers) + self.lm_head = nn.Linear(hidden, vocab, bias=False) + + def forward(self, input_ids=None, labels=None, **_kw): + return _lm_forward( + self.transformer(input_ids), self.lm_head, input_ids, labels + ) + + +def _peft_wrap(base): + """Minimal PEFT-style wrapper exposing ``get_base_model``.""" + + class _Peft: + def __init__(self, b): + self._b = b + + def get_base_model(self): + return self._b + + return _Peft(base) + + +class _PeftCallable: + """Callable PEFT-style wrapper: delegates forward + train/eval to base. + + Unlike ``_peft_wrap`` (locate-only), this supports the full kernel path — + ``model(**inputs)`` / ``model.training`` / ``eval()`` / ``train()`` — so a + GPT-2 base can be edited end-to-end while wrapped (L1). + """ + + def __init__(self, base): + self._b = base + + def get_base_model(self): + return self._b + + def __call__(self, *a, **k): + return self._b(*a, **k) + + @property + def training(self): + return self._b.training + + def eval(self): + self._b.eval() + return self + + def train(self, mode=True): + self._b.train(mode) + return self + + +_SUBJECT = "Paris is the capital of" +_TARGET = "Lyon" +_CORPUS = [ + "the quick brown fox jumps over the lazy dog", + "lorem ipsum dolor sit amet consectetur", + "machine learning models edit facts surgically", +] + + +# =========================================================================== +# #251 — GPT-2 transformer.h / mlp.c_proj support +# =========================================================================== + + +class TestLocateDecoderLayersGpt2: + def test_locates_transformer_h(self): + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + model = _FakeGPT2LM(layers=4) + layers = _locate_decoder_layers(model) + assert len(layers) == 4 + + def test_peft_wrapped_gpt2(self): + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + model = _peft_wrap(_FakeGPT2LM(layers=3)) + layers = _locate_decoder_layers(model) + assert len(layers) == 3 + + def test_peft_wrapped_llama_still_works(self): + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + model = _peft_wrap(_FakeLlamaLM(layers=2)) + layers = _locate_decoder_layers(model) + assert len(layers) == 2 + + def test_unknown_arch_raises(self): + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + with pytest.raises(ValueError, match="decoder layers"): + _locate_decoder_layers(nn.Linear(2, 2)) + + def test_get_base_model_raises_swallowed(self): + """A PEFT wrapper whose get_base_model() blows up must fall through to + the clear ValueError (DEBUG-logged, not masked) — review L2.""" + from soup_cli.utils.edit_kernels import _locate_decoder_layers + + class _BadPeft: + def get_base_model(self): + raise RuntimeError("boom") + + with pytest.raises(ValueError, match="decoder layers"): + _locate_decoder_layers(_BadPeft()) + + +class TestDownProjGpt2: + def test_returns_c_proj(self): + from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers + + model = _FakeGPT2LM() + down = _down_proj(_locate_decoder_layers(model), 1) + assert hasattr(down, "weight") + # Conv1D weight is [in, out] = [inter, hidden] = [16, 8]. + assert tuple(down.weight.shape) == (16, 8) + assert down.nf == 8 + + def test_llama_down_proj_unchanged(self): + from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers + + model = _FakeLlamaLM() + down = _down_proj(_locate_decoder_layers(model), 0) + # nn.Linear weight is [out, in] = [hidden, inter] = [8, 16]. + assert tuple(down.weight.shape) == (8, 16) + + def test_out_of_range(self): + from soup_cli.utils.edit_kernels import _down_proj, _locate_decoder_layers + + layers = _locate_decoder_layers(_FakeGPT2LM()) + with pytest.raises(ValueError, match="out of range"): + _down_proj(layers, 99) + + +class TestProjHelpers: + def test_is_transposed_proj_conv1d(self): + from soup_cli.utils.edit_kernels import _is_transposed_proj + + assert _is_transposed_proj(_Conv1D(16, 8)) is True + + def test_is_transposed_proj_linear(self): + from soup_cli.utils.edit_kernels import _is_transposed_proj + + assert _is_transposed_proj(nn.Linear(16, 8, bias=False)) is False + + def test_proj_out_dim_conv1d_uses_nf(self): + from soup_cli.utils.edit_kernels import _proj_out_dim + + # Conv1D nf = output (hidden) dim, NOT weight.shape[0] (= in dim). + assert _proj_out_dim(_Conv1D(16, 8)) == 8 + + def test_proj_out_dim_linear_uses_shape0(self): + from soup_cli.utils.edit_kernels import _proj_out_dim + + assert _proj_out_dim(nn.Linear(16, 8, bias=False)) == 8 + + def test_is_transposed_proj_rejects_bool_nf(self): + """``nf=True`` (bool, a subclass of int) must NOT be treated as Conv1D + — review L3.""" + from soup_cli.utils.edit_kernels import _is_transposed_proj + + assert _is_transposed_proj(SimpleNamespace(nf=True)) is False + + +class TestRank1UpdateTransposed: + def test_conv1d_post_condition(self): + """Conv1D: ``key @ W`` must gain exactly ``delta`` after the update.""" + from soup_cli.utils.edit_kernels import _rank1_update + + conv = _Conv1D(16, 8) + key = torch.ones(16) + delta = torch.full((8,), 0.5) + before = key @ conv.weight # [out] = [8] + norm = _rank1_update(conv, key, delta) + after = key @ conv.weight + assert norm > 0 + assert torch.allclose(after - before, delta, atol=1e-4) + + def test_linear_post_condition_regression(self): + from soup_cli.utils.edit_kernels import _rank1_update + + lin = nn.Linear(16, 8, bias=False) + key = torch.ones(16) + delta = torch.full((8,), 0.5) + before = lin.weight @ key + norm = _rank1_update(lin, key, delta) + after = lin.weight @ key + assert norm > 0 + assert torch.allclose(after - before, delta, atol=1e-4) + + def test_conv1d_zero_key_rejected(self): + from soup_cli.utils.edit_kernels import _rank1_update + + with pytest.raises(ValueError, match="zero norm"): + _rank1_update(_Conv1D(16, 8), torch.zeros(16), torch.ones(8)) + + +class TestAlphaEditProjectTransposed: + def test_conv1d_shape_and_determinism(self): + from soup_cli.utils.edit_kernels import _alphaedit_project + + conv = _Conv1D(16, 8) + # Logical [out, in] update. + upd = torch.full((8, 16), 0.3) + p1 = _alphaedit_project(conv, upd) + p2 = _alphaedit_project(conv, upd) + assert tuple(p1.shape) == (8, 16) + assert torch.allclose(p1, p2) + + def test_conv1d_projection_idempotent(self): + from soup_cli.utils.edit_kernels import _alphaedit_project + + conv = _Conv1D(16, 8) + upd = torch.randn(8, 16) + once = _alphaedit_project(conv, upd) + twice = _alphaedit_project(conv, once) + # P is a projection: P(P(u)) == P(u). + assert torch.allclose(once, twice, atol=1e-4) + + +class TestApplyKernelsGpt2EndToEnd: + def _run(self, method, model): + from soup_cli.utils.edit_kernels import measure_target_prob, run_edit_kernel + + tok = _FakeTok() + before = measure_target_prob( + model, tok, subject=_SUBJECT, target=_TARGET, device="cpu" + ) + result = run_edit_kernel( + model, tok, method=method, subject=_SUBJECT, target=_TARGET, + layer=1, device="cpu", + ) + after = measure_target_prob( + model, tok, subject=_SUBJECT, target=_TARGET, device="cpu" + ) + return before, result, after + + def test_rome_gpt2_changes_target(self): + torch.manual_seed(0) + before, result, after = self._run("rome", _FakeGPT2LM(layers=3)) + assert result.method == "rome" + assert result.layers_edited == (1,) + assert result.norm_delta > 0 + assert 0.0 <= after <= 1.0 + assert after > before # the fact was edited + + def test_rome_llama_regression(self): + torch.manual_seed(0) + before, result, after = self._run("rome", _FakeLlamaLM(layers=3)) + assert result.norm_delta > 0 + assert after > before + + def test_memit_gpt2_runs(self): + torch.manual_seed(1) + _before, result, after = self._run("memit", _FakeGPT2LM(layers=3)) + assert result.method == "memit" + assert len(result.layers_edited) >= 1 + assert result.norm_delta > 0 + assert 0.0 <= after <= 1.0 + + def test_alphaedit_gpt2_runs(self): + torch.manual_seed(2) + _before, result, after = self._run("alphaedit", _FakeGPT2LM(layers=3)) + assert result.method == "alphaedit" + assert result.layers_edited == (1,) + assert result.norm_delta > 0 + assert 0.0 <= after <= 1.0 + + def test_memit_gpt2_edits_full_band(self): + """The #251 ``_proj_out_dim`` Conv1D fix makes the MEMIT band dim-check + MATCH across uniform-width GPT-2 layers, so the whole band is edited.""" + from soup_cli.utils.edit_kernels import apply_memit_edit + + torch.manual_seed(1) + result = apply_memit_edit( + _FakeGPT2LM(layers=3), _FakeTok(), + subject=_SUBJECT, target=_TARGET, layer=2, device="cpu", + ) + # _MEMIT_BAND=3, layer=2 → band [0, 1, 2]; all same width → all edited. + assert result.layers_edited == (0, 1, 2) + + def test_memit_raises_when_no_layer_editable(self, monkeypatch): + """Defensive band-skip → raise branch (#251 Conv1D dim-check). + + A call-counted ``_proj_out_dim`` returns the real dim for the residual + sizing (first call) then a mismatching dim for every band check, so the + sole band layer is skipped and the empty-edit guard fires. + """ + import soup_cli.utils.edit_kernels as ek + + real = ek._proj_out_dim + state = {"n": 0} + + def fake(module): + state["n"] += 1 + return real(module) if state["n"] == 1 else real(module) + 1 + + monkeypatch.setattr(ek, "_proj_out_dim", fake) + with pytest.raises(ValueError, match="could not edit any layer"): + ek.apply_memit_edit( + _FakeGPT2LM(layers=1), _FakeTok(), + subject=_SUBJECT, target=_TARGET, layer=0, device="cpu", + ) + + def test_alphaedit_conv1d_weight_orientation(self): + """AlphaEdit's ``.t()`` apply keeps the Conv1D weight in [in, out] + layout — an un-transposed [out, in] apply would shape-error in add_ + (review H3).""" + from soup_cli.utils.edit_kernels import ( + _down_proj, + _locate_decoder_layers, + apply_alphaedit_edit, + ) + + torch.manual_seed(3) + model = _FakeGPT2LM(layers=3) + down = _down_proj(_locate_decoder_layers(model), 1) + w0 = down.weight.detach().clone() # [in, out] = [16, 8] + apply_alphaedit_edit( + model, _FakeTok(), + subject=_SUBJECT, target=_TARGET, layer=1, device="cpu", + ) + delta_w = down.weight.detach() - w0 + assert tuple(delta_w.shape) == (16, 8) + assert torch.isfinite(delta_w).all() + assert float(torch.linalg.norm(delta_w)) > 0 + + def test_rome_peft_wrapped_gpt2(self): + """A PEFT-wrapped GPT-2 base is editable end-to-end through the kernel + (review L1 — PEFT fallback at the kernel level, not just locate).""" + from soup_cli.utils.edit_kernels import run_edit_kernel + + torch.manual_seed(0) + model = _PeftCallable(_FakeGPT2LM(layers=3)) + result = run_edit_kernel( + model, _FakeTok(), method="rome", + subject=_SUBJECT, target=_TARGET, layer=1, device="cpu", + ) + assert result.method == "rome" + assert result.norm_delta > 0 + + +# =========================================================================== +# #252 — atomic EditGovernor edit-count increment +# =========================================================================== + + +class TestGovernorAtomicIncrement: + def test_baseline_set_on_fresh(self): + from soup_cli.utils.edit_governor import EditGovernor + + gov = EditGovernor(base_model="m") + assert gov._persisted_edit_count == 0 + + def test_baseline_set_on_loaded(self, tmp_path, monkeypatch): + from soup_cli.utils.edit_governor import ( + EditGovernor, + EditGovernorStore, + load_governor, + save_governor, + ) + + monkeypatch.chdir(tmp_path) + db = str(tmp_path / "g.db") + with EditGovernorStore(db) as store: + gov = EditGovernor(base_model="m") + gov.record_edit(method="rome", norm_delta=0.1) + gov.record_edit(method="rome", norm_delta=0.1) + save_governor(store, gov) + with EditGovernorStore(db) as store2: + restored = load_governor(store2, "m") + # Baseline mirrors the loaded count so a further edit merges as +1. + assert restored.edit_count == 2 + assert restored._persisted_edit_count == 2 + + def test_concurrent_save_merges_increments(self, tmp_path, monkeypatch): + """Two governors loaded from the same state both record + save. + + With the pre-#252 absolute-write behaviour the second save would + clobber the first to 1. The baseline-delta merge keeps both → 2. + """ + from soup_cli.utils.edit_governor import ( + EditGovernorStore, + load_governor, + save_governor, + ) + + monkeypatch.chdir(tmp_path) + db = str(tmp_path / "g.db") + with EditGovernorStore(db) as store: + gov_a = load_governor(store, "m") + gov_b = load_governor(store, "m") + gov_a.record_edit(method="rome", norm_delta=0.1) + gov_b.record_edit(method="rome", norm_delta=0.2) + save_governor(store, gov_a) # persists 1 + save_governor(store, gov_b) # MERGES → persists 2 (not clobber) + final = load_governor(store, "m") + assert final.edit_count == 2 + + def test_concurrent_save_merges_multi_increments(self, tmp_path, monkeypatch): + """Two governors record MULTIPLE edits each → merged is the sum of the + deltas (3 + 2 = 5), proving the baseline-delta merge (not a naive +1 + per save) — review M4.""" + from soup_cli.utils.edit_governor import ( + EditGovernorStore, + load_governor, + save_governor, + ) + + monkeypatch.chdir(tmp_path) + db = str(tmp_path / "g.db") + with EditGovernorStore(db) as store: + a = load_governor(store, "m") + b = load_governor(store, "m") + for _ in range(3): + a.record_edit(method="rome", norm_delta=0.1) + for _ in range(2): + b.record_edit(method="rome", norm_delta=0.1) + save_governor(store, a) # persists 3 + save_governor(store, b) # merges +2 → 5 (NOT clobber to 2, NOT +1) + assert load_governor(store, "m").edit_count == 5 + + def test_merge_onto_existing_row(self, tmp_path, monkeypatch): + from soup_cli.utils.edit_governor import ( + EditGovernorStore, + load_governor, + save_governor, + ) + + monkeypatch.chdir(tmp_path) + db = str(tmp_path / "g.db") + with EditGovernorStore(db) as store: + # Seed a persisted count of 5. + g0 = load_governor(store, "m") + for _ in range(5): + g0.record_edit(method="rome", norm_delta=0.1) + save_governor(store, g0) + with EditGovernorStore(db) as store2: + g1 = load_governor(store2, "m") # baseline 5 + g1.record_edit(method="rome", norm_delta=0.2) # → 6 + save_governor(store2, g1) + final = load_governor(store2, "m") + assert final.edit_count == 6 + + def test_save_updates_in_memory_count(self, tmp_path, monkeypatch): + """After an atomic merge, ``governor.edit_count`` reflects the merged + value so the CLI summary shows the real count.""" + from soup_cli.utils.edit_governor import ( + EditGovernorStore, + load_governor, + save_governor, + ) + + monkeypatch.chdir(tmp_path) + db = str(tmp_path / "g.db") + with EditGovernorStore(db) as store: + a = load_governor(store, "m") + b = load_governor(store, "m") + a.record_edit(method="rome", norm_delta=0.1) + b.record_edit(method="rome", norm_delta=0.2) + save_governor(store, a) + save_governor(store, b) + # b merged onto a's persisted 1 → 2; b.edit_count must reflect it. + assert b.edit_count == 2 + # Re-saving b must NOT double-count (baseline now 2). + save_governor(store, b) + assert load_governor(store, "m").edit_count == 2 + + def test_save_state_atomic_uses_lock(self): + """Regression: the get+insert in save_state runs under the lock.""" + import inspect + + from soup_cli.utils.edit_governor import EditGovernorStore + + src = inspect.getsource(EditGovernorStore.save_state) + assert "_cross_process_lock" in src + assert "get_state" in src # read inside the lock + + +# =========================================================================== +# #250 — covariance-preconditioned ROME +# =========================================================================== + + +class TestEstimateKeyCovariance: + def test_shape_and_spd(self): + from soup_cli.utils.edit_kernels import ( + _down_proj, + _locate_decoder_layers, + estimate_key_covariance, + ) + + model = _FakeGPT2LM() + down = _down_proj(_locate_decoder_layers(model), 1) + cov = estimate_key_covariance( + model, _FakeTok(), down, _CORPUS, device="cpu", + ) + # Covariance dim = down-proj INPUT dim (intermediate) = 16. + assert tuple(cov.shape) == (16, 16) + # Symmetric. + assert torch.allclose(cov, cov.t(), atol=1e-5) + # SPD (ridge guarantees positive eigenvalues). + eigvals = torch.linalg.eigvalsh(cov) + assert float(eigvals.min()) > 0.0 + + def test_empty_corpus_rejected(self): + from soup_cli.utils.edit_kernels import ( + _down_proj, + _locate_decoder_layers, + estimate_key_covariance, + ) + + model = _FakeGPT2LM() + down = _down_proj(_locate_decoder_layers(model), 0) + with pytest.raises(ValueError, match="corpus"): + estimate_key_covariance(model, _FakeTok(), down, [], device="cpu") + + def test_bad_caps_rejected(self): + from soup_cli.utils.edit_kernels import ( + _down_proj, + _locate_decoder_layers, + estimate_key_covariance, + ) + + model = _FakeGPT2LM() + down = _down_proj(_locate_decoder_layers(model), 0) + with pytest.raises(ValueError, match="max_prompts"): + estimate_key_covariance( + model, _FakeTok(), down, _CORPUS, device="cpu", max_prompts=0 + ) + with pytest.raises(ValueError, match="max_tokens"): + estimate_key_covariance( + model, _FakeTok(), down, _CORPUS, device="cpu", max_tokens=-1 + ) + with pytest.raises(ValueError, match="ridge"): + estimate_key_covariance( + model, _FakeTok(), down, _CORPUS, device="cpu", ridge=-1.0 + ) + + def test_all_blank_corpus_rejected(self): + """A corpus where every entry is blank / non-str captures no keys → + the runtime ``count == 0`` branch raises (review H4).""" + from soup_cli.utils.edit_kernels import ( + _down_proj, + _locate_decoder_layers, + estimate_key_covariance, + ) + + model = _FakeGPT2LM() + down = _down_proj(_locate_decoder_layers(model), 0) + with pytest.raises(ValueError, match="no key vectors"): + estimate_key_covariance( + model, _FakeTok(), down, ["", " ", 123], device="cpu" + ) + + +class TestRank1UpdatePreconditioned: + def test_post_condition_preserved_linear(self): + """With C != I the ROME post-condition ``down(key*) += delta`` still + holds exactly — the covariance only redistributes the update mass.""" + from soup_cli.utils.edit_kernels import _rank1_update + + lin = nn.Linear(16, 8, bias=False) + key = torch.randn(16) + delta = torch.randn(8) + # Arbitrary SPD covariance. + a = torch.randn(16, 16) + cov = a @ a.t() + torch.eye(16) + before = lin.weight @ key + norm = _rank1_update(lin, key, delta, cov=cov) + after = lin.weight @ key + assert norm > 0 + assert torch.allclose(after - before, delta, atol=1e-3) + + def test_post_condition_preserved_conv1d(self): + from soup_cli.utils.edit_kernels import _rank1_update + + conv = _Conv1D(16, 8) + key = torch.randn(16) + delta = torch.randn(8) + a = torch.randn(16, 16) + cov = a @ a.t() + torch.eye(16) + before = key @ conv.weight + norm = _rank1_update(conv, key, delta, cov=cov) + after = key @ conv.weight + assert norm > 0 + assert torch.allclose(after - before, delta, atol=1e-3) + + def test_non_finite_cov_rejected(self): + """A non-finite covariance must raise (not silently corrupt weights). + + Review-fix: ``denom = NaN`` would slip past the bare ``<= 0`` guard. + """ + from soup_cli.utils.edit_kernels import _rank1_update + + lin = nn.Linear(16, 8, bias=False) + bad_cov = torch.full((16, 16), float("nan")) + with pytest.raises(ValueError): + _rank1_update(lin, torch.ones(16), torch.ones(8), cov=bad_cov) + + def test_singular_cov_rejected(self): + """A singular (rank-deficient) covariance makes the solve fail / + produce a non-finite denom → clean ValueError (review M3).""" + from soup_cli.utils.edit_kernels import _rank1_update + + lin = nn.Linear(16, 8, bias=False) + with pytest.raises(ValueError, match="covariance solve failed|degenerate"): + _rank1_update( + lin, torch.ones(16), torch.ones(8), cov=torch.zeros(16, 16) + ) + + def test_cov_changes_update_direction(self): + """A non-identity covariance produces a different update than C=I.""" + from soup_cli.utils.edit_kernels import _rank1_update + + key = torch.randn(16) + delta = torch.randn(8) + a = torch.randn(16, 16) + cov = a @ a.t() + torch.eye(16) + + lin_iden = nn.Linear(16, 8, bias=False) + w0 = lin_iden.weight.detach().clone() + _rank1_update(lin_iden, key, delta) + upd_iden = lin_iden.weight.detach() - w0 + + lin_cov = nn.Linear(16, 8, bias=False) + with torch.no_grad(): + lin_cov.weight.copy_(w0) + _rank1_update(lin_cov, key, delta, cov=cov) + upd_cov = lin_cov.weight.detach() - w0 + + assert not torch.allclose(upd_iden, upd_cov, atol=1e-3) + + +class TestApplyRomeWithCovCorpus: + def test_apply_edit_rome_cov(self, monkeypatch): + import soup_cli.utils.live_eval as live_eval + from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan + + torch.manual_seed(0) + model = _FakeGPT2LM(layers=3) + tok = _FakeTok() + monkeypatch.setattr( + live_eval, "load_model_and_tokenizer", + lambda *a, **k: (model, tok, "cpu"), + ) + plan = build_edit_plan( + base="b", method="rome", subject=_SUBJECT, target=_TARGET, layer=1, + ) + result = apply_edit(plan, cov_corpus=_CORPUS) + assert result.method == "rome" + assert result.norm_delta > 0 + assert result.target_prob_after > result.target_prob_before + + def test_cov_corpus_rejected_for_memit(self, monkeypatch): + import soup_cli.utils.live_eval as live_eval + from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan + + # load_model_and_tokenizer must NOT be reached — the reject is before it. + monkeypatch.setattr( + live_eval, "load_model_and_tokenizer", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")), + ) + plan = build_edit_plan( + base="b", method="memit", subject="s", target="t", + ) + with pytest.raises(ValueError, match="cov-corpus"): + apply_edit(plan, cov_corpus=_CORPUS) + + def test_cov_corpus_rejected_for_alphaedit(self, monkeypatch): + import soup_cli.utils.live_eval as live_eval + from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan + + monkeypatch.setattr( + live_eval, "load_model_and_tokenizer", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")), + ) + plan = build_edit_plan( + base="b", method="alphaedit", subject="s", target="t", + ) + with pytest.raises(ValueError, match="cov-corpus"): + apply_edit(plan, cov_corpus=_CORPUS) + + def test_cov_corpus_rejected_for_grace(self): + """grace takes a different code path (codebook sidecar) but the cov + reject still fires first — before any import / model load (review H1).""" + from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan + + plan = build_edit_plan(base="b", method="grace", subject="s", target="t") + with pytest.raises(ValueError, match="cov-corpus"): + apply_edit(plan, cov_corpus=_CORPUS) + + def test_cov_reject_runs_before_governor(self, monkeypatch): + """Order matters: a non-ROME method + cov_corpus + a governor that + WOULD refuse must report the cov error, not the governance refusal + (the cov check precedes governor.check_can_edit) — review M5.""" + import soup_cli.utils.live_eval as live_eval + from soup_cli.utils.edit_governor import EditGovernor + from soup_cli.utils.knowledge_edit import apply_edit, build_edit_plan + + monkeypatch.setattr( + live_eval, "load_model_and_tokenizer", + lambda *a, **k: (_ for _ in ()).throw(AssertionError("loaded")), + ) + gov = EditGovernor(base_model="b") + gov.record_edit(method="rome", norm_delta=100.0) # → BLOWUP: would refuse + plan = build_edit_plan(base="b", method="memit", subject="s", target="t") + with pytest.raises(ValueError, match="cov-corpus"): + apply_edit(plan, governor=gov, cov_corpus=_CORPUS) + + +class TestLoadCovCorpus: + def test_parses_jsonl_text_field(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + f = tmp_path / "corpus.jsonl" + f.write_text( + json.dumps({"text": "row one"}) + "\n" + + json.dumps({"prompt": "row two"}) + "\n" + + json.dumps({"content": "row three"}) + "\n", + encoding="utf-8", + ) + rows = _load_cov_corpus("corpus.jsonl") + assert rows == ["row one", "row two", "row three"] + + def test_parses_raw_text_lines(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + f = tmp_path / "corpus.txt" + f.write_text("plain line one\nplain line two\n", encoding="utf-8") + rows = _load_cov_corpus("corpus.txt") + assert rows == ["plain line one", "plain line two"] + + def test_outside_cwd_rejected(self): + from soup_cli.commands.edit import _load_cov_corpus + + with pytest.raises(ValueError, match="cwd"): + _load_cov_corpus("/etc/passwd") + + def test_null_byte_rejected(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + with pytest.raises(ValueError, match="null"): + _load_cov_corpus("a\x00b.jsonl") + + def test_missing_file(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + with pytest.raises(FileNotFoundError): + _load_cov_corpus("nope.jsonl") + + def test_directory_rejected(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + (tmp_path / "adir").mkdir() + # POSIX: os.open succeeds, fstat → not S_ISREG → "regular file". + # Windows: os.open on a dir raises PermissionError → "not readable". + with pytest.raises(ValueError, match="regular file|not readable"): + _load_cov_corpus("adir") + + def test_empty_corpus_rejected(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + f = tmp_path / "empty.jsonl" + f.write_text("\n \n", encoding="utf-8") + with pytest.raises(ValueError, match="no usable"): + _load_cov_corpus("empty.jsonl") + + def test_jsonl_dict_without_usable_field_skipped(self, tmp_path, monkeypatch): + """A JSONL object with no text/prompt/content field is silently dropped + (NOT appended as raw JSON) — review L5.""" + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + f = tmp_path / "c.jsonl" + f.write_text( + json.dumps({"other": "x"}) + "\n" + json.dumps({"text": "y"}) + "\n", + encoding="utf-8", + ) + assert _load_cov_corpus("c.jsonl") == ["y"] + + def test_oversize_rejected(self, tmp_path, monkeypatch): + """File larger than the byte cap is rejected before any read — L4.""" + import soup_cli.commands.edit as edit_mod + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(edit_mod, "_MAX_COV_CORPUS_BYTES", 8) + f = tmp_path / "big.txt" + f.write_text("this is definitely more than eight bytes\n", encoding="utf-8") + with pytest.raises(ValueError, match="too large"): + _load_cov_corpus("big.txt") + + def test_line_cap_truncates(self, tmp_path, monkeypatch): + """Reading stops at the line cap rather than consuming the whole file — L4.""" + import soup_cli.commands.edit as edit_mod + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(edit_mod, "_MAX_COV_CORPUS_LINES", 2) + f = tmp_path / "many.txt" + f.write_text("a\nb\nc\nd\n", encoding="utf-8") + assert _load_cov_corpus("many.txt") == ["a", "b"] + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") + def test_symlink_rejected(self, tmp_path, monkeypatch): + from soup_cli.commands.edit import _load_cov_corpus + + monkeypatch.chdir(tmp_path) + real = tmp_path / "real.jsonl" + real.write_text(json.dumps({"text": "x"}) + "\n", encoding="utf-8") + os.symlink(real, tmp_path / "link.jsonl") + with pytest.raises(ValueError, match="symlink"): + _load_cov_corpus("link.jsonl") + + +# =========================================================================== +# #147 — Mixtral in the LongLoRA allowlist +# =========================================================================== + + +class TestIsMixtralModel: + def test_basic(self): + from soup_cli.utils.longlora import is_mixtral_model + + assert is_mixtral_model("mistralai/Mixtral-8x7B-v0.1") is True + assert is_mixtral_model("mistralai/Mixtral-8x22B-Instruct-v0.1") is True + + def test_bare_token(self): + """The start-of-string anchor matches a lone ``mixtral-...`` id (M1).""" + from soup_cli.utils.longlora import is_mixtral_model + + assert is_mixtral_model("Mixtral-8x7B-v0.1") is True + + def test_not_plain_mistral(self): + from soup_cli.utils.longlora import is_mixtral_model + + assert is_mixtral_model("mistralai/Mistral-7B-v0.1") is False + + def test_word_boundary(self): + from soup_cli.utils.longlora import is_mixtral_model + + assert is_mixtral_model("my-mixtralish-finetune") is False + assert is_mixtral_model("unmixtral-7b") is False + + def test_input_guards(self): + from soup_cli.utils.longlora import is_mixtral_model + + assert is_mixtral_model("") is False + with pytest.raises(TypeError): + is_mixtral_model(None) # type: ignore[arg-type] + with pytest.raises(ValueError): + is_mixtral_model("a\x00b") + assert is_mixtral_model("a" * 1024) is False + + +class TestMixtralStillNotMistral: + def test_is_mistral_model_excludes_mixtral(self): + from soup_cli.utils.longlora import is_mistral_model + + # Regression — is_mistral_model stays narrow; Mixtral is detected by + # the dedicated is_mixtral_model helper. + assert is_mistral_model("mistralai/Mixtral-8x7B-v0.1") is False + + +class TestMixtralInAllowlist: + def test_supported(self): + from soup_cli.utils.longlora import is_supported_longlora_arch + + assert is_supported_longlora_arch("mistralai/Mixtral-8x7B-v0.1") is True + assert ( + is_supported_longlora_arch("mistralai/Mixtral-8x22B-Instruct-v0.1") + is True + ) + + def test_unsupported_unchanged(self): + from soup_cli.utils.longlora import is_supported_longlora_arch + + assert is_supported_longlora_arch("google/gemma-2-9b") is False + assert is_supported_longlora_arch("databricks/dbrx-base") is False + + def test_separate_qkv_families_includes_mixtral(self): + from soup_cli.utils.longlora import _SEPARATE_QKV_FAMILIES + + assert "Mixtral" in _SEPARATE_QKV_FAMILIES + + def test_defensive_surface(self): + """The rewired ``or``-chain must still swallow non-str / null-byte input + (returns False, never raises) — review M2.""" + from soup_cli.utils.longlora import is_supported_longlora_arch + + assert is_supported_longlora_arch(None) is False + assert is_supported_longlora_arch(123) is False + assert is_supported_longlora_arch("a\x00b") is False + + +class TestValidateLongloraCompatMixtral: + def test_accepts_mixtral(self, monkeypatch): + from soup_cli.utils import longlora + + monkeypatch.setattr( + "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False + ) + # Should not raise. + longlora.validate_longlora_compat( + model_name="mistralai/Mixtral-8x7B-v0.1", + task="sft", + backend="transformers", + use_ring_attention=False, + ) + + def test_error_message_lists_mixtral(self, monkeypatch): + from soup_cli.utils import longlora + + monkeypatch.setattr( + "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False + ) + with pytest.raises(ValueError) as exc: + longlora.validate_longlora_compat( + model_name="google/gemma-2-9b", + task="sft", + backend="transformers", + use_ring_attention=False, + ) + assert "Mixtral" in str(exc.value) + + def test_mixtral_ring_attention_rejected(self, monkeypatch): + """Now that Mixtral passes the arch gate, the downstream ring-attention + exclusivity becomes reachable for it — confirm it still fires (M6).""" + from soup_cli.utils import longlora + + monkeypatch.setattr( + "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False + ) + with pytest.raises(ValueError, match="ring"): + longlora.validate_longlora_compat( + model_name="mistralai/Mixtral-8x7B-v0.1", + task="sft", + backend="transformers", + use_ring_attention=True, + ) + + +class TestMixtralForwardOverride: + def _fake_attn_model(self, cls_name): + # The override matches on the attention module's CLASS NAME, so the + # instance's class must literally be named e.g. ``MixtralAttention``. + def _init(self): + nn.Module.__init__(self) + self.head_dim = 4 + self.num_heads = 2 + self.num_key_value_heads = 2 + self.q_proj = nn.Linear(8, 8, bias=False) + self.k_proj = nn.Linear(8, 8, bias=False) + self.v_proj = nn.Linear(8, 8, bias=False) + + attn_cls = type(cls_name, (nn.Module,), {"__init__": _init}) + model = nn.Module() + model.attn = attn_cls() # registers as a submodule + return model + + def test_patches_mixtral_attention(self): + from soup_cli.utils.longlora import LongLoRAForwardOverride + + model = self._fake_attn_model("MixtralAttention") + with LongLoRAForwardOverride(model, group_size=4): + assert getattr( + model.attn.q_proj.forward, "_soup_longlora_patched", False + ) + assert getattr( + model.attn.k_proj.forward, "_soup_longlora_patched", False + ) + # Restored on exit. + assert not getattr( + model.attn.q_proj.forward, "_soup_longlora_patched", False + ) + + def test_llama_attention_still_patched(self): + from soup_cli.utils.longlora import LongLoRAForwardOverride + + model = self._fake_attn_model("LlamaAttention") + with LongLoRAForwardOverride(model, group_size=4): + assert getattr( + model.attn.q_proj.forward, "_soup_longlora_patched", False + ) + + +class TestMixtralSchemaGate: + def test_schema_accepts_mixtral(self, monkeypatch): + from soup_cli.config.loader import load_config_from_string + + monkeypatch.setattr( + "soup_cli.utils.flash_attn.is_flash_attn_v3_available", lambda: False + ) + cfg = load_config_from_string( + "base: mistralai/Mixtral-8x7B-v0.1\n" + "task: sft\n" + "data:\n" + " train: data.jsonl\n" + "training:\n" + " use_longlora: true\n" + ) + assert cfg.training.use_longlora is True + + +# =========================================================================== +# CLI plumbing +# =========================================================================== + + +class TestCliCovCorpus: + def test_edit_set_help_has_cov_corpus(self): + import re + + from typer.testing import CliRunner + + from soup_cli.cli import app + + result = CliRunner().invoke(app, ["edit", "set", "--help"]) + assert result.exit_code == 0, result.output + clean = re.sub(r"\x1b\[[0-9;]*m", "", result.output) + assert "cov-corpus" in clean + + +# =========================================================================== +# Patch invariants +# =========================================================================== + + +class TestPatchInvariants: + def test_version_bumped(self): + import soup_cli + + parts = soup_cli.__version__.split(".") + assert (int(parts[0]), int(parts[1]), int(parts[2])) >= (0, 71, 16) + + @pytest.mark.parametrize( + "module", + [ + "soup_cli.utils.edit_kernels", + "soup_cli.utils.edit_governor", + "soup_cli.utils.longlora", + "soup_cli.commands.edit", + ], + ) + def test_no_top_level_torch(self, module): + import importlib + + mod = importlib.import_module(module) + with open(mod.__file__, encoding="utf-8") as fh: + src = fh.read() + for line in src.splitlines(): + assert not line.startswith("import torch"), module + assert not line.startswith("from torch"), module + + def test_edit_kernels_no_top_level_safetensors(self): + import importlib + + mod = importlib.import_module("soup_cli.utils.edit_kernels") + with open(mod.__file__, encoding="utf-8") as fh: + src = fh.read() + assert "estimate_key_covariance" in src