Add from scratch pretrain converter

README.md

@@ -184,12 +184,12 @@ cover:
 
 Verification lives in the integration tests at
 `tests/integration/model_bridge/test_mamba_adapter.py` and
-`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total). The
-`verify_models` benchmark suite does not currently cover SSM architectures — the
-benchmark has transformer-shaped assumptions (hook path patterns, layer norm
-folding, block submodule dispatch) that would need a dedicated refactor. SSM
-benchmark coverage will be revisited if hybrid architectures like Jamba or
-Falcon-H1 become priority or a user explicitly requests it.
+`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total), and the
+`verify_models` benchmark suite now covers the SSM and hybrid families. Mamba-1,
+Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next), NemotronH, and GraniteMoeHybrid
+all declare `applicable_phases = [1, 2, 3, 4]`, so their forward parity (P1, vs raw
+HF), hook/cache coverage (P2/P3, which skip the HookedTransformer comparison SSMs
+lack), and generation quality (P4) are benchmarked like any transformer.
 
 ## Credits
 

docs/source/content/ssm_interpretability.md

@@ -0,0 +1,117 @@
+# SSM / Recurrent-Model Interpretability
+
+TransformerBridge exposes a family-agnostic interpretability surface for state-space
+and linear-attention models — Mamba-1, Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next),
+and the SSM hybrids (NemotronH, GraniteMoeHybrid). The same three `ActivationCache`
+methods work across every family, dispatching to each block's SSM mixer regardless of
+which slot it occupies (`.mixer` for Mamba, `.linear_attn` for gated-delta-net) and
+skipping a hybrid's passthrough (attention / MLP / MoE) layers automatically.
+
+> These features are **bridge-only** — there is no `HookedTransformer` counterpart for
+> SSM models. Load with `TransformerBridge.boot_transformers(...)`.
+
+---
+
+## Discovery
+
+```python
+bridge = TransformerBridge.boot_transformers("state-spaces/mamba-130m-hf")
+_, cache = bridge.run_with_cache(tokens, use_cache=False)
+
+cache.ssm_layers()          # -> [0, 1, 2, ...]  block indices whose mixer is recurrent
+```
+
+`ssm_layers()` is purely structural (no dependence on `cfg.layers_block_type`): a
+hybrid's attention/MLP layers are excluded, so on NemotronH you get only the Mamba
+layers.
+
+## Read-only analysis (no forward re-run)
+
+Both methods reconstruct their quantities post-hoc from cached hooks. They return a
+single tensor when *every* block is an SSM layer, else a `{layer_idx: tensor}` dict
+over the SSM layers.
+
+```python
+# Effective ("hidden") attention  M = L ⊙ (C Bᵀ)   [batch, heads, seq, seq]
+M = cache.compute_ssm_effective_attention()          # all SSM layers
+M0 = cache.compute_ssm_effective_attention(layer=0)  # one layer
+
+# Recurrent state trajectory S_t
+S = cache.compute_ssm_state()                        # all SSM layers
+S5 = cache.compute_ssm_state(layer=5, time_step=-1)  # one layer, final step (memory-bounded)
+```
+
+State shape is family-specific (Mamba-1 is per-channel; Mamba-2 and gated-delta-net
+are per-head), so `compute_ssm_state()` returns a dict except when one mixer type
+covers every block.
+
+> **`use_cache=False` is required for gated-delta-net.** Its interior hooks
+> (`hook_q/k/v`, `hook_beta`, `hook_log_decay`) fire only on the hooked prefill path;
+> the default cached path exposes only `hook_in`/`hook_out`, and the reconstruction
+> methods then raise. Mamba reconstructs from `in_proj`/`conv1d` hooks either way, but
+> passing `use_cache=False` uniformly is safe.
+
+## Canonical hook vocabulary
+
+Every family exposes the same state-mutation names (defined once, in
+`SSMStateHookMixin`), so interp tooling can address the same quantity across families:
+
+| Hook | Meaning | Kind |
+|---|---|---|
+| `hook_ssm_state` | post-scan recurrent-state trajectory `S_t` | real HookPoint (fires on the eager-scan path) |
+| `hook_ssm_write` | per-step write influence | **real** for Mamba-1/2 (`dt·(x⊗B)`); **alias → `hook_beta`** for gated-delta-net (state-dependent write) |
+
+Additional canonical aliases resolve per family where the quantity exists:
+`hook_ssm_out`, `hook_ssm_B`, `hook_ssm_C`, `hook_ssm_decay`, `hook_ssm_dt`.
+
+---
+
+## Causal intervention: the `eager_scan` opt-in
+
+By default the mixer runs HF's fused recurrence kernel, which is opaque — the state
+trajectory is never materialized, so it cannot be patched. Setting `eager_scan = True`
+on a realized SSM mixer swaps the kernel for a readable Python scan that fires
+`hook_ssm_state` (and, for the input-linear families, `hook_ssm_write`), so you can
+read and edit the state mid-recurrence.
+
+```python
+from transformer_lens.model_bridge.generalized_components.ssm_protocol import find_ssm_mixer
+
+# Enable on every realized SSM mixer (skips hybrid passthrough slots).
+mixers = [m for b in bridge.blocks if (m := find_ssm_mixer(b)) is not None]
+for m in mixers:
+    m.eager_scan = True
+
+# hook path uses the mixer's slot name: `.mixer` (Mamba) or `.linear_attn` (gated-delta-net)
+STATE = "blocks.0.mixer.hook_ssm_state"
+
+def zero_state(state, hook):
+    return torch.zeros_like(state)   # ablate the whole recurrent state
+
+logits = bridge.run_with_hooks(tokens, use_cache=False, fwd_hooks=[(STATE, zero_state)])
+
+for m in mixers:
+    m.eager_scan = False             # restore the default (bit-identical) fused path
+```
+
+**Two intervention semantics** (mirroring Mamba-Knockout and state-patching):
+
+- Patching **`hook_ssm_state`** changes only the *same-position* readout
+  (`y_t = C_t·S_t` or `S_t^T q_t`); it does not alter the forward recurrence.
+- Patching **`hook_ssm_write`** re-runs the recurrence, so the edit **propagates** to
+  every later state. For gated-delta-net this is the write-strength gate
+  (`hook_beta`), so zeroing it suppresses all writes.
+
+### Caveats
+
+| | |
+|---|---|
+| **Prefill only** | The eager scan runs when `cache_params is None` (i.e. `use_cache=False`, no generation step). Autoregressive decode falls back to the fused kernel. |
+| **Numerical** | The Python scan matches the fused kernel only to floating-point tolerance (≈1e-6 fp32), never bit-for-bit. |
+| **Cost** | O(seq) Python with an O(batch·seq·…) state tensor — orders of magnitude slower/heavier than the kernel. Use short sequences. |
+| **Default untouched** | With `eager_scan = False` (the default), `run_with_cache` is bit-identical to HF and `hook_ssm_state` does not fire. |
+
+> See [`ssm2_mixer.py`](../../../transformer_lens/model_bridge/generalized_components/ssm2_mixer.py),
+> [`ssm_mixer.py`](../../../transformer_lens/model_bridge/generalized_components/ssm_mixer.py), and
+> [`gated_delta_net.py`](../../../transformer_lens/model_bridge/generalized_components/gated_delta_net.py)
+> for the per-family recurrences and reconstruction identities.

