feat(edit): GPT-2 Conv1D edits, covariance ROME, atomic governor, Mixtral LongLoRA (v0.71.16)

Knowledge edit depth — 4-issue patch (validated on real gpt2 + SmolLM2-135M, RTX 3050).

- #251 GPT-2 (transformer.h / mlp.c_proj Conv1D) support in ROME/MEMIT/AlphaEdit:
  transpose-aware _rank1_update + _alphaedit_project + MEMIT dim-check + PEFT unwrap.
- #250 covariance-preconditioned ROME via --cov-corpus (estimate_key_covariance +
  C^-1 k* solve; fail-loud reject for non-rome; cwd-contained O_NOFOLLOW loader).
- #252 atomic EditGovernor increment: save_state merges this run's delta under the
  cross-process lock (baseline-delta, mirrors namespace_pin).
- #147 Mixtral in the LongLoRA allowlist (is_mixtral_model + MixtralAttention
  forward override + _SEPARATE_QKV_FAMILIES).

Tests: 13511 -> 13595 (+81 in tests/test_v07116.py). Full suite green; ruff clean.

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 <jsonl|txt>` 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

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
 ```
 

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 <path>`** — persists the measured energy/CO2 so
-  `soup bom emit --energy <path>` 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 <jsonl|txt>` 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) &middot; [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).
 

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 <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 <jsonl|txt>` 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 <id>` attaches the edited model (or GRACE codebook) into the Registry lineage.
 
 
 ## Activation Steering (`soup steer`)

docs/commands.md

@@ -171,7 +171,7 @@ soup can pack --entry-id <id> --out r.can --attest <statement.json>  Embed in-to
 soup audit-log tail / rotate  Tail / rotate the per-command HIPAA/SOC2 audit log (~/.soup/audit.jsonl)
 soup --no-audit-log <cmd> / SOUP_NO_AUDIT_LOG=1  Opt out of the per-command audit line
 soup eval unlearning <run-id> --benchmark tofu|muse|wmdp  Forget Quality + Model Utility + PrivLeak verdict
-soup edit set --base <m> --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output <dir>] [--device cpu] [--governor/--no-governor] [--registry-id <id>]  Live surgical knowledge edit (--plan-only available)
+soup edit set --base <m> --method rome|memit|alphaedit|grace --subject "..." --target "..." [--output <dir>] [--device cpu] [--governor/--no-governor] [--registry-id <id>] [--cov-corpus <jsonl|txt>]  Live surgical knowledge edit (GPT-2 Conv1D + Llama; --cov-corpus = covariance-preconditioned ROME, rome-only; --plan-only available)
 soup edit diff <before-run> <after-run> --probes p.jsonl [--before-model <m> --after-model <m>]  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

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

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"

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"

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(

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:

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(

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,

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):