Fix gated Qwen W_Q analysis weights

tests/integration/model_bridge/test_qwen3_gated_query_weights.py

@@ -0,0 +1,180 @@
+"""Download-free integration tests for Qwen3 gated query projections."""
+
+import copy
+from typing import NamedTuple
+
+import pytest
+import torch
+from transformers import (
+    Qwen3_5ForCausalLM,
+    Qwen3_5TextConfig,
+    Qwen3Config,
+    Qwen3ForCausalLM,
+    Qwen3NextConfig,
+    Qwen3NextForCausalLM,
+)
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.sources import build_bridge_from_module
+
+N_HEADS = 2
+D_HEAD = 8
+D_MODEL = 16
+N_LAYERS = 2
+VOCAB_SIZE = 32
+
+
+class GatedQueryCase(NamedTuple):
+    bridge: TransformerBridge
+    reference_logits: torch.Tensor
+    tokens: torch.Tensor
+    raw_q_weights: torch.Tensor
+
+
+def _tiny_hybrid_config(architecture: str):
+    common = dict(
+        hidden_size=D_MODEL,
+        num_hidden_layers=N_LAYERS,
+        num_attention_heads=N_HEADS,
+        num_key_value_heads=1,
+        head_dim=D_HEAD,
+        intermediate_size=32,
+        vocab_size=VOCAB_SIZE,
+        rms_norm_eps=1e-6,
+        hidden_act="silu",
+        full_attention_interval=1,
+        linear_conv_kernel_dim=4,
+        linear_key_head_dim=8,
+        linear_value_head_dim=8,
+        linear_num_key_heads=N_HEADS,
+        linear_num_value_heads=N_HEADS,
+        rope_parameters={
+            "rope_theta": 10000.0,
+            "partial_rotary_factor": 0.25,
+            "rope_type": "default",
+        },
+    )
+    if architecture == "Qwen3_5ForCausalLM":
+        return Qwen3_5TextConfig(**common)
+    return Qwen3NextConfig(
+        **common,
+        num_experts=2,
+        num_experts_per_tok=1,
+        moe_intermediate_size=16,
+        shared_expert_intermediate_size=16,
+        decoder_sparse_step=1,
+        mlp_only_layers=[],
+    )
+
+
+@pytest.fixture(
+    scope="module",
+    params=[
+        ("Qwen3_5ForCausalLM", Qwen3_5ForCausalLM),
+        ("Qwen3NextForCausalLM", Qwen3NextForCausalLM),
+    ],
+    ids=["qwen3_5", "qwen3_next"],
+)
+def gated_query_case(request: pytest.FixtureRequest) -> GatedQueryCase:
+    architecture, model_cls = request.param
+    torch.manual_seed(0)
+    cfg = _tiny_hybrid_config(architecture)
+    hf_model = model_cls(cfg).eval()
+    with torch.no_grad():
+        for layer_index, layer in enumerate(hf_model.model.layers):
+            q_weight = layer.self_attn.q_proj.weight
+            values = torch.arange(q_weight.numel(), dtype=q_weight.dtype).reshape_as(q_weight)
+            q_weight.copy_((values + layer_index * q_weight.numel()) / q_weight.numel())
+    raw_q_weights = torch.stack(
+        [layer.self_attn.q_proj.weight.detach().clone() for layer in hf_model.model.layers]
+    )
+
+    tokens = torch.arange(4).unsqueeze(0)
+    with torch.no_grad():
+        reference_logits = hf_model(tokens).logits
+
+    bridge = build_bridge_from_module(
+        hf_model,
+        architecture,
+        hf_config=copy.deepcopy(cfg),
+        tokenizer=None,
+        device="cpu",
+    ).eval()
+    bridge.process_weights(
+        fold_ln=False,
+        center_writing_weights=False,
+        center_unembed=False,
+        fold_value_biases=False,
+    )
+    return GatedQueryCase(
+        bridge=bridge,
+        reference_logits=reference_logits,
+        tokens=tokens,
+        raw_q_weights=raw_q_weights,
+    )
+
+
+def test_gated_w_q_exposes_query_rows_only(gated_query_case: GatedQueryCase) -> None:
+    expected = gated_query_case.raw_q_weights.view(N_LAYERS, N_HEADS, D_HEAD * 2, D_MODEL)[
+        :, :, :D_HEAD, :
+    ].transpose(-1, -2)
+
+    assert gated_query_case.bridge.W_Q.shape == (N_LAYERS, N_HEADS, D_MODEL, D_HEAD)
+    torch.testing.assert_close(gated_query_case.bridge.W_Q, expected)
+
+
+def test_gated_w_q_access_preserves_forward_and_gate_hook(
+    gated_query_case: GatedQueryCase,
+) -> None:
+    captured_gate: list[torch.Tensor] = []
+    _ = gated_query_case.bridge.W_Q
+
+    with torch.no_grad():
+        bridge_logits = gated_query_case.bridge.run_with_hooks(
+            gated_query_case.tokens,
+            fwd_hooks=[
+                (
+                    "blocks.0.attn.hook_q_gate",
+                    lambda gate, hook: captured_gate.append(gate.detach().clone()),
+                )
+            ],
+        )
+
+    torch.testing.assert_close(
+        bridge_logits, gated_query_case.reference_logits, atol=1e-5, rtol=1e-5
+    )
+    assert len(captured_gate) == 1
+    assert captured_gate[0].shape == (1, gated_query_case.tokens.shape[1], N_HEADS * D_HEAD)
+    live_q_weights = torch.stack(
+        [
+            gated_query_case.bridge.state_dict()[f"blocks.{layer}.attn.q.weight"]
+            for layer in range(N_LAYERS)
+        ]
+    )
+    torch.testing.assert_close(live_q_weights, gated_query_case.raw_q_weights)
+
+
+def test_standard_qwen3_w_q_is_unchanged() -> None:
+    cfg = Qwen3Config(
+        hidden_size=D_MODEL,
+        num_hidden_layers=1,
+        num_attention_heads=N_HEADS,
+        num_key_value_heads=1,
+        head_dim=D_HEAD,
+        intermediate_size=32,
+        vocab_size=VOCAB_SIZE,
+        max_position_embeddings=32,
+    )
+    hf_model = Qwen3ForCausalLM(cfg).eval()
+    bridge = build_bridge_from_module(
+        hf_model,
+        "Qwen3ForCausalLM",
+        hf_config=copy.deepcopy(cfg),
+        tokenizer=None,
+        device="cpu",
+    ).eval()
+    raw_q_weight = hf_model.model.layers[0].self_attn.q_proj.weight.detach()
+    expected = raw_q_weight.view(N_HEADS, D_HEAD, D_MODEL).transpose(-1, -2).unsqueeze(0)
+
+    assert bridge.W_Q.shape == (1, N_HEADS, D_MODEL, D_HEAD)
+    torch.testing.assert_close(bridge.W_Q, expected)

tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py

@@ -335,117 +335,6 @@ def test_n_key_value_heads_not_set_when_absent(self, qwen3_5_dependency_availabl
         )
 
 
-class TestQwen3_5PreprocessWeights:
-    """q_proj rows are interleaved per-head (query, gate, query, gate, ...) — naive first-half slice is wrong."""
-
-    N_HEADS = 4
-    D_HEAD = 8
-    HIDDEN_SIZE = 32
-
-    @pytest.fixture
-    def adapter(self, qwen3_5_dependency_available):
-        from transformer_lens.model_bridge.supported_architectures.qwen3_5 import (
-            Qwen3_5ArchitectureAdapter,
-        )
-
-        cfg = _make_bridge_cfg(
-            n_heads=self.N_HEADS,
-            d_head=self.D_HEAD,
-            d_model=self.HIDDEN_SIZE,
-            n_key_value_heads=self.N_HEADS,
-        )
-        return Qwen3_5ArchitectureAdapter(cfg)
-
-    def _make_q_proj_weight(self):
-        import torch
-
-        total_rows = self.N_HEADS * self.D_HEAD * 2
-        w = torch.zeros(total_rows, self.HIDDEN_SIZE)
-        for row_idx in range(total_rows):
-            w[row_idx] = float(row_idx)
-        return w
-
-    def test_q_proj_output_shape(self, adapter):
-        import torch
-
-        w = self._make_q_proj_weight()
-        state_dict = {"model.layers.3.self_attn.q_proj.weight": w}
-        result = adapter.preprocess_weights(state_dict)
-        out = result["model.layers.3.self_attn.q_proj.weight"]
-        assert out.shape == (self.N_HEADS * self.D_HEAD, self.HIDDEN_SIZE)
-
-    def test_q_proj_selects_query_rows_not_naive_first_half(self, adapter):
-        import torch
-
-        w = self._make_q_proj_weight()
-        state_dict = {"model.layers.0.self_attn.q_proj.weight": w}
-        result = adapter.preprocess_weights(state_dict)
-        out = result["model.layers.0.self_attn.q_proj.weight"]
-
-        for head_idx in range(self.N_HEADS):
-            out_rows = out[head_idx * self.D_HEAD : (head_idx + 1) * self.D_HEAD]
-            expected_start = head_idx * self.D_HEAD * 2
-            expected_rows = w[expected_start : expected_start + self.D_HEAD]
-            assert torch.equal(out_rows, expected_rows), (
-                f"Head {head_idx}: output rows do not match expected query rows. "
-                f"Got row values starting at {out_rows[0, 0].item()}, "
-                f"expected starting at {expected_rows[0, 0].item()}"
-            )
-
-    def test_naive_slice_would_be_wrong(self, adapter):
-        import torch
-
-        w = self._make_q_proj_weight()
-        state_dict = {"model.layers.0.self_attn.q_proj.weight": w}
-        result = adapter.preprocess_weights(state_dict)
-        correct_out = result["model.layers.0.self_attn.q_proj.weight"]
-        naive_out = w[: self.N_HEADS * self.D_HEAD]
-
-        if self.N_HEADS > 1:
-            assert not torch.equal(correct_out, naive_out), (
-                "Naive first-half slice gave the same result as per-head slice — "
-                "test setup may be wrong"
-            )
-
-    def test_non_q_proj_weights_unchanged(self, adapter):
-        import torch
-
-        k_proj = torch.randn(self.N_HEADS * self.D_HEAD, self.HIDDEN_SIZE)
-        down_proj = torch.randn(self.HIDDEN_SIZE, self.N_HEADS * self.D_HEAD)
-        state_dict = {
-            "model.layers.0.self_attn.k_proj.weight": k_proj.clone(),
-            "model.layers.0.mlp.down_proj.weight": down_proj.clone(),
-        }
-        result = adapter.preprocess_weights(state_dict)
-        assert torch.equal(result["model.layers.0.self_attn.k_proj.weight"], k_proj)
-        assert torch.equal(result["model.layers.0.mlp.down_proj.weight"], down_proj)
-
-    def test_multiple_layers_all_processed(self, adapter):
-        import torch
-
-        w0 = self._make_q_proj_weight()
-        w3 = self._make_q_proj_weight() * 2
-        state_dict = {
-            "model.layers.0.self_attn.q_proj.weight": w0,
-            "model.layers.3.self_attn.q_proj.weight": w3,
-        }
-        result = adapter.preprocess_weights(state_dict)
-        expected_shape = (self.N_HEADS * self.D_HEAD, self.HIDDEN_SIZE)
-        assert result["model.layers.0.self_attn.q_proj.weight"].shape == expected_shape
-        assert result["model.layers.3.self_attn.q_proj.weight"].shape == expected_shape
-
-    def test_empty_state_dict_returns_empty(self, adapter):
-        assert adapter.preprocess_weights({}) == {}
-
-    def test_state_dict_without_q_proj_unchanged(self, adapter):
-        import torch
-
-        state_dict = {"model.embed_tokens.weight": torch.randn(100, self.HIDDEN_SIZE)}
-        original_keys = set(state_dict.keys())
-        result = adapter.preprocess_weights(state_dict)
-        assert set(result.keys()) == original_keys
-
-
 @pytest.mark.skipif(
     not _QWEN3_5_AVAILABLE,
     reason="Qwen3_5TextConfig / Qwen3_5ForCausalLM not available in installed transformers",
@@ -508,7 +397,7 @@ def adapter(self):
         return Qwen3_5ArchitectureAdapter(_make_bridge_cfg())
 
     def test_gated_q_proj_flag_set(self, adapter):
-        """Flag drives preprocess_weights to slice the gated half of q_proj."""
+        """Flag drives the query-only W_Q analysis view and gate hook path."""
         assert getattr(adapter.cfg, "gated_q_proj", False) is True
 
 

tests/unit/model_bridge/supported_architectures/test_qwen3_5_multimodal_adapter.py

@@ -9,7 +9,6 @@
 
 import torch
 
-from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig
 from transformer_lens.model_bridge.generalized_components import (
     LinearBridge,
@@ -110,17 +109,26 @@ def test_vision_tower_decomposed(self):
             assert isinstance(block.submodules[comp].submodules[sub], LinearBridge)
 
 
-def test_gated_q_proj_query_half_is_sliced_under_nested_path():
-    """preprocess_weights slices the query half from the 2x-wide gated q_proj, matching the
-    nested model.language_model.* key."""
+def test_gated_q_proj_exposes_query_only_w_q_without_mutating_live_weight():
+    """The multimodal adapter shares the non-mutating gated W_Q analysis view."""
     adapter = Qwen3_5MultimodalArchitectureAdapter(_make_cfg())
+    attention = adapter.component_mapping["blocks"].submodules["attn"]
     n_heads, d_head, hidden = adapter.cfg.n_heads, adapter.cfg.d_head, adapter.cfg.d_model
-    key = "model.language_model.layers.1.self_attn.q_proj.weight"
-    # Per head: rows [query(d_head), gate(d_head)] -> 2*d_head wide.
-    full = torch.randn(n_heads * d_head * 2, hidden)
-    out = adapter.preprocess_weights({key: full.clone()})
-    assert out[key].shape == (n_heads * d_head, hidden)
-    expected = full.view(n_heads, d_head * 2, hidden)[:, :d_head, :].reshape(
-        n_heads * d_head, hidden
+    q_proj = torch.nn.Linear(hidden, n_heads * d_head * 2, bias=False)
+    with torch.no_grad():
+        values = torch.arange(q_proj.weight.numel(), dtype=q_proj.weight.dtype)
+        q_proj.weight.copy_(values.reshape_as(q_proj.weight))
+    original_weight = q_proj.weight.detach().clone()
+    query = attention.submodules["q"]
+    query.set_original_component(q_proj)
+    attention.add_module("q", query)
+
+    expected = (
+        original_weight.view(n_heads, d_head * 2, hidden)[:, :d_head, :]
+        .transpose(-1, -2)
+        .contiguous()
     )
-    assert torch.equal(out[key], expected)
+
+    assert attention.W_Q.shape == (n_heads, hidden, d_head)
+    torch.testing.assert_close(attention.W_Q, expected)
+    torch.testing.assert_close(q_proj.weight, original_weight)