diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 50ea07b16..4cb68a524 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -97,10 +97,19 @@ jobs: ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2 ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM + ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3 ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2 ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM - key: ${{ runner.os }}-huggingface-models-v5 + ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5 + key: ${{ runner.os }}-huggingface-models-v6 restore-keys: | + ${{ runner.os }}-huggingface-models-v5 ${{ runner.os }}-huggingface-models-v4 ${{ runner.os }}-huggingface-models-v2 ${{ runner.os }}-huggingface-models-v1 @@ -143,7 +152,22 @@ jobs: with: path: | ~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M* - key: ${{ runner.os }}-huggingface-models-mps-v1 + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2 + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM + ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3 + ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2 + ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM + ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5 + key: ${{ runner.os }}-huggingface-models-mps-v2 + restore-keys: | + ${{ runner.os }}-huggingface-models-mps-v1 - name: Install dependencies run: | uv lock --check @@ -297,10 +321,19 @@ jobs: ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2 ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM + ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM + ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3 ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2 ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM - key: ${{ runner.os }}-huggingface-models-v5 + ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5 + key: ${{ runner.os }}-huggingface-models-v6 restore-keys: | + ${{ runner.os }}-huggingface-models-v5 ${{ runner.os }}-huggingface-models-v4 ${{ runner.os }}-huggingface-models-v2 ${{ runner.os }}-huggingface-models-v1 @@ -350,12 +383,10 @@ jobs: with: path: | ~/.cache/huggingface/hub/models--gpt2 - key: ${{ runner.os }}-huggingface-models-v5 - restore-keys: | - ${{ runner.os }}-huggingface-models-v4 - ${{ runner.os }}-huggingface-models-v2 - ${{ runner.os }}-huggingface-models-v1 - ${{ runner.os }}-huggingface-models + # Distinct key: this job needs only gpt2, and sharing the full-suite + # key let whichever job saved first pin the cache contents for all + # three (a one-path save then serves the suites that need the rest). + key: ${{ runner.os }}-huggingface-models-benchmark-v1 - name: Install dependencies run: | uv lock --check diff --git a/docs/source/content/adapter_development/adapter-specification.md b/docs/source/content/adapter_development/adapter-specification.md index 3b2756c5a..afb68bf59 100644 --- a/docs/source/content/adapter_development/adapter-specification.md +++ b/docs/source/content/adapter_development/adapter-specification.md @@ -276,6 +276,32 @@ Standard positional embeddings + LayerNorm + standard MLP + combined QKV. Used b Similar to Llama-like but with `MoEBridge` replacing the MLP. Used by: Mixtral, GraniteMoE, OLMoE. +**Interleaved / dense-prefix stacks.** When the config can make some layers a plain +gated MLP under the same attribute name (`first_k_dense_replace`, +`mlp_only_layers`/`decoder_sparse_step`, `moe_layer_start_index`, +`num_dense_layers`, `mlp_layer_types`), also declare the dense projections as +optional submodules named `dense_gate` / `dense_in` / `dense_out`: + +```python +"mlp": MoEBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": MoERouterBridge(name="gate", optional=True), + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), + }, +), +``` + +`MoEBridge` detects that binding per layer and gives the dense layers the +gated-MLP hook set (`hook_pre`, `hook_pre_linear`, `hook_post`) and weight +accessors (`W_gate`/`W_in`/`W_out`), while sparse layers keep MoE semantics. +Without the declaration those layers report `d_model` boundary tensors under +neuron-hook names (#1645). The keys are `dense_*` rather than `gate`/`in`/`out` +because `gate` already means the *router* on the sparse layers of the same model. + ### Pattern 4: Multimodal Extends a text-only pattern with vision encoder and projection bridges. Used by: LLaVA, LLaVA-Next, Gemma3 Multimodal. diff --git a/tests/integration/model_bridge/test_bridge_vs_hf_eager_parity.py b/tests/integration/model_bridge/test_bridge_vs_hf_eager_parity.py index dcecd45ee..b407338bc 100644 --- a/tests/integration/model_bridge/test_bridge_vs_hf_eager_parity.py +++ b/tests/integration/model_bridge/test_bridge_vs_hf_eager_parity.py @@ -5,21 +5,17 @@ SDPA, which reorders ops in a fused kernel. Bridge vs HF *eager* matches to fp32-noise. """ -import platform from typing import Callable import pytest import torch from transformers import AutoModelForCausalLM, AutoTokenizer +from tests.tiny_checkpoints import FP32_NOISE_TOL, assert_tiny_parity, parity_params from transformer_lens.model_bridge import TransformerBridge MODEL_NAME = "EleutherAI/pythia-70m" -# Wider fp32 op-order noise floor on GH Actions macOS-arm64; ~3e-3 at output. -_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64" -FP32_NOISE_TOL = 1e-2 if _MACOS_ARM64 else 1e-5 - @pytest.fixture(scope="module") def tokenizer(): @@ -123,3 +119,28 @@ def test_bridge_attention_reconstruction_actually_runs(bridge, tokenize): "blocks.0.attn.hook_attn_scores did not fire — bridge no longer runs its " "own attention reconstruction, making the parity tests tautological." ) + + +@pytest.mark.parametrize("model_name", parity_params()) +def test_tiny_bridge_logits_match_hf_eager(model_name: str) -> None: + """The hooked forward must reproduce an independent HF eager load's logits.""" + bridge = TransformerBridge.boot_transformers(model_name, device="cpu", dtype=torch.float32) + tokens = bridge.to_tokens("The quick brown fox jumps") + hf_eager = AutoModelForCausalLM.from_pretrained( + model_name, torch_dtype=torch.float32, attn_implementation="eager" + ).eval() + + with torch.inference_mode(): + bridge_logits = bridge(tokens) + hf_logits = hf_eager(tokens).logits + assert_tiny_parity(bridge_logits, hf_logits, model_name) + + # Non-tautology guard: on archs with a reconstructed attention path, its + # scores hook must fire, or the parity above proves nothing. + if "blocks.0.attn.hook_attn_scores" in bridge.hook_dict: + fired: list[bool] = [] + bridge.run_with_hooks( + tokens, + fwd_hooks=[("blocks.0.attn.hook_attn_scores", lambda v, hook: fired.append(True))], + ) + assert fired, f"{model_name!r}: hook_attn_scores did not fire — reconstruction bypassed." diff --git a/tests/integration/model_bridge/test_deepseek_adapter.py b/tests/integration/model_bridge/test_deepseek_adapter.py index 9409299fd..d4ab64640 100644 --- a/tests/integration/model_bridge/test_deepseek_adapter.py +++ b/tests/integration/model_bridge/test_deepseek_adapter.py @@ -79,12 +79,18 @@ def test_forward_matches_hf(self, tiny_deepseek_bridge): class TestDeepSeekDenseVsMoELayers: - def test_dense_layer_has_no_moe_hooks(self, tiny_deepseek_bridge): + def test_dense_layer_binds_gated_mlp_hooks(self, tiny_deepseek_bridge): + """Dense-prefix layers expose the gate PROJECTION under dense_gate + (d_mlp neuron basis) — `gate` stays the router name, so one hook name + never means two things across layers (#1645).""" tokens = torch.tensor([[1, 2, 3, 4]]) _, cache = tiny_deepseek_bridge.run_with_cache(tokens) - cache_keys = set(cache.keys()) - assert not any("blocks.0.mlp.gate" in k for k in cache_keys) - assert not any("blocks.0.mlp.shared_experts" in k for k in cache_keys) + d_mlp = tiny_deepseek_bridge.original_model.config.intermediate_size + assert cache["blocks.0.mlp.dense_gate.hook_out"].shape[-1] == d_mlp + torch.testing.assert_close( + cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"] + ) + assert not any("blocks.0.mlp.shared_experts" in k for k in cache.keys()) def test_moe_layer_has_gate_hooks(self, tiny_deepseek_bridge): tokens = torch.tensor([[1, 2, 3, 4]]) diff --git a/tests/integration/model_bridge/test_deepseek_v2_adapter.py b/tests/integration/model_bridge/test_deepseek_v2_adapter.py index 0ac05c208..dce8a1a94 100644 --- a/tests/integration/model_bridge/test_deepseek_v2_adapter.py +++ b/tests/integration/model_bridge/test_deepseek_v2_adapter.py @@ -113,9 +113,16 @@ def test_forward_matches_hf(self, tiny_deepseek_v2_bridge): class TestDeepSeekV2DenseVsMoELayers: - def test_dense_layer_has_no_moe_hooks(self, tiny_deepseek_v2_bridge): + def test_dense_layer_binds_gated_mlp_hooks(self, tiny_deepseek_v2_bridge): + """Dense-prefix layers expose the gate PROJECTION under dense_gate + (d_mlp neuron basis) — `gate` stays the router name, so one hook name + never means two things across layers (#1645).""" _, cache = tiny_deepseek_v2_bridge.run_with_cache(_tokens()) - assert not any("blocks.0.mlp.gate" in k for k in cache) + d_mlp = tiny_deepseek_v2_bridge.original_model.config.intermediate_size + assert cache["blocks.0.mlp.dense_gate.hook_out"].shape[-1] == d_mlp + torch.testing.assert_close( + cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"] + ) assert not any("blocks.0.mlp.shared_experts" in k for k in cache) def test_moe_layer_has_router_and_shared_expert_hooks(self, tiny_deepseek_v2_bridge): diff --git a/tests/integration/model_bridge/test_glm4_moe_lite_adapter.py b/tests/integration/model_bridge/test_glm4_moe_lite_adapter.py index 4c69a621d..eef5824bb 100644 --- a/tests/integration/model_bridge/test_glm4_moe_lite_adapter.py +++ b/tests/integration/model_bridge/test_glm4_moe_lite_adapter.py @@ -28,12 +28,22 @@ def test_adapter_selected(self, glm_bridge): assert isinstance(glm_bridge.adapter, Glm4MoeLiteArchitectureAdapter) def test_dense_sparse_layer_mix(self, glm_bridge): - """The tiny checkpoint declares layer 0 dense, layer 1 sparse.""" + """The tiny checkpoint declares layer 0 dense, layer 1 sparse; the mlp + template binds each accordingly (#1645).""" hf_model = glm_bridge.original_model assert hf_model.config.mlp_layer_types == ["dense", "sparse"] - assert not hasattr(hf_model.model.layers[0].mlp, "gate") - assert hasattr(hf_model.model.layers[1].mlp, "gate") - assert hasattr(hf_model.model.layers[1].mlp, "shared_experts") + dense_mlp = glm_bridge.blocks[0].mlp + sparse_mlp = glm_bridge.blocks[1].mlp + # Dense binding: gated-MLP neuron hooks via the dense_* projections. + assert not hasattr(dense_mlp.original_component, "experts") + assert dense_mlp._bound_dense is True + assert dense_mlp.hook_aliases["hook_pre"] == "dense_gate.hook_out" + assert not hasattr(dense_mlp, "hook_router_scores") + # Sparse binding: router + shared experts survive setup on the BRIDGE + # (a silently-popped optional router would otherwise stay green here). + assert hasattr(sparse_mlp.original_component, "experts") + assert {"gate", "shared_experts"} <= set(sparse_mlp.submodules) + assert hasattr(sparse_mlp, "hook_router_scores") class TestGlm4MoeLiteForwardEquivalence: diff --git a/tests/integration/model_bridge/test_moe_dense_dispatch.py b/tests/integration/model_bridge/test_moe_dense_dispatch.py new file mode 100644 index 000000000..16c2eb0d8 --- /dev/null +++ b/tests/integration/model_bridge/test_moe_dense_dispatch.py @@ -0,0 +1,113 @@ +"""Dense-prefix MoE dispatch on a real mixed checkpoint (#1645). + +katuni4ka/tiny-random-deepseek-v3 has first_k_dense_replace=1 over two layers: +layer 0 is a plain gated MLP, layer 1 a sparse MoE block — one boot exercises +both bindings of the DenseOrMoEBridge template. +""" + +import torch +from transformers import AutoModelForCausalLM + +from tests.tiny_checkpoints import MIXED_DENSE_SPARSE_MOE, assert_tiny_parity +from transformer_lens.model_bridge import TransformerBridge + +MODEL_NAME = MIXED_DENSE_SPARSE_MOE + + +def test_mixed_checkpoint_dense_and_sparse_mlp_hooks() -> None: + bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + tokens = bridge.to_tokens("The quick brown fox jumps") + d_model = bridge.cfg.d_model + d_mlp = bridge.original_model.config.intermediate_size + + with torch.no_grad(): + logits, cache = bridge.run_with_cache(tokens) + + # Dense layer 0: neuron-basis hooks with gated-MLP semantics. + assert cache["blocks.0.mlp.hook_pre"].shape[-1] == d_mlp + torch.testing.assert_close( + cache["blocks.0.mlp.hook_pre"], cache["blocks.0.mlp.dense_gate.hook_out"] + ) + torch.testing.assert_close( + cache["blocks.0.mlp.hook_pre_linear"], cache["blocks.0.mlp.dense_in.hook_out"] + ) + torch.testing.assert_close( + cache["blocks.0.mlp.hook_post"], cache["blocks.0.mlp.dense_out.hook_in"] + ) + + # Sparse layer 1: MoE boundary semantics unchanged. + assert cache["blocks.1.mlp.hook_pre"].shape[-1] == d_model + torch.testing.assert_close(cache["blocks.1.mlp.hook_pre"], cache["blocks.1.mlp.hook_in"]) + + # The dispatch must not perturb the forward itself. + hf_eager = AutoModelForCausalLM.from_pretrained( + MODEL_NAME, torch_dtype=torch.float32, attn_implementation="eager" + ).eval() + with torch.no_grad(): + hf_logits = hf_eager(tokens).logits + assert_tiny_parity(logits, hf_logits, MODEL_NAME) + + +def test_dense_layer_neuron_intervention_targets_gate_output() -> None: + """A write on the dense layer's hook_pre must be a write on the gate + projection's output — the #1645 complaint was precisely that no + neuron-basis intervention point existed on dense layers. (Pre-fix, + hook_pre aliased the MLP *input*, so zeroing it also moved logits — + equivalence with the gate target is what pins the semantics.)""" + bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + tokens = bridge.to_tokens("The quick brown fox jumps") + + def zero(tensor: torch.Tensor, hook) -> torch.Tensor: + return tensor * 0.0 + + with torch.no_grad(): + base_logits = bridge(tokens) + via_alias = bridge.run_with_hooks(tokens, fwd_hooks=[("blocks.0.mlp.hook_pre", zero)]) + via_target = bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.0.mlp.dense_gate.hook_out", zero)] + ) + torch.testing.assert_close(via_alias, via_target) + assert not torch.allclose(base_logits, via_alias) + + +def test_dense_layer_weights_reach_get_params() -> None: + """The neuron hooks advertise a neuron basis; the weight API must back it. + get_params() keyed on submodules named `in`/`out`, so dense MoE layers were + silently zero-filled — real weights on the component, zeros in the dict.""" + bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + params = bridge.get_params() + + dense = bridge.blocks[0].mlp + assert dense.bound_dense is True + # Anchor on the wrapped HF module, NOT on dense.W_in/W_gate/W_out: + # get_params_util builds these entries by reading those very properties, so + # comparing the two would be the same expression evaluated twice and could + # not detect a projection being read from the wrong submodule. + hf_mlp = dense.original_component + torch.testing.assert_close(params["blocks.0.mlp.W_in"], hf_mlp.up_proj.weight.T) + torch.testing.assert_close(params["blocks.0.mlp.W_gate"], hf_mlp.gate_proj.weight.T) + torch.testing.assert_close(params["blocks.0.mlp.W_out"], hf_mlp.down_proj.weight.T) + assert (params["blocks.0.mlp.W_in"] != 0).any() + # Negative control: gate and up are same-shaped, so this checkpoint can tell + # them apart only if their weights actually differ. + assert not torch.equal(hf_mlp.gate_proj.weight, hf_mlp.up_proj.weight) + + # Sparse layers legitimately have no single W_* and keep the placeholder. + assert bridge.blocks[1].mlp.bound_dense is False + assert (params["blocks.1.mlp.W_in"] == 0).all() + + +def test_hook_dict_follows_a_dense_to_sparse_rebind() -> None: + """The alias cache must not outlive the binding it described: a rebind can + leave the hook registry exactly the same size while changing what the + aliases point at.""" + bridge = TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + _ = bridge.hook_dict # populate the cache while layer 0 is dense + assert bridge._collect_block_instance_aliases()["blocks.0.mlp.hook_pre"].endswith( + "dense_gate.hook_out" + ) + + bridge.blocks[0].mlp.set_original_component(bridge.blocks[1].mlp.original_component) + + assert bridge.blocks[0].mlp.bound_dense is False + assert bridge._collect_block_instance_aliases()["blocks.0.mlp.hook_pre"].endswith("mlp.hook_in") diff --git a/tests/integration/model_bridge/test_residual_decomposition_identities.py b/tests/integration/model_bridge/test_residual_decomposition_identities.py index e95d6a655..5a75d9277 100644 --- a/tests/integration/model_bridge/test_residual_decomposition_identities.py +++ b/tests/integration/model_bridge/test_residual_decomposition_identities.py @@ -16,6 +16,12 @@ - olmo2: post-norm inside the residual branch — RMSNorm applies to the sublayer output before the add, so the contributions are the norm outputs (the #1648 case) - mpt: residual added inside the HF MLP module only (attention adds at block level) +- granite: the scaled-residual block wiring end to end. NOTE: the hub tiny ships + residual_multiplier=1.0, so this entry does NOT exercise the multiplier — + test_granite_hook_semantics.py builds a fixture with a real multiplier and is + the coverage for the #1648 Granite case. +- gemma2: sandwich norms — a post-sublayer norm inside each residual branch, so + the contributions are the post-norm outputs Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope here: they have no ``hook_resid_mid``. Their two-term identity @@ -26,15 +32,10 @@ import pytest import torch +from tests.tiny_checkpoints import sequential_residual_params from transformer_lens.model_bridge import TransformerBridge -SEQUENTIAL_RESIDUAL_MODELS = [ - pytest.param("hf-internal-testing/tiny-random-gpt2", id="gpt2"), - pytest.param("trl-internal-testing/tiny-MistralForCausalLM-0.2", id="mistral"), - pytest.param("trl-internal-testing/tiny-BloomForCausalLM", id="bloom"), - pytest.param("hf-internal-testing/tiny-random-Olmo2ForCausalLM", id="olmo2"), - pytest.param("hf-internal-testing/tiny-random-MptForCausalLM", id="mpt"), -] +SEQUENTIAL_RESIDUAL_MODELS = sequential_residual_params() @pytest.mark.parametrize("model_name", SEQUENTIAL_RESIDUAL_MODELS) diff --git a/tests/tiny_checkpoints.py b/tests/tiny_checkpoints.py new file mode 100644 index 000000000..a412088f4 --- /dev/null +++ b/tests/tiny_checkpoints.py @@ -0,0 +1,127 @@ +"""Shared roster of tiny checkpoints and the parity tolerance policy. + +Single source for the bridge-vs-HF suites so the roster does not drift between +files (it must still be mirrored into the CI model caches in +``.github/workflows/checks.yml`` — those lists are the other half of this +contract). +""" + +from __future__ import annotations + +import platform + +import pytest + +# Wider fp32 op-order noise floor on GH Actions macOS-arm64. +_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64" + +# Tolerance for the REAL-model (pythia-70m) parity tests. The macOS-arm64 value +# is an op-order noise floor measured on that checkpoint (~3e-3 observed). +# Do not reuse it for the tiny roster — see TINY_PARITY_REL_TOL. +FP32_NOISE_TOL = 1e-2 if _MACOS_ARM64 else 1e-5 + +# Tolerance for the TINY roster, expressed RELATIVE to each checkpoint's logit +# scale. Measured on macOS-arm64 (the noisier of the two CI platforms) across +# all 12 roster entries: worst absolute drift 2.2e-07, worst relative 2.0e-06, +# and 9 of 12 bit-identical. A 10% error injected into the reconstructed +# attention scale produces ~0.35 relative on mistral, so this threshold sits +# ~50x above real noise and ~3500x below the effect it must catch. +# +# It is relative because the roster's logit spreads differ by ~700x (qwen2 std +# 0.0103 vs olmo 5.6540); a single absolute number cannot be both non-flaky on +# the widest checkpoint and discriminating on the narrowest — which is exactly +# how the previous platform-keyed 1e-2 became unable to fail. +TINY_PARITY_REL_TOL = 1e-4 + + +def assert_tiny_parity(bridge_logits, hf_logits, model_name: str) -> None: + """Assert bridge-vs-HF logit parity scaled to this checkpoint's magnitude.""" + max_diff = (bridge_logits - hf_logits).abs().max().item() + scale = max(1.0, hf_logits.abs().max().item()) + limit = TINY_PARITY_REL_TOL * scale + assert max_diff < limit, ( + f"{model_name!r} bridge vs HF eager drift={max_diff:.3e} exceeds " + f"{limit:.3e} (rel tol {TINY_PARITY_REL_TOL:.0e} x logit scale " + f"{scale:.3f}) — a reconstructed term (scale, norm, clamp, sink, " + "routing) may have been dropped." + ) + + +# One tiny checkpoint per reconstruction variant the bridge re-derives; each +# carries a term a regression could silently drop (the #1618 failure class): +# scale flags, GQA rope, ALiBi + in-module residual, flat qk-norm, fused QKV + +# softmax_scale/clip machinery, logit softcapping, sandwich norms + per-head +# qk-norm, residual_multiplier, clip_qkv, MoE routing, and MLA latent attention. +# +# ``sequential_residual`` marks the blocks whose attn/mlp contributions are added +# to a single residual stream in sequence — the ones for which +# ``resid_pre + attn_out == resid_mid`` is meaningful (parallel-residual and +# MoE/MLA-only stacks are excluded). +TINY_CHECKPOINTS: dict[str, dict[str, object]] = { + "gpt2": { + "name": "hf-internal-testing/tiny-random-gpt2", + "sequential_residual": True, + }, + "mistral": { + "name": "trl-internal-testing/tiny-MistralForCausalLM-0.2", + "sequential_residual": True, + }, + "bloom": { + "name": "trl-internal-testing/tiny-BloomForCausalLM", + "sequential_residual": True, + }, + "olmo2": { + "name": "hf-internal-testing/tiny-random-Olmo2ForCausalLM", + "sequential_residual": True, + }, + "mpt": { + "name": "hf-internal-testing/tiny-random-MptForCausalLM", + "sequential_residual": True, + }, + "granite": { + "name": "hf-internal-testing/tiny-random-GraniteForCausalLM", + "sequential_residual": True, + }, + "gemma2": { + "name": "hf-internal-testing/tiny-random-Gemma2ForCausalLM", + "sequential_residual": True, + }, + "gemma3": { + "name": "hf-internal-testing/tiny-random-Gemma3ForCausalLM", + "sequential_residual": False, + }, + "olmo": { + "name": "hf-internal-testing/tiny-random-OlmoForCausalLM", + "sequential_residual": False, + }, + "olmoe": { + "name": "hf-internal-testing/tiny-random-OlmoeForCausalLM", + "sequential_residual": False, + }, + "qwen2": { + "name": "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", + "sequential_residual": False, + }, + "deepseek_v3": { + "name": "hf-internal-testing/tiny-random-DeepseekV3ForCausalLM", + "sequential_residual": False, + }, +} + +# Mixed dense+sparse MoE stack (first_k_dense_replace=1 over two layers): one +# boot exercises both bindings of a single MoEBridge template. +MIXED_DENSE_SPARSE_MOE = "katuni4ka/tiny-random-deepseek-v3" + + +def parity_params() -> list: + """pytest params for every tiny checkpoint.""" + return [pytest.param(spec["name"], id=key) for key, spec in TINY_CHECKPOINTS.items()] + + +def sequential_residual_params() -> list: + """pytest params for checkpoints with a sequential residual decomposition.""" + return [ + pytest.param(spec["name"], id=key) + for key, spec in TINY_CHECKPOINTS.items() + if spec["sequential_residual"] + ] 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 new file mode 100644 index 000000000..f999749ae --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_moe_dense_dispatch.py @@ -0,0 +1,438 @@ +"""Unit tests for MoEBridge's per-layer dense/sparse dispatch (#1645). + +Interleaved and dense-prefix MoE architectures build a plain gated MLP on some +layers under the same attribute name as the sparse block. A single uniform +template then exposed MoE boundary aliases there — ``hook_pre``/``hook_post`` +delivering d_model residual tensors under neuron-hook names, no hooks on the +dense projections, and no weight accessors. MoEBridge now detects the dense +binding and adopts gated-MLP semantics for that layer only. +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +import transformer_lens.model_bridge.supported_architectures +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 +from transformer_lens.model_bridge.generalized_components import LinearBridge, MoEBridge + +D_MODEL, D_MLP = 8, 16 + + +class _DenseMLP(nn.Module): + """Standard SwiGLU gated MLP — the dense-prefix layer shape.""" + + def __init__(self) -> None: + super().__init__() + self.gate_proj = nn.Linear(D_MODEL, D_MLP, bias=False) + self.up_proj = nn.Linear(D_MODEL, D_MLP, bias=False) + self.down_proj = nn.Linear(D_MLP, D_MODEL, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down_proj( + torch.nn.functional.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states) + ) + + +class _SparseMoE(nn.Module): + """Minimal sparse block: a router plus experts, no dense projections.""" + + def __init__(self) -> None: + super().__init__() + self.gate = nn.Linear(D_MODEL, 4, bias=False) + self.experts = nn.ModuleList([nn.Linear(D_MODEL, D_MODEL) for _ in range(4)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + 0.0 * self.gate(hidden_states).sum(-1, keepdim=True) + + +def _adapter(): + """A real adapter (beartype enforces the type on setup_submodules), whose + mlp template is the interleaved-MoE mapping under test.""" + return ArchitectureAdapterFactory.select_architecture_adapter( + make_bridge_cfg("DeepseekV3ForCausalLM", d_model=D_MODEL, n_heads=2, d_head=4) + ) + + +def _template() -> MoEBridge: + """The adapter's own mlp template, as block setup deepcopies it per layer.""" + template = _adapter().component_mapping["blocks"].submodules["mlp"] + assert isinstance(template, MoEBridge) + return template + + +def _bind(module: nn.Module) -> MoEBridge: + """Bind through the real component-setup path, as boot does.""" + adapter = _adapter() + bridge = copy.deepcopy(adapter.component_mapping["blocks"].submodules["mlp"]) + assert isinstance(bridge, MoEBridge) + bridge.set_original_component(module) + setup_submodules(bridge, adapter, module) + return bridge + + +class TestDenseBinding: + def test_dense_layer_adopts_gated_mlp_aliases(self) -> None: + bridge = _bind(_DenseMLP()) + assert bridge._bound_dense is True + assert bridge.hook_aliases == { + "hook_pre": "dense_gate.hook_out", + "hook_pre_linear": "dense_in.hook_out", + "hook_post": "dense_out.hook_in", + } + + def test_dense_layer_exposes_neuron_basis_weights(self) -> None: + """The #1645 fix advertises neuron hooks; the weight API must back them + (ActivationCache neuron-result stacking and weight collection read these, + and a missing accessor is silently zero-filled by get_params_util).""" + module = _DenseMLP() + bridge = _bind(module) + # Pin each accessor to the projection it must read. Shape alone cannot: + # gate_proj and up_proj are both [d_model, d_mlp], so a W_gate that + # returned the UP projection would satisfy any shape-only assertion. + torch.testing.assert_close(bridge.W_gate, module.gate_proj.weight.T) + torch.testing.assert_close(bridge.W_in, module.up_proj.weight.T) + torch.testing.assert_close(bridge.W_out, module.down_proj.weight.T) + # Negative control: the two same-shaped projections are distinguishable + # in this fixture, so the assertions above are not trivially satisfiable. + assert not torch.equal(module.gate_proj.weight, module.up_proj.weight) + + def test_dense_layer_drops_the_router_hook(self) -> None: + """A dense layer has no router; advertising hook_router_scores would be a + hook that can never fire.""" + assert not hasattr(_bind(_DenseMLP()), "hook_router_scores") + + def test_dense_gate_hook_carries_the_gate_projection_output(self) -> None: + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + module = _DenseMLP() + x = torch.randn(2, 3, D_MODEL) + expected = module.gate_proj(x).detach() + bridge = _bind(module) + + seen: list[torch.Tensor] = [] + bridge.submodules["dense_gate"].hook_out.add_hook( + lambda tensor, hook: seen.append(tensor.clone()) + ) + try: + with torch.no_grad(): + out = bridge(x) + finally: + bridge.submodules["dense_gate"].hook_out.remove_hooks() + + assert len(seen) == 1 + assert seen[0].shape == (2, 3, D_MLP) # neuron basis, not d_model + torch.testing.assert_close(seen[0], expected) + assert out.shape == (2, 3, D_MODEL) + + +class TestSparseBinding: + def test_sparse_layer_keeps_moe_semantics(self) -> None: + bridge = _bind(_SparseMoE()) + assert bridge._bound_dense is False + # Pin the concrete targets rather than comparing to + # MoEBridge.hook_aliases: that is the class constant this bind is + # supposed to reproduce, so the comparison holds even if the constant + # itself is wrong. hook_pre/hook_post are the MoE block boundaries. + assert bridge.hook_aliases == {"hook_pre": "hook_in", "hook_post": "hook_out"} + assert hasattr(bridge, "hook_router_scores") + + def test_sparse_layer_has_no_dense_weight_accessors(self) -> None: + """Sparse layers have per-expert weights and no single W_*; hasattr must + stay False so weight-collection helpers skip them as before.""" + bridge = _bind(_SparseMoE()) + assert not hasattr(bridge, "W_in") + assert not hasattr(bridge, "W_gate") + assert not hasattr(bridge, "W_out") + + def test_gate_name_is_not_overloaded_across_layers(self) -> None: + """blocks.N.mlp.gate must not mean the router on one layer and a d_mlp + gate projection on another — that per-layer semantic flip is the #1645 + flaw class this dispatch removes.""" + dense, sparse = _bind(_DenseMLP()), _bind(_SparseMoE()) + # The two layers must not resolve the same hook name to different + # KINDS of tensor. Compare the resolved targets directly. + assert dense.hook_aliases["hook_pre"] == "dense_gate.hook_out" + assert sparse.hook_aliases["hook_pre"] == "hook_in" + # And `gate` must never be a dense projection: on the dense layer the + # key is absent entirely, so `blocks.N.mlp.gate.hook_out` means the + # router on every layer that has it. + assert "gate" not in dense.submodules + assert "gate" in sparse.submodules + + +class TestDispatchRobustness: + def test_one_template_serves_both_layer_types(self) -> None: + """Block setup deepcopies a single template per layer; each copy must + bind independently and leave the template untouched.""" + template = _template() + dense = copy.deepcopy(template) + sparse = copy.deepcopy(template) + dense.set_original_component(_DenseMLP()) + sparse.set_original_component(_SparseMoE()) + assert dense.hook_aliases["hook_pre"] == "dense_gate.hook_out" + assert sparse.hook_aliases == {"hook_pre": "hook_in", "hook_post": "hook_out"} + assert template.hook_aliases == {"hook_pre": "hook_in", "hook_post": "hook_out"} + + def test_rebinding_sparse_after_dense_restores_moe_state(self) -> None: + """The morph is symmetric: a rebinding harness must not leave a chimera + with dense aliases and no router hook on a sparse layer.""" + bridge = _bind(_DenseMLP()) + assert bridge._bound_dense is True + bridge.set_original_component(_SparseMoE()) + assert bridge._bound_dense is False + assert bridge.hook_aliases == {"hook_pre": "hook_in", "hook_post": "hook_out"} + assert hasattr(bridge, "hook_router_scores") + + def test_alias_rebind_survives_the_attribute_passthrough(self) -> None: + """GeneralizedComponent.__setattr__ forwards unknown attributes to the + wrapped module; hook_aliases must be exempt or the rebind vanishes + whenever the HF module happens to expose that attribute.""" + module = _DenseMLP() + module.hook_aliases = {"decoy": "value"} # type: ignore[assignment] + bridge = _bind(module) + assert bridge.hook_aliases["hook_pre"] == "dense_gate.hook_out" + + def test_undeclared_dense_projections_leave_moe_mapping(self) -> None: + """Detection is positive-only: without declared dense_* submodules the + bridge must not guess a layer is dense and strip its MoE hooks.""" + bridge = MoEBridge(name="mlp", submodules={}) + bridge.set_original_component(_DenseMLP()) + assert bridge._bound_dense is False + assert bridge.hook_aliases == {"hook_pre": "hook_in", "hook_post": "hook_out"} + + +# Every adapter declaring dense_* projections. Kept in sync with +# `grep -l '"dense_in"' transformer_lens/model_bridge/supported_architectures/` +# by test_roster_covers_every_dense_declaring_adapter below, so a new adapter +# cannot join the dispatch without also joining these guards. +DENSE_AWARE_ARCHS = [ + "DeepseekV2ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", + "GlmMoeDsaForCausalLM", + "Ernie4_5_MoeForCausalLM", + "AfmoeForCausalLM", + "Qwen2MoeForCausalLM", + "Qwen3MoeForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3VLMoeForConditionalGeneration", + "LLaDA2MoeModelLM", + "LagunaForCausalLM", + "Llama4ForConditionalGeneration", +] + + +def test_roster_covers_every_dense_declaring_adapter() -> None: + """The roster parametrizes the guards below; an adapter that declares dense_* + but is missing here would be the one place the guards do not reach.""" + adapters_dir = Path(transformer_lens.model_bridge.supported_architectures.__file__).parent + declaring = { + path.stem for path in adapters_dir.glob("*.py") if '"dense_in"' in path.read_text() + } + # Adapters whose dense mapping is reached through a different template shape + # (a per-config builder or an encoder-decoder block list) rather than + # blocks.mlp, so the blocks-based parametrization cannot construct them. + # jamba builds its MoEBridge only when num_experts > 1 (a per-config + # builder) and switch_transformers maps encoder/decoder block lists, so + # neither is reachable through a blocks.mlp template here. + NOT_BLOCKS_MLP = {"jamba", "switch_transformers"} + # Walk the MRO: an arch string can resolve to a subclass in another module + # (Llama4ForConditionalGeneration -> the multimodal adapter, which inherits + # llama4's block mapping), and the declaration lives on the base. + covered_modules = { + klass.__module__.rsplit(".", 1)[-1] + for arch in DENSE_AWARE_ARCHS + for klass in type( + ArchitectureAdapterFactory.select_architecture_adapter(make_bridge_cfg(arch, d_head=8)) + ).__mro__ + } + missing = declaring - covered_modules - NOT_BLOCKS_MLP + assert not missing, f"adapters declare dense_* but are outside the guards: {sorted(missing)}" + + +@pytest.mark.parametrize("architecture", DENSE_AWARE_ARCHS) +def test_adapter_templates_bind_dense_layers_as_gated_mlps(architecture: str) -> None: + """Every interleaved/dense-prefix MoE adapter must declare the dense + projections, so its dense layers get neuron-basis hooks (#1645).""" + 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"] + instance = copy.deepcopy(template) + instance.set_original_component(_DenseMLP()) + assert instance._bound_dense is True, f"{architecture} did not declare dense projections" + assert instance.hook_aliases["hook_pre"] == "dense_gate.hook_out" + + +class _UngatedDenseFF(nn.Module): + """Switch-style ungated dense feed-forward (wi/wo, no gate projection).""" + + def __init__(self) -> None: + super().__init__() + self.wi = nn.Linear(D_MODEL, D_MLP, bias=False) + self.wo = nn.Linear(D_MLP, D_MODEL, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.wo(torch.relu(self.wi(hidden_states))) + + +class _RenamedRouterSparseMoE(nn.Module): + """Sparse block whose router HF renamed out from under the adapter.""" + + def __init__(self) -> None: + super().__init__() + self.router = nn.Linear(D_MODEL, 4, bias=False) + self.experts = nn.ModuleList([nn.Linear(D_MODEL, D_MODEL) for _ in range(4)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +class TestUngatedDenseBinding: + """Not every dense feed-forward is gated — Switch's is wi/wo.""" + + def test_ungated_dense_binds_mlp_alias_set(self) -> None: + bridge = MoEBridge( + name="mlp", + submodules={ + "dense_in": LinearBridge(name="wi", optional=True), + "dense_out": LinearBridge(name="wo", optional=True), + }, + ) + module = _UngatedDenseFF() + bridge.set_original_component(module) + for key, attr in (("dense_in", "wi"), ("dense_out", "wo")): + sub = bridge.submodules[key] + sub.set_original_component(getattr(module, attr)) + bridge.add_module(key, sub) + + assert bridge.bound_dense is True + assert bridge.hook_aliases == { + "hook_pre": "dense_in.hook_out", + "hook_post": "dense_out.hook_in", + } + # No gate projection exists, so no hook_pre_linear and no W_gate. + assert "hook_pre_linear" not in bridge.hook_aliases + assert not hasattr(bridge, "W_gate") + assert bridge.W_in.shape == (D_MODEL, D_MLP) + + +class TestSparseRequiredGuard: + """`optional` must not mean a renamed HF router silently loses its hooks.""" + + def test_renamed_router_on_sparse_layer_raises(self) -> None: + adapter = _adapter() + bridge = copy.deepcopy(adapter.component_mapping["blocks"].submodules["mlp"]) + module = _RenamedRouterSparseMoE() + bridge.set_original_component(module) + with pytest.raises(ValueError, match="required submodule"): + setup_submodules(bridge, adapter, module) + + def test_correctly_named_sparse_router_binds_clean(self) -> None: + bridge = _bind(_SparseMoE()) + assert bridge.bound_dense is False + assert "gate" in bridge.submodules + + def test_dense_layer_may_legitimately_lack_the_router(self) -> None: + """Router absence is expected on dense layers — must not raise.""" + bridge = _bind(_DenseMLP()) + assert bridge.bound_dense is True + + def test_undeclared_sparse_required_key_raises_at_construction(self) -> None: + """An unvalidated opt-in string would silently disable the guard.""" + with pytest.raises(ValueError, match="not declared submodules"): + MoEBridge( + name="mlp", + submodules={"gate": LinearBridge(name="gate", optional=True)}, + sparse_required=("rooter",), + ) + + def test_opting_in_is_what_makes_a_missing_router_loud(self) -> None: + """Differential on the opt-in alone: the SAME template shape and the SAME + module bind silently without `sparse_required` and raise with it. + + Driven through the real setup_submodules so the skipped set is computed + rather than hand-fed — hand-feeding it would exercise the guard's body + while skipping the machinery that decides when the guard applies. + """ + adapter = _adapter() + module = _RenamedRouterSparseMoE() + + def build(**kwargs) -> MoEBridge: + return MoEBridge( + name="mlp", + submodules={"gate": LinearBridge(name="gate", optional=True)}, + **kwargs, + ) + + # jamba/switch-style: no opt-in, so a skipped optional stays silent. + opted_out = build() + opted_out.set_original_component(module) + setup_submodules(opted_out, adapter, module) + assert "gate" not in opted_out.submodules # precondition: it WAS skipped + + # Same template + same module, only the opt-in differs. + opted_in = build(sparse_required=("gate",)) + opted_in.set_original_component(module) + with pytest.raises(ValueError, match="required submodule"): + setup_submodules(opted_in, adapter, module) + + +@pytest.mark.parametrize("architecture", DENSE_AWARE_ARCHS) +def test_dense_aware_adapters_declare_a_sparse_required_router(architecture: str) -> None: + """Every dense-aware MoE adapter must opt into the loud-on-rename guard, so a + new adapter cannot be added with a silently-skippable router.""" + 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"] + assert template._sparse_required, f"{architecture} declares no sparse_required router" + assert set(template._sparse_required) <= set(template.submodules) + + +class TestDeclaredGateMustResolve: + """A declared dense_gate that does not resolve is a RENAME, not evidence the + MLP is ungated — binding it ungated aliases hook_pre to the up projection.""" + + def test_renamed_gate_raises_instead_of_binding_ungated(self) -> None: + class _RenamedGateDenseMLP(nn.Module): + def __init__(self) -> None: + super().__init__() + self.w_gate = nn.Linear(D_MODEL, D_MLP, bias=False) + self.up_proj = nn.Linear(D_MODEL, D_MLP, bias=False) + self.down_proj = nn.Linear(D_MLP, D_MODEL, bias=False) + + bridge = copy.deepcopy(_template()) + with pytest.raises(ValueError, match="dense MLPs are gated"): + bridge.set_original_component(_RenamedGateDenseMLP()) + + def test_undeclared_gate_still_binds_ungated(self) -> None: + """Switch-style: no dense_gate declared, so ungated is the truth.""" + bridge = MoEBridge( + name="mlp", + submodules={ + "dense_in": LinearBridge(name="wi", optional=True), + "dense_out": LinearBridge(name="wo", optional=True), + }, + ) + + class _Ungated(nn.Module): + def __init__(self) -> None: + super().__init__() + self.wi = nn.Linear(D_MODEL, D_MLP) + self.wo = nn.Linear(D_MLP, D_MODEL) + + bridge.set_original_component(_Ungated()) + assert bridge.bound_dense is True + assert bridge.hook_aliases["hook_pre"] == "dense_in.hook_out" 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 4ec2bf327..967c84b2c 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 @@ -215,7 +215,13 @@ class TestDeepSeekV3AdapterMoE: 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"} + assert set(mlp.submodules.keys()) == { + "gate", + "shared_experts", + "dense_gate", + "dense_in", + "dense_out", + } 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 3d5e29911..c95ed05a0 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 @@ -55,7 +55,13 @@ 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"} + assert set(mlp.submodules) == { + "gate", + "shared_experts", + "dense_gate", + "dense_in", + "dense_out", + } 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_gemma2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gemma2_adapter.py index eb83b6707..f13969f6f 100644 --- a/tests/unit/model_bridge/supported_architectures/test_gemma2_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_gemma2_adapter.py @@ -7,7 +7,6 @@ of R * sqrt(d_model)) — the dev-4.x vLLM-comparison investigation surfaced this. """ -import os import pytest import torch @@ -47,7 +46,6 @@ def test_adapter_does_not_override_setup_hook_compatibility( assert "setup_hook_compatibility" not in vars(type(adapter)) -@pytest.mark.skipif(bool(os.getenv("CI")), reason="Network/disk fetch of tiny Gemma2 — skip in CI") def test_gemma2_embed_hook_out_magnitude_matches_sqrt_d_model_scaling(): """End-to-end regression for the embed double-scale bug. 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 7fae03e61..fccc5daa3 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 @@ -191,7 +191,7 @@ 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"} + assert set(mlp.submodules.keys()) == {"gate", "dense_gate", "dense_in", "dense_out"} 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 5939c77ae..c93b9a960 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 @@ -87,7 +87,13 @@ 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"} + assert set(mlp.submodules.keys()) == { + "gate", + "shared_experts", + "dense_gate", + "dense_in", + "dense_out", + } 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 30169f23d..f71386fc7 100644 --- a/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_jamba_adapter.py @@ -204,10 +204,14 @@ def test_moe_bridge_when_num_experts_gt_one(self) -> None: mlp = adapter.get_component_mapping()["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) assert mlp.name == "feed_forward" - assert set(mlp.submodules) == {"gate", "in", "out", "router"} + # 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["gate"].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",) class TestJambaFactoryRegistration: 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 a990eca37..d84308002 100644 --- a/tests/unit/model_bridge/supported_architectures/test_llama4_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_llama4_adapter.py @@ -48,12 +48,21 @@ def test_attention_stays_native(self, adapter): assert attn.name == "self_attn" def test_moe_with_optional_shared_expert(self, adapter): - """The router returns a tuple so it stays unwrapped; non-MoE layers - hold a dense gated MLP under the same feed_forward name.""" + """The router is mapped so a rename fails loudly; non-MoE layers hold a + dense gated MLP under the same feed_forward name.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, _Llama4MoEBridge) assert mlp.name == "feed_forward" - assert set(mlp.submodules) == {"shared_expert", "dense_gate", "dense_in", "dense_out"} + assert set(mlp.submodules) == { + "router", + "shared_expert", + "dense_gate", + "dense_in", + "dense_out", + } + # 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"] 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 f405b2fd1..025a06509 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 @@ -116,6 +116,9 @@ def test_moe_submodules(self, adapter: Qwen2MoeArchitectureAdapter) -> None: "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_5_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_moe_adapter.py index cd9ffd6f5..1c68de0f1 100644 --- a/tests/unit/model_bridge/supported_architectures/test_qwen3_5_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_qwen3_5_moe_adapter.py @@ -172,7 +172,17 @@ def test_mlp_is_moe_not_dense(self, adapter): mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) assert not isinstance(mlp, GatedMLPBridge) - assert set(mlp.submodules) == {"gate", "experts", "shared_expert", "shared_expert_gate"} + # No dense fallback exists for this arch — HF builds + # Qwen3_5MoeSparseMoeBlock unconditionally (modeling_qwen3_5_moe.py:849), + # so there are no dense_* projections to declare and all four MoE + # submodules are required rather than optional. + assert set(mlp.submodules) == { + "gate", + "experts", + "shared_expert", + "shared_expert_gate", + } + assert all(not sub.optional for sub in mlp.submodules.values()) assert isinstance(mlp.submodules["gate"], MoERouterBridge) def test_gated_q_proj_flag_set(self, adapter): 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 6f88369ae..31c4f74ef 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 @@ -220,10 +220,14 @@ def test_mlp_is_moe_not_gated_mlp(self, adapter: Qwen3MoeArchitectureAdapter) -> assert isinstance(mlp, MoEBridge) assert not isinstance(mlp, GatedMLPBridge) - def test_mlp_has_only_gate_submodule(self, adapter: Qwen3MoeArchitectureAdapter) -> None: - """Experts are batched 3D tensors inside the MoE block — only the router is mapped.""" + 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.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] - assert set(mlp.submodules.keys()) == {"gate"} + assert set(mlp.submodules.keys()) == {"gate", "dense_gate", "dense_in", "dense_out"} 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 65ab7602f..803d950fc 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 @@ -82,10 +82,20 @@ def test_mlp_path(self, adapter): submodules = adapter.component_mapping["blocks"].submodules assert submodules["mlp"].name == "mlp" - def test_mlp_has_no_submodules(self, adapter): - """Qwen3NextSparseMoeBlock has a non-Linear router and 3D batched experts; MoEBridge delegates to HF forward, so no internal subs are mapped.""" + 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.""" mlp = adapter.component_mapping["blocks"].submodules["mlp"] - assert mlp.submodules == {} + assert set(mlp.submodules) == { + "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.""" 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 3c19a1836..d23a7142e 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 @@ -46,7 +46,7 @@ def test_moe_mlp(self, adapter): mlp = adapter.component_mapping["blocks"].submodules["mlp"] assert isinstance(mlp, MoEBridge) - assert set(mlp.submodules) == {"gate", "experts"} + assert set(mlp.submodules) == {"gate", "experts", "dense_gate", "dense_in", "dense_out"} 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/model_bridge/test_alias_collection_traversal.py b/tests/unit/model_bridge/test_alias_collection_traversal.py new file mode 100644 index 000000000..829d1a8b8 --- /dev/null +++ b/tests/unit/model_bridge/test_alias_collection_traversal.py @@ -0,0 +1,146 @@ +"""Traversal safety for the two alias collectors on the hook_dict path. + +Both walk an arbitrary object graph (component templates and, per layer, bound +block subtrees). A component reachable from itself must be cut rather than +recursed forever, and a component legitimately shared under two names must +still contribute aliases at both — a globally-visited set would silently drop +the second path. +""" + +from __future__ import annotations + +import re +import sys + +import pytest +import torch + +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components import LinearBridge, MLPBridge + +MODEL_NAME = "hf-internal-testing/tiny-random-gpt2" + + +@pytest.fixture(scope="module") +def bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers(MODEL_NAME, device="cpu", dtype=torch.float32) + + +@pytest.fixture +def block_template(bridge: TransformerBridge): + return bridge.adapter.component_mapping["blocks"] + + +def _shared_component() -> MLPBridge: + """A component that actually declares hook_aliases (LinearBridge declares none, + so sharing one would make a diamond test vacuous).""" + return MLPBridge( + name="c_fc", + submodules={"in": LinearBridge(name="c_fc"), "out": LinearBridge(name="c_proj")}, + ) + + +class TestTemplateCollectorTraversal: + """_collect_component_aliases runs first on every hook_dict access.""" + + def test_self_cycle_is_cut(self, bridge: TransformerBridge, block_template) -> None: + mlp = block_template.submodules["mlp"] + mlp.submodules["self_cycle"] = mlp + original_limit = sys.getrecursionlimit() + sys.setrecursionlimit(300) + try: + bridge._collect_component_aliases(bridge.adapter.component_mapping) + finally: + sys.setrecursionlimit(original_limit) + del mlp.submodules["self_cycle"] + + def test_mutual_cycle_is_cut(self, bridge: TransformerBridge, block_template) -> None: + mlp = block_template.submodules["mlp"] + attn = block_template.submodules["attn"] + mlp.submodules["to_attn"] = attn + attn.submodules["to_mlp"] = mlp + original_limit = sys.getrecursionlimit() + sys.setrecursionlimit(300) + try: + bridge._collect_component_aliases(bridge.adapter.component_mapping) + finally: + sys.setrecursionlimit(original_limit) + del mlp.submodules["to_attn"] + del attn.submodules["to_mlp"] + + def test_shared_component_aliases_at_every_path( + self, bridge: TransformerBridge, block_template + ) -> None: + """Diamond, not cycle: the guard must be path-scoped.""" + mlp = block_template.submodules["mlp"] + shared = _shared_component() + mlp.submodules["alpha"] = shared + mlp.submodules["beta"] = shared + try: + aliases = bridge._collect_component_aliases(bridge.adapter.component_mapping) + finally: + del mlp.submodules["alpha"] + del mlp.submodules["beta"] + assert "blocks.mlp.alpha.hook_pre" in aliases + assert "blocks.mlp.beta.hook_pre" in aliases + + def test_list_valued_target_does_not_reach_the_cached_consumer( + self, bridge: TransformerBridge, block_template + ) -> None: + """_compute_hook_aliases_cached reverse-matches with str.endswith and is + lru_cached, so a list target is neither matchable nor hashable.""" + attn = block_template.submodules["attn"] + saved = attn.hook_aliases + attn.hook_aliases = {**dict(saved), "hook_fallback": ["hook_out", "hook_in"]} + try: + bridge._hook_registry_initialized = False + bridge._block_alias_cache = None + bridge._initialize_hook_registry() + bridge._collect_hook_aliases_from_registry() # must not raise TypeError + finally: + attn.hook_aliases = saved + bridge._hook_registry_initialized = False + bridge._block_alias_cache = None + bridge._initialize_hook_registry() + + +class TestBoundBlockWalkTraversal: + """_collect_block_instance_aliases walks each bound block's subtree.""" + + def test_self_cycle_in_a_bound_block_is_cut(self, bridge: TransformerBridge) -> None: + mlp = bridge.blocks[0].mlp + mlp.submodules["self_cycle"] = mlp + bridge._block_alias_cache = None + try: + bridge._collect_block_instance_aliases() + finally: + del mlp.submodules["self_cycle"] + bridge._block_alias_cache = None + + def test_shared_submodule_aliases_at_every_path(self, bridge: TransformerBridge) -> None: + """A component reachable under two names must be walked at BOTH — a + globally-visited guard would silently skip the second path. + + The shared component's alias targets are not in the hook registry (it is + attached for this test only), so the observable evidence that both paths + were walked is the unresolved-alias report: one entry per alias PER + PATH. Asserting on the resolved map instead would be vacuous — it is + empty for both paths, so any comparison between them holds trivially. + """ + mlp = bridge.blocks[0].mlp + shared = _shared_component() + mlp.submodules["alpha"] = shared + mlp.submodules["beta"] = shared + bridge._block_alias_cache = None + try: + with pytest.warns(UserWarning, match="did not resolve") as record: + bridge._collect_block_instance_aliases() + finally: + del mlp.submodules["alpha"] + del mlp.submodules["beta"] + bridge._block_alias_cache = None + + reported = " ".join(str(w.message) for w in record) + count = int(re.search(r"^(\d+) block hook alias", reported).group(1)) + # _shared_component declares 2 aliases; both paths walked => 2 per path. + assert count == 4, f"expected 2 aliases x 2 paths, got {count}: {reported}" diff --git a/tests/unit/model_bridge/test_get_params_util.py b/tests/unit/model_bridge/test_get_params_util.py index 97362dd2d..e6ff03b7b 100644 --- a/tests/unit/model_bridge/test_get_params_util.py +++ b/tests/unit/model_bridge/test_get_params_util.py @@ -73,15 +73,19 @@ def test_get_bridge_params_attention_reshaping(self): w_v = params[f"blocks.{layer_idx}.attn.W_V"] w_o = params[f"blocks.{layer_idx}.attn.W_O"] - # Should be reshaped to [n_heads, d_model, d_head] format - expected_shape = (12, 768, 64) # n_heads=12, d_model=768, d_head=64 - assert w_q.shape == expected_shape - assert w_k.shape == expected_shape - assert w_v.shape == expected_shape - - # Output should be [n_heads, d_head, d_model] - expected_o_shape = (12, 64, 768) - assert w_o.shape == expected_o_shape + # Shape alone cannot catch a reshape that scrambles the elements, + # nor Q/K/V being read from the wrong projection: pin the VALUES + # against the source weights the mock block exposes. + block = mock_bridge.blocks[layer_idx] + n_heads, d_model, d_head = 12, 768, 64 + assert w_q.shape == (n_heads, d_model, d_head) + assert w_o.shape == (n_heads, d_head, d_model) + torch.testing.assert_close(w_q, block.attn.q.weight.reshape(n_heads, d_model, d_head)) + torch.testing.assert_close(w_o, block.attn.o.weight.reshape(n_heads, d_head, d_model)) + # Negative control: Q and K must be distinguishable in this fixture, + # or reading either one would satisfy the assertions above. + assert not torch.equal(w_q, w_k) + assert not torch.equal(w_k, w_v) def test_get_bridge_params_bias_handling(self): """Test that biases are handled correctly, including None biases.""" diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py index 2e89f497d..8934bb98a 100644 --- a/transformer_lens/model_bridge/bridge.py +++ b/transformer_lens/model_bridge/bridge.py @@ -16,6 +16,7 @@ Any, Callable, Dict, + FrozenSet, Iterator, List, Literal, @@ -43,6 +44,7 @@ from transformer_lens.model_bridge.exceptions import StopAtLayerException from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, + alias_generation, ) from transformer_lens.model_bridge.generalized_components.block import ( _BLOCK_INTERNAL_MODULES, @@ -191,6 +193,7 @@ def __init__(self, model: nn.Module, adapter: ArchitectureAdapter, tokenizer: An self._hook_registry: Dict[str, HookPoint] = {} self._hook_registry_initialized = False self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {} + self._block_alias_cache: Optional[Tuple[Tuple[int, int], Dict[str, str]]] = None self._property_alias_registry: Dict[str, str] = {} # real_components maps TL keys to (remote_path, actual_instance) tuples # For list components, actual_instance will be a list of component instances @@ -523,23 +526,44 @@ def _initialize_hook_registry(self) -> None: self._scan_existing_hooks(self, "") self._hook_registry_initialized = True - def _collect_component_aliases(self, component_mapping, prefix=""): - """Recursively collect aliases from components.""" - aliases = {} + def _collect_component_aliases(self, component_mapping, prefix="", _ancestors=frozenset()): + """Recursively collect aliases from the architecture's component templates. + + ``_ancestors`` holds the ids on the current path, cutting a component + reachable from itself (directly or mutually): the mapping is an + arbitrary object graph, and without the guard a cycle recurses until + RecursionError during hook-registry construction. Tracking the path + rather than every visited node keeps a component legitimately shared + under two names (a diamond) contributing aliases at both. + """ + aliases: Dict[str, str] = {} + if id(component_mapping) in _ancestors: + return aliases + _ancestors = _ancestors | {id(component_mapping)} if isinstance(component_mapping, dict): for name, component in component_mapping.items(): sub_prefix = f"{prefix}.{name}" if prefix else name - aliases.update(self._collect_component_aliases(component, sub_prefix)) + aliases.update(self._collect_component_aliases(component, sub_prefix, _ancestors)) else: if hasattr(component_mapping, "hook_aliases") and component_mapping.hook_aliases: for alias_name, target in component_mapping.hook_aliases.items(): + # Fallback-list targets are skipped here: the consumer + # (_compute_hook_aliases_cached) reverse-matches with + # str.endswith and is lru_cached, so a list is neither + # matchable nor hashable. Component-level lists resolve + # through _collect_block_instance_aliases (block subtrees) + # and _add_aliases_to_hooks (bridge level) instead. + if not isinstance(target, str): + continue full_alias = f"{prefix}.{alias_name}" if prefix else alias_name full_target = f"{prefix}.{target}" if prefix else target aliases[full_alias] = full_target if hasattr(component_mapping, "submodules") and component_mapping.submodules: for sub_name, sub_component in component_mapping.submodules.items(): sub_prefix = f"{prefix}.{sub_name}" if prefix else sub_name - aliases.update(self._collect_component_aliases(sub_component, sub_prefix)) + aliases.update( + self._collect_component_aliases(sub_component, sub_prefix, _ancestors) + ) return aliases @staticmethod @@ -591,8 +615,20 @@ def _collect_block_instance_aliases(self) -> Dict[str, str]: The template collection above reads ``adapter.component_mapping`` and so cannot see aliases a block rebinds per layer at bind time (heterogeneous - architectures like OlmoHybrid) or prunes for absent optional submodules. + architectures like OlmoHybrid, MoEBridge's dense/sparse dispatch) or + prunes for absent optional submodules. + + Memoized against (hook-registry size, alias generation): this walks + every block's submodule tree and hook_dict reads it on each access. + Size alone is not a sufficient key — a dense<->sparse rebind changes + what the aliases point at while leaving the registry the same size, so + the generation counter (bumped on every hook_aliases assignment) is + what makes a post-boot rebind visible here. """ + cache_key = (len(self._hook_registry), alias_generation()) + cached = self._block_alias_cache + if cached is not None and cached[0] == cache_key: + return cached[1] aliases: Dict[str, str] = {} unresolved: List[str] = [] for bl_name in _BLOCK_LIST_ATTRS: @@ -600,31 +636,51 @@ def _collect_block_instance_aliases(self) -> Dict[str, str]: if block_list is None: continue for i, block in enumerate(block_list): - block_aliases = getattr(block, "hook_aliases", None) - if not block_aliases: - continue # A block with no registered hooks means the registry hasn't # scanned it yet — unresolved aliases there are timing, not drops. block_prefix = f"{bl_name}.{i}." if f"{block_prefix}hook_in" not in self._hook_registry: continue - for alias_name, target in block_aliases.items(): - targets = target if isinstance(target, list) else [target] - for single_target in targets: - full_target = f"{block_prefix}{single_target}" - if full_target in self._hook_registry: - aliases[f"{block_prefix}{alias_name}"] = full_target - break - else: - unresolved.append(f"{block_prefix}{alias_name}") + # Walk the block and its submodule tree: components rebind + # aliases per layer at bind time either at block level + # (OlmoHybrid) or one level down (MoEBridge's dense/sparse + # dispatch). id()-seen guards against shared/cyclic submodule + # references, which would otherwise hang boot. + # (prefix, component, ids-on-this-path): path-scoped rather than + # globally-visited so a cycle is cut while a component shared + # under two names still contributes aliases at both. + stack: List[Tuple[str, Any, FrozenSet[int]]] = [("", block, frozenset())] + while stack: + sub_prefix, component, ancestors = stack.pop() + if id(component) in ancestors: + continue + ancestors = ancestors | {id(component)} + component_aliases = getattr(component, "hook_aliases", None) + if component_aliases: + for alias_name, target in component_aliases.items(): + targets = target if isinstance(target, list) else [target] + for single_target in targets: + full_target = f"{block_prefix}{sub_prefix}{single_target}" + if full_target in self._hook_registry: + aliases[f"{block_prefix}{sub_prefix}{alias_name}"] = full_target + break + else: + unresolved.append(f"{block_prefix}{sub_prefix}{alias_name}") + for nested_name, nested in ( + getattr(component, "submodules", None) or {} + ).items(): + stack.append((f"{sub_prefix}{nested_name}.", nested, ancestors)) if unresolved: # Surface drops instead of silently swallowing, mirroring # GeneralizedComponent._register_aliases. warnings.warn( f"{len(unresolved)} block hook alias(es) did not resolve to a " - f"registered hook (e.g. '{unresolved[0]}').", + f"registered hook (e.g. '{unresolved[0]}'). Any such alias falls " + "back to the architecture template's mapping, which for a " + "per-layer rebind is the wrong tensor for this layer.", stacklevel=2, ) + self._block_alias_cache = (cache_key, aliases) return aliases def _add_aliases_to_hooks(self, hooks: Dict[str, HookPoint]) -> None: diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index 6294247b6..c628c643d 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -152,6 +152,12 @@ def setup_submodules( component.submodules.pop(name, None) if skipped_optional: _prune_hook_aliases_for_skipped(component, skipped_optional) + # Components whose submodules are optional only in some configurations + # validate here: which optionals were actually skipped is not knowable + # until every submodule has been resolved. + validate_after_setup = getattr(component, "validate_after_setup", None) + if callable(validate_after_setup): + validate_after_setup(skipped_optional) def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None: diff --git a/transformer_lens/model_bridge/generalized_components/base.py b/transformer_lens/model_bridge/generalized_components/base.py index a58f62df5..af8087282 100644 --- a/transformer_lens/model_bridge/generalized_components/base.py +++ b/transformer_lens/model_bridge/generalized_components/base.py @@ -30,6 +30,23 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: return out +# Bumped whenever any component rebinds its hook_aliases (MoEBridge's dense/ +# sparse dispatch, OlmoHybrid's per-layer selection). Alias caches key on it: +# a rebind can leave the hook REGISTRY unchanged while changing what the +# aliases point at, so size-based cache keys cannot see it. +_ALIAS_GENERATION = 0 + + +def alias_generation() -> int: + """Current global alias-rebind generation.""" + return _ALIAS_GENERATION + + +def _bump_alias_generation() -> None: + global _ALIAS_GENERATION + _ALIAS_GENERATION += 1 + + class GeneralizedComponent(nn.Module): """Base class for generalized transformer components. @@ -386,6 +403,9 @@ def __setattr__(self, name: str, value: Any) -> None: self._register_hook(name, value) super().__setattr__(name, value) return + if name == "hook_aliases": + # Any alias rebind invalidates alias caches downstream. + _bump_alias_generation() if name.startswith("_") or name in [ "name", "config", @@ -394,6 +414,12 @@ def __setattr__(self, name: str, value: Any) -> None: "compatibility_mode", "disable_warnings", "optional", + # Components rebind these per layer at bind time (MoEBridge's + # dense/sparse dispatch). Without the carve-out the assignment is + # forwarded to the wrapped HF module whenever it happens to expose + # the attribute — the rebind then silently vanishes. + "hook_aliases", + "property_aliases", # train()/eval() set self.training; redirecting it to the original # component leaves the wrapper stuck in training mode (dropout at # inference). Recursion still reaches the original via _modules. diff --git a/transformer_lens/model_bridge/generalized_components/mlp.py b/transformer_lens/model_bridge/generalized_components/mlp.py index bb0059b28..a56e40cf1 100644 --- a/transformer_lens/model_bridge/generalized_components/mlp.py +++ b/transformer_lens/model_bridge/generalized_components/mlp.py @@ -11,6 +11,58 @@ ) +def weight_layout_in_out(proj: Any) -> Optional[bool]: + """Whether proj's wrapped module stores its weight as [in, out]. + + Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores + [out_features, in_features]. Returns None when the wrapped module is + neither, so callers fall back to in_features/out_features or a shape heuristic. + """ + from transformers.pytorch_utils import Conv1D + + component = getattr(proj, "original_component", None) + if isinstance(component, Conv1D): + return True + if isinstance(component, torch.nn.Linear): + return False + return None + + +def normalize_mlp_weight( + weight: torch.Tensor, layout: Optional[bool], proj: Any, pattern: str = "in" +) -> torch.Tensor: + """Normalize an MLP projection weight to TransformerLens orientation. + + Args: + weight: 2D weight tensor from the projection + layout: True if [in, out] (Conv1D), False if [out, in] (nn.Linear), + None falls back to in_features/out_features or shape heuristic. + proj: The projection module (LinearBridge) for metadata fallback. + pattern: "in" for W_in/W_gate [d_model, d_mlp], "out" for W_out [d_mlp, d_model] + + Note: the shape heuristic assumes d_model < d_mlp, which fails for + architectures like GIDD's ScaledLinear where d_mlp < d_model. When layout is + None, in_features/out_features are checked first; the heuristic is last resort. + """ + if layout is None: + # Try in_features/out_features from the wrapped module (works for bare nn.Module) + component = getattr(proj, "original_component", None) + in_f = getattr(component, "in_features", None) + out_f = getattr(component, "out_features", None) + if in_f is not None and out_f is not None: + # Module declares its orientation; weight[0] == in_f means [in, out] layout + layout = weight.shape[0] == in_f + else: + # Last resort: shape heuristic. WARNING: assumes d_model < d_mlp. + if pattern == "in": + layout = weight.shape[0] < weight.shape[1] + else: + layout = weight.shape[0] > weight.shape[1] + if layout: + return weight # Conv1D-style: already in TL orientation + return weight.T # nn.Linear-style: transpose to TL orientation + + class MLPBridge(GeneralizedComponent): """Bridge component for MLP layers. @@ -101,55 +153,14 @@ def forward(self, *args, **kwargs) -> Any: return output def _weight_layout_in_out(self, proj: Any) -> Optional[bool]: - """Whether proj's wrapped module stores its weight as [in, out]. - - Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores - [out_features, in_features]. Returns None when the wrapped module is - neither, so callers can fall back to in_features/out_features or shape heuristic. - """ - from transformers.pytorch_utils import Conv1D - - component = getattr(proj, "original_component", None) - if isinstance(component, Conv1D): - return True - if isinstance(component, torch.nn.Linear): - return False - return None + """Whether proj's wrapped module stores its weight as [in, out].""" + return weight_layout_in_out(proj) def _normalize_mlp_weight( self, weight: torch.Tensor, layout: Optional[bool], proj: Any, pattern: str = "in" ) -> torch.Tensor: - """Normalize MLP weight to TL orientation. - - Args: - weight: 2D weight tensor from the projection - layout: True if [in, out] (Conv1D), False if [out, in] (nn.Linear), - None falls back to in_features/out_features or shape heuristic. - proj: The projection module (LinearBridge) for metadata fallback. - pattern: "in" for W_in/W_gate [d_model, d_mlp], "out" for W_out [d_mlp, d_model] - - Note: Shape heuristic assumes d_model < d_mlp, which fails for architectures - like GIDD's ScaledLinear where d_mlp < d_model. When layout is None, we first - check in_features/out_features; shape heuristic is last resort. - """ - if layout is None: - # Try in_features/out_features from the wrapped module (works for bare nn.Module) - component = getattr(proj, "original_component", None) - in_f = getattr(component, "in_features", None) - out_f = getattr(component, "out_features", None) - if in_f is not None and out_f is not None: - # Module declares its orientation; weight[0] == out_f means [out, in] layout - layout = weight.shape[0] == in_f # True if [in, out], False if [out, in] - else: - # Last resort: shape heuristic. WARNING: assumes d_model < d_mlp. - # This will mis-orient architectures where d_mlp < d_model. - if pattern == "in": - layout = weight.shape[0] < weight.shape[1] - else: - layout = weight.shape[0] > weight.shape[1] - if layout: - return weight # Conv1D-style: already in TL orientation - return weight.T # nn.Linear-style: transpose to TL orientation + """Normalize MLP weight to TL orientation.""" + return normalize_mlp_weight(weight, layout, proj, pattern=pattern) @property def W_in(self) -> torch.Tensor: diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index a01235bcb..4ae3c243e 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/transformer_lens/model_bridge/generalized_components/moe.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import Any, Dict, Mapping, Optional +from typing import Any, Dict, Mapping, Optional, Tuple import torch @@ -14,6 +14,10 @@ GeneralizedComponent, ) from transformer_lens.model_bridge.generalized_components.linear import LinearBridge +from transformer_lens.model_bridge.generalized_components.mlp import ( + normalize_mlp_weight, + weight_layout_in_out, +) class MoEBridge(GeneralizedComponent): @@ -30,12 +34,39 @@ class MoEBridge(GeneralizedComponent): hook_aliases = {"hook_pre": "hook_in", "hook_post": "hook_out"} + # Dense-layer alias set, adopted per layer at bind time (see + # set_original_component). The keys are deliberately NOT gate/in/out: on a + # sparse layer of the same model ``gate`` is the ROUTER, so reusing it would + # make blocks.N.mlp.gate.hook_out mean the router on some layers and a d_mlp + # gate projection on others — the per-layer semantic flip this dispatch + # exists to remove. ``dense_*`` matches what llada2_moe/llama4/laguna already + # declare. dense_in/dense_out are what make a layer dense; dense_gate is + # present only for gated (SwiGLU) MLPs — ungated dense feed-forwards exist + # too (Switch Transformers' wi/wo), so the alias set follows the bound shape. + DENSE_SUBMODULE_KEYS = ("dense_in", "dense_out") + DENSE_GATE_KEY = "dense_gate" + _DENSE_HOOK_ALIASES = { + "hook_pre": "dense_in.hook_out", + "hook_post": "dense_out.hook_in", + } + _DENSE_GATED_HOOK_ALIASES = { + "hook_pre": "dense_gate.hook_out", + "hook_pre_linear": "dense_in.hook_out", + "hook_post": "dense_out.hook_in", + } + _DENSE_PROPERTY_ALIASES = { + "b_gate": "dense_gate.bias", + "b_in": "dense_in.bias", + "b_out": "dense_out.bias", + } + def __init__( self, name: str, config: Optional[Any] = None, submodules: Optional[Dict[str, GeneralizedComponent]] = {}, optional: bool = False, + sparse_required: Tuple[str, ...] = (), ): """Initialize the MoE bridge. @@ -44,9 +75,180 @@ def __init__( config: Optional configuration (unused for MoEBridge) submodules: Dictionary of GeneralizedComponent submodules to register optional: If True, setup skips this subtree when absent (dense layers) + sparse_required: Submodule keys that must be declared ``optional`` + (dense layers of an interleaved stack do not have them) but whose + absence on a SPARSE layer is an error rather than a silent skip. + Routers belong here: HF creates them unconditionally on sparse + blocks, so a skip means the attribute was renamed or moved, and + plain ``optional`` would drop their hooks without a word. """ super().__init__(name, config, submodules=submodules, optional=optional) self.hook_router_scores = HookPoint() + self._bound_dense = False + self._bound_dense_gate = False + # A misspelled key would silently disable the very guard that exists to + # stop silent degradation, so the opt-in is validated on construction. + unknown = set(sparse_required) - set(submodules or {}) + if unknown: + raise ValueError( + f"{name}: sparse_required {sorted(unknown)} are not declared " + f"submodules (declared: {sorted(submodules or {})})" + ) + self._sparse_required = sparse_required + + def _binds_dense_projections(self, component: torch.nn.Module) -> bool: + """Whether this layer is the dense variant of an interleaved MoE stack. + + Positive detection only: the adapter must have declared the dense + projections AND the wrapped module must actually expose them. Guessing + from the absence of ``experts`` would risk silently stripping hooks off + a sparse block whose experts are named differently. + """ + declared = [self.submodules.get(key) for key in self.DENSE_SUBMODULE_KEYS] + if not all(declared): + return False + return all( + sub is not None and sub.name is not None and hasattr(component, sub.name) + for sub in declared + ) + + def _binds_dense_gate(self, component: torch.nn.Module) -> bool: + """Whether the bound dense MLP is gated (SwiGLU) rather than plain. + + Declaring ``dense_gate`` is the adapter asserting this architecture's + dense layers ARE gated, so a declared-but-unresolvable gate means the HF + attribute was renamed — not that the MLP is ungated. Binding it as + ungated would silently alias hook_pre to the UP projection, which is the + #1645 defect class in a subtler form; raise instead. + """ + gate = self.submodules.get(self.DENSE_GATE_KEY) + if gate is None or gate.name is None: + return False # never declared: a genuinely ungated dense MLP + if hasattr(component, gate.name): + return True + raise ValueError( + f"{self.name}: dense layer wrapped {type(component).__name__} which has " + f"no {gate.name!r}, but this adapter declares {self.DENSE_GATE_KEY!r} — " + "so its dense MLPs are gated and the attribute was renamed or moved. " + "Binding it as ungated would alias hook_pre to the up projection. " + "Update the adapter's submodule name, or drop the declaration if this " + "architecture's dense layers really are ungated." + ) + + def set_original_component(self, component: torch.nn.Module) -> None: + """Bind the layer, adopting gated-MLP semantics on dense layers. + + Interleaved/dense-prefix MoE architectures build a plain gated MLP on + some layers under the same attribute name as the sparse block. Those + layers get the neuron-basis hooks (and weight accessors) every all-dense + architecture exposes, instead of MoE boundary tensors under + neuron-hook names (#1645). + """ + super().set_original_component(component) + is_dense = self._binds_dense_projections(component) + if is_dense == self._bound_dense: + # First sparse bind, or an idempotent rebind: nothing to morph. + if not is_dense: + return + self._bound_dense = is_dense + if is_dense: + self._bound_dense_gate = self._binds_dense_gate(component) + self.hook_aliases = dict( + self._DENSE_GATED_HOOK_ALIASES + if self._bound_dense_gate + else self._DENSE_HOOK_ALIASES + ) + dense_props = dict(self._DENSE_PROPERTY_ALIASES) + if not self._bound_dense_gate: + dense_props.pop("b_gate", None) + self.property_aliases = {**self.property_aliases, **dense_props} + # A dense layer has no router; leaving the hook in hook_dict would + # advertise an intervention point that can never fire. + if hasattr(self, "hook_router_scores"): + self._hook_registry.pop("hook_router_scores", None) + del self.hook_router_scores + else: + # Symmetric restore so a rebinding harness cannot leave a chimera. + self.hook_aliases = dict(type(self).hook_aliases) + self.property_aliases = { + key: value + for key, value in self.property_aliases.items() + if key not in self._DENSE_PROPERTY_ALIASES + } + self._bound_dense_gate = False + if not hasattr(self, "hook_router_scores"): + self.hook_router_scores = HookPoint() + + def validate_after_setup(self, skipped_optional: list[str]) -> None: + """Fail loudly when a sparse layer is missing a submodule only dense + layers may lack (see ``sparse_required``). + + Called by setup_submodules once every submodule has resolved — the + skipped set is not knowable at bind time. + """ + if self._bound_dense or not self._sparse_required: + return + missing = [key for key in self._sparse_required if key in skipped_optional] + if missing: + component = type(self.original_component).__name__ + raise ValueError( + f"{self.name}: sparse MoE layer wrapped {component} but required " + f"submodule(s) {missing} were absent. These are optional only so " + "dense layers of an interleaved stack can skip them; on a sparse " + "layer their absence means the HF attribute was renamed or moved. " + "Update the adapter's submodule name(s) rather than losing the hooks." + ) + + @property + def bound_dense(self) -> bool: + """Whether this layer bound the dense variant of an interleaved MoE stack. + + Public so weight-collection helpers can find the projections under + ``DENSE_SUBMODULE_KEYS`` instead of the sparse ``gate`` (the router). + """ + return self._bound_dense + + def _dense_projection(self, key: str) -> Any: + """Return a bound dense projection. + + Raises AttributeError on sparse layers, which have per-expert weights + and no single W_*: that keeps ``hasattr`` False so weight-collection + helpers skip the layer as they always have. (The base __getattr__ + substitutes its own message for property AttributeErrors, so this text + is a code comment in practice, not user-facing.) + """ + if not self._bound_dense: + raise AttributeError(f"{self.name}: {key} exists only on dense layers") + module = getattr(self, key, None) + if module is None: + raise AttributeError(f"{self.name}: dense projection {key!r} is not bound") + return module + + @property + def W_gate(self) -> torch.Tensor: + """Gated dense layer's gate weight in TL orientation [d_model, d_mlp].""" + if not self._bound_dense_gate: + raise AttributeError(f"{self.name}: this dense layer is ungated (no W_gate)") + module = self._dense_projection("dense_gate") + return normalize_mlp_weight( + module.weight, weight_layout_in_out(module), module, pattern="in" + ) + + @property + def W_in(self) -> torch.Tensor: + """Dense-layer input weight in TL orientation [d_model, d_mlp].""" + module = self._dense_projection("dense_in") + return normalize_mlp_weight( + module.weight, weight_layout_in_out(module), module, pattern="in" + ) + + @property + def W_out(self) -> torch.Tensor: + """Dense-layer output weight in TL orientation [d_mlp, d_model].""" + module = self._dense_projection("dense_out") + return normalize_mlp_weight( + module.weight, weight_layout_in_out(module), module, pattern="out" + ) def get_random_inputs( self, @@ -138,7 +340,10 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: router_scores = next( (t for t in router_scores if isinstance(t, torch.Tensor)), None ) - if isinstance(router_scores, torch.Tensor): + # The hook is removed on dense binds (no router exists there), so + # a dense layer whose wrapped module still returns a tuple must + # pass the extras through untouched rather than raise. + if isinstance(router_scores, torch.Tensor) and hasattr(self, "hook_router_scores"): self.hook_router_scores(router_scores) hidden_states = self.hook_out(hidden_states) return (hidden_states,) + output[1:] diff --git a/transformer_lens/model_bridge/get_params_util.py b/transformer_lens/model_bridge/get_params_util.py index adf9de034..a1a242ce7 100644 --- a/transformer_lens/model_bridge/get_params_util.py +++ b/transformer_lens/model_bridge/get_params_util.py @@ -122,7 +122,21 @@ def _get_device_dtype(): e, ) try: - mlp_in = getattr(block.mlp, "in", None) or getattr(block.mlp, "input", None) + # Dense layers of an interleaved MoE stack keep their projections + # under dense_* — `gate` there is the sparse layers' ROUTER, so the + # standard names would either miss the weights (silently zero-filling + # a real dense MLP) or read the router as a gate projection. + # `is True`, not truthiness: auto-vivifying stand-ins (Mock blocks in + # this module's own tests) return a truthy object for any attribute + # and would take the dense branch with non-tensor projections. + if getattr(block.mlp, "bound_dense", False) is True: + mlp_in = getattr(block.mlp, "dense_in", None) + mlp_out = getattr(block.mlp, "dense_out", None) + mlp_gate = getattr(block.mlp, "dense_gate", None) + else: + mlp_in = getattr(block.mlp, "in", None) or getattr(block.mlp, "input", None) + mlp_out = getattr(block.mlp, "out", None) + mlp_gate = getattr(block.mlp, "gate", None) if mlp_in is None: raise AttributeError("MLP has no 'in' or 'input' attribute") # Use normalized accessors for consistent TL orientation @@ -137,7 +151,7 @@ def _get_device_dtype(): params_dict[f"blocks.{layer_idx}.mlp.b_in"] = torch.zeros( d_mlp, device=device, dtype=dtype ) - mlp_out_bias = block.mlp.out.bias + mlp_out_bias = mlp_out.bias if mlp_out is not None else None if mlp_out_bias is not None: params_dict[f"blocks.{layer_idx}.mlp.b_out"] = mlp_out_bias else: @@ -145,13 +159,24 @@ def _get_device_dtype(): params_dict[f"blocks.{layer_idx}.mlp.b_out"] = torch.zeros( bridge.cfg.d_model, device=device, dtype=dtype ) - if hasattr(block.mlp, "gate") and hasattr(block.mlp.gate, "weight"): + if mlp_gate is not None and hasattr(mlp_gate, "weight"): w_gate = block.mlp.W_gate if w_gate is not None: params_dict[f"blocks.{layer_idx}.mlp.W_gate"] = w_gate - if hasattr(block.mlp.gate, "bias") and block.mlp.gate.bias is not None: - params_dict[f"blocks.{layer_idx}.mlp.b_gate"] = block.mlp.gate.bias - except AttributeError: + if getattr(mlp_gate, "bias", None) is not None: + params_dict[f"blocks.{layer_idx}.mlp.b_gate"] = mlp_gate.bias + except AttributeError as e: + # Zero-filling a real MLP silently yields wrong numbers downstream + # (SVD/weight analyses decompose zeros). Say so — the fill stays for + # architectures that genuinely have no MLP under this name. + logger.warning( + "Block %d MLP weights could not be extracted (%s) — emitting " + "ZEROS for blocks.%d.mlp.W_in/W_out/b_in/b_out. Any weight-space " + "analysis of this layer will be meaningless.", + layer_idx, + e, + layer_idx, + ) device, dtype = _get_device_dtype() d_mlp = bridge.cfg.d_mlp if bridge.cfg.d_mlp is not None else 4 * bridge.cfg.d_model params_dict[f"blocks.{layer_idx}.mlp.W_in"] = torch.zeros( diff --git a/transformer_lens/model_bridge/supported_architectures/AGENTS.md b/transformer_lens/model_bridge/supported_architectures/AGENTS.md index 5f915ad52..af6b569c9 100644 --- a/transformer_lens/model_bridge/supported_architectures/AGENTS.md +++ b/transformer_lens/model_bridge/supported_architectures/AGENTS.md @@ -242,6 +242,7 @@ Failure message names the missing set. (`INTENTIONAL_EXCLUDES` in the test handl | Combined QKV (`c_attn`) | `gpt2.py`, `bloom.py` | `QKVSplitRearrangeConversion` to split + rearrange | | Split Q/K/V (standard) | `llama.py`, `mistral.py`, most modern | `self._qkvo_weight_conversions()` helper | | MoE routing | `mixtral.py`, `deepseek_v3.py`, `qwen3_moe.py`, `granite_moe.py` | `MoEBridge` with `gate` + batched expert submodules | +| MoE with dense layers (`first_k_dense_replace`, `mlp_only_layers`, `mlp_layer_types`) | `deepseek_v3.py`, `glm4_moe.py`, `qwen3_moe.py` | Add optional `dense_gate`/`dense_in`/`dense_out` LinearBridges — `MoEBridge` binds gated-MLP neuron hooks on those layers (#1645) | | Missing biases (RMSNorm has no `b`; Llama has no attn/MLP biases) | `llama.py` (documented in docstring) | Weight processing handles `None` via `ProcessWeights._safe_get_tensor()` | | KV cache layout | All (implicit) | Adapter delegates; HF module manages internally | diff --git a/transformer_lens/model_bridge/supported_architectures/afmoe.py b/transformer_lens/model_bridge/supported_architectures/afmoe.py index ba9a99969..d13aa91e6 100644 --- a/transformer_lens/model_bridge/supported_architectures/afmoe.py +++ b/transformer_lens/model_bridge/supported_architectures/afmoe.py @@ -64,15 +64,24 @@ def __init__(self, cfg: Any) -> None: requires_attention_mask=True, ), # Dense layers (< num_dense_layers) hold a plain gated MLP - # under the same name; router and shared experts are - # optional. The tuple-returning router stays unwrapped — - # only its inner gate Linear is hookable. + # under the same name; the dense_* projections below make + # MoEBridge bind gated-MLP neuron hooks there. The + # tuple-returning router stays unwrapped — only its inner + # gate Linear is hookable. "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("router_gate",), submodules={ "router_gate": LinearBridge(name="router.gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/deepseek_v2.py b/transformer_lens/model_bridge/supported_architectures/deepseek_v2.py index c94e8e1c3..d2606e356 100644 --- a/transformer_lens/model_bridge/supported_architectures/deepseek_v2.py +++ b/transformer_lens/model_bridge/supported_architectures/deepseek_v2.py @@ -29,16 +29,6 @@ ) -class _DeepSeekV2MLPBridge(MoEBridge): - """Expose neuron-basis hooks on dense-prefix layers and MoE boundaries elsewhere.""" - - hook_aliases = { - "hook_pre": ["dense_gate.hook_out", "hook_in"], - "hook_pre_linear": "dense_in.hook_out", - "hook_post": ["dense_out.hook_in", "hook_out"], - } - - class DeepSeekV2ArchitectureAdapter(ArchitectureAdapter): """Architecture adapter for DeepSeek V2 / V2-Lite / Coder-V2 models. @@ -98,9 +88,9 @@ def __init__(self, cfg: Any) -> None: "o": LinearBridge(name="o_proj"), }, ), - # On dense layers (idx < first_k_dense_replace), gate and - # shared_experts are absent — marked optional so setup gracefully - # skips them when the layer is DeepseekV2MLP instead of MoE. + # Layers before first_k_dense_replace are DeepseekV2MLP: + # the MoE parts are skipped and the dense_* projections bind, + # which is what gives those layers gated-MLP neuron hooks. "mlp": self._build_mlp_bridge(), }, ), @@ -109,14 +99,22 @@ def __init__(self, cfg: Any) -> None: } def _build_mlp_bridge(self): - """Routed MoE with optional shared experts; Youtu (all-dense) overrides.""" - return _DeepSeekV2MLPBridge( + """Routed MoE with optional shared experts; Youtu (all-dense) overrides. + + Dense-prefix layers (idx < first_k_dense_replace) bind as gated MLPs + with neuron-basis hook_pre/hook_pre_linear/hook_post (#1645). + """ + return MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ # Router is a custom Module, not nn.Linear. "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on layers before + # first_k_dense_replace); their presence is what makes + # MoEBridge bind gated-MLP neuron hooks there (#1645). "dense_gate": LinearBridge(name="gate_proj", optional=True), "dense_in": LinearBridge(name="up_proj", optional=True), "dense_out": LinearBridge(name="down_proj", optional=True), diff --git a/transformer_lens/model_bridge/supported_architectures/deepseek_v3.py b/transformer_lens/model_bridge/supported_architectures/deepseek_v3.py index acd8955ac..b97a91ff0 100644 --- a/transformer_lens/model_bridge/supported_architectures/deepseek_v3.py +++ b/transformer_lens/model_bridge/supported_architectures/deepseek_v3.py @@ -75,16 +75,24 @@ def __init__(self, cfg: Any) -> None: "o": LinearBridge(name="o_proj"), }, ), - # On dense layers (idx < first_k_dense_replace), gate and - # shared_experts are marked optional so setup gracefully - # skips them when the layer is DeepseekV3MLP instead of MoE. + # Dense-prefix layers (idx < first_k_dense_replace) bind as + # gated MLPs with neuron-basis hooks; sparse layers keep the + # MoE mapping with its optional router/shared experts (#1645). "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ # Router is a custom Module, not nn.Linear "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/ernie4_5_moe.py b/transformer_lens/model_bridge/supported_architectures/ernie4_5_moe.py index f42d8e797..ab60e54c8 100644 --- a/transformer_lens/model_bridge/supported_architectures/ernie4_5_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/ernie4_5_moe.py @@ -67,14 +67,23 @@ def __init__(self, cfg: Any) -> None: requires_position_embeddings=True, ), # Layers before moe_layer_start_index hold a plain gated - # MLP; router and shared experts are absent there. + # MLP; the dense_* projections bind there and carry the + # neuron hooks, while router/shared experts are skipped. "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ # Raw-Parameter router; tuple-safe hook via base. "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/glm4_moe.py b/transformer_lens/model_bridge/supported_architectures/glm4_moe.py index ff545da3b..b6d563a5d 100644 --- a/transformer_lens/model_bridge/supported_architectures/glm4_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/glm4_moe.py @@ -84,13 +84,22 @@ def __init__(self, cfg: Any) -> None: requires_attention_mask=True, requires_position_embeddings=True, ), - # Dense prefix layers expose `mlp` but no router; mark gate optional - # for the dense-MoE boundary. + # Dense prefix layers (idx < first_k_dense_replace) expose + # `mlp` with no router; the dense_* projections bind there + # and carry the gated-MLP neuron hooks. "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": Glm4MoeRouterBridge(name="gate", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/glm4_moe_lite.py b/transformer_lens/model_bridge/supported_architectures/glm4_moe_lite.py index 9ea6626c9..42a854639 100644 --- a/transformer_lens/model_bridge/supported_architectures/glm4_moe_lite.py +++ b/transformer_lens/model_bridge/supported_architectures/glm4_moe_lite.py @@ -80,14 +80,23 @@ def __init__(self, cfg: Any) -> None: "o": LinearBridge(name="o_proj"), }, ), - # Layers marked "dense" in mlp_layer_types hold a plain gated MLP: - # router and shared expert absent, so both are optional. + # Layers marked "dense" in mlp_layer_types hold a plain gated + # MLP: router/shared expert are skipped and the dense_* + # projections bind, carrying the neuron hooks. "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": Glm4MoeRouterBridge(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/glm_moe_dsa.py b/transformer_lens/model_bridge/supported_architectures/glm_moe_dsa.py index 2522d53ce..542e83054 100644 --- a/transformer_lens/model_bridge/supported_architectures/glm_moe_dsa.py +++ b/transformer_lens/model_bridge/supported_architectures/glm_moe_dsa.py @@ -65,9 +65,17 @@ def __init__(self, cfg: Any) -> None: "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), + # Dense-layer projections (present only on the + # dense layers of this interleaved stack); their + # presence is what makes MoEBridge bind gated-MLP + # neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/jamba.py b/transformer_lens/model_bridge/supported_architectures/jamba.py index 9d663385e..d6810af32 100644 --- a/transformer_lens/model_bridge/supported_architectures/jamba.py +++ b/transformer_lens/model_bridge/supported_architectures/jamba.py @@ -200,10 +200,13 @@ def _build_ffn_bridge(self, num_experts: int) -> GatedMLPBridge | MoEBridge: return MoEBridge( name="feed_forward", config=self.cfg, + sparse_required=("router",), submodules={ - "gate": LinearBridge(name="gate_proj", optional=True), - "in": LinearBridge(name="up_proj", optional=True), - "out": LinearBridge(name="down_proj", optional=True), + # Dense (JambaMLP) layers: dense_* is what makes MoEBridge + # bind gated-MLP neuron hooks there (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), "router": LinearBridge(name="router", optional=True), }, ) diff --git a/transformer_lens/model_bridge/supported_architectures/laguna.py b/transformer_lens/model_bridge/supported_architectures/laguna.py index e98c8bc84..f964273de 100644 --- a/transformer_lens/model_bridge/supported_architectures/laguna.py +++ b/transformer_lens/model_bridge/supported_architectures/laguna.py @@ -77,6 +77,7 @@ def __init__(self, cfg: Any) -> None: "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), diff --git a/transformer_lens/model_bridge/supported_architectures/llada2_moe.py b/transformer_lens/model_bridge/supported_architectures/llada2_moe.py index 5f97aa62b..6763e0e7f 100644 --- a/transformer_lens/model_bridge/supported_architectures/llada2_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/llada2_moe.py @@ -110,6 +110,7 @@ def __init__(self, cfg: Any) -> None: "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": GeneralizedComponent(name="gate", optional=True), "shared_experts": self._gated_mlp(name="shared_experts", optional=True), diff --git a/transformer_lens/model_bridge/supported_architectures/llama4.py b/transformer_lens/model_bridge/supported_architectures/llama4.py index 080e82af8..f9e23d5a2 100644 --- a/transformer_lens/model_bridge/supported_architectures/llama4.py +++ b/transformer_lens/model_bridge/supported_architectures/llama4.py @@ -25,6 +25,7 @@ ) from transformer_lens.model_bridge.generalized_components.base import ( CloneOutputUnderGradMixin, + GeneralizedComponent, ) @@ -105,7 +106,12 @@ def __init__(self, cfg: Any) -> None: "mlp": _Llama4MoEBridge( name="feed_forward", config=self.cfg, + sparse_required=("router",), submodules={ + # HF creates the router unconditionally on MoE + # layers; mapping it makes a rename loud and gives + # the layer real router observability. + "router": GeneralizedComponent(name="router", optional=True), # Dense-layer projections (absent on MoE layers). "dense_gate": LinearBridge(name="gate_proj", optional=True), "dense_in": LinearBridge(name="up_proj", optional=True), diff --git a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py index 6eb77015e..2d52f8e3e 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen2_moe.py @@ -36,10 +36,18 @@ def __init__(self, cfg: Any) -> None: blocks.submodules["mlp"] = MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate", "experts"), submodules={ - "gate": Qwen2MoeRouterBridge(name="gate"), - "experts": MoEBridge(name="experts", config=self.cfg), - "shared_expert": self._gated_mlp(name="shared_expert"), - "shared_expert_gate": LinearBridge(name="shared_expert_gate"), + # mlp_only_layers / decoder_sparse_step let HF build a dense + # Qwen2MoeMLP on some layers: the MoE parts are optional there, + # and the dense projections below make those layers bind + # gated-MLP neuron hooks instead of MoE boundary tensors (#1645). + "gate": Qwen2MoeRouterBridge(name="gate", optional=True), + "experts": MoEBridge(name="experts", config=self.cfg, optional=True), + "shared_expert": self._gated_mlp(name="shared_expert", optional=True), + "shared_expert_gate": LinearBridge(name="shared_expert_gate", optional=True), + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ) diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen3_moe.py index 99ad984c3..147604756 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_moe.py @@ -91,8 +91,15 @@ def __init__(self, cfg: Any) -> None: "mlp": MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate",), submodules={ - "gate": MoERouterBridge(name="gate"), + # Dense fallback layers (mlp_only_layers / + # decoder_sparse_step) have no router; their + # projections bind gated-MLP neuron hooks (#1645). + "gate": MoERouterBridge(name="gate", optional=True), + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ), }, diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_next.py b/transformer_lens/model_bridge/supported_architectures/qwen3_next.py index 31e1be3cd..11e7ed69f 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_next.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_next.py @@ -9,7 +9,11 @@ import torch -from transformer_lens.model_bridge.generalized_components import MoEBridge +from transformer_lens.model_bridge.generalized_components import ( + LinearBridge, + MoEBridge, + MoERouterBridge, +) from transformer_lens.model_bridge.supported_architectures.qwen3 import ( Qwen3ArchitectureAdapter, ) @@ -27,7 +31,22 @@ def __init__(self, cfg: Any) -> None: def _build_mlp_bridge(self): """Sparse MoE MLP (router + batched experts + shared expert).""" - return MoEBridge(name="mlp", config=self.cfg) + return MoEBridge( + name="mlp", + config=self.cfg, + sparse_required=("gate",), + submodules={ + # Plain-tensor SparseMoeBlock: router observability comes from + # the gate submodule. Dense fallback layers (mlp_only_layers / + # decoder_sparse_step) bind the projections below (#1645). + "gate": MoERouterBridge(name="gate", optional=True), + "shared_expert": self._gated_mlp(name="shared_expert", optional=True), + "shared_expert_gate": LinearBridge(name="shared_expert_gate", optional=True), + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), + }, + ) def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Slice query half from gated q_proj.weight for weight-space analysis.""" diff --git a/transformer_lens/model_bridge/supported_architectures/qwen3_vl_moe.py b/transformer_lens/model_bridge/supported_architectures/qwen3_vl_moe.py index 0b46fb2fd..8d9a73948 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen3_vl_moe.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen3_vl_moe.py @@ -11,6 +11,7 @@ from typing import Any from transformer_lens.model_bridge.generalized_components import ( + LinearBridge, MoEBridge, MoERouterBridge, ) @@ -27,8 +28,13 @@ def _build_mlp_bridge(self) -> Any: return MoEBridge( name="mlp", config=self.cfg, + sparse_required=("gate", "experts"), submodules={ "gate": MoERouterBridge(name="gate", optional=True), "experts": MoEBridge(name="experts", config=self.cfg, optional=True), + # Dense layers listed in mlp_only_layers (#1645). + "dense_gate": LinearBridge(name="gate_proj", optional=True), + "dense_in": LinearBridge(name="up_proj", optional=True), + "dense_out": LinearBridge(name="down_proj", optional=True), }, ) diff --git a/transformer_lens/model_bridge/supported_architectures/switch_transformers.py b/transformer_lens/model_bridge/supported_architectures/switch_transformers.py index befe6f260..2493d7269 100644 --- a/transformer_lens/model_bridge/supported_architectures/switch_transformers.py +++ b/transformer_lens/model_bridge/supported_architectures/switch_transformers.py @@ -90,12 +90,14 @@ def _build_ff_bridge(self, layer_prefix: str) -> MoEBridge: return MoEBridge( name=f"{layer_prefix}.mlp", config=self.cfg, + sparse_required=("gate",), submodules={ "gate": GeneralizedComponent(name="router", optional=True), - # Dense-layer projections (absent on sparse layers); in/out per - # T5-family convention (nonstandard names dodge the component - # prober's down-projection skip and get crash-probed). - "in": LinearBridge(name="wi", optional=True), - "out": LinearBridge(name="wo", optional=True), + # Dense-layer projections (absent on sparse layers). The dense + # feed-forward here is UNGATED (wi/wo), so MoEBridge binds the + # two-key dense shape: hook_pre/hook_post over d_ff, no + # hook_pre_linear (#1645). + "dense_in": LinearBridge(name="wi", optional=True), + "dense_out": LinearBridge(name="wo", optional=True), }, ) diff --git a/transformer_lens/model_bridge/supported_architectures/youtu.py b/transformer_lens/model_bridge/supported_architectures/youtu.py index e917af257..436cf51c1 100644 --- a/transformer_lens/model_bridge/supported_architectures/youtu.py +++ b/transformer_lens/model_bridge/supported_architectures/youtu.py @@ -5,8 +5,8 @@ DeepSeek-V2's MLA (q LoRA + compressed KV with decoupled rope) with every layer dense. Module names match DeepSeek-V2 exactly and the MLA bridge already handles both q-projection variants, so this is a pure subclass: -the MoE router/shared-expert submodules are optional and simply never -bind on the dense MLPs. +the MoE router/shared-expert submodules are optional and never bind, so +every layer takes MoEBridge's dense path and exposes gated-MLP neuron hooks. """ from transformer_lens.model_bridge.supported_architectures.deepseek_v2 import (