From f1a3503e9f64331022d459ed1ac9ca47607d66d9 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:10:50 +0800 Subject: [PATCH] Fix gated Qwen W_Q analysis weights --- .../test_qwen3_gated_query_weights.py | 180 ++++++++++++++++++ .../test_qwen3_5_adapter.py | 113 +---------- .../test_qwen3_5_multimodal_adapter.py | 32 ++-- .../test_qwen3_adapter.py | 51 ----- .../test_qwen3_next_adapter.py | 109 ----------- .../generalized_components/attention.py | 29 ++- .../supported_architectures/qwen3.py | 18 -- .../supported_architectures/qwen3_5.py | 13 +- .../qwen3_5_multimodal.py | 6 - .../supported_architectures/qwen3_next.py | 6 - 10 files changed, 227 insertions(+), 330 deletions(-) create mode 100644 tests/integration/model_bridge/test_qwen3_gated_query_weights.py diff --git a/tests/integration/model_bridge/test_qwen3_gated_query_weights.py b/tests/integration/model_bridge/test_qwen3_gated_query_weights.py new file mode 100644 index 0000000000..c490eab77f --- /dev/null +++ b/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) diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py index 91236361a2..30fb9aaf69 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_adapter.py +++ b/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 diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_multimodal_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_multimodal_adapter.py index 544c7f6bda..f76cfbf7a5 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_multimodal_adapter.py +++ b/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) diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_adapter.py index f5508623f0..9c937b91de 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_adapter.py @@ -4,14 +4,12 @@ - Config attributes - Component mapping structure and HF module names (incl. q_norm/k_norm) - Weight conversion keys/types (GQA: k/v use n_key_value_heads) -- _preprocess_gated_q_proj static helper (gated q_proj slicing) - Factory registration """ from types import SimpleNamespace from typing import Any import pytest -import torch from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg from transformer_lens.config import TransformerBridgeConfig @@ -140,55 +138,6 @@ def test_no_linear_attn_when_dense(self, adapter: Qwen3ArchitectureAdapter) -> N assert "linear_attn" not in blocks.submodules -class TestPreprocessGatedQProj: - """Numerical correctness of the _preprocess_gated_q_proj static helper - on synthetic interleaved [query, gate] rows: asserts query-half slicing, - that unrelated state-dict keys are untouched, and that the rewrite - applies across all matching layers.""" - - def test_slices_query_half(self) -> None: - """Interleaved [query, gate] rows per head must be reduced to query-only.""" - n_heads, d_head, d_model = 4, 8, 16 - # Build q_proj.weight as (n_heads, d_head*2, d_model): query=1.0, gate=9.0 - w = torch.empty(n_heads, d_head * 2, d_model) - w[:, :d_head, :] = 1.0 - w[:, d_head:, :] = 9.0 - w_flat = w.reshape(n_heads * d_head * 2, d_model) - - state_dict = {"model.layers.0.self_attn.q_proj.weight": w_flat.clone()} - out = Qwen3ArchitectureAdapter._preprocess_gated_q_proj(state_dict, n_heads, d_head) - - result = out["model.layers.0.self_attn.q_proj.weight"] - assert result.shape == (n_heads * d_head, d_model) - assert torch.all(result == 1.0), "gate rows must be dropped" - - def test_only_q_proj_keys_modified(self) -> None: - n_heads, d_head, d_model = 2, 4, 8 - q_w = torch.ones(n_heads * d_head * 2, d_model) - other = torch.full((d_model, d_model), 7.0) - state_dict = { - "model.layers.0.self_attn.q_proj.weight": q_w, - "model.layers.0.self_attn.k_proj.weight": other.clone(), - "model.layers.0.mlp.gate_proj.weight": other.clone(), - } - out = Qwen3ArchitectureAdapter._preprocess_gated_q_proj(state_dict, n_heads, d_head) - assert torch.equal(out["model.layers.0.self_attn.k_proj.weight"], other) - assert torch.equal(out["model.layers.0.mlp.gate_proj.weight"], other) - - def test_multiple_layers(self) -> None: - n_heads, d_head, d_model = 2, 4, 8 - state_dict = { - f"model.layers.{i}.self_attn.q_proj.weight": torch.ones(n_heads * d_head * 2, d_model) - for i in range(3) - } - out = Qwen3ArchitectureAdapter._preprocess_gated_q_proj(state_dict, n_heads, d_head) - for i in range(3): - assert out[f"model.layers.{i}.self_attn.q_proj.weight"].shape == ( - n_heads * d_head, - d_model, - ) - - class TestQwen3HybridConstructor: """The hybrid=True constructor branch on the base class. The Qwen3_5 / Qwen3Next subclasses exercise this path transitively; pinning it here diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_next_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_next_adapter.py index 2b2ed53326..640f53e9fb 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_next_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_next_adapter.py @@ -123,115 +123,6 @@ def test_rotary_emb_bridge_type(self, adapter): assert isinstance(adapter.component_mapping["rotary_emb"], RotaryEmbeddingBridge) -class TestQwen3NextWeightConversions: - """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): - from transformer_lens.model_bridge.supported_architectures.qwen3_next import ( - Qwen3NextArchitectureAdapter, - ) - - 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 Qwen3NextArchitectureAdapter(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_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): - result = adapter.preprocess_weights({}) - assert result == {} - - 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 - - class TestQwen3NextConfigAttributes: """cfg attributes set by the adapter.""" diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index 9a79c3e2a3..bc9aa4a861 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -766,12 +766,33 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: @property def W_Q(self) -> torch.Tensor: - """Get W_Q in 3D format [n_heads, d_model, d_head].""" + """Get W_Q in 3D format [n_heads, d_model, d_head]. + + Gated query projections retain their live query-and-gate parameter; + this analysis view selects the query rows interleaved within each head. + """ weight = self.q.weight if weight.ndim == 2 and self.config is not None: - return self._reshape_weight_to_3d( - weight, self._get_n_heads(), in_out_layout=self._weight_layout_in_out(self.q) - ) + n_heads = self._get_n_heads() + in_out_layout = self._weight_layout_in_out(self.q) + if getattr(self.config, "gated_q_proj", False): + d_head = int(self.config.d_head) + gated_width = n_heads * d_head * 2 + if in_out_layout is True: + output_first_weight = weight.T + elif in_out_layout is False or weight.shape[0] == gated_width: + output_first_weight = weight + elif weight.shape[1] == gated_width: + output_first_weight = weight.T + else: + output_first_weight = None + + if output_first_weight is not None and output_first_weight.shape[0] == gated_width: + # Preserve the live query-gate projection; W_Q is an analysis-only query view. + per_head_weight = output_first_weight.reshape(n_heads, d_head * 2, -1) + return per_head_weight[:, :d_head, :].transpose(-1, -2) + + return self._reshape_weight_to_3d(weight, n_heads, in_out_layout=in_out_layout) return weight @property diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3.py b/transformer_lens/model_bridge/supported_architectures/qwen3.py index 9551c4ebc0..a4481dfcc0 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3.py @@ -7,8 +7,6 @@ from typing import Any -import torch - from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( AttentionBridge, @@ -93,19 +91,3 @@ def _build_component_mapping(self, *, hybrid: bool = False, lm_prefix: str = "mo "ln_final": RMSNormalizationBridge(name=f"{lm_prefix}.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head"), } - - @staticmethod - def _preprocess_gated_q_proj( - state_dict: dict[str, torch.Tensor], n_heads: int, d_head: int - ) -> dict[str, torch.Tensor]: - """Slice query half from gated q_proj.weight (interleaved per-head layout). - - q_proj.weight has shape (n_heads * d_head * 2, hidden_size) with - interleaved [query, gate] rows per head. Extracts query-only half. - """ - keys_to_update = [k for k in state_dict if k.endswith(".self_attn.q_proj.weight")] - for key in keys_to_update: - w = state_dict[key] - w = w.view(n_heads, d_head * 2, -1) - state_dict[key] = w[:, :d_head, :].reshape(n_heads * d_head, -1) - return state_dict diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_5.py b/transformer_lens/model_bridge/supported_architectures/qwen3_5.py index 9b3a98bc32..a155055579 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_5.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_5.py @@ -7,8 +7,6 @@ from typing import Any -import torch - from transformer_lens.model_bridge.supported_architectures.qwen3 import ( Qwen3ArchitectureAdapter, ) @@ -19,7 +17,7 @@ class Qwen3_5ArchitectureAdapter(Qwen3ArchitectureAdapter): Inherits Qwen3 config/attention/MLP structure. Differences: - Attention + linear_attn are optional (per-layer type) - - Gated q_proj (2x wide) sliced by preprocess_weights for weight analysis + - Gated q_proj (2x wide); AttentionBridge exposes a query-only W_Q view """ # Multimodal wrapper architecture this text-only adapter rejects; the MoE @@ -58,12 +56,3 @@ def prepare_model(self, hf_model: Any) -> None: f"TransformerBridge.boot_transformers(...) so {multimodal_arch} " f"checkpoints route to the multimodal adapter automatically." ) - - def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Slice query half from gated q_proj.weight for weight-space analysis. - - In processed mode, W_Q is the pure query projection (for composition - scores, logit lens). Gate signal available in unprocessed mode on - full-attention layers via blocks.N.attn.hook_q_gate. - """ - return self._preprocess_gated_q_proj(state_dict, self.cfg.n_heads, self.cfg.d_head) diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_5_multimodal.py b/transformer_lens/model_bridge/supported_architectures/qwen3_5_multimodal.py index 6e2aa3b277..e5cbd8975b 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_5_multimodal.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_5_multimodal.py @@ -7,8 +7,6 @@ from typing import Any -import torch - from transformer_lens.model_bridge.generalized_components import VisionProjectionBridge from transformer_lens.model_bridge.generalized_components.qwen3_5_vision_encoder import ( Qwen3_5VisionEncoderBridge, @@ -39,7 +37,3 @@ def __init__(self, cfg: Any) -> None: name="model.visual", config=self.cfg ) self.components["vision_projector"] = VisionProjectionBridge(name="model.visual.merger") - - def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Slice query half from gated q_proj.weight (matcher is path-prefix-agnostic).""" - return self._preprocess_gated_q_proj(state_dict, self.cfg.n_heads, self.cfg.d_head) diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_next.py b/transformer_lens/model_bridge/supported_architectures/qwen3_next.py index 31e1be3cdc..d0649ace9d 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_next.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_next.py @@ -7,8 +7,6 @@ from typing import Any -import torch - from transformer_lens.model_bridge.generalized_components import MoEBridge from transformer_lens.model_bridge.supported_architectures.qwen3 import ( Qwen3ArchitectureAdapter, @@ -28,7 +26,3 @@ def __init__(self, cfg: Any) -> None: def _build_mlp_bridge(self): """Sparse MoE MLP (router + batched experts + shared expert).""" return MoEBridge(name="mlp", config=self.cfg) - - def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Slice query half from gated q_proj.weight for weight-space analysis.""" - return self._preprocess_gated_q_proj(state_dict, self.cfg.n_heads, self.cfg.d_head)