docs/source/index.md

@@ -56,6 +56,7 @@ content/citation
 content/contributing
 content/hook_system
 content/compatibility_mode
+content/ssm_interpretability
 content/debugging_numerical_divergence
 generated/demos/Main_Demo
 generated/demos/Exploratory_Analysis_Demo

tests/QUARANTINES.md

@@ -46,6 +46,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip
 | [`unit/model_bridge/supported_architectures/test_gemma2_adapter.py`:49](unit/model_bridge/supported_architectures/test_gemma2_adapter.py) | "Network/disk fetch of tiny Gemma2" |
 | [`integration/model_bridge/test_bridge_integration.py`:801](integration/model_bridge/test_bridge_integration.py) | "Skip Gemma2 in CI to avoid timeout" |
 | [`acceptance/model_bridge/compatibility/test_hook_completeness.py`:156](acceptance/model_bridge/compatibility/test_hook_completeness.py) | "Gemma2 too large for CI" |
+| [`integration/model_bridge/test_ouro_adapter.py`:35](integration/model_bridge/test_ouro_adapter.py) module-level (+ `slow` marker) | "ByteDance/Ouro-1.4B: 2.8GB download + ~11GB RAM, too large for CI" |
 
 **Un-skip:** locally with `HF_TOKEN` sourced.
 

tests/integration/model_bridge/test_arcee_adapter.py

@@ -0,0 +1,119 @@
+"""Integration tests for the Arcee architecture adapter.
+
+Loads the tiny-random ArceeForCausalLM checkpoint and asserts logit parity with
+HF plus Arcee-specific structure: an ungated ReLU^2 MLP whose post-activation
+neurons are inspectable via ``blocks.{i}.mlp.hook_post``.
+
+Runs at fp32 + eager attention (the adapter forces eager so output_attentions —
+and therefore hook_attn_scores / hook_pattern — work).
+"""
+import pytest
+import torch
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+MODEL = "optimum-intel-internal-testing/tiny-random-ArceeForCausalLM"
+
+
+@pytest.fixture(scope="module")
+def arcee_bridge():
+    # The adapter forces eager attention via cfg.attn_implementation, so we do not
+    # (and cannot) pass attn_implementation to boot_transformers.
+    return TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32)
+
+
+class TestArceeBridgeCreation:
+    def test_block_count(self, arcee_bridge):
+        assert len(arcee_bridge.blocks) == 2
+
+    def test_has_core_components(self, arcee_bridge):
+        assert hasattr(arcee_bridge, "embed")
+        assert hasattr(arcee_bridge, "unembed")
+        assert hasattr(arcee_bridge, "ln_final")
+
+    def test_config_flags(self, arcee_bridge):
+        cfg = arcee_bridge.cfg
+        assert cfg.normalization_type == "RMS"
+        assert cfg.positional_embedding_type == "rotary"
+        assert cfg.gated_mlp is False  # ungated ReLU^2 MLP
+        assert cfg.act_fn == "relu2"
+
+    def test_gqa_config(self, arcee_bridge):
+        """tiny-random config: 4 attention heads, 2 KV heads."""
+        assert arcee_bridge.cfg.n_key_value_heads == 2
+
+
+class TestArceeForwardEquivalence:
+    def test_forward_returns_logits(self, arcee_bridge):
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        with torch.no_grad():
+            output = arcee_bridge(tokens)
+        assert output.shape[0] == 1
+        assert output.shape[1] == 4
+        assert not torch.isnan(output).any()
+
+    def test_forward_matches_hf(self, arcee_bridge):
+        """Bridge delegates to HF native forward — output should match at fp32."""
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        hf_model = arcee_bridge.original_model
+        with torch.no_grad():
+            bridge_out = arcee_bridge(tokens)
+            hf_out = hf_model(tokens).logits
+        max_diff = (bridge_out - hf_out).abs().max().item()
+        assert max_diff < 1e-4, f"Bridge vs HF max diff = {max_diff}"
+
+
+class TestArceeHFDelegation:
+    def test_mlp_delegates_to_hf_relu2_mlp(self, arcee_bridge):
+        """The ungated MLP bridge wraps HF's ArceeMLP, whose activation is HF's
+        ReLUSquaredActivation (ReLU^2) — i.e. the bridge delegates the MLP forward
+        to the real squared-ReLU implementation rather than a gated substitute."""
+        hf_mlp = arcee_bridge.blocks[0].mlp.original_component
+        assert type(hf_mlp).__name__ == "ArceeMLP"
+        assert type(hf_mlp.act_fn).__name__ == "ReLUSquaredActivation"
+
+
+class TestArceeHookShapes:
+    def test_attn_and_mlp_hooks_fire(self, arcee_bridge):
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        _, cache = arcee_bridge.run_with_cache(tokens)
+        for i in range(2):
+            assert f"blocks.{i}.attn.hook_in" in cache
+            assert f"blocks.{i}.attn.hook_out" in cache
+            assert f"blocks.{i}.mlp.hook_in" in cache
+            assert f"blocks.{i}.mlp.hook_out" in cache
+
+    def test_residual_hooks_fire(self, arcee_bridge):
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        _, cache = arcee_bridge.run_with_cache(tokens)
+        for i in range(2):
+            assert f"blocks.{i}.hook_resid_pre" in cache
+            assert f"blocks.{i}.hook_resid_post" in cache
+
+
+class TestArceeReLU2MLP:
+    """Arcee's distinguishing quirk: the post-activation MLP neurons (ReLU^2
+    output, input to down_proj) are exposed via hook_post and are non-negative
+    (ReLU^2 >= 0), which is the sparse-activation structure the issue targets."""
+
+    def test_mlp_post_activation_hook_fires(self, arcee_bridge):
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        _, cache = arcee_bridge.run_with_cache(tokens)
+        for i in range(2):
+            assert f"blocks.{i}.mlp.hook_post" in cache
+
+    def test_mlp_post_activation_shape(self, arcee_bridge):
+        """hook_post has the intermediate (d_mlp) width."""
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        _, cache = arcee_bridge.run_with_cache(tokens)
+        post = cache["blocks.0.mlp.hook_post"]
+        assert post.shape[0] == 1
+        assert post.shape[1] == 4
+        assert post.shape[-1] == arcee_bridge.cfg.d_mlp
+
+    def test_mlp_post_activation_nonnegative(self, arcee_bridge):
+        """ReLU^2 output is always >= 0."""
+        tokens = torch.tensor([[1, 2, 3, 4]])
+        _, cache = arcee_bridge.run_with_cache(tokens)
+        post = cache["blocks.0.mlp.hook_post"]
+        assert (post >= 0).all(), "ReLU^2 activations must be non-negative"