fix: apply BitNet's ffn_sub_norm on the processed-weights MLP path

tests/integration/model_bridge/test_bitnet_compatibility_mode.py

@@ -0,0 +1,93 @@
+"""BitNet in compatibility mode, on a tiny random model built without hub access.
+
+BitNet applies an RMSNorm (attn_sub_norm / ffn_sub_norm) right before each output
+projection. The processed-weights path rebuilds attention and the MLP from weights,
+so it has to apply both norms to match HF.
+"""
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_config_from_hf,
+)
+
+TOKENS = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]])
+
+
+def _hf_bitnet():
+    from transformers import AutoModelForCausalLM, BitNetConfig
+
+    torch.manual_seed(0)
+    c = BitNetConfig(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        max_position_embeddings=64,
+    )
+    c.architectures = ["BitNetForCausalLM"]
+    hf = AutoModelForCausalLM.from_config(c, attn_implementation="eager").eval()
+    with torch.no_grad():
+        # Non-trivial sub-norm weights, so skipping a norm changes the output.
+        for name, param in hf.named_parameters():
+            if "norm" in name:
+                param.normal_(1.0, 0.3)
+    return hf
+
+
+def _bridge(hf) -> TransformerBridge:
+    from transformer_lens.factories.architecture_adapter_factory import (
+        ArchitectureAdapterFactory,
+    )
+
+    cfg = build_bridge_config_from_hf(hf.config, "BitNetForCausalLM", "tiny", torch.float32)
+    adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+    return TransformerBridge(model=hf, adapter=adapter, tokenizer=MagicMock())
+
+
+@pytest.fixture(scope="module")
+def reference():
+    hf = _hf_bitnet()
+    with torch.no_grad():
+        logits = hf(TOKENS).logits
+    _, cache = _bridge(hf).run_with_cache(TOKENS)
+    return logits, cache
+
+
+@pytest.fixture(scope="module")
+def compat():
+    bridge = _bridge(_hf_bitnet())
+    bridge.enable_compatibility_mode()
+    with torch.no_grad():
+        logits, cache = bridge.run_with_cache(TOKENS)
+    return bridge, logits, cache
+
+
+def test_compatibility_mode_matches_hf(reference, compat):
+    ref_logits, _ = reference
+    _, logits, _ = compat
+    torch.testing.assert_close(
+        logits.log_softmax(-1), ref_logits.log_softmax(-1), atol=1e-5, rtol=0
+    )
+
+
+@pytest.mark.parametrize(
+    "hook", ["blocks.0.mlp.hook_post", "blocks.0.hook_mlp_out", "blocks.1.hook_resid_post"]
+)
+def test_mlp_hooks_match_the_unprocessed_bridge(reference, compat, hook):
+    """hook_post is the down_proj input, after ffn_sub_norm, on both paths."""
+    _, ref_cache = reference
+    _, _, cache = compat
+    torch.testing.assert_close(cache[hook], ref_cache[hook], atol=1e-5, rtol=1e-5)
+
+
+def test_neuron_decomposition_sums_to_residual(compat):
+    bridge, _, cache = compat
+    final = cache[f"blocks.{bridge.cfg.n_layers - 1}.hook_resid_post"]
+    stack = cache.get_full_resid_decomposition(expand_neurons=True)
+    torch.testing.assert_close(stack.sum(0), final, atol=1e-5, rtol=0)

tests/unit/model_bridge/supported_architectures/test_bitnet_adapter.py

@@ -17,6 +17,7 @@
 from transformer_lens.model_bridge.supported_architectures.bitnet import (
     BitNetArchitectureAdapter,
     _BitNetAttentionBridge,
+    _BitNetMLPBridge,
 )
 
 
@@ -43,11 +44,18 @@ def adapter() -> BitNetArchitectureAdapter:
 class TestBitNetComponentMapping:
     def test_attention_uses_sub_norm_bridge(self, adapter):
         """BitNet's distinguishing feature: RMSNorms before both output
-        projections, applied via the adapter-local attention bridge and the
-        delegated HF MLP."""
+        projections, applied via the adapter-local attention and MLP bridges."""
         submodules = adapter.component_mapping["blocks"].submodules
         assert isinstance(submodules["attn"], _BitNetAttentionBridge)
+        assert isinstance(submodules["mlp"], _BitNetMLPBridge)
         assert isinstance(submodules["mlp"], GatedMLPBridge)
+        mlp = submodules["mlp"]
+        assert mlp.name == "mlp"
+        assert {k: v.name for k, v in mlp.submodules.items()} == {
+            "gate": "gate_proj",
+            "in": "up_proj",
+            "out": "down_proj",
+        }
 
     def test_sub_norms_disable_folding(self, adapter):
         assert adapter.supports_fold_ln is False
@@ -73,6 +81,28 @@ def __init__(self):
         out = attn._pre_output_projection(x)
         assert torch.equal(out, x * 2.0)
 
+    def test_mlp_reconstruction_applies_sub_norm(self, adapter):
+        """The generic processed-weights MLP goes straight from act(gate) * up to
+        down_proj; BitNet's ffn_sub_norm must be applied in between."""
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+
+        class _FakeNorm(torch.nn.Module):
+            def forward(self, x):
+                return x * 2.0
+
+        class _FakeMLP(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.ffn_sub_norm = _FakeNorm()
+
+        mlp._modules["_original_component"] = _FakeMLP()
+        x = torch.ones(1, 2, 4)
+        assert torch.equal(mlp._pre_output_projection(x), x * 2.0)
+
+    def test_generic_gated_mlp_pre_output_projection_is_identity(self):
+        x = torch.randn(1, 2, 4)
+        assert GatedMLPBridge(name="mlp")._pre_output_projection(x) is x
+
 
 class TestBitNetPackedCheckpointGuard:
     """prepare_model must refuse packed 1.58-bit checkpoints: weight-space reads

transformer_lens/model_bridge/generalized_components/gated_mlp.py

@@ -120,7 +120,7 @@ def forward(self, *args, **kwargs) -> torch.Tensor:
                 linear_output = in_module.hook_out(linear_output)  # type: ignore[misc]
             act_fn = resolve_activation_fn(self.config)
             activated = act_fn(gate_output)
-            hidden = activated * linear_output
+            hidden = self._pre_output_projection(activated * linear_output)
             if hasattr(self, "out") and hasattr(self.out, "hook_in"):
                 hidden = self.out.hook_in(hidden)
             output = torch.nn.functional.linear(hidden, W_out, b_out)
@@ -141,6 +141,15 @@ def forward(self, *args, **kwargs) -> torch.Tensor:
         output = self.hook_out(output)
         return output
 
+    def _pre_output_projection(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Transform the gated activation before ``out`` on the processed-weights path.
+
+        Identity by default. Override when the HF MLP applies more between the gating
+        product and ``down_proj`` (BitNet's ``ffn_sub_norm``); ``hook_post`` sees the result,
+        as it does on the HF path, where it is ``out.hook_in``.
+        """
+        return hidden
+
     def _processed_projection(self, name: str) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Current weight and bias of a projection, still attached to autograd so they train."""
         projection = getattr(self, name)

transformer_lens/model_bridge/supported_architectures/bitnet.py

@@ -1,11 +1,13 @@
 """BitNet b1.58 (``BitNetForCausalLM``) adapter: llama layout plus attn/ffn
-sub-layer RMSNorms (attn_sub_norm reapplied by an adapter-local attention bridge)."""
+sub-layer RMSNorms (attn_sub_norm and ffn_sub_norm reapplied by adapter-local
+attention and MLP bridges on the processed-weights path)."""
 
 from typing import Any
 
 import torch
 
 from transformer_lens.model_bridge.generalized_components import (
+    GatedMLPBridge,
     PositionEmbeddingsAttentionBridge,
 )
 from transformer_lens.model_bridge.supported_architectures.llama import (
@@ -29,6 +31,21 @@ def _pre_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor:
         return attn_output
 
 
+class _BitNetMLPBridge(GatedMLPBridge):
+    """Applies BitNet's ffn_sub_norm before the down projection.
+
+    The generic processed-weights path goes straight from ``act(gate) * up`` to
+    down_proj; BitNet inserts an RMSNorm in between.
+    """
+
+    def _pre_output_projection(self, hidden: torch.Tensor) -> torch.Tensor:
+        oc = self.original_component
+        sub_norm = getattr(oc, "ffn_sub_norm", None) if oc is not None else None
+        if isinstance(sub_norm, torch.nn.Module):
+            hidden = sub_norm(hidden)
+        return hidden
+
+
 class BitNetArchitectureAdapter(LlamaArchitectureAdapter):
     """Architecture adapter for BitNetForCausalLM models."""
 
@@ -47,6 +64,11 @@ def __init__(self, cfg: Any) -> None:
         # standard LN folding and W_O centering do not model them.
         self.supports_fold_ln = False
         self.supports_center_writing_weights = False
+        blocks = self.components["blocks"]
+        mlp = blocks.submodules["mlp"]
+        blocks.submodules["mlp"] = _BitNetMLPBridge(
+            name=mlp.name, config=self.cfg, submodules=mlp.submodules
+        )
 
     def prepare_model(self, hf_model: Any) -> None:
         """Refuse packed 1.58-bit checkpoints, which need BitNet dequant kernels.