Refuse quantized weights in weight-space code paths instead of reading them as matrices

tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py

@@ -275,6 +275,58 @@ def test_adapter_templates_bind_dense_layers_as_gated_mlps(architecture: str) ->
     assert instance.hook_aliases["hook_pre"] == "dense_gate.hook_out"
 
 
+@pytest.mark.parametrize("architecture", DENSE_AWARE_ARCHS)
+def test_dense_keys_read_the_projection_they_name(architecture: str) -> None:
+    """`dense_gate` must be the gate projection, not the up projection.
+
+    Both are [d_model, d_mlp] and both bind without complaint, so a
+    dense_gate/dense_in swap in an adapter survives every shape check, key-set
+    check and the binding guard above — while making `hook_pre` report the
+    wrong tensor under the right name. That is #1645's own confusion one level
+    in, so it is checked here, once, for every adapter, rather than in the two
+    that happened to have bespoke assertions.
+
+    Hooks the concrete targets rather than the aliases; that the aliases point
+    here is asserted by TestDenseBinding, and the two compose.
+    """
+    cfg = make_bridge_cfg(architecture, d_head=8)
+    adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+    blocks = adapter.component_mapping["blocks"]
+    template = blocks.submodules.get("mlp") or blocks.submodules["feed_forward"]
+    bridge = copy.deepcopy(template)
+    module = _DenseMLP()
+    bridge.set_original_component(module)
+    setup_submodules(bridge, adapter, module)
+
+    with torch.random.fork_rng(devices=[]):
+        torch.manual_seed(0)
+        x = torch.randn(1, 3, D_MODEL)
+
+    captured: dict = {}
+    for key in ("dense_gate", "dense_in"):
+        getattr(bridge, key).hook_out.add_hook(
+            lambda t, hook, key=key: captured.__setitem__(key, t.clone())
+        )
+    with torch.no_grad():
+        bridge(x)
+        expected_gate = module.gate_proj(x)
+        expected_in = module.up_proj(x)
+
+    torch.testing.assert_close(
+        captured["dense_gate"],
+        expected_gate,
+        msg=lambda m: f"{architecture}: dense_gate is not the gate projection\n{m}",
+    )
+    torch.testing.assert_close(
+        captured["dense_in"],
+        expected_in,
+        msg=lambda m: f"{architecture}: dense_in is not the up projection\n{m}",
+    )
+    # The fixture must be able to tell them apart, or neither assertion means
+    # anything (equal weights would satisfy both under a swap).
+    assert not torch.allclose(expected_gate, expected_in)
+
+
 class _UngatedDenseFF(nn.Module):
     """Switch-style ungated dense feed-forward (wi/wo, no gate projection)."""
 

tests/unit/model_bridge/supported_architectures/helpers.py

@@ -9,6 +9,14 @@
 
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.component_setup import setup_submodules
+from transformer_lens.model_bridge.generalized_components import MoEBridge
+
+# The keys MoEBridge binds a dense layer's projections under. Adapter key-set
+# assertions subtract these and check only the sparse-side keys: that every
+# dense-aware adapter declares dense_* AND that each key reads the projection it
+# names is covered behaviorally, for all 14 of them, by the roster in
+# tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py.
+DENSE_KEYS = frozenset({*MoEBridge.DENSE_SUBMODULE_KEYS, MoEBridge.DENSE_GATE_KEY})
 
 
 def make_bridge_cfg(architecture: str, **overrides) -> TransformerBridgeConfig:

tests/unit/model_bridge/supported_architectures/test_bitnet_adapter.py

@@ -6,6 +6,7 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg
 from transformer_lens.config import TransformerBridgeConfig
@@ -73,6 +74,68 @@ def __init__(self):
         assert torch.equal(out, x * 2.0)
 
 
+class TestBitNetPackedCheckpointGuard:
+    """`prepare_model` must refuse packed 1.58-bit checkpoints.
+
+    The flagship microsoft/bitnet-b1.58-2B-4T stores `weight` as packed uint8
+    with a collapsed first dim plus a separate weight_scale, so every
+    weight-space read reshapes it into a wrong-but-plausible matrix instead of
+    failing — the registry records that checkpoint at 0% on the forward phase.
+    """
+
+    @staticmethod
+    def _model(dtype, packed_shape=(8, 1)):
+        """Mirrors BitNet's real layout: a FLOAT embedding, then packed linears.
+
+        The embedding is load-bearing. BitNet leaves it unquantized and it sorts
+        first in named_modules(), so a guard that samples only the first
+        weight-bearing module inspects the one weight that is never packed. A
+        fixture whose only weight is the quantized one passes such a guard for
+        the wrong reason — which is exactly how the `break` survived review.
+        """
+
+        class _Tiny(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.embed_tokens = torch.nn.Embedding(8, 4)
+                self.q_proj = torch.nn.Linear(4, 4, bias=False)
+                self.q_proj.weight = torch.nn.Parameter(
+                    torch.zeros(*packed_shape, dtype=dtype), requires_grad=False
+                )
+
+        model = _Tiny()
+        # prepare_model's base implementation reads cfg.attn_implementation.
+        model.config = SimpleNamespace(attn_implementation="eager")
+        return model
+
+    def test_fixture_orders_the_float_embedding_first(self, adapter):
+        """Pins the property that makes the tests above meaningful: if the
+        packed weight were seen first, they would pass even with the guard
+        sampling a single module."""
+        weighted = [
+            (name, module.weight.dtype)
+            for name, module in self._model(torch.uint8).named_modules()
+            if getattr(module, "weight", None) is not None
+        ]
+        assert weighted[0][1].is_floating_point, weighted
+        assert any(not dtype.is_floating_point for _, dtype in weighted), weighted
+
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.int8])
+    def test_packed_weights_are_refused(self, adapter, dtype):
+        with pytest.raises(NotImplementedError, match="packed weights"):
+            adapter.prepare_model(self._model(dtype))
+
+    def test_error_names_the_dequantized_sibling(self, adapter):
+        with pytest.raises(NotImplementedError, match="bitnet-b1.58-2B-4T-bf16"):
+            adapter.prepare_model(self._model(torch.uint8))
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16])
+    def test_dequantized_checkpoints_still_load(self, adapter, dtype):
+        """The positive control: the bf16 sibling this error points users at
+        must pass, or the guard would make BitNet unusable entirely."""
+        adapter.prepare_model(self._model(dtype, packed_shape=(4, 4)))
+
+
 class TestBitNetRegistration:
     def test_factory_lookup(self):
         assert SUPPORTED_ARCHITECTURES["BitNetForCausalLM"] is BitNetArchitectureAdapter

tests/unit/model_bridge/supported_architectures/test_deepseek_v2_adapter.py

@@ -12,6 +12,7 @@
 
 import pytest
 
+from tests.unit.model_bridge.supported_architectures.helpers import DENSE_KEYS
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.generalized_components import (
     EmbeddingBridge,
@@ -214,30 +215,14 @@ def test_kv_path_submodules_are_required(self, adapter: DeepSeekV2ArchitectureAd
 
 
 class TestDeepSeekV2AdapterMoE:
-    """Tests the MoE submodule mapping and its dense-layer fallback."""
+    """Tests the MoE submodule mapping."""
 
     def test_moe_submodule_keys(self, adapter: DeepSeekV2ArchitectureAdapter) -> None:
         """Gate and shared experts are bridged; DeepseekV2Moe.forward calls self.gate,
         so its routing logits are hookable."""
         mlp = adapter.component_mapping["blocks"].submodules["mlp"]
-        assert set(mlp.submodules.keys()) == {
-            "gate",
-            "shared_experts",
-            "dense_gate",
-            "dense_in",
-            "dense_out",
-        }
-
-    def test_dense_projections_are_optional(self, adapter: DeepSeekV2ArchitectureAdapter) -> None:
-        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
-        for key, path in {
-            "dense_gate": "gate_proj",
-            "dense_in": "up_proj",
-            "dense_out": "down_proj",
-        }.items():
-            assert isinstance(mlp.submodules[key], LinearBridge)
-            assert mlp.submodules[key].name == path
-            assert mlp.submodules[key].optional is True
+        # dense_* are covered by the roster in test_moe_dense_dispatch.py.
+        assert set(mlp.submodules) - DENSE_KEYS == {"gate", "shared_experts"}
 
     def test_shared_experts_is_optional(self, adapter: DeepSeekV2ArchitectureAdapter) -> None:
         """Dense layers (idx < first_k_dense_replace) have no shared_experts."""

tests/unit/model_bridge/supported_architectures/test_deepseek_v3_adapter.py

@@ -14,6 +14,7 @@
 
 import pytest
 
+from tests.unit.model_bridge.supported_architectures.helpers import DENSE_KEYS
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.generalized_components import (
     EmbeddingBridge,
@@ -210,18 +211,13 @@ def test_all_projections_required(self, adapter: DeepSeekV3ArchitectureAdapter)
 
 
 class TestDeepSeekV3AdapterMoE:
-    """Tests the MoE submodule mapping and its dense-layer fallback."""
+    """Tests the MoE submodule mapping."""
 
     def test_moe_submodule_keys(self, adapter: DeepSeekV3ArchitectureAdapter) -> None:
         """V3 bridges the router gate (a custom Module), unlike V2."""
         mlp = adapter.component_mapping["blocks"].submodules["mlp"]
-        assert set(mlp.submodules.keys()) == {
-            "gate",
-            "shared_experts",
-            "dense_gate",
-            "dense_in",
-            "dense_out",
-        }
+        # dense_* are covered by the roster in test_moe_dense_dispatch.py.
+        assert set(mlp.submodules) - DENSE_KEYS == {"gate", "shared_experts"}
 
     def test_gate_is_optional_plain_component(self, adapter: DeepSeekV3ArchitectureAdapter) -> None:
         """The router gate is a custom Module (not nn.Linear) and absent on dense layers."""

tests/unit/model_bridge/supported_architectures/test_ernie4_5_moe_adapter.py

@@ -7,7 +7,10 @@
 
 import pytest
 
-from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg
+from tests.unit.model_bridge.supported_architectures.helpers import (
+    DENSE_KEYS,
+    make_bridge_cfg,
+)
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.factories.architecture_adapter_factory import (
     SUPPORTED_ARCHITECTURES,
@@ -55,13 +58,8 @@ def test_moe_with_optional_shared_experts(self, adapter):
         sigmoid router is fully delegated."""
         mlp = adapter.component_mapping["blocks"].submodules["mlp"]
         assert isinstance(mlp, MoEBridge)
-        assert set(mlp.submodules) == {
-            "gate",
-            "shared_experts",
-            "dense_gate",
-            "dense_in",
-            "dense_out",
-        }
+        # dense_* are covered by the roster in test_moe_dense_dispatch.py.
+        assert set(mlp.submodules) - DENSE_KEYS == {"gate", "shared_experts"}
         assert mlp.submodules["gate"].optional is True
         shared = mlp.submodules["shared_experts"]
         assert isinstance(shared, GatedMLPBridge)