diff --git a/tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py b/tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py index f999749ae..653d54364 100644 --- a/tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py +++ b/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).""" diff --git a/tests/unit/model_bridge/supported_architectures/helpers.py b/tests/unit/model_bridge/supported_architectures/helpers.py index 70f2aefd6..eb598e235 100644 --- a/tests/unit/model_bridge/supported_architectures/helpers.py +++ b/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: diff --git a/tests/unit/model_bridge/supported_architectures/test_bitnet_adapter.py b/tests/unit/model_bridge/supported_architectures/test_bitnet_adapter.py index 107f38dd4..fc15d3387 100644 --- a/tests/unit/model_bridge/supported_architectures/test_bitnet_adapter.py +++ b/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 diff --git a/tests/unit/model_bridge/supported_architectures/test_deepseek_v2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_deepseek_v2_adapter.py index 56c2af2a1..b7b8d8ddd 100644 --- a/tests/unit/model_bridge/supported_architectures/test_deepseek_v2_adapter.py +++ b/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.""" diff --git a/tests/unit/model_bridge/supported_architectures/test_deepseek_v3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_deepseek_v3_adapter.py index 967c84b2c..0410ca4ef 100644 --- a/tests/unit/model_bridge/supported_architectures/test_deepseek_v3_adapter.py +++ b/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.""" diff --git a/tests/unit/model_bridge/supported_architectures/test_ernie4_5_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_ernie4_5_moe_adapter.py index c95ed05a0..7b4cdae57 100644 --- a/tests/unit/model_bridge/supported_architectures/test_ernie4_5_moe_adapter.py +++ b/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) diff --git a/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py index 09a4d12fc..4b154c9a9 100644 --- a/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py @@ -447,3 +447,155 @@ def test_bridge_matches_hf_with_key_multiplier( attention_mask=attention_mask, )[0] assert not torch.allclose(unscaled_out, hf_out, atol=1e-5) + + +KEY_MULTIPLIER = 1.5 + + +class TestFalconH1ScaledHookK: + """`hook_k` must carry the tensor that reaches attention, not the raw projection. + + Falcon-H1 scales K between the projection and RoPE. While `hook_k` was an + alias for `k.hook_out` it reported the unscaled tensor, and a value written + there was silently multiplied by key_multiplier on its way in. This mirrors + the split Granite's residual_multiplier established: the TL-semantic name + carries the scaled value, the module-shaped `k.hook_out` stays raw. + """ + + @staticmethod + def _wire(adapter, cfg): + import torch + from transformers import FalconH1Config + from transformers.models.falcon_h1.modeling_falcon_h1 import FalconH1Attention + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + hf_config = FalconH1Config( + hidden_size=cfg.d_model, + num_attention_heads=cfg.n_heads, + num_key_value_heads=cfg.n_key_value_heads, + key_multiplier=KEY_MULTIPLIER, + ) + hf_config._attn_implementation = "eager" + hf_attn = FalconH1Attention(hf_config, layer_idx=0) + batch, seq = 2, 5 + hidden_states = torch.randn(batch, seq, cfg.d_model) + + bridge = wire_attention_bridge( + adapter, hf_attn, expected_type=PositionEmbeddingsAttentionBridge + ) + return bridge, hf_attn, hidden_states, batch, seq + + @staticmethod + def _run(bridge, hidden_states, seq, d_head): + import torch + + batch = hidden_states.shape[0] + with torch.no_grad(): + bridge( + hidden_states=hidden_states, + position_embeddings=identity_rope(seq, d_head), + attention_mask=make_additive_causal_mask(batch, seq), + ) + + def test_hook_k_is_scaled_and_k_hook_out_stays_raw(self, adapter, cfg) -> None: + """Both legs are asserted: checking only the ratio would still pass if + `k.hook_out` had been scaled too, which is the wrong fix.""" + import torch + + bridge, hf_attn, hidden_states, batch, seq = self._wire(adapter, cfg) + + captured: dict = {} + bridge.hook_k.add_hook(lambda t, hook: captured.__setitem__("scaled", t.clone())) + bridge.k.hook_out.add_hook(lambda t, hook: captured.__setitem__("raw", t.clone())) + self._run(bridge, hidden_states, seq, cfg.d_head) + + with torch.no_grad(): + expected_raw = hf_attn.k_proj(hidden_states).view( + batch, seq, cfg.n_key_value_heads, cfg.d_head + ) + torch.testing.assert_close(captured["raw"], expected_raw) + torch.testing.assert_close(captured["scaled"], expected_raw * KEY_MULTIPLIER) + + def test_hook_k_matches_the_tensor_entering_rope(self, adapter, cfg) -> None: + """Under identity RoPE, hook_rot_k is what attention consumes; hook_k + must already agree with it.""" + import torch + + bridge, _, hidden_states, _, seq = self._wire(adapter, cfg) + + captured: dict = {} + bridge.hook_k.add_hook(lambda t, hook: captured.__setitem__("k", t.clone())) + bridge.hook_rot_k.add_hook(lambda t, hook: captured.__setitem__("rot_k", t.clone())) + self._run(bridge, hidden_states, seq, cfg.d_head) + + torch.testing.assert_close(captured["k"], captured["rot_k"]) + + def test_a_value_written_at_hook_k_is_not_rescaled(self, adapter, cfg) -> None: + """The write path is the half users hit when patching. Previously a + write landed as value * key_multiplier.""" + import torch + + bridge, _, hidden_states, batch, seq = self._wire(adapter, cfg) + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(7) + written = torch.randn(batch, seq, cfg.n_key_value_heads, cfg.d_head) + + captured: dict = {} + bridge.hook_k.add_hook(lambda t, hook: written) + bridge.hook_rot_k.add_hook(lambda t, hook: captured.__setitem__("rot_k", t.clone())) + self._run(bridge, hidden_states, seq, cfg.d_head) + + torch.testing.assert_close(captured["rot_k"], written) + + def test_hook_k_keeps_the_per_head_cache_shape(self, adapter, cfg) -> None: + """A real HookPoint has to inherit k's ReshapeForAttentionHeads, or the + de-aliasing silently changes hook_k's shape from 4D to flat 3D.""" + bridge, _, hidden_states, batch, seq = self._wire(adapter, cfg) + + captured: dict = {} + bridge.hook_k.add_hook(lambda t, hook: captured.__setitem__("k", t)) + self._run(bridge, hidden_states, seq, cfg.d_head) + + assert captured["k"].shape == (batch, seq, cfg.n_key_value_heads, cfg.d_head) + + def test_unscaled_architectures_keep_the_alias(self, adapter, cfg) -> None: + """The de-aliasing is opt-in on the HF module carrying key_multiplier; + every other architecture must still resolve hook_k to k.hook_out.""" + import torch + from transformers import LlamaConfig + from transformers.models.llama.modeling_llama import LlamaAttention + + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + + llama_cfg = TransformerBridgeConfig( + d_model=cfg.d_model, + d_head=cfg.d_head, + n_heads=cfg.n_heads, + n_key_value_heads=cfg.n_key_value_heads, + n_layers=1, + n_ctx=16, + d_vocab=16, + architecture="LlamaForCausalLM", + original_architecture="LlamaForCausalLM", + ) + llama_adapter = ArchitectureAdapterFactory.select_architecture_adapter(llama_cfg) + hf_attn = LlamaAttention( + LlamaConfig( + hidden_size=cfg.d_model, + num_attention_heads=cfg.n_heads, + num_key_value_heads=cfg.n_key_value_heads, + head_dim=cfg.d_head, + ), + layer_idx=0, + ) + assert getattr(hf_attn, "key_multiplier", None) is None + + bridge = wire_attention_bridge( + llama_adapter, hf_attn, expected_type=PositionEmbeddingsAttentionBridge + ) + assert bridge.hook_aliases.get("hook_k") == "k.hook_out" + assert not isinstance(bridge.__dict__.get("_modules", {}).get("hook_k"), torch.nn.Module) diff --git a/tests/unit/model_bridge/supported_architectures/test_glm4_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_glm4_moe_adapter.py index fccc5daa3..a974237d9 100644 --- a/tests/unit/model_bridge/supported_architectures/test_glm4_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_glm4_moe_adapter.py @@ -4,7 +4,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.conversion_utils.conversion_steps import RearrangeTensorConversion from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import ( @@ -191,7 +194,8 @@ def test_gate_submodule_is_optional_for_dense_prefix_layers( gate = mlp.submodules["gate"] assert isinstance(gate, LinearBridge) assert getattr(gate, "optional", False) is True - assert set(mlp.submodules.keys()) == {"gate", "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"} class TestGlm4MoeComponentTypes: diff --git a/tests/unit/model_bridge/supported_architectures/test_glm4_moe_lite_adapter.py b/tests/unit/model_bridge/supported_architectures/test_glm4_moe_lite_adapter.py index c93b9a960..36b6fb8dd 100644 --- a/tests/unit/model_bridge/supported_architectures/test_glm4_moe_lite_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_glm4_moe_lite_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, @@ -87,13 +90,8 @@ def test_moe_with_optional_router_and_shared_expert(self, adapter): """Dense layers in mlp_layer_types have neither router nor shared expert.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) - 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"} assert mlp.submodules["gate"].optional is True shared = mlp.submodules["shared_experts"] assert isinstance(shared, GatedMLPBridge) diff --git a/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py b/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py index f71386fc7..55124c3c4 100644 --- a/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py @@ -207,9 +207,6 @@ def test_moe_bridge_when_num_experts_gt_one(self) -> None: # Dense (JambaMLP) layers bind under dense_* so they get gated-MLP # neuron hooks; `router` stays the sparse layers' router (#1645). assert set(mlp.submodules) == {"dense_gate", "dense_in", "dense_out", "router"} - assert mlp.submodules["router"].name == "router" - assert mlp.submodules["router"].optional is True - assert mlp.submodules["dense_gate"].optional is True # A renamed router on a sparse layer must fail loudly, not bind silently. assert mlp._sparse_required == ("router",) diff --git a/tests/unit/model_bridge/supported_architectures/test_llama4_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llama4_adapter.py index d84308002..b6f9f8568 100644 --- a/tests/unit/model_bridge/supported_architectures/test_llama4_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_llama4_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, @@ -48,23 +51,15 @@ def test_attention_stays_native(self, adapter): assert attn.name == "self_attn" def test_moe_with_optional_shared_expert(self, adapter): - """The router is mapped so a rename fails loudly; non-MoE layers hold a - dense gated MLP under the same feed_forward name.""" + """The router is mapped so a rename fails loudly.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, _Llama4MoEBridge) assert mlp.name == "feed_forward" - assert set(mlp.submodules) == { - "router", - "shared_expert", - "dense_gate", - "dense_in", - "dense_out", - } + # dense_* are covered by the roster in test_moe_dense_dispatch.py. + assert set(mlp.submodules) - DENSE_KEYS == {"router", "shared_expert"} # Optional so dense layers may skip it, required on sparse ones. assert mlp.submodules["router"].optional is True assert mlp._sparse_required == ("router",) - for key in ("dense_gate", "dense_in", "dense_out"): - assert mlp.submodules[key].optional is True shared = mlp.submodules["shared_expert"] assert isinstance(shared, _Llama4SharedExpertBridge) assert shared.optional is True diff --git a/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py index e2cb42527..8b1d7eb41 100644 --- a/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py @@ -400,3 +400,50 @@ def test_derived_geometry_unchanged(self) -> None: N_KV_HEADS * D_HEAD, N_KV_HEADS * D_HEAD, ] + + +class TestPhi3FusedSplitRefusesQuantizedWeights: + """The splitters Phi-3, GLM and GLM-4V actually install must refuse packed + or scale-separated weights. + + A guard was added to the *default* splitters in JointQKVAttentionBridge / + JointGateUpMLPBridge, but every in-tree user overrides those defaults with + these two methods — so the guard had no reachable caller for this family. + FP8 is the case that matters: `tensor_split` and `nn.Parameter` both accept + it, so the split silently yields scale-less halves. + """ + + QUANTIZED = [ + pytest.param(torch.int8, "packed integer storage", id="int8"), + pytest.param(torch.uint8, "packed integer storage", id="uint8"), + pytest.param(torch.float8_e4m3fn, "narrow float", id="fp8-e4m3fn"), + ] + + @pytest.mark.parametrize("dtype,reason", QUANTIZED) + def test_qkv_split_refuses(self, dtype, reason) -> None: + adapter = Phi3ArchitectureAdapter(_make_cfg(d_model=16, n_heads=3, n_kv_heads=3, d_head=4)) + fake = _FakeAttention(d_model=16, qkv_rows=36, bias=False) + fake.qkv_proj.weight = torch.nn.Parameter( + fake.qkv_proj.weight.detach().to(dtype), requires_grad=False + ) + with pytest.raises(NotImplementedError, match=reason): + adapter._split_phi3_qkv(fake) + + @pytest.mark.parametrize("dtype,reason", QUANTIZED) + def test_gate_up_split_refuses(self, dtype, reason) -> None: + class _FakeMLP(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.gate_up_proj = torch.nn.Linear(16, 32, bias=False) + self.gate_up_proj.weight = torch.nn.Parameter( + torch.zeros(32, 16, dtype=dtype), requires_grad=False + ) + + with pytest.raises(NotImplementedError, match=reason): + Phi3ArchitectureAdapter._split_gate_up(_FakeMLP()) + + def test_float_weights_still_split(self) -> None: + """Positive control: the guard must not break ordinary Phi-3 boot.""" + adapter = Phi3ArchitectureAdapter(_make_cfg(d_model=16, n_heads=3, n_kv_heads=3, d_head=4)) + q, k, v = adapter._split_phi3_qkv(_FakeAttention(d_model=16, qkv_rows=36, bias=False)) + assert q.weight.shape == (12, 16) diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py index 025a06509..b0cd3bbbe 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen2_moe_adapter.py @@ -4,6 +4,7 @@ import torch from transformers import Qwen2MoeConfig +from tests.unit.model_bridge.supported_architectures.helpers import DENSE_KEYS from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import ( RearrangeTensorConversion, ) @@ -111,14 +112,12 @@ def test_mlp_is_moe_bridge(self, adapter: Qwen2MoeArchitectureAdapter) -> None: def test_moe_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: mlp = adapter.component_mapping["blocks"].submodules["mlp"] - assert set(mlp.submodules.keys()) == { + # dense_* are covered by the roster in test_moe_dense_dispatch.py. + assert set(mlp.submodules) - DENSE_KEYS == { "gate", "experts", "shared_expert", "shared_expert_gate", - "dense_gate", - "dense_in", - "dense_out", } assert isinstance(mlp.submodules["gate"], LinearBridge) assert isinstance(mlp.submodules["experts"], MoEBridge) diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_moe_adapter.py index 31c4f74ef..18d1b2694 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_moe_adapter.py @@ -2,7 +2,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.conversion_utils.conversion_steps.rearrange_tensor_conversion import ( RearrangeTensorConversion, @@ -224,10 +227,10 @@ def test_mlp_maps_router_and_dense_projections( self, adapter: Qwen3MoeArchitectureAdapter ) -> None: """Experts are batched 3D tensors inside the MoE block, so only the - router is mapped for sparse layers; the dense_* projections carry the - neuron hooks on mlp_only_layers / decoder_sparse_step dense layers.""" + router is mapped for sparse layers.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] - assert set(mlp.submodules.keys()) == {"gate", "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"} class TestQwen3MoeArchitectureGuards: 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 803d950fc..686174669 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 @@ -7,7 +7,9 @@ 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 MoEBridge def _make_bridge_cfg(**overrides): @@ -84,23 +86,18 @@ def test_mlp_path(self, adapter): def test_mlp_maps_router_shared_expert_and_dense_projections(self, adapter): """Qwen3NextSparseMoeBlock has 3D batched experts (delegated to HF), but - its router and shared expert are hookable, and the dense_* projections - carry neuron hooks on dense mlp_only_layers / decoder_sparse_step layers.""" + its router and shared expert are hookable.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] - assert set(mlp.submodules) == { + # dense_* are covered by the roster in test_moe_dense_dispatch.py. + assert set(mlp.submodules) - DENSE_KEYS == { "gate", "shared_expert", "shared_expert_gate", - "dense_gate", - "dense_in", - "dense_out", } assert all(sub.optional for sub in mlp.submodules.values()) def test_mlp_bridge_type(self, adapter): """Every real checkpoint is sparse MoE.""" - from transformer_lens.model_bridge.generalized_components import MoEBridge - mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen3_vl_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_vl_moe_adapter.py index d23a7142e..d7086dcbc 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_vl_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_vl_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, @@ -46,7 +49,8 @@ def test_moe_mlp(self, adapter): mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) - assert set(mlp.submodules) == {"gate", "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", "experts"} assert isinstance(mlp.submodules["gate"], MoERouterBridge) assert mlp.submodules["gate"].optional is True assert mlp.submodules["experts"].optional is True diff --git a/tests/unit/pretrained_weight_conversions/test_mixtral_conversion.py b/tests/unit/pretrained_weight_conversions/test_mixtral_conversion.py index e394f9297..4240c0f2c 100644 --- a/tests/unit/pretrained_weight_conversions/test_mixtral_conversion.py +++ b/tests/unit/pretrained_weight_conversions/test_mixtral_conversion.py @@ -136,16 +136,24 @@ def test_router_weights_come_from_the_moe_gate(hf_model, state_dict) -> None: ) -def test_quantized_expert_weights_are_refused(hf_model, tl_cfg) -> None: - """Slicing packed/scaled expert weights would silently drop their scales, so - the converter must refuse rather than emit plausible garbage.""" +@pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.uint8, "packed integer storage"), + # float8 reports is_floating_point=True, so a plain float check admits + # it — and it is the one family that slices without complaint. + (torch.float8_e4m3fn, "narrow float"), + ], +) +def test_quantized_expert_weights_are_refused(hf_model, tl_cfg, dtype, reason) -> None: + """Slicing packed or scale-separated expert weights would silently drop the + scales, so the converter must refuse rather than emit plausible garbage.""" experts = hf_model.model.layers[0].mlp.experts original = experts.gate_up_proj try: - experts.gate_up_proj = torch.nn.Parameter( - original.detach().to(torch.int8), requires_grad=False - ) - with pytest.raises(NotImplementedError, match="floating-point"): + experts.gate_up_proj = torch.nn.Parameter(original.detach().to(dtype), requires_grad=False) + with pytest.raises(NotImplementedError, match=reason): convert_mixtral_weights(hf_model, tl_cfg) finally: experts.gate_up_proj = original @@ -213,3 +221,27 @@ def test_routing_renormalization_matches_hf(hf_model, tl_cfg, state_dict) -> Non hf_out = hf_out[0] if isinstance(hf_out, tuple) else hf_out torch.testing.assert_close(tl_out, hf_out, atol=1e-5, rtol=1e-4) + + +@pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.float8_e4m3fn, "narrow float"), + ], +) +def test_quantized_router_weight_is_refused(hf_model, tl_cfg, dtype, reason) -> None: + """The router sits next to the guarded expert reads and had no guard. + + load_state_dict does not save it: a SAME-SHAPE int8/FP8 tensor is accepted + and silently cast to float32, so the router would land as garbage with no + error anywhere. Only the shape-changing 4-bit case is caught downstream. + """ + moe = hf_model.model.layers[0].mlp + original = moe.gate.weight + try: + moe.gate.weight = torch.nn.Parameter(original.detach().to(dtype), requires_grad=False) + with pytest.raises(NotImplementedError, match=reason): + convert_mixtral_weights(hf_model, tl_cfg) + finally: + moe.gate.weight = original diff --git a/tests/unit/pretrained_weight_conversions/test_olmoe.py b/tests/unit/pretrained_weight_conversions/test_olmoe.py new file mode 100644 index 000000000..9d98667fe --- /dev/null +++ b/tests/unit/pretrained_weight_conversions/test_olmoe.py @@ -0,0 +1,150 @@ +"""convert_olmoe_weights: the batched-expert split and its quantization guard. + +OLMoE stores all experts in two batched Parameters and the converter slices +them, so a packed or scale-separated tensor would be reshaped into a +plausible-looking matrix rather than raising. These guards existed but nothing +exercised them — removing both left the whole suite green. + +The model is built from a tiny config in memory: no download, no hub access. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F +from transformers import OlmoeConfig, OlmoeForCausalLM + +from transformer_lens.config import HookedTransformerConfig +from transformer_lens.pretrained.weight_conversions import convert_olmoe_weights + +D_MODEL, D_MLP, N_EXPERTS, N_LAYERS = 8, 16, 4, 1 +EXPERTS_PER_TOKEN = 2 + + +@pytest.fixture(scope="module") +def hf_model() -> OlmoeForCausalLM: + """Tiny OLMoE with AMPLIFIED weights. + + Same reason as the Mixtral fixture: SiLU is near-linear at small magnitudes, + so `act(gate) * up ~= act(up) * gate` and a swapped gate/up mapping reads as + noise at HF's default init. Do not lower without re-measuring the negative + control in test_expert_weights_reproduce_hf_expert_output. + """ + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + model = OlmoeForCausalLM( + OlmoeConfig( + hidden_size=D_MODEL, + intermediate_size=D_MLP, + num_hidden_layers=N_LAYERS, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=32, + num_experts=N_EXPERTS, + num_experts_per_tok=EXPERTS_PER_TOKEN, + max_position_embeddings=32, + ) + ).eval() + for param in model.parameters(): + torch.nn.init.normal_(param, std=0.3) + return model + + +@pytest.fixture(scope="module") +def tl_cfg() -> HookedTransformerConfig: + return HookedTransformerConfig( + d_model=D_MODEL, + d_head=4, + n_heads=2, + n_key_value_heads=1, + n_layers=N_LAYERS, + n_ctx=32, + d_vocab=32, + d_mlp=D_MLP, + num_experts=N_EXPERTS, + experts_per_token=EXPERTS_PER_TOKEN, + act_fn="silu", + normalization_type="RMS", + positional_embedding_type="rotary", + use_qk_norm=True, + ) + + +def test_expert_weights_reproduce_hf_expert_output(hf_model, tl_cfg) -> None: + """The decisive check on the fused-projection split. + + Shapes cannot catch a swapped gate/up — both halves are [d_mlp, d_model] — + so this compares numerically and carries a negative control proving the + swapped assignment would differ. + """ + state_dict = convert_olmoe_weights(hf_model, tl_cfg) + + with torch.random.fork_rng(devices=[]): + torch.manual_seed(1) + x = torch.randn(1, D_MODEL) + + experts = hf_model.model.layers[0].mlp.experts + for expert in range(N_EXPERTS): + gate_hf, up_hf = F.linear(x, experts.gate_up_proj[expert]).chunk(2, dim=-1) + expected = F.linear(F.silu(gate_hf) * up_hf, experts.down_proj[expert]) + + w_gate = state_dict[f"blocks.0.mlp.experts.{expert}.W_gate.weight"] + w_in = state_dict[f"blocks.0.mlp.experts.{expert}.W_in.weight"] + w_out = state_dict[f"blocks.0.mlp.experts.{expert}.W_out.weight"] + actual = F.linear(F.silu(F.linear(x, w_gate)) * F.linear(x, w_in), w_out) + torch.testing.assert_close(actual, expected) + + swapped = F.linear(F.silu(F.linear(x, w_in)) * F.linear(x, w_gate), w_out) + assert not torch.allclose(swapped, expected, atol=1e-6), ( + f"expert {expert}: gate and up are interchangeable in this fixture, " + "so the assertion above cannot detect a swapped mapping" + ) + + +@pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.uint8, "packed integer storage"), + # float8 reports is_floating_point=True, so a plain float check admits + # it — and it is the one family that slices without complaint. + (torch.float8_e4m3fn, "narrow float"), + (torch.float8_e5m2, "narrow float"), + ], +) +@pytest.mark.parametrize("tensor_name", ["gate_up_proj", "down_proj"]) +def test_quantized_expert_weights_are_refused(hf_model, tl_cfg, dtype, reason, tensor_name) -> None: + """Both batched expert tensors are sliced, so both must be guarded — + slicing either would silently drop the scales stored beside it.""" + experts = hf_model.model.layers[0].mlp.experts + original = getattr(experts, tensor_name) + try: + setattr( + experts, + tensor_name, + torch.nn.Parameter(original.detach().to(dtype), requires_grad=False), + ) + with pytest.raises(NotImplementedError, match=reason): + convert_olmoe_weights(hf_model, tl_cfg) + finally: + setattr(experts, tensor_name, original) + + +@pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.float8_e4m3fn, "narrow float"), + ], +) +def test_quantized_router_weight_is_refused(hf_model, tl_cfg, dtype, reason) -> None: + """Same unguarded router read as Mixtral's, caught by the same guard.""" + moe = hf_model.model.layers[0].mlp + original = moe.gate.weight + try: + moe.gate.weight = torch.nn.Parameter(original.detach().to(dtype), requires_grad=False) + with pytest.raises(NotImplementedError, match=reason): + convert_olmoe_weights(hf_model, tl_cfg) + finally: + moe.gate.weight = original diff --git a/tests/unit/pretrained_weight_conversions/test_openai.py b/tests/unit/pretrained_weight_conversions/test_openai.py index 25e0536a0..b0dc60b0a 100644 --- a/tests/unit/pretrained_weight_conversions/test_openai.py +++ b/tests/unit/pretrained_weight_conversions/test_openai.py @@ -220,6 +220,33 @@ def test_packed_mxfp4_experts_raise_legible_error(): convert_gpt_oss_weights(model, cfg) +@pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.uint8, "packed integer storage"), + # The one that a plain isinstance/is_floating_point check admits: FP8 + # reports is_floating_point=True and slices without complaint, so this + # converter used to emit scale-less garbage for it silently. + (torch.float8_e4m3fn, "narrow float"), + (torch.float8_e5m2, "narrow float"), + ], +) +@pytest.mark.parametrize("tensor_name", ["gate_up_proj", "down_proj"]) +def test_quantized_expert_weights_are_refused(dtype, reason, tensor_name): + """Both fused expert tensors are sliced, so both must be guarded — a packed + down_proj would silently drop its scales just as gate_up_proj would.""" + cfg = make_cfg() + model = make_mock_model(cfg) + experts = model.model.layers[0].mlp.experts + setattr(experts, tensor_name, getattr(experts, tensor_name).to(dtype)) + + with pytest.raises(NotImplementedError, match=reason) as excinfo: + convert_gpt_oss_weights(model, cfg) + # The gpt-oss remedy must survive the move onto the shared helper. + assert "Mxfp4Config(dequantize=True)" in str(excinfo.value) + + def test_sinks_are_converted(): """The learned attention-sink logits must reach the state dict (#1619 follow-up).""" cfg = make_cfg() diff --git a/tests/unit/test_loading_from_pretrained_utilities.py b/tests/unit/test_loading_from_pretrained_utilities.py index b0c948f94..c37c3a2f4 100644 --- a/tests/unit/test_loading_from_pretrained_utilities.py +++ b/tests/unit/test_loading_from_pretrained_utilities.py @@ -2,6 +2,7 @@ from unittest import mock import pytest +import torch from transformer_lens import HookedTransformer from transformer_lens.config import HookedTransformerConfig @@ -18,7 +19,9 @@ def get_default_config(): ) -def _config_with_architecture(architecture: str) -> HookedTransformerConfig: +def _config_with_architecture( + architecture: str, quantization_method: str | None = None +) -> HookedTransformerConfig: return HookedTransformerConfig( d_model=128, d_head=8, @@ -28,50 +31,254 @@ def _config_with_architecture(architecture: str) -> HookedTransformerConfig: d_vocab=50257, attn_only=True, original_architecture=architecture, + quantization_method=quantization_method, ) class TestMxfp4DequantizeConfig: - """Packed-MXFP4 gpt-oss checkpoints must load dequantized so the weight - converter sees plain torch.Tensors instead of triton-kernels wrappers (#1619).""" + """Packed-MXFP4 checkpoints must load dequantized so the weight converter + sees plain torch.Tensors instead of triton-kernels wrappers (#1619).""" - @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") - def test_mxfp4_gpt_oss_gets_dequantize_config(self, mock_auto_config: mock.MagicMock): - mock_auto_config.from_pretrained.return_value = SimpleNamespace( - quantization_config={"quant_method": "mxfp4"} - ) - result = _mxfp4_dequantize_config( - "openai/gpt-oss-20b", _config_with_architecture("GptOssForCausalLM"), None - ) + def test_mxfp4_gets_dequantize_config(self): + result = _mxfp4_dequantize_config(_config_with_architecture("GptOssForCausalLM", "mxfp4")) + assert result is not None + assert result.dequantize is True + + def test_applies_beyond_gpt_oss(self): + """The architecture gate this used to carry existed only to dodge a + second AutoConfig fetch. MXFP4 Qwen3-MoE checkpoints are in the registry + and pack their experts the same way, so they must dequantize too.""" + result = _mxfp4_dequantize_config(_config_with_architecture("Qwen3MoeForCausalLM", "mxfp4")) assert result is not None assert result.dequantize is True + def test_unquantized_finetune_untouched(self): + assert _mxfp4_dequantize_config(_config_with_architecture("GptOssForCausalLM")) is None + + def test_other_quantizations_are_not_dequantized_here(self): + """Only MXFP4 has a dequantize-on-load path; the rest are refused later + by the converter guards, which give a better-targeted message.""" + for method in ("bitsandbytes", "gptq", "awq", "fp8"): + assert ( + _mxfp4_dequantize_config(_config_with_architecture("LlamaForCausalLM", method)) + is None + ) + @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") - def test_object_style_quantization_config(self, mock_auto_config: mock.MagicMock): - mock_auto_config.from_pretrained.return_value = SimpleNamespace( - quantization_config=SimpleNamespace(quant_method="mxfp4") + def test_costs_no_hub_round_trip(self, mock_auto_config: mock.MagicMock): + """The whole point of carrying the method on the cfg: this runs on every + load, and a fetch here would be a Hub HEAD request per model load.""" + _mxfp4_dequantize_config(_config_with_architecture("GptOssForCausalLM", "mxfp4")) + mock_auto_config.from_pretrained.assert_not_called() + + +class TestUnsupportedQuantizationRefusal: + """The ordinary `from_pretrained("")` path must refuse quantized + weights too. + + The older refusal lived inside `if hf_model is not None:`, so it fired only + when a user handed TL a pre-loaded model. On the internal-load path nothing + stopped a quantized checkpoint reaching converters that slice `.weight` + directly — and same-shape int8/FP8 survives both the converter *and* + `load_state_dict`, which casts it to float32, so an int8 code of 107 lands + as the weight 107.0 with no error anywhere. + """ + + @staticmethod + def _model(dtype=None, quant_method="gptq"): + class _Tiny(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q_proj = torch.nn.Linear(4, 4, bias=False) + if dtype is not None: + self.q_proj.weight = torch.nn.Parameter( + torch.zeros(4, 4, dtype=dtype), requires_grad=False + ) + + model = _Tiny() + quantization_config = {"quant_method": quant_method} if quant_method is not None else None + model.config = SimpleNamespace(quantization_config=quantization_config) + return model + + @pytest.mark.parametrize( + "dtype", + [torch.int8, torch.uint8, torch.float8_e4m3fn], + ) + def test_quantized_storage_is_refused(self, dtype): + from transformer_lens.loading_from_pretrained import ( + _refuse_unsupported_quantization, ) - result = _mxfp4_dequantize_config( - "openai/gpt-oss-20b", _config_with_architecture("GptOssForCausalLM"), None + + cfg = _config_with_architecture("LlamaForCausalLM") + with pytest.raises(NotImplementedError, match="gptq"): + _refuse_unsupported_quantization(cfg, self._model(dtype)) + + def test_unquantized_model_passes(self): + """The positive control — every ordinary load goes through this.""" + from transformer_lens.loading_from_pretrained import ( + _refuse_unsupported_quantization, ) - assert result is not None - assert result.dequantize is True + + cfg = _config_with_architecture("LlamaForCausalLM") + _refuse_unsupported_quantization(cfg, self._model(quant_method=None)) + + def test_dequantized_checkpoint_still_loads(self): + """A model loaded with dequantize=True keeps advertising its original + quant_method while holding real bf16 tensors. Refusing on the + declaration alone would break the MXFP4 auto-dequantize path.""" + from transformer_lens.loading_from_pretrained import ( + _refuse_unsupported_quantization, + ) + + cfg = _config_with_architecture("GptOssForCausalLM", "mxfp4") + _refuse_unsupported_quantization(cfg, self._model(quant_method="mxfp4")) + + def test_bitsandbytes_4bit_llama_flow_is_preserved(self): + """The one supported quantized HT flow: weight conversion and + abstract_attention's matmul_4bit both depend on it.""" + from transformer_lens.loading_from_pretrained import ( + _refuse_unsupported_quantization, + ) + + cfg = _config_with_architecture("LlamaForCausalLM") + cfg.load_in_4bit = True + _refuse_unsupported_quantization(cfg, self._model(torch.uint8, quant_method="bitsandbytes")) + + @mock.patch("transformer_lens.loading_from_pretrained.AutoModelForCausalLM") + def test_refusal_is_wired_into_the_internal_load_path(self, mock_auto_model: mock.MagicMock): + """The wiring, not just the helper. + + This is the whole point of the finding: the refusal must fire on + `from_pretrained("")`, where TL loads the model itself and the + caller never sees an hf_model to be checked against. + """ + from transformer_lens.loading_from_pretrained import get_pretrained_state_dict + + mock_auto_model.from_pretrained.return_value = self._model(torch.int8) + + cfg = _config_with_architecture("LlamaForCausalLM") + with pytest.raises(NotImplementedError, match="gptq"): + get_pretrained_state_dict("meta-llama/Llama-2-7b-hf", cfg) + + +class TestQuantizationMethodCapture: + """`convert_hf_model_config` must record the checkpoint's quant_method while + the HF config is in hand — that capture is what lets the loader act on the + quantization without refetching.""" + + @pytest.mark.parametrize( + "quantization_config", + [ + {"quant_method": "mxfp4"}, + SimpleNamespace(quant_method="mxfp4"), + ], + ids=["dict-style", "object-style"], + ) + def test_capture_handles_both_config_shapes(self, quantization_config): + from transformer_lens.utilities.quantization import quantization_method + + assert ( + quantization_method(SimpleNamespace(quantization_config=quantization_config)) == "mxfp4" + ) + + def test_unquantized_config_captures_none(self): + from transformer_lens.utilities.quantization import quantization_method + + assert quantization_method(SimpleNamespace()) is None + assert quantization_method(None) is None + + @staticmethod + def _gpt2_shaped_config(**extra): + """The attributes convert_hf_model_config's GPT2 branch reads, so the + real function can run without touching the Hub.""" + return SimpleNamespace( + architectures=["GPT2LMHeadModel"], + n_embd=128, + n_head=4, + n_layer=2, + n_ctx=1024, + layer_norm_epsilon=1e-5, + vocab_size=50257, + activation_function="gelu_new", + scale_attn_by_inverse_layer_idx=False, + **extra, + ) + + @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") + def test_convert_hf_model_config_records_the_method(self, mock_auto_config: mock.MagicMock): + """The wiring, not just the extractor: without this the loader would + silently stop dequantizing MXFP4 and the converter would raise instead.""" + from transformer_lens.loading_from_pretrained import convert_hf_model_config + + mock_auto_config.from_pretrained.return_value = self._gpt2_shaped_config( + quantization_config={"quant_method": "mxfp4"} + ) + assert convert_hf_model_config("gpt2")["quantization_method"] == "mxfp4" + + @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") + def test_convert_hf_model_config_records_none_when_unquantized( + self, mock_auto_config: mock.MagicMock + ): + from transformer_lens.loading_from_pretrained import convert_hf_model_config + + mock_auto_config.from_pretrained.return_value = self._gpt2_shaped_config() + assert convert_hf_model_config("gpt2")["quantization_method"] is None @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") - def test_unquantized_gpt_oss_finetune_untouched(self, mock_auto_config: mock.MagicMock): - mock_auto_config.from_pretrained.return_value = SimpleNamespace() - result = _mxfp4_dequantize_config( - "someone/gpt-oss-finetune-bf16", _config_with_architecture("GptOssForCausalLM"), None + def test_user_supplied_hf_model_config_wins(self, mock_auto_config: mock.MagicMock): + """A user passing hf_model= says how the weights in hand are stored, + which can differ from the Hub repo's declaration (they may have loaded + it dequantized, or quantized one that ships in bf16).""" + mock_auto_config.from_pretrained.return_value = self._gpt2_shaped_config() + + cfg = get_pretrained_model_config( + "gpt2", hf_cfg={"quantization_config": {"quant_method": "bitsandbytes"}} ) - assert result is None + assert cfg.quantization_method == "bitsandbytes" @mock.patch("transformer_lens.loading_from_pretrained.AutoConfig") - def test_non_gpt_oss_skips_config_fetch(self, mock_auto_config: mock.MagicMock): - result = _mxfp4_dequantize_config( - "gpt2", _config_with_architecture("GPT2LMHeadModel"), None + def test_unquantized_hf_model_does_not_erase_the_repo_method( + self, mock_auto_config: mock.MagicMock + ): + """hf_cfg is often present but carries no quantization_config at all; + that absence must not wipe what the repo config declared.""" + mock_auto_config.from_pretrained.return_value = self._gpt2_shaped_config( + quantization_config={"quant_method": "mxfp4"} ) - assert result is None - mock_auto_config.from_pretrained.assert_not_called() + + cfg = get_pretrained_model_config("gpt2", hf_cfg={"vocab_size": 50257}) + assert cfg.quantization_method == "mxfp4" + + @mock.patch("transformer_lens.loading_from_pretrained.convert_neel_model_config") + def test_config_builders_that_never_see_an_hf_config(self, mock_neel: mock.MagicMock): + """Not every cfg_dict comes from convert_hf_model_config. + + `convert_neel_model_config` builds one from the model name alone, so the + quantization_method key is simply absent — reading it with [] raised + KeyError and broke every NeelNanda load. + """ + mock_neel.return_value = { + "d_model": 128, + "d_head": 8, + "n_heads": 16, + "n_ctx": 128, + "n_layers": 1, + "d_vocab": 50257, + "attn_only": True, + "original_architecture": "neel", + } + + cfg = get_pretrained_model_config("NeelNanda/SoLU_2L512W_C4_Code", hf_cfg={}) + assert cfg.quantization_method is None + + def test_name_based_architecture_branches_capture_none(self): + """Llama/gemma names never fetch a config, so nothing is there to read. + Pinned because the field must be present-and-None rather than missing — + HookedTransformerConfig.from_dict passes cfg_dict through unfiltered.""" + from transformer_lens.loading_from_pretrained import convert_hf_model_config + + cfg_dict = convert_hf_model_config("llama-7b-hf") + assert cfg_dict["quantization_method"] is None # Successes diff --git a/tests/unit/utilities/test_quantization_guards.py b/tests/unit/utilities/test_quantization_guards.py new file mode 100644 index 000000000..bfe318534 --- /dev/null +++ b/tests/unit/utilities/test_quantization_guards.py @@ -0,0 +1,245 @@ +"""Guards against reading quantized weights as if they were plain matrices. + +TransformerLens supports quantized *forward* passes — the wrapped HF module +dequantizes internally. What it cannot do is arithmetic on the stored weight: +reshaping it per head, slicing a fused projection, folding LayerNorm into it. +Those reads must fail loudly rather than return packed bytes or scale-less +narrow floats, which look like ordinary tensors and produce wrong numbers. + +Everything here is built in memory: bitsandbytes is an optional extra and is not +installed in CI, so the quantized shapes are simulated. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from transformer_lens.utilities.quantization import ( + describe_quantization, + require_readable_weight, + unreadable_weight_reason, +) + +# The shapes a quantized weight actually takes in this stack, with the reason +# fragment each should produce. float8 is the important one: it reports +# is_floating_point=True, so a plain float check admits it — and it is the only +# family that survives both slicing and nn.Parameter() without complaint. +UNREADABLE = [ + pytest.param(torch.zeros(4, 4, dtype=torch.int8), "packed integer storage", id="bnb-int8"), + pytest.param(torch.zeros(8, 1, dtype=torch.uint8), "packed integer storage", id="packed-uint8"), + pytest.param(torch.zeros(4, 4, dtype=torch.int32), "packed integer storage", id="gptq-int32"), + pytest.param(torch.zeros(4, 4, dtype=torch.float8_e4m3fn), "narrow float", id="fp8-e4m3fn"), + pytest.param(torch.zeros(4, 4, dtype=torch.float8_e5m2), "narrow float", id="fp8-e5m2"), +] + +# Legitimate weights, including the dtypes a quantized *forward* pass uses for +# its compute. A guard that fires on any of these is a regression. +READABLE = [ + pytest.param(torch.zeros(4, 4, dtype=torch.float32), id="fp32"), + pytest.param(torch.zeros(4, 4, dtype=torch.bfloat16), id="bf16"), + pytest.param(torch.zeros(4, 4, dtype=torch.float16), id="fp16"), + pytest.param(torch.zeros(4, 4, dtype=torch.float64), id="fp64"), +] + + +class _PackedParams(torch.Tensor): + """Stands in for bitsandbytes Params4bit: a Tensor SUBCLASS holding packed + bytes, so isinstance(w, Tensor) is True and only the dtype gives it away.""" + + +class _TritonWrapper: + """Stands in for the MXFP4 triton-kernels object — named Tensor upstream, + which is what makes the resulting TypeError look like a torch bug.""" + + dtype = "packed" + + +@pytest.mark.parametrize("weight,reason", UNREADABLE) +def test_quantized_storage_is_rejected(weight, reason) -> None: + assert reason in (unreadable_weight_reason(weight) or "") + with pytest.raises(NotImplementedError, match=reason): + require_readable_weight(weight, operation="read W_Q") + + +@pytest.mark.parametrize("weight", READABLE) +def test_real_float_weights_pass(weight) -> None: + """The positive control: these are what a working model (including a + quantized forward pass's compute dtype) presents.""" + assert unreadable_weight_reason(weight) is None + assert require_readable_weight(weight, operation="read W_Q") is weight + + +def test_tensor_subclass_with_packed_storage_is_rejected() -> None: + """bitsandbytes params pass isinstance(w, torch.Tensor), so the dtype branch + is the only thing standing between them and a silent reshape.""" + packed = _PackedParams(torch.zeros(8, 1, dtype=torch.uint8)) + assert isinstance(packed, torch.Tensor) + with pytest.raises(NotImplementedError, match="packed integer storage"): + require_readable_weight(packed, operation="read W_Q") + + +def test_non_tensor_wrapper_is_rejected() -> None: + with pytest.raises(NotImplementedError, match="not a torch.Tensor"): + require_readable_weight(_TritonWrapper(), operation="read W_Q") + + +def test_meta_device_is_reported_as_a_load_problem_not_quantization() -> None: + """A meta tensor is float and full-width, so it passes the dtype checks; it + must be named as an unmaterialized load rather than sent after a + quantization that is not there.""" + with pytest.raises(NotImplementedError) as excinfo: + require_readable_weight(torch.zeros(4, 4, device="meta"), operation="read W_Q") + assert "meta device" in str(excinfo.value) + assert "quantization:" not in str(excinfo.value) + + +def test_error_names_the_quantization_method() -> None: + class _Owner: + class config: + class quantization_config: + quant_method = "bitsandbytes" + + with pytest.raises(NotImplementedError, match="bitsandbytes"): + require_readable_weight( + torch.zeros(4, 4, dtype=torch.int8), operation="read W_Q", owner=_Owner() + ) + assert describe_quantization(_Owner()) == "bitsandbytes" + + +def test_quantization_method_falls_back_when_unknowable() -> None: + assert describe_quantization(nn.Linear(2, 2)) == "an unknown quantization" + + +class TestBridgeAccessorsAreGuarded: + """The accessors are the most dangerous site: they returned packed bytes + with no error, so downstream analyses decomposed garbage.""" + + @staticmethod + def _bridge(): + from tests.unit.model_bridge.supported_architectures.helpers import ( + make_bridge_cfg, + ) + from transformer_lens.factories.architecture_adapter_factory import ( + ArchitectureAdapterFactory, + ) + from transformer_lens.model_bridge.component_setup import setup_submodules + + cfg = make_bridge_cfg("LlamaForCausalLM", d_model=8, n_heads=2, d_head=4) + adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg) + return adapter, setup_submodules + + @pytest.mark.parametrize("dtype", [torch.int8, torch.uint8, torch.float8_e4m3fn]) + def test_attention_accessor_refuses_quantized_weight(self, dtype) -> None: + import copy + + adapter, setup = self._bridge() + attn_template = adapter.component_mapping["blocks"].submodules["attn"] + + class _QuantAttn(nn.Module): + def __init__(self) -> None: + super().__init__() + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + linear = nn.Linear(8, 8, bias=False) + linear.weight = nn.Parameter( + torch.zeros(8, 8, dtype=dtype), requires_grad=False + ) + setattr(self, name, linear) + + module = _QuantAttn() + bridge = copy.deepcopy(attn_template) + bridge.set_original_component(module) + setup(bridge, adapter, module) + + with pytest.raises(NotImplementedError, match="W_Q"): + _ = bridge.W_Q + + +class TestBitsandbytesParameterSubclassIsNamed: + """bitsandbytes' Params4bit and Int8Params are nn.Parameter SUBCLASSES. + + The class-name fallback used to be gated behind `not isinstance(weight, + nn.Parameter)`, which excluded exactly the classes it was written to + identify. transformers keys off the same two class names. + """ + + class _Params4bit(torch.nn.Parameter): + """Stands in for bitsandbytes; the real package is an optional extra.""" + + def test_parameter_subclass_is_named(self) -> None: + owner = nn.Linear(2, 2) + owner.weight = self._Params4bit(torch.zeros(2, 2)) + assert describe_quantization(owner) == "_Params4bit" + + def test_plain_parameter_stays_unknown(self) -> None: + """The negative control: an ordinary weight must not be reported as a + quantization method named 'Parameter'.""" + assert describe_quantization(nn.Linear(2, 2)) == "an unknown quantization" + + def test_hf_config_still_wins(self) -> None: + class _Owner: + class config: + class quantization_config: + quant_method = "bitsandbytes" + + owner = _Owner() + owner.weight = self._Params4bit(torch.zeros(2, 2)) + assert describe_quantization(owner) == "bitsandbytes" + + +class TestProcessWeightsScansBatchedMoEParameters: + """`process_weights` must see batched-MoE expert tensors. + + Its guard filtered `named_parameters()` with `name.endswith(".weight")`, + but on transformers 5.x every batched-MoE family (Mixtral, OLMoE, gpt-oss) + stores experts as Parameters named `mlp.experts.gate_up_proj` / + `down_proj` — no `.weight` suffix — so the tensors the converters were + guarded for were the exact ones this filter skipped. + + Calls the unbound method against a stub carrying only `original_model`: + the guard is the first statement in the body, so the raise happens before + anything else is touched. The no-raise direction needs no stub — every + bridge test in the suite runs process_weights on float weights. + """ + + @staticmethod + def _stub(dtype): + class _Experts(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + # Batched expert layout: a Parameter that is NOT named *.weight. + self.gate_up_proj = torch.nn.Parameter( + torch.zeros(2, 8, 4, dtype=dtype), requires_grad=False + ) + + class _Stub: + original_model = _Experts() + + return _Stub() + + @pytest.mark.parametrize( + "dtype,reason", + [ + (torch.int8, "packed integer storage"), + (torch.uint8, "packed integer storage"), + (torch.float8_e4m3fn, "narrow float"), + ], + ) + def test_batched_expert_parameters_are_scanned(self, dtype, reason) -> None: + from transformer_lens.model_bridge.bridge import TransformerBridge + + with pytest.raises(NotImplementedError, match=reason): + TransformerBridge.process_weights(self._stub(dtype)) + + def test_float_batched_experts_pass_the_guard(self) -> None: + """Positive control on the guard itself: a float batched Parameter must + not trip it (it then fails later, on the real processing this stub + cannot supply — which is why only the guard's verdict is asserted).""" + from transformer_lens.model_bridge.bridge import TransformerBridge + + with pytest.raises(Exception) as excinfo: + TransformerBridge.process_weights(self._stub(torch.float32)) + assert not isinstance(excinfo.value, NotImplementedError) or "cannot" not in str( + excinfo.value + ) diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py index f46b697bd..7184426b5 100644 --- a/transformer_lens/HookedTransformer.py +++ b/transformer_lens/HookedTransformer.py @@ -1369,6 +1369,18 @@ def from_pretrained( assert ( qc.get("quant_method", "") == "bitsandbytes" ), "Only bitsandbytes quantization is supported" + elif quant_method: + # Anything other than the supported bitsandbytes 4-bit Llama + # flow reaches the converters, which slice `.weight` directly: + # packed or scale-separated storage yields wrong numbers rather + # than an error. Refuse instead. + raise NotImplementedError( + f"HookedTransformer cannot convert a {quant_method!r}-quantized " + "checkpoint: its weight converters read weights directly, and " + "packed or scale-separated storage would silently produce wrong " + "values. Load the model dequantized, or use TransformerBridge " + "for a quantized forward pass." + ) else: hf_cfg = {} diff --git a/transformer_lens/config/hooked_transformer_config.py b/transformer_lens/config/hooked_transformer_config.py index 26756cb67..75bee9156 100644 --- a/transformer_lens/config/hooked_transformer_config.py +++ b/transformer_lens/config/hooked_transformer_config.py @@ -173,6 +173,11 @@ class HookedTransformerConfig(TransformerLensConfig): We need this information to dynamically control bos prepending. load_in_4bit(bool): If this flag is set, then it's assumed that parameters are 4-bit quantized with bitsandbytes. Currently only supported for Llama. + quantization_method (str, *optional*): the ``quant_method`` declared by the checkpoint's + HF config ("mxfp4", "bitsandbytes", "gptq", ...), captured while that config is already + in hand so later load steps need not refetch it. None when unquantized, and also when + the config was never fetched (the llama/gemma name-based branches of + ``convert_hf_model_config`` infer the architecture from the model name instead). n_key_value_heads (int, *optional*): The number of groups of heads that use the same key and value matrix. Only for models that use Grouped Query Attention. post_embedding_ln (bool): Whether to apply layer normalization after embedding the tokens. Defaults @@ -287,6 +292,7 @@ class HookedTransformerConfig(TransformerLensConfig): trust_remote_code: bool = False rotary_adjacent_pairs: bool = False load_in_4bit: bool = False + quantization_method: Optional[str] = None num_experts: Optional[int] = None experts_per_token: Optional[int] = None relative_attention_max_distance: Optional[int] = None diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py index 22127cb2a..4281200f5 100644 --- a/transformer_lens/loading_from_pretrained.py +++ b/transformer_lens/loading_from_pretrained.py @@ -59,6 +59,10 @@ ) from transformer_lens.supported_models import MODEL_ALIASES, OFFICIAL_MODEL_NAMES from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config +from transformer_lens.utilities.quantization import ( + quantization_method, + unreadable_weight_reason, +) NON_HF_HOSTED_MODEL_NAMES = [ "llama-7b-hf", @@ -139,7 +143,9 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]: else: official_model_name = get_official_model_name(model_name) - # Load HuggingFace model config + # Load HuggingFace model config. Stays None on the name-based branches + # below, which infer the architecture from the model name and never fetch. + hf_config: Any = None if "llama" in official_model_name.lower(): architecture = "LlamaForCausalLM" elif "gemma-3" in official_model_name.lower() or "medgemma" in official_model_name.lower(): @@ -1592,6 +1598,9 @@ def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]: raise NotImplementedError(f"{architecture} is not currently supported.") # All of these models use LayerNorm cfg_dict["original_architecture"] = architecture + # Carried on the cfg so the loader can act on the quantization without a + # second AutoConfig fetch (which would be a Hub round trip per load). + cfg_dict["quantization_method"] = quantization_method(hf_config) # The name such that AutoTokenizer.from_pretrained works cfg_dict["tokenizer_name"] = official_model_name if kwargs.get("trust_remote_code", False): @@ -1788,6 +1797,13 @@ def get_pretrained_model_config( if hf_cfg is not None: cfg_dict["load_in_4bit"] = hf_cfg.get("quantization_config", {}).get("load_in_4bit", False) + # A user-supplied hf_model is the more authoritative source: it says how + # the weights in hand are actually stored, not how the Hub repo declares + # them. .get, not []: convert_neel_model_config builds cfg_dict without + # ever seeing an HF config, so the key need not be there. + cfg_dict["quantization_method"] = quantization_method(hf_cfg) or cfg_dict.get( + "quantization_method" + ) cfg_dict["d_vocab"] = hf_cfg.get("vocab_size", cfg_dict["d_vocab"]) if cfg_dict["original_architecture"] == "Qwen2ForCausalLM": rope_params = hf_cfg.get("rope_parameters", {}) or {} @@ -1875,31 +1891,79 @@ def get_checkpoint_labels(model_name: str, **kwargs: Any) -> tuple[list[int], st # %% Loading state dicts -def _mxfp4_dequantize_config( - official_model_name: str, - cfg: HookedTransformerConfig, - token: str | None, -) -> Any | None: - """Return ``Mxfp4Config(dequantize=True)`` for packed-MXFP4 gpt-oss checkpoints. - - The gpt-oss weight converter slices expert tensors, which only works on - materialized torch.Tensors — packed MXFP4 weights stay wrapped in - triton-kernels objects. Returns None for anything else, including - already-dequantized gpt-oss finetunes. +def _mxfp4_dequantize_config(cfg: HookedTransformerConfig) -> Any | None: + """Return ``Mxfp4Config(dequantize=True)`` for packed-MXFP4 checkpoints. + + Weight converters slice expert tensors, which only works on materialized + torch.Tensors — packed MXFP4 weights stay wrapped in triton-kernels objects. + Returns None for anything else, including already-dequantized finetunes. + + Reads the method off ``cfg``, captured when the HF config was already loaded, + so this costs nothing. It used to short-circuit on + ``original_architecture == "GptOssForCausalLM"`` purely to avoid a second + AutoConfig fetch; dropping that reaches MXFP4 checkpoints of other + architectures, and Qwen3-MoE ones are in the registry. + + One blind spot, by construction: ``convert_hf_model_config`` infers llama + and gemma from the model *name* and never fetches a config for them, so + ``cfg.quantization_method`` is always None there and an MXFP4 checkpoint + under such a name will not auto-dequantize. It is refused rather than + mis-read — ``_refuse_unsupported_quantization`` inspects the loaded model + instead of the cfg for exactly this reason. """ - if cfg.original_architecture != "GptOssForCausalLM": - return None - hf_cfg = AutoConfig.from_pretrained(official_model_name, token=token) - quant_cfg = getattr(hf_cfg, "quantization_config", None) - if isinstance(quant_cfg, dict): - quant_method = quant_cfg.get("quant_method") - else: - quant_method = getattr(quant_cfg, "quant_method", None) - if quant_method != "mxfp4": + if cfg.quantization_method != "mxfp4": return None return Mxfp4Config(dequantize=True) +def _refuse_unsupported_quantization(cfg: HookedTransformerConfig, hf_model: Any) -> None: + """Refuse a quantized checkpoint before the weight converters read it. + + The converters slice ``.weight`` directly, and only three of them guard the + read. Same-shape storage is the dangerous case: an int8 or FP8 weight + converts silently *and* survives ``load_state_dict``, which casts it to + float32 — an int8 code of 107 lands as the weight 107.0. Packed 4-bit is + caught late by a shape mismatch, and GPTQ/AWQ happen to fail loudly only + because their ``QuantLinear`` has no ``.weight`` at all. + + Reads the *loaded model's* config rather than ``cfg.quantization_method``: + the latter is populated from ``convert_hf_model_config``, which infers + llama and gemma from the model name and never fetches a config for them — + exactly the family where bitsandbytes is most common. + """ + if hf_model is None: + return + method = quantization_method(getattr(hf_model, "config", None)) + if method is None: + return + # The one supported quantized HookedTransformer flow (weight conversion and + # abstract_attention's matmul_4bit both handle it). + if cfg.load_in_4bit and method == "bitsandbytes": + return + # Refuse on the stored weights, not the declaration: a checkpoint loaded + # with dequantize=True still advertises its original quant_method while + # holding perfectly readable bf16 tensors. Meta params are skipped — an + # offloaded load is a different problem with a different message. + offender = next( + ( + (name, unreadable_weight_reason(param)) + for name, param in hf_model.named_parameters() + if param.device.type != "meta" and unreadable_weight_reason(param) is not None + ), + None, + ) + if offender is None: + return + name, reason = offender + raise NotImplementedError( + f"HookedTransformer cannot convert this {method!r}-quantized checkpoint: " + f"{name} cannot be read because {reason}. The weight converters read " + "weights directly, so packed or scale-separated storage silently " + "produces wrong values. Load the model dequantized, or use " + "TransformerBridge for a quantized forward pass." + ) + + def get_pretrained_state_dict( official_model_name: str, cfg: HookedTransformerConfig, @@ -2021,11 +2085,7 @@ def get_pretrained_state_dict( ) else: if "quantization_config" not in kwargs: - mxfp4_dequantize = _mxfp4_dequantize_config( - official_model_name, - cfg, - huggingface_token if len(huggingface_token) > 0 else None, - ) + mxfp4_dequantize = _mxfp4_dequantize_config(cfg) if mxfp4_dequantize is not None: kwargs = {**kwargs, "quantization_config": mxfp4_dequantize} # Older models may lack pad_token_id (required in newer transformers) @@ -2058,6 +2118,8 @@ def get_pretrained_state_dict( for param in hf_model.parameters(): param.requires_grad = False + _refuse_unsupported_quantization(cfg, hf_model) + if cfg.original_architecture == "GPT2LMHeadModel": state_dict = convert_gpt2_weights(hf_model, cfg) elif cfg.original_architecture == "GPTNeoForCausalLM": diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 8934bb98a..cfb9184a5 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -57,6 +57,7 @@ from transformer_lens.utilities.aliases import resolve_alias from transformer_lens.utilities.devices import move_to_and_update_config from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss +from transformer_lens.utilities.quantization import require_readable_weight if TYPE_CHECKING: from transformer_lens.ActivationCache import ActivationCache @@ -1039,6 +1040,28 @@ def process_weights( fold_value_biases: Fold value biases into output bias. Default: True refactor_factored_attn_matrices: Experimental QK/OV factorization. Default: False """ + # Folding and centering do arithmetic on raw weights, so packed or + # scale-separated storage would produce silent garbage. The forward + # path stays usable when quantized; only this transformation does not. + for name, param in self.original_model.named_parameters(): + # No name filter: modern batched-MoE experts are Parameters that do + # NOT end in .weight (mlp.experts.gate_up_proj), and those are the + # very tensors the converters had to guard. unreadable_weight_reason + # is dtype-driven, so every full-width float parameter still passes. + # Routed through the shared helper so meta gets its own "load with + # real weights" message instead of being silently skipped — folding + # on meta tensors yields meta tensors, which is the same + # silent-garbage failure this guard exists to stop. + require_readable_weight( + param, + operation=f"process weights ({name})", + owner=self.original_model, + remedy=( + "Load the model dequantized, or use the bridge without weight " + "processing (enable_compatibility_mode(no_processing=True))." + ), + ) + # A failed or partial processing attempt is no longer guaranteed to retain # the raw HuggingFace basis, so invalidate that contract before any work. self._weights_processed = True diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index 9a79c3e2a..babc92d1c 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -22,6 +22,7 @@ GeneralizedComponent, ) from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config +from transformer_lens.utilities.quantization import require_readable_weight class AttentionBridge(GeneralizedComponent): @@ -331,6 +332,12 @@ def revert(self, input_value, *full_context): if hasattr(self, "k") and self.k is not None and hasattr(self.k, "hook_out"): k_reshape = ReshapeForAttentionHeads(n_kv_heads, d_head) self.k.hook_out.hook_conversion = k_reshape + # Subclasses that de-alias hook_k onto their own HookPoint (K-scaling + # architectures) need the same conversion, whichever order runs first. + if "hook_k" not in self.hook_aliases and isinstance( + getattr(self, "hook_k", None), HookPoint + ): + self.hook_k.hook_conversion = k_reshape if hasattr(self, "v") and self.v is not None and hasattr(self.v, "hook_out"): v_reshape = ReshapeForAttentionHeads(n_kv_heads, d_head) self.v.hook_out.hook_conversion = v_reshape @@ -616,7 +623,11 @@ def _project_per_head_qkv( """ component = linear_bridge.original_component assert component is not None, "LinearBridge.original_component not set" - weight = component.weight + weight = require_readable_weight( + component.weight, + operation="project per head (use_split_qkv_input / use_attn_in)", + owner=component, + ) bias = component.bias w3d = einops.rearrange( weight, @@ -767,7 +778,9 @@ 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].""" - weight = self.q.weight + weight = require_readable_weight( + self.q.weight, operation=f"read W_Q from {self.name}", owner=self.q + ) 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) @@ -777,7 +790,9 @@ def W_Q(self) -> torch.Tensor: @property def W_K(self) -> torch.Tensor: """Get W_K in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA).""" - weight = self.k.weight + weight = require_readable_weight( + self.k.weight, operation=f"read W_K from {self.name}", owner=self.k + ) if weight.ndim == 2 and self.config is not None: return self._reshape_weight_to_3d( weight, @@ -789,7 +804,9 @@ def W_K(self) -> torch.Tensor: @property def W_V(self) -> torch.Tensor: """Get W_V in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA).""" - weight = self.v.weight + weight = require_readable_weight( + self.v.weight, operation=f"read W_V from {self.name}", owner=self.v + ) if weight.ndim == 2 and self.config is not None: return self._reshape_weight_to_3d( weight, @@ -801,7 +818,9 @@ def W_V(self) -> torch.Tensor: @property def W_O(self) -> torch.Tensor: """Get W_O in 3D format [n_heads, d_head, d_model].""" - weight = self.o.weight + weight = require_readable_weight( + self.o.weight, operation=f"read W_O from {self.name}", owner=self.o + ) if weight.ndim == 2 and self.config is not None: return self._reshape_weight_to_3d( weight, diff --git a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py index 79573e7a4..898ed5c72 100644 --- a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py @@ -14,6 +14,7 @@ resolve_activation_fn, ) from transformer_lens.model_bridge.generalized_components.linear import LinearBridge +from transformer_lens.utilities.quantization import require_readable_weight class JointGateUpMLPBridge(GatedMLPBridge): @@ -78,7 +79,13 @@ def _default_split_gate_up( original_mlp_component: Any, ) -> tuple[torch.nn.Module, torch.nn.Module]: """Split gate_up_proj [2*d_mlp, d_model] into (gate, up) nn.Linear modules.""" - fused_weight = original_mlp_component.gate_up_proj.weight + # Guard before the split: float8 slices and survives nn.Parameter() + # silently, producing scale-less projections with no error anywhere. + fused_weight = require_readable_weight( + original_mlp_component.gate_up_proj.weight, + operation="split a fused gate/up projection at boot", + owner=original_mlp_component.gate_up_proj, + ) gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0) d_model = fused_weight.shape[1] d_mlp = gate_w.shape[0] diff --git a/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py b/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py index 2490c3ee3..438437474 100644 --- a/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py +++ b/transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py @@ -18,6 +18,7 @@ GeneralizedComponent, ) from transformer_lens.model_bridge.generalized_components.linear import LinearBridge +from transformer_lens.utilities.quantization import require_readable_weight class JointQKVAttentionBridge(AttentionBridge): @@ -230,8 +231,14 @@ def _default_split_qkv_matrix( qkv_component = getattr(original_attention_component, qkv_name) - qkv_weights = qkv_component.weight - assert isinstance(qkv_weights, torch.Tensor) + # Before the tensor_split/nn.Parameter below: int8 and uint8 would die + # there with an opaque "Only Tensors of floating point ... can require + # gradients", and float8 would split SILENTLY into scale-less pieces. + qkv_weights = require_readable_weight( + qkv_component.weight, + operation="split a fused QKV projection at boot", + owner=qkv_component, + ) # Original qkv_weights shape: [d_model, 3 * d_model] # Split into three equal parts along dimension 1 to get Q, K, V weights diff --git a/transformer_lens/model_bridge/generalized_components/mlp.py b/transformer_lens/model_bridge/generalized_components/mlp.py index a56e40cf1..387888b6a 100644 --- a/transformer_lens/model_bridge/generalized_components/mlp.py +++ b/transformer_lens/model_bridge/generalized_components/mlp.py @@ -9,6 +9,7 @@ from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) +from transformer_lens.utilities.quantization import require_readable_weight def weight_layout_in_out(proj: Any) -> Optional[bool]: @@ -168,7 +169,9 @@ def W_in(self) -> torch.Tensor: in_module = getattr(self, "in", None) if in_module is None: raise AttributeError("No 'in' submodule on this MLP bridge") - weight = in_module.weight + weight = require_readable_weight( + in_module.weight, operation=f"read W_in from {self.name}", owner=in_module + ) layout = self._weight_layout_in_out(in_module) return self._normalize_mlp_weight(weight, layout, in_module, pattern="in") @@ -178,7 +181,9 @@ def W_gate(self) -> Optional[torch.Tensor]: gate_module = getattr(self, "gate", None) if gate_module is None: return None - weight = gate_module.weight + weight = require_readable_weight( + gate_module.weight, operation=f"read W_gate from {self.name}", owner=gate_module + ) layout = self._weight_layout_in_out(gate_module) return self._normalize_mlp_weight(weight, layout, gate_module, pattern="in") @@ -188,6 +193,8 @@ def W_out(self) -> torch.Tensor: out_module = getattr(self, "out", None) if out_module is None: raise AttributeError("No 'out' submodule on this MLP bridge") - weight = out_module.weight + weight = require_readable_weight( + out_module.weight, operation=f"read W_out from {self.name}", owner=out_module + ) layout = self._weight_layout_in_out(out_module) return self._normalize_mlp_weight(weight, layout, out_module, pattern="out") diff --git a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py index ec6503be5..827d1a5c4 100644 --- a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py +++ b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py @@ -189,6 +189,45 @@ def set_original_component(self, component: torch.nn.Module) -> None: _setup_eager_attention_hook_wrapper() self._validate_submodule_declarations(component) self._qk_norm_phase = self._decide_qk_norm_phase(component) + self._own_scaled_hook_k(component) + + def _own_scaled_hook_k(self, hf_attn: torch.nn.Module) -> None: + """Replace the ``hook_k`` alias with a real HookPoint when K is scaled. + + Falcon-H1 multiplies K by a learned mup scalar between the projection + and RoPE, so the aliased ``hook_k`` (= ``k.hook_out``) reports a tensor + that never reaches attention, and a value written there is silently + rescaled on the way in. Same split Granite's residual_multiplier + established: the TL-semantic name carries the scaled tensor, the + module-shaped ``k.hook_out`` stays the raw projection. + + No-op for every other architecture, which keeps the alias. + """ + if getattr(hf_attn, "key_multiplier", None) is None: + return + if self.hook_aliases is type(self).hook_aliases: + self.hook_aliases = dict(self.hook_aliases) + self.hook_aliases.pop("hook_k", None) + # The per-head hook_conversion is attached later, by + # _setup_qkv_hook_reshaping — component binding always precedes hook + # compatibility setup (bridge.py wires components, then calls it). + self.hook_k = HookPoint() + + def _fire_scaled_hook_k(self, key_states: torch.Tensor) -> torch.Tensor: + """Fire an owned ``hook_k`` on the flat 3D tensor, preserving input rank. + + The split-qkv path arrives 4D; the hook_conversion presents 4D to the + user either way, and its revert only fires on 4D returns, so a 4D input + has to be flattened first or a hook that edits the tensor would hand + back a shape the RoPE call below cannot use. + """ + if "hook_k" in self.hook_aliases or not hasattr(self, "hook_k"): + return key_states + if key_states.dim() == 4: + b, s, n_h, d_h = key_states.shape + flat = self.hook_k(key_states.reshape(b, s, n_h * d_h)) + return flat.reshape(b, s, n_h, d_h) + return self.hook_k(key_states) def _validate_submodule_declarations(self, hf_attn: torch.nn.Module) -> None: """Raise if adapter omits q/k/v/o or a QK-norm the HF module has.""" @@ -399,9 +438,12 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: query_states = query_states.reshape(*input_shape, -1) # Falcon-H1 scales K by a learned mup scalar between projection and RoPE. + # hook_k fires after the scale (see _own_scaled_hook_k) so it carries the + # tensor that actually reaches attention; k.hook_out kept the raw one. key_multiplier = getattr(hf_attn, "key_multiplier", None) if key_multiplier is not None: key_states = key_states * key_multiplier + key_states = self._fire_scaled_hook_k(key_states) has_q_norm = "q_norm" in self.submodules has_k_norm = "k_norm" in self.submodules diff --git a/transformer_lens/model_bridge/supported_architectures/bitnet.py b/transformer_lens/model_bridge/supported_architectures/bitnet.py index 2f5e158b5..0b933d9d4 100644 --- a/transformer_lens/model_bridge/supported_architectures/bitnet.py +++ b/transformer_lens/model_bridge/supported_architectures/bitnet.py @@ -11,6 +11,7 @@ from transformer_lens.model_bridge.supported_architectures.llama import ( LlamaArchitectureAdapter, ) +from transformer_lens.utilities.quantization import unreadable_weight_reason class _BitNetAttentionBridge(PositionEmbeddingsAttentionBridge): @@ -46,3 +47,27 @@ 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 + + def prepare_model(self, hf_model: Any) -> None: + """Refuse packed 1.58-bit checkpoints, which need BitNet dequant kernels. + + The flagship microsoft/bitnet-b1.58-2B-4T stores `weight` as packed + uint8 with a collapsed first dim (out_features // 4) plus a separate + weight_scale, so every weight-space read reshapes it into a + wrong-but-plausible matrix rather than failing. The registry records + this checkpoint at 0% on the forward-pass phase for exactly that reason. + """ + super().prepare_model(hf_model) + # Every weight-bearing module, not just the first: BitNet leaves the + # embedding unquantized, and it sorts first in named_modules(), so + # sampling one module inspected the one weight that is never packed. + for name, module in hf_model.named_modules(): + weight = getattr(module, "weight", None) + if weight is None: + continue + if unreadable_weight_reason(weight) is not None: + raise NotImplementedError( + f"BitNet checkpoint stores packed weights ({name}); " + "TransformerLens needs the dequantized sibling — use " + "microsoft/bitnet-b1.58-2B-4T-bf16." + ) diff --git a/transformer_lens/model_bridge/supported_architectures/phi3.py b/transformer_lens/model_bridge/supported_architectures/phi3.py index bd9f73d2c..05d2fe573 100644 --- a/transformer_lens/model_bridge/supported_architectures/phi3.py +++ b/transformer_lens/model_bridge/supported_architectures/phi3.py @@ -26,6 +26,7 @@ RotaryEmbeddingBridge, UnembeddingBridge, ) +from transformer_lens.utilities.quantization import require_readable_weight class _SizedSplitConversion(BaseTensorConversion): @@ -134,7 +135,14 @@ def _split_gate_up( original_mlp_component: Any, ) -> tuple[torch.nn.Module, torch.nn.Module]: """Split Phi-3's fused gate_up_proj into separate gate and up Linear modules.""" - fused_weight = original_mlp_component.gate_up_proj.weight + # This override, not the guarded default, is what Phi-3, GLM and GLM-4V + # install — so the guard has to be here too. FP8 is the case that needs + # it: tensor_split and nn.Parameter both accept it without complaint. + fused_weight = require_readable_weight( + original_mlp_component.gate_up_proj.weight, + operation="split a fused gate/up projection at boot", + owner=original_mlp_component.gate_up_proj, + ) gate_w, up_w = torch.tensor_split(fused_weight, 2, dim=0) d_model = fused_weight.shape[1] d_mlp = gate_w.shape[0] @@ -166,7 +174,11 @@ def _split_phi3_qkv( self, original_attention_component: Any ) -> tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]: """Split Phi-3's fused qkv_proj into separate Q, K, V linear modules.""" - qkv_weight = original_attention_component.qkv_proj.weight + qkv_weight = require_readable_weight( + original_attention_component.qkv_proj.weight, + operation="split a fused QKV projection at boot", + owner=original_attention_component.qkv_proj, + ) d_model = qkv_weight.shape[1] # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head each. diff --git a/transformer_lens/pretrained/weight_conversions/mixtral.py b/transformer_lens/pretrained/weight_conversions/mixtral.py index 8587241ea..be58c5fb9 100644 --- a/transformer_lens/pretrained/weight_conversions/mixtral.py +++ b/transformer_lens/pretrained/weight_conversions/mixtral.py @@ -2,6 +2,7 @@ import torch from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig +from transformer_lens.utilities.quantization import require_readable_weight def convert_mixtral_weights(mixtral, cfg: HookedTransformerConfig): @@ -51,20 +52,24 @@ def convert_mixtral_weights(mixtral, cfg: HookedTransformerConfig): # gate_up_proj: [num_experts, 2 * d_mlp, d_model] (gate fused above up) # down_proj: [num_experts, d_model, d_mlp] moe = mixtral.model.layers[l].mlp - state_dict[f"blocks.{l}.mlp.W_gate.weight"] = moe.gate.weight + # Guarded like the experts below: load_state_dict accepts a SAME-SHAPE + # int8/FP8 router and silently casts it to float32, so nothing + # downstream catches it. + state_dict[f"blocks.{l}.mlp.W_gate.weight"] = require_readable_weight( + moe.gate.weight, operation="convert the Mixtral router weight", owner=mixtral + ) experts = moe.experts - gate_up = experts.gate_up_proj - down = experts.down_proj - if not isinstance(gate_up, torch.Tensor) or not gate_up.dtype.is_floating_point: - # Quantized checkpoints wrap or pack these; slicing them silently - # drops the scales instead of failing (cf. the MXFP4 gpt-oss case). - raise NotImplementedError( - "convert_mixtral_weights needs plain floating-point expert " - f"weights; got {type(gate_up).__name__} with dtype " - f"{getattr(gate_up, 'dtype', None)}. Load the checkpoint " - "dequantized (e.g. without a quantization_config)." - ) + gate_up = require_readable_weight( + experts.gate_up_proj, + operation="convert Mixtral expert weights (gate_up_proj)", + owner=mixtral, + ) + down = require_readable_weight( + experts.down_proj, + operation="convert Mixtral expert weights (down_proj)", + owner=mixtral, + ) # MixtralExperts.forward does # gate, up = F.linear(x, gate_up_proj[e]).chunk(2, dim=-1) diff --git a/transformer_lens/pretrained/weight_conversions/olmoe.py b/transformer_lens/pretrained/weight_conversions/olmoe.py index a38ea758f..e71f57c56 100644 --- a/transformer_lens/pretrained/weight_conversions/olmoe.py +++ b/transformer_lens/pretrained/weight_conversions/olmoe.py @@ -2,6 +2,7 @@ import torch from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig +from transformer_lens.utilities.quantization import require_readable_weight def convert_olmoe_weights(olmoe, cfg: HookedTransformerConfig): @@ -45,15 +46,27 @@ def convert_olmoe_weights(olmoe, cfg: HookedTransformerConfig): state_dict[f"blocks.{l}.ln2.w"] = olmoe_layer.post_attention_layernorm.weight - state_dict[f"blocks.{l}.mlp.W_gate.weight"] = olmoe_layer.mlp.gate.weight + state_dict[f"blocks.{l}.mlp.W_gate.weight"] = require_readable_weight( + olmoe_layer.mlp.gate.weight, + operation="convert the OLMoE router weight", + owner=olmoe, + ) # HF OLMoE uses batched expert weights: # gate_up_proj: [num_experts, 2 * intermediate_size, hidden_size] # down_proj: [num_experts, hidden_size, intermediate_size] # The gate_up_proj fuses gate and up projections along dim 1. experts = olmoe_layer.mlp.experts - gate_up = experts.gate_up_proj # [num_experts, 2*d_mlp, d_model] - down = experts.down_proj # [num_experts, d_model, d_mlp] + gate_up = require_readable_weight( + experts.gate_up_proj, + operation="convert OLMoE expert weights (gate_up_proj)", + owner=olmoe, + ) # [num_experts, 2*d_mlp, d_model] + down = require_readable_weight( + experts.down_proj, + operation="convert OLMoE expert weights (down_proj)", + owner=olmoe, + ) # [num_experts, d_model, d_mlp] for e in range(cfg.num_experts): # Split fused gate_up into gate and up projections diff --git a/transformer_lens/pretrained/weight_conversions/openai.py b/transformer_lens/pretrained/weight_conversions/openai.py index e4d3ede46..f6b89e3f1 100644 --- a/transformer_lens/pretrained/weight_conversions/openai.py +++ b/transformer_lens/pretrained/weight_conversions/openai.py @@ -11,6 +11,17 @@ import torch from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig +from transformer_lens.utilities.quantization import require_readable_weight + +# Phrased to hold for any quantization: the guard catches int8 and FP8 too, and +# cannot know which one it caught, so it must not assert this *is* MXFP4. +_GPT_OSS_REMEDY = ( + "If this is a packed-MXFP4 checkpoint, load it dequantized so the converter " + "sees plain tensors: pass hf_model=AutoModelForCausalLM.from_pretrained(name, " + "quantization_config=Mxfp4Config(dequantize=True)), or load by model name and " + "TransformerLens dequantizes automatically. Otherwise reload without a " + "quantization_config. Quantized *forward* passes remain supported." +) def convert_gpt_oss_weights(gpt_oss, cfg: HookedTransformerConfig): @@ -75,7 +86,12 @@ def convert_gpt_oss_weights(gpt_oss, cfg: HookedTransformerConfig): ) # MoE - Router (GPT-OSS uses 'router' with bias) - state_dict[f"blocks.{l}.mlp.W_gate.weight"] = layer.mlp.router.weight + state_dict[f"blocks.{l}.mlp.W_gate.weight"] = require_readable_weight( + layer.mlp.router.weight, + operation="convert the gpt-oss router weight", + owner=gpt_oss, + remedy=_GPT_OSS_REMEDY, + ) state_dict[f"blocks.{l}.mlp.W_gate.bias"] = layer.mlp.router.bias # MoE - Experts @@ -84,22 +100,35 @@ def convert_gpt_oss_weights(gpt_oss, cfg: HookedTransformerConfig): # down_proj: (num_experts, expert_dim, hidden_size) experts = layer.mlp.experts gate_up_proj = experts.gate_up_proj # (num_experts, hidden_size, 2*expert_dim) - gate_up_bias = experts.gate_up_proj_bias # (num_experts, 2*expert_dim) + gate_up_bias = require_readable_weight( + experts.gate_up_proj_bias, + operation=f"convert gpt-oss expert biases (blocks.{l}.mlp.experts.gate_up_proj_bias)", + owner=gpt_oss, + remedy=_GPT_OSS_REMEDY, + ) # (num_experts, 2*expert_dim) down_proj = experts.down_proj # (num_experts, expert_dim, hidden_size) - down_bias = experts.down_proj_bias # (num_experts, hidden_size) - - if not isinstance(gate_up_proj, torch.Tensor): - # Packed MXFP4 checkpoints wrap expert weights in a triton-kernels - # object (confusingly also named "Tensor") that cannot be sliced. - raise NotImplementedError( - f"blocks.{l}.mlp.experts.gate_up_proj is a " - f"{type(gate_up_proj).__module__}.{type(gate_up_proj).__name__}, not a " - "torch.Tensor — this gpt-oss checkpoint has packed MXFP4 expert weights. " - "Load it dequantized so the converter sees plain tensors: pass " - "hf_model=AutoModelForCausalLM.from_pretrained(name, " - "quantization_config=Mxfp4Config(dequantize=True)), or load by model " - "name and TransformerLens dequantizes automatically." - ) + down_bias = require_readable_weight( + experts.down_proj_bias, + operation=f"convert gpt-oss expert biases (blocks.{l}.mlp.experts.down_proj_bias)", + owner=gpt_oss, + remedy=_GPT_OSS_REMEDY, + ) # (num_experts, hidden_size) + + # Packed MXFP4 wraps these in a triton-kernels object (confusingly also + # named "Tensor"), but int8 and FP8 gpt-oss finetunes slice without + # complaint and would emit plausible garbage, so check the dtype too. + gate_up_proj = require_readable_weight( + gate_up_proj, + operation=f"convert gpt-oss expert weights (blocks.{l}.mlp.experts.gate_up_proj)", + owner=gpt_oss, + remedy=_GPT_OSS_REMEDY, + ) + down_proj = require_readable_weight( + down_proj, + operation=f"convert gpt-oss expert weights (blocks.{l}.mlp.experts.down_proj)", + owner=gpt_oss, + remedy=_GPT_OSS_REMEDY, + ) for e in range(cfg.num_experts): # Split interleaved gate_up_proj into separate gate and up (in) projections diff --git a/transformer_lens/utilities/quantization.py b/transformer_lens/utilities/quantization.py new file mode 100644 index 000000000..ad9e39d8a --- /dev/null +++ b/transformer_lens/utilities/quantization.py @@ -0,0 +1,164 @@ +"""Guards for weight-space code paths that cannot read quantized weights. + +TransformerLens supports quantized *forward* passes: the wrapped HF module +dequantizes internally, and the bridge's forward paths deliberately skip +non-floating-point parameters when picking a compute dtype. Those paths must +keep working. + +What does not work is reading a quantized ``.weight`` and doing arithmetic on it +directly — reshaping it into per-head matrices, slicing a fused projection, +folding LayerNorm into it. There the storage is packed (bitsandbytes 4-bit keeps +a ``[N, 1]`` uint8 buffer), split from its scales (FP8 keeps a separate +``weight_scale_inv``), or not a tensor at all (MXFP4 wraps a triton-kernels +object). Slicing those yields plausible-looking garbage rather than an error, +which is the failure mode this module exists to prevent. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + + +def unreadable_weight_reason(weight: Any) -> Optional[str]: + """Why ``weight`` cannot be read as a plain weight matrix, or None if it can. + + Readable means exactly the four full-width float dtypes (float16, bfloat16, + float32, float64); every 1-byte dtype in torch is either integer storage or + a narrow float that carries its scales separately. + + Scope this to ``.weight`` reads. Complex dtypes also report + ``is_floating_point == False`` and would be described as integer storage — + they appear in this codebase only as rotary buffers, never as weights. + + Returns a fragment that completes "... cannot be read because ", or + None when the weight is usable. Meta-device tensors are reported here too, + but callers should surface them as a load problem, not a quantization one + (see :func:`require_readable_weight`). + """ + if not isinstance(weight, torch.Tensor): + # MXFP4 and similar hand back a wrapper object holding the packed + # payload plus scales. It is often *named* Tensor, which makes the + # resulting TypeError look like a torch bug — hence the qualified name, + # which is the only thing distinguishing it from the real one. + cls = type(weight) + return f"it is a {cls.__module__}.{cls.__qualname__}, not a torch.Tensor" + if weight.device.type == "meta": + return "it is on the meta device, so no values have been materialized" + if not weight.dtype.is_floating_point: + # bitsandbytes int8/4-bit, GPTQ/AWQ int32, BitNet packed uint8. Note + # that bnb's Params4bit IS a torch.Tensor subclass, so only the dtype + # distinguishes it — and its shape is the packed [N, 1], not [out, in]. + return f"its dtype is {weight.dtype}, which is packed integer storage" + if weight.dtype.itemsize < 2: + # float8_e4m3fn and friends report is_floating_point=True, so a plain + # float check lets them through. They are the only family that slices + # and even survives nn.Parameter() silently, so omitting this branch + # leaves the worst case uncaught. + return ( + f"its dtype is {weight.dtype}, a narrow float whose values are " + "meaningless without the separate scale tensor stored beside them " + "(e.g. weight_scale_inv)" + ) + return None + + +def _is_meta(weight: Any) -> bool: + return isinstance(weight, torch.Tensor) and weight.device.type == "meta" + + +def quantization_method(config: Any) -> Optional[str]: + """The ``quant_method`` declared on an HF config, or None if unquantized. + + Accepts a ``PretrainedConfig`` or its ``to_dict()`` form, and tolerates the + nested ``quantization_config`` being either shape — both appear in the wild, + depending on whether the config was loaded or round-tripped through JSON. + """ + if config is None: + return None + if isinstance(config, dict): + quant_config = config.get("quantization_config") + else: + quant_config = getattr(config, "quantization_config", None) + if quant_config is None: + return None + if isinstance(quant_config, dict): + method = quant_config.get("quant_method") + else: + method = getattr(quant_config, "quant_method", None) + # HF stores this as a str or a str-valued QuantizationMethod enum; require + # that rather than str()-ing whatever turned up, or a stub object's repr + # ends up quoted back at the user as a method name. + method = getattr(method, "value", method) + return method if isinstance(method, str) else None + + +def describe_quantization(owner: Any) -> str: + """Best-effort name for how ``owner``'s weights are quantized. + + Resolution order: the HF config's declared ``quant_method``, then the + weight's class name (bitsandbytes subclasses are identifiable that way), + then a generic fallback. + """ + for holder in (owner, getattr(owner, "original_component", None)): + method = quantization_method(getattr(holder, "config", None)) + if method is not None: + return method + weight = getattr(owner, "weight", None) + if weight is not None: + # Exclude the two uninformative names rather than nn.Parameter itself: + # bitsandbytes' Params4bit and Int8Params ARE Parameter subclasses, so + # an isinstance test excluded exactly the classes this identifies. + # transformers keys off the same class names (integrations/bitsandbytes). + name = type(weight).__name__ + if name not in ("Tensor", "Parameter"): + return name + return "an unknown quantization" + + +_GENERIC_REMEDY = ( + "Weight-space operations need dequantized weights — reload the model " + "without a quantization_config (or with dequantize enabled). Quantized " + "*forward* passes remain supported." +) + + +def require_readable_weight( + weight: Any, *, operation: str, owner: Any = None, remedy: Optional[str] = None +) -> torch.Tensor: + """Return ``weight`` if it can be read as a plain matrix, else raise. + + Args: + weight: the tensor (or packed stand-in) about to be read. + operation: what the caller is trying to do, phrased to complete + "TransformerLens cannot because ...". + owner: the component or module holding the weight, used to name the + quantization method in the error. + remedy: architecture-specific advice replacing the generic "reload + without a quantization_config". Pass this where a concrete recipe + exists, e.g. gpt-oss's ``Mxfp4Config(dequantize=True)``. Phrase it + so it holds for any quantization, since the caller cannot know + which one it caught. + + Raises: + NotImplementedError: naming the quantization and how to proceed. This is + deliberately loud: every alternative silently produces wrong numbers. + """ + reason = unreadable_weight_reason(weight) + if reason is None: + assert isinstance(weight, torch.Tensor) + return weight + if _is_meta(weight): + # Not a quantization problem — naming one here would send the reader + # after the wrong cause. This is an offloaded or never-materialized load. + raise NotImplementedError( + f"TransformerLens cannot {operation} because {reason}. Load the " + "model with real weights (no meta-device or disk offload) before " + "reading them." + ) + method = describe_quantization(owner) if owner is not None else "an unknown quantization" + raise NotImplementedError( + f"TransformerLens cannot {operation} because {reason} " + f"(quantization: {method}). {remedy or _GENERIC_REMEDY}" + )