Generalize dense-layer MLP hooks across interleaved MoE architectures

.github/workflows/checks.yml

@@ -97,10 +97,19 @@ jobs:
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
             ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
             ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
-          key: ${{ runner.os }}-huggingface-models-v5
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-v6
           restore-keys: |
+            ${{ runner.os }}-huggingface-models-v5
             ${{ runner.os }}-huggingface-models-v4
             ${{ runner.os }}-huggingface-models-v2
             ${{ runner.os }}-huggingface-models-v1
@@ -143,7 +152,22 @@ jobs:
         with:
           path: |
             ~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M*
-          key: ${{ runner.os }}-huggingface-models-mps-v1
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-mps-v2
+          restore-keys: |
+            ${{ runner.os }}-huggingface-models-mps-v1
       - name: Install dependencies
         run: |
           uv lock --check
@@ -297,10 +321,19 @@ jobs:
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
             ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
             ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
             ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
-          key: ${{ runner.os }}-huggingface-models-v5
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-v6
           restore-keys: |
+            ${{ runner.os }}-huggingface-models-v5
             ${{ runner.os }}-huggingface-models-v4
             ${{ runner.os }}-huggingface-models-v2
             ${{ runner.os }}-huggingface-models-v1
@@ -350,12 +383,10 @@ jobs:
         with:
           path: |
             ~/.cache/huggingface/hub/models--gpt2
-          key: ${{ runner.os }}-huggingface-models-v5
-          restore-keys: |
-            ${{ runner.os }}-huggingface-models-v4
-            ${{ runner.os }}-huggingface-models-v2
-            ${{ runner.os }}-huggingface-models-v1
-            ${{ runner.os }}-huggingface-models
+          # Distinct key: this job needs only gpt2, and sharing the full-suite
+          # key let whichever job saved first pin the cache contents for all
+          # three (a one-path save then serves the suites that need the rest).
+          key: ${{ runner.os }}-huggingface-models-benchmark-v1
       - name: Install dependencies
         run: |
           uv lock --check

docs/source/content/adapter_development/adapter-specification.md

@@ -276,6 +276,32 @@ Standard positional embeddings + LayerNorm + standard MLP + combined QKV. Used b
 
 Similar to Llama-like but with `MoEBridge` replacing the MLP. Used by: Mixtral, GraniteMoE, OLMoE.
 
+**Interleaved / dense-prefix stacks.** When the config can make some layers a plain
+gated MLP under the same attribute name (`first_k_dense_replace`,
+`mlp_only_layers`/`decoder_sparse_step`, `moe_layer_start_index`,
+`num_dense_layers`, `mlp_layer_types`), also declare the dense projections as
+optional submodules named `dense_gate` / `dense_in` / `dense_out`:
+
+```python
+"mlp": MoEBridge(
+    name="mlp",
+    config=self.cfg,
+    submodules={
+        "gate": MoERouterBridge(name="gate", optional=True),
+        "dense_gate": LinearBridge(name="gate_proj", optional=True),
+        "dense_in": LinearBridge(name="up_proj", optional=True),
+        "dense_out": LinearBridge(name="down_proj", optional=True),
+    },
+),
+```
+
+`MoEBridge` detects that binding per layer and gives the dense layers the
+gated-MLP hook set (`hook_pre`, `hook_pre_linear`, `hook_post`) and weight
+accessors (`W_gate`/`W_in`/`W_out`), while sparse layers keep MoE semantics.
+Without the declaration those layers report `d_model` boundary tensors under
+neuron-hook names (#1645). The keys are `dense_*` rather than `gate`/`in`/`out`
+because `gate` already means the *router* on the sparse layers of the same model.
+
 ### Pattern 4: Multimodal
 
 Extends a text-only pattern with vision encoder and projection bridges. Used by: LLaVA, LLaVA-Next, Gemma3 Multimodal.

tests/integration/model_bridge/test_bridge_vs_hf_eager_parity.py

@@ -5,21 +5,17 @@
 SDPA, which reorders ops in a fused kernel. Bridge vs HF *eager* matches to fp32-noise.
 """
 
-import platform
 from typing import Callable
 
 import pytest
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from tests.tiny_checkpoints import FP32_NOISE_TOL, assert_tiny_parity, parity_params
 from transformer_lens.model_bridge import TransformerBridge
 
 MODEL_NAME = "EleutherAI/pythia-70m"
 
-# Wider fp32 op-order noise floor on GH Actions macOS-arm64; ~3e-3 at output.
-_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64"
-FP32_NOISE_TOL = 1e-2 if _MACOS_ARM64 else 1e-5
-
 
 @pytest.fixture(scope="module")
 def tokenizer():
@@ -123,3 +119,28 @@ def test_bridge_attention_reconstruction_actually_runs(bridge, tokenize):
         "blocks.0.attn.hook_attn_scores did not fire — bridge no longer runs its "
         "own attention reconstruction, making the parity tests tautological."
     )
+
+
+@pytest.mark.parametrize("model_name", parity_params())
+def test_tiny_bridge_logits_match_hf_eager(model_name: str) -> None:
+    """The hooked forward must reproduce an independent HF eager load's logits."""
+    bridge = TransformerBridge.boot_transformers(model_name, device="cpu", dtype=torch.float32)
+    tokens = bridge.to_tokens("The quick brown fox jumps")
+    hf_eager = AutoModelForCausalLM.from_pretrained(
+        model_name, torch_dtype=torch.float32, attn_implementation="eager"
+    ).eval()
+
+    with torch.inference_mode():
+        bridge_logits = bridge(tokens)
+        hf_logits = hf_eager(tokens).logits
+    assert_tiny_parity(bridge_logits, hf_logits, model_name)
+
+    # Non-tautology guard: on archs with a reconstructed attention path, its
+    # scores hook must fire, or the parity above proves nothing.
+    if "blocks.0.attn.hook_attn_scores" in bridge.hook_dict:
+        fired: list[bool] = []
+        bridge.run_with_hooks(
+            tokens,
+            fwd_hooks=[("blocks.0.attn.hook_attn_scores", lambda v, hook: fired.append(True))],
+        )
+        assert fired, f"{model_name!r}: hook_attn_scores did not fire — reconstruction bypassed."

tests/integration/model_bridge/test_deepseek_adapter.py

@@ -79,12 +79,18 @@ def test_forward_matches_hf(self, tiny_deepseek_bridge):
 
 
 class TestDeepSeekDenseVsMoELayers:
-    def test_dense_layer_has_no_moe_hooks(self, tiny_deepseek_bridge):
+    def test_dense_layer_binds_gated_mlp_hooks(self, tiny_deepseek_bridge):
+        """Dense-prefix layers expose the gate PROJECTION under dense_gate
+        (d_mlp neuron basis) — `gate` stays the router name, so one hook name
+        never means two things across layers (#1645)."""
         tokens = torch.tensor([[1, 2, 3, 4]])
         _, cache = tiny_deepseek_bridge.run_with_cache(tokens)
-        cache_keys = set(cache.keys())
-        assert not any("blocks.0.mlp.gate" in k for k in cache_keys)
-        assert not any("blocks.0.mlp.shared_experts" in k for k in cache_keys)
+        d_mlp = tiny_deepseek_bridge.original_model.config.intermediate_size
+        assert cache["blocks.0.mlp.dense_gate.hook_out"].shape[-1] == d_mlp
+        torch.testing.assert_close(
+            cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"]
+        )
+        assert not any("blocks.0.mlp.shared_experts" in k for k in cache.keys())
 
     def test_moe_layer_has_gate_hooks(self, tiny_deepseek_bridge):
         tokens = torch.tensor([[1, 2, 3, 4]])

tests/integration/model_bridge/test_deepseek_v2_adapter.py

@@ -113,9 +113,16 @@ def test_forward_matches_hf(self, tiny_deepseek_v2_bridge):
 
 
 class TestDeepSeekV2DenseVsMoELayers:
-    def test_dense_layer_has_no_moe_hooks(self, tiny_deepseek_v2_bridge):
+    def test_dense_layer_binds_gated_mlp_hooks(self, tiny_deepseek_v2_bridge):
+        """Dense-prefix layers expose the gate PROJECTION under dense_gate
+        (d_mlp neuron basis) — `gate` stays the router name, so one hook name
+        never means two things across layers (#1645)."""
         _, cache = tiny_deepseek_v2_bridge.run_with_cache(_tokens())
-        assert not any("blocks.0.mlp.gate" in k for k in cache)
+        d_mlp = tiny_deepseek_v2_bridge.original_model.config.intermediate_size
+        assert cache["blocks.0.mlp.dense_gate.hook_out"].shape[-1] == d_mlp
+        torch.testing.assert_close(
+            cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"]
+        )
         assert not any("blocks.0.mlp.shared_experts" in k for k in cache)
 
     def test_moe_layer_has_router_and_shared_expert_hooks(self, tiny_deepseek_v2_bridge):

tests/integration/model_bridge/test_glm4_moe_lite_adapter.py

@@ -28,12 +28,22 @@ def test_adapter_selected(self, glm_bridge):
         assert isinstance(glm_bridge.adapter, Glm4MoeLiteArchitectureAdapter)
 
     def test_dense_sparse_layer_mix(self, glm_bridge):
-        """The tiny checkpoint declares layer 0 dense, layer 1 sparse."""
+        """The tiny checkpoint declares layer 0 dense, layer 1 sparse; the mlp
+        template binds each accordingly (#1645)."""
         hf_model = glm_bridge.original_model
         assert hf_model.config.mlp_layer_types == ["dense", "sparse"]
-        assert not hasattr(hf_model.model.layers[0].mlp, "gate")
-        assert hasattr(hf_model.model.layers[1].mlp, "gate")
-        assert hasattr(hf_model.model.layers[1].mlp, "shared_experts")
+        dense_mlp = glm_bridge.blocks[0].mlp
+        sparse_mlp = glm_bridge.blocks[1].mlp
+        # Dense binding: gated-MLP neuron hooks via the dense_* projections.
+        assert not hasattr(dense_mlp.original_component, "experts")
+        assert dense_mlp._bound_dense is True
+        assert dense_mlp.hook_aliases["hook_pre"] == "dense_gate.hook_out"
+        assert not hasattr(dense_mlp, "hook_router_scores")
+        # Sparse binding: router + shared experts survive setup on the BRIDGE
+        # (a silently-popped optional router would otherwise stay green here).
+        assert hasattr(sparse_mlp.original_component, "experts")
+        assert {"gate", "shared_experts"} <= set(sparse_mlp.submodules)
+        assert hasattr(sparse_mlp, "hook_router_scores")
 
 
 class TestGlm4MoeLiteForwardEquivalence:

tests/integration/model_bridge/test_moe_dense_dispatch.py

@@ -0,0 +1,113 @@
+"""Dense-prefix MoE dispatch on a real mixed checkpoint (#1645).
+
+katuni4ka/tiny-random-deepseek-v3 has first_k_dense_replace=1 over two layers:
+layer 0 is a plain gated MLP, layer 1 a sparse MoE block — one boot exercises
+both bindings of the DenseOrMoEBridge template.
+"""
+
+import torch
+from transformers import AutoModelForCausalLM
+
+from tests.tiny_checkpoints import MIXED_DENSE_SPARSE_MOE, assert_tiny_parity
+from transformer_lens.model_bridge import TransformerBridge
+
+MODEL_NAME = MIXED_DENSE_SPARSE_MOE
+
+
+def test_mixed_checkpoint_dense_and_sparse_mlp_hooks() -> None:
+    bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32)
+    tokens = bridge.to_tokens("The quick brown fox jumps")
+    d_model = bridge.cfg.d_model
+    d_mlp = bridge.original_model.config.intermediate_size
+
+    with torch.no_grad():
+        logits, cache = bridge.run_with_cache(tokens)
+
+    # Dense layer 0: neuron-basis hooks with gated-MLP semantics.
+    assert cache["blocks.0.mlp.hook_pre"].shape[-1] == d_mlp
+    torch.testing.assert_close(
+        cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"]
+    )
+    torch.testing.assert_close(
+        cache["blocks.0.mlp.hook_pre_linear"], cache["blocks.0.mlp.dense_in.hook_out"]
+    )
+    torch.testing.assert_close(
+        cache["blocks.0.mlp.hook_post"], cache["blocks.0.mlp.dense_out.hook_in"]
+    )
+
+    # Sparse layer 1: MoE boundary semantics unchanged.
+    assert cache["blocks.1.mlp.hook_pre"].shape[-1] == d_model
+    torch.testing.assert_close(cache["blocks.1.mlp.hook_pre"], cache["blocks.1.mlp.hook_in"])
+
+    # The dispatch must not perturb the forward itself.
+    hf_eager = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, torch_dtype=torch.float32, attn_implementation="eager"
+    ).eval()
+    with torch.no_grad():
+        hf_logits = hf_eager(tokens).logits
+    assert_tiny_parity(logits, hf_logits, MODEL_NAME)
+
+
+def test_dense_layer_neuron_intervention_targets_gate_output() -> None:
+    """A write on the dense layer's hook_pre must be a write on the gate
+    projection's output — the #1645 complaint was precisely that no
+    neuron-basis intervention point existed on dense layers. (Pre-fix,
+    hook_pre aliased the MLP *input*, so zeroing it also moved logits —
+    equivalence with the gate target is what pins the semantics.)"""
+    bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32)
+    tokens = bridge.to_tokens("The quick brown fox jumps")
+
+    def zero(tensor: torch.Tensor, hook) -> torch.Tensor:
+        return tensor * 0.0
+
+    with torch.no_grad():
+        base_logits = bridge(tokens)
+        via_alias = bridge.run_with_hooks(tokens, fwd_hooks=[("blocks.0.mlp.hook_pre", zero)])
+        via_target = bridge.run_with_hooks(
+            tokens, fwd_hooks=[("blocks.0.mlp.dense_gate.hook_out", zero)]
+        )
+    torch.testing.assert_close(via_alias, via_target)
+    assert not torch.allclose(base_logits, via_alias)
+
+
+def test_dense_layer_weights_reach_get_params() -> None:
+    """The neuron hooks advertise a neuron basis; the weight API must back it.
+    get_params() keyed on submodules named `in`/`out`, so dense MoE layers were
+    silently zero-filled — real weights on the component, zeros in the dict."""
+    bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32)
+    params = bridge.get_params()
+
+    dense = bridge.blocks[0].mlp
+    assert dense.bound_dense is True
+    # Anchor on the wrapped HF module, NOT on dense.W_in/W_gate/W_out:
+    # get_params_util builds these entries by reading those very properties, so
+    # comparing the two would be the same expression evaluated twice and could
+    # not detect a projection being read from the wrong submodule.
+    hf_mlp = dense.original_component
+    torch.testing.assert_close(params["blocks.0.mlp.W_in"], hf_mlp.up_proj.weight.T)
+    torch.testing.assert_close(params["blocks.0.mlp.W_gate"], hf_mlp.gate_proj.weight.T)
+    torch.testing.assert_close(params["blocks.0.mlp.W_out"], hf_mlp.down_proj.weight.T)
+    assert (params["blocks.0.mlp.W_in"] != 0).any()
+    # Negative control: gate and up are same-shaped, so this checkpoint can tell
+    # them apart only if their weights actually differ.
+    assert not torch.equal(hf_mlp.gate_proj.weight, hf_mlp.up_proj.weight)
+
+    # Sparse layers legitimately have no single W_* and keep the placeholder.
+    assert bridge.blocks[1].mlp.bound_dense is False
+    assert (params["blocks.1.mlp.W_in"] == 0).all()
+
+
+def test_hook_dict_follows_a_dense_to_sparse_rebind() -> None:
+    """The alias cache must not outlive the binding it described: a rebind can
+    leave the hook registry exactly the same size while changing what the
+    aliases point at."""
+    bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32)
+    _ = bridge.hook_dict  # populate the cache while layer 0 is dense
+    assert bridge._collect_block_instance_aliases()["blocks.0.mlp.hook_pre"].endswith(
+        "dense_gate.hook_out"
+    )
+
+    bridge.blocks[0].mlp.set_original_component(bridge.blocks[1].mlp.original_component)
+
+    assert bridge.blocks[0].mlp.bound_dense is False
+    assert bridge._collect_block_instance_aliases()["blocks.0.mlp.hook_pre"].endswith("mlp.hook_in")

tests/integration/model_bridge/test_residual_decomposition_identities.py

@@ -16,6 +16,12 @@
 - olmo2: post-norm inside the residual branch — RMSNorm applies to the sublayer
   output before the add, so the contributions are the norm outputs (the #1648 case)
 - mpt: residual added inside the HF MLP module only (attention adds at block level)
+- granite: the scaled-residual block wiring end to end. NOTE: the hub tiny ships
+  residual_multiplier=1.0, so this entry does NOT exercise the multiplier —
+  test_granite_hook_semantics.py builds a fixture with a real multiplier and is
+  the coverage for the #1648 Granite case.
+- gemma2: sandwich norms — a post-sublayer norm inside each residual branch, so
+  the contributions are the post-norm outputs
 
 Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope
 here: they have no ``hook_resid_mid``. Their two-term identity
@@ -26,15 +32,10 @@
 import pytest
 import torch
 
+from tests.tiny_checkpoints import sequential_residual_params
 from transformer_lens.model_bridge import TransformerBridge
 
-SEQUENTIAL_RESIDUAL_MODELS = [
-    pytest.param("hf-internal-testing/tiny-random-gpt2", id="gpt2"),
-    pytest.param("trl-internal-testing/tiny-MistralForCausalLM-0.2", id="mistral"),
-    pytest.param("trl-internal-testing/tiny-BloomForCausalLM", id="bloom"),
-    pytest.param("hf-internal-testing/tiny-random-Olmo2ForCausalLM", id="olmo2"),
-    pytest.param("hf-internal-testing/tiny-random-MptForCausalLM", id="mpt"),
-]
+SEQUENTIAL_RESIDUAL_MODELS = sequential_residual_params()
 
 
 @pytest.mark.parametrize("model_name", SEQUENTIAL_RESIDUAL_MODELS)