diff --git a/tests/integration/model_bridge/test_attention_score_sentinel.py b/tests/integration/model_bridge/test_attention_score_sentinel.py index 424198df5..dce95b060 100644 --- a/tests/integration/model_bridge/test_attention_score_sentinel.py +++ b/tests/integration/model_bridge/test_attention_score_sentinel.py @@ -7,19 +7,27 @@ def test_gpt2_compatibility_scores_use_negative_infinity( - gpt2_bridge_compat, gpt2_hooked_processed + gpt2_bridge_compat, gpt2_goldens_processed ) -> None: - """GPT-2's direct HF mask is normalized before the compatibility hook.""" - tokens = gpt2_hooked_processed.to_tokens("The capital of France is") + """GPT-2's direct HF mask is normalized before the compatibility hook. + + Anchored on the frozen HookedTransformer goldens rather than a live + HookedTransformer, matching the rest of the compatibility suite. + """ + golden = gpt2_goldens_processed + tokens = golden.scalars["short_prompt"] _, bridge_cache = gpt2_bridge_compat.run_with_cache(tokens, names_filter=[SCORES]) - _, hooked_cache = gpt2_hooked_processed.run_with_cache(tokens, names_filter=[SCORES]) + hooked_cache = golden.tensors("activations") bridge_scores, hooked_scores = bridge_cache[SCORES], hooked_cache[SCORES] causal_mask = torch.isneginf(hooked_scores) assert causal_mask.any() assert torch.isneginf(bridge_scores[causal_mask]).all() + # The goldens were captured on different hardware, so the unmasked scores + # agree to fp32 accumulation noise rather than bit-exactly. Same tolerance + # the sibling golden comparison uses for this hook. torch.testing.assert_close( - bridge_scores[~causal_mask], hooked_scores[~causal_mask], rtol=0, atol=0 + bridge_scores[~causal_mask], hooked_scores[~causal_mask], rtol=1e-4, atol=1e-4 ) diff --git a/tests/integration/model_bridge/test_audio_frame_entry.py b/tests/integration/model_bridge/test_audio_frame_entry.py new file mode 100644 index 000000000..c9f199ef9 --- /dev/null +++ b/tests/integration/model_bridge/test_audio_frame_entry.py @@ -0,0 +1,100 @@ +"""Audio frame entry: run the encoder from precomputed frames. + +Mirrors HookedAudioEncoder.encoder_output, the audio-path analogue of +start_at_layer. Deletion evidence for that method: without it the bridge can +only enter at the waveform, so injecting frames means re-running the conv front +end. start_at_layer stays refused for audio — this is a separate entry point. +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +import torch + +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "facebook/hubert-base-ls960" +SAMPLE_RATE = 16000 +FRAMES_HOOK = "feat_proj.hook_out" + + +@pytest.fixture(scope="module") +def audio_bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers(MODEL, device="cpu") + + +@pytest.fixture(scope="module") +def waveform() -> torch.Tensor: + t = np.linspace(0, 1.0, SAMPLE_RATE, endpoint=False, dtype=np.float32) + return torch.tensor(0.1 * np.sin(2 * math.pi * 440.0 * t))[None, :] + + +@pytest.fixture(scope="module") +def full_run(audio_bridge, waveform): + last = f"blocks.{audio_bridge.cfg.n_layers - 1}.hook_out" + _, cache = audio_bridge.run_with_cache( + waveform, names_filter=[FRAMES_HOOK, "blocks.0.hook_out", last] + ) + return cache, last + + +def test_frame_entry_matches_the_full_waveform_run(audio_bridge, full_run): + """Re-entering at the frames reproduces the encoder output exactly.""" + cache, last = full_run + resid = audio_bridge.encoder_output(cache[FRAMES_HOOK]) + torch.testing.assert_close(resid, cache[last], atol=0.0, rtol=0.0) + + +def test_hooks_fire_from_frame_entry(audio_bridge, full_run): + """Block hooks fire on the frame path, so caching composes with it.""" + cache, last = full_run + wanted = {"blocks.0.hook_out", last} + cached, fwd_hooks, _ = audio_bridge.get_caching_hooks(names_filter=lambda name: name in wanted) + with audio_bridge.hooks(fwd_hooks=fwd_hooks): + audio_bridge.encoder_output(cache[FRAMES_HOOK]) + + assert set(cached) == wanted + for name in wanted: + torch.testing.assert_close(cached[name], cache[name], atol=0.0, rtol=0.0) + + +def test_padding_mask_changes_the_encoding(audio_bridge, full_run): + """The mask is applied, not ignored.""" + cache, last = full_run + frames = cache[FRAMES_HOOK] + mask = torch.ones(frames.shape[:2], dtype=torch.long) + mask[:, -10:] = 0 + + masked = audio_bridge.encoder_output(frames, one_zero_attention_mask=mask) + assert not torch.allclose(masked, cache[last]) + + +def test_waveform_shaped_input_is_rejected(audio_bridge, waveform): + """A 2D waveform is not frames; say so instead of silently mis-running.""" + with pytest.raises(ValueError, match=r"\[batch, frames, d_model\]"): + audio_bridge.encoder_output(waveform) + + +def test_start_at_layer_remains_refused_for_audio(audio_bridge, waveform): + """Frame entry is a separate API; the residual-injection guard is untouched.""" + with pytest.raises(NotImplementedError, match="audio models"): + audio_bridge(waveform, start_at_layer=1) + + +def test_text_models_reject_the_audio_frame_entry(): + """Non-audio bridges have no conv-frame stage to re-enter.""" + bridge = TransformerBridge.boot_transformers("gpt2", device="cpu") + with pytest.raises(NotImplementedError, match="not an audio model"): + bridge.encoder_output(torch.zeros(1, 4, bridge.cfg.d_model)) + + +def test_spectrogram_encoders_reject_the_frame_entry(): + """AST has no conv feature extractor, so there is no frame stage to bypass.""" + bridge = TransformerBridge.boot_transformers( + "MIT/ast-finetuned-audioset-10-10-0.4593", device="cpu" + ) + with pytest.raises(NotImplementedError, match="convolutional front end"): + bridge.encoder_output(torch.zeros(1, 4, bridge.cfg.d_model)) diff --git a/tests/integration/model_bridge/test_bert_pooler_hook.py b/tests/integration/model_bridge/test_bert_pooler_hook.py new file mode 100644 index 000000000..afef28fed --- /dev/null +++ b/tests/integration/model_bridge/test_bert_pooler_hook.py @@ -0,0 +1,69 @@ +"""The BERT [CLS] pooler is observable through a named bridge hook. + +Mirrors HookedEncoder's BertPooler, whose hook_pooler_out carries the +post-tanh pooled [CLS]. Deletion evidence for that component: without a named +bridge hook the pooled [CLS] is only reachable coincidentally, via the NSP +head's unembed.hook_in. +""" + +from __future__ import annotations + +import pytest +import torch +from transformers import BertForMaskedLM, BertForNextSentencePrediction + +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "google-bert/bert-base-cased" + + +@pytest.fixture(scope="module") +def nsp_bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers( + MODEL, device="cpu", model_class=BertForNextSentencePrediction + ) + + +def _tokens(bridge: TransformerBridge) -> torch.Tensor: + return bridge.tokenizer("Hello there my friend.", return_tensors="pt")["input_ids"] + + +def test_pooler_hook_matches_huggingfaces_own_pooler(nsp_bridge): + """hook_out is the pooled [CLS], checked against HF's pooler directly.""" + tokens = _tokens(nsp_bridge) + _, cache = nsp_bridge.run_with_cache(tokens) + + hf = nsp_bridge.original_model + with torch.no_grad(): + expected = hf.bert.pooler(hf.bert(tokens).last_hidden_state) + + torch.testing.assert_close(cache["pooler.hook_out"], expected, atol=0.0, rtol=0.0) + + +def test_pooler_hook_is_post_activation(nsp_bridge): + """HookedEncoder fires hook_pooler_out after tanh; the projection is separate.""" + tokens = _tokens(nsp_bridge) + _, cache = nsp_bridge.run_with_cache(tokens) + + pre_activation = cache["pooler.dense.hook_out"] + pooled = cache["pooler.hook_out"] + assert not torch.allclose(pre_activation, pooled) + torch.testing.assert_close(torch.tanh(pre_activation), pooled) + + +def test_hooked_encoder_hook_name_is_aliased(nsp_bridge): + """Code migrated from HookedEncoder asks for hook_pooler_out.""" + tokens = _tokens(nsp_bridge) + _, cache = nsp_bridge.run_with_cache(tokens) + + assert torch.equal(cache["pooler.hook_pooler_out"], cache["pooler.hook_out"]) + + +def test_masked_lm_checkpoint_without_a_pooler_still_boots(): + """BertForMaskedLM leaves bert.pooler as None; the mapping must skip it.""" + bridge = TransformerBridge.boot_transformers(MODEL, device="cpu", model_class=BertForMaskedLM) + tokens = _tokens(bridge) + logits, cache = bridge.run_with_cache(tokens) + + assert logits.shape[0] == 1 + assert not [name for name in cache if "pooler" in name] diff --git a/tests/integration/model_bridge/test_encdec_weight_stacking_parity.py b/tests/integration/model_bridge/test_encdec_weight_stacking_parity.py new file mode 100644 index 000000000..a84b3d00e --- /dev/null +++ b/tests/integration/model_bridge/test_encdec_weight_stacking_parity.py @@ -0,0 +1,41 @@ +"""Stacked enc-dec weights on the bridge match HookedEncoderDecoder. + +Deletion evidence for HookedEncoderDecoder's stacked-weight properties: the +bridge must produce the same tensors over chain(encoder, decoder), and the same +head labels, before the legacy class can go. +""" + +from __future__ import annotations + +import pytest +import torch + +from transformer_lens import HookedEncoderDecoder +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "google-t5/t5-small" +STACKED = ["W_Q", "W_K", "W_V", "W_O", "W_in", "W_out"] + + +@pytest.fixture(scope="module") +def hooked() -> HookedEncoderDecoder: + return HookedEncoderDecoder.from_pretrained(MODEL, device="cpu") + + +@pytest.fixture(scope="module") +def bridge() -> TransformerBridge: + return TransformerBridge.boot_transformers(MODEL, device="cpu") + + +@pytest.mark.parametrize("name", STACKED) +def test_stacked_weights_match_hooked_encoder_decoder(name, hooked, bridge): + """HookedEncoderDecoder does no weight processing, so these are directly comparable.""" + expected = getattr(hooked, name) + actual = getattr(bridge, name) + assert actual.shape == expected.shape + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + + +def test_head_labels_match_hooked_encoder_decoder(hooked, bridge): + """all_head_labels is a property on the bridge; HT exposes it as a method.""" + assert bridge.all_head_labels == hooked.all_head_labels() diff --git a/tests/integration/model_bridge/test_nsp_sentence_pair_helper.py b/tests/integration/model_bridge/test_nsp_sentence_pair_helper.py new file mode 100644 index 000000000..fa364707b --- /dev/null +++ b/tests/integration/model_bridge/test_nsp_sentence_pair_helper.py @@ -0,0 +1,82 @@ +"""Next-sentence prediction from strings on the bridge. + +Mirrors BertNextSentencePrediction's string interface, which cannot be adapted +onto a bridge (its forward reaches for encoder_output/pooler/nsp_head). This +helper is where that ergonomics survives the Hooked* removal. +""" + +from __future__ import annotations + +import pytest +import torch +from transformers import AutoTokenizer, BertForNextSentencePrediction + +from transformer_lens.model_bridge.bridge import TransformerBridge + +MODEL = "google-bert/bert-base-cased" +SENTENCE_A = "A man walked into a grocery store." +SEQUENTIAL_B = "He bought an apple." +UNRELATED_B = "The Eiffel Tower is in Paris." + + +@pytest.fixture(scope="module") +def nsp_bridge() -> TransformerBridge: + bridge = TransformerBridge.boot_transformers( + MODEL, device="cpu", model_class=BertForNextSentencePrediction + ) + bridge.enable_compatibility_mode() + return bridge + + +@pytest.fixture(scope="module") +def hf_tokenizer(): + return AutoTokenizer.from_pretrained(MODEL) + + +def test_pair_tokenization_matches_huggingface(nsp_bridge, hf_tokenizer): + """[CLS] a [SEP] b [SEP] with segment ids, identical to tokenizer(a, b).""" + tokens = nsp_bridge.to_sentence_pair_tokens(SENTENCE_A, SEQUENTIAL_B) + expected = hf_tokenizer(SENTENCE_A, SEQUENTIAL_B, return_tensors="pt") + + assert torch.equal(tokens["input_ids"], expected["input_ids"]) + assert torch.equal(tokens["token_type_ids"], expected["token_type_ids"]) + assert tokens["token_type_ids"].unique().tolist() == [0, 1] + + +def test_logits_match_a_direct_huggingface_nsp_forward(nsp_bridge, hf_tokenizer): + encodings = hf_tokenizer(SENTENCE_A, SEQUENTIAL_B, return_tensors="pt") + with torch.no_grad(): + expected = nsp_bridge.original_model(**encodings).logits + + actual = nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B, return_type="logits") + torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0) + + +def test_predictions_distinguish_sequential_from_unrelated(nsp_bridge): + assert nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B) == ( + "The sentences are sequential" + ) + assert nsp_bridge.predict_next_sentence(SENTENCE_A, UNRELATED_B) == ( + "The sentences are NOT sequential" + ) + + +def test_segment_ids_are_load_bearing(nsp_bridge): + """Dropping token_type_ids collapses the NSP margin — why the helper owns them.""" + tokens = nsp_bridge.to_sentence_pair_tokens(SENTENCE_A, SEQUENTIAL_B) + with_segments = nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B, return_type="logits") + without_segments = nsp_bridge( + tokens["input_ids"], + attention_mask=tokens["attention_mask"], + return_type="logits", + ) + + margin = lambda logits: float((logits[0, 0] - logits[0, 1]).abs()) + assert margin(with_segments) > 2 * margin(without_segments) + + +def test_helper_rejects_a_model_without_an_nsp_head(): + """An MLM-headed bridge has no 2-class output; say so instead of decoding noise.""" + bridge = TransformerBridge.boot_transformers(MODEL, device="cpu") + with pytest.raises(ValueError, match="next-sentence-prediction head"): + bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B) diff --git a/tests/integration/model_bridge/test_parent_module_traversal.py b/tests/integration/model_bridge/test_parent_module_traversal.py index 5f60c7930..3d80d25d1 100644 --- a/tests/integration/model_bridge/test_parent_module_traversal.py +++ b/tests/integration/model_bridge/test_parent_module_traversal.py @@ -370,7 +370,11 @@ def test_direct_assign_load_stays_current_after_apply( @pytest.mark.parametrize( ("case_name", "key_fragment", "expected_keys"), ( - ("bert-nsp", "pooler", {"pooler.weight", "pooler.bias"}), + # BERT wraps the pooler module (not its inner Linear) so hook_out carries + # the post-tanh pooled [CLS], which nests the weights one level deeper. + # Still two keys, so nothing is re-expanded — only renamed. ViT wraps + # pooler.dense and keeps the flat names. + ("bert-nsp", "pooler", {"pooler.dense.weight", "pooler.dense.bias"}), ("vit-bare-pooler", "pooler", {"pooler.weight", "pooler.bias"}), ( "ast-audio-classifier", diff --git a/tests/integration/model_bridge/test_qwen2_moe_bridge.py b/tests/integration/model_bridge/test_qwen2_moe_bridge.py index 05ee6a351..2f2ca75fb 100644 --- a/tests/integration/model_bridge/test_qwen2_moe_bridge.py +++ b/tests/integration/model_bridge/test_qwen2_moe_bridge.py @@ -117,3 +117,13 @@ def test_run_with_cache_captures_moe_hooks(self) -> None: router_scores_key = f"blocks.{layer_idx}.mlp.hook_router_scores" assert router_scores_key not in cache + + # Routing observables mirror HookedTransformer: weights at full + # expert width, indices at top-k. + weights_key = f"blocks.{layer_idx}.mlp.gate.hook_expert_weights" + assert weights_key in cache, f"Missing cache key: {weights_key}" + assert cache[weights_key].shape == (flat_tokens, num_experts) + + indices_key = f"blocks.{layer_idx}.mlp.gate.hook_expert_indices" + assert indices_key in cache, f"Missing cache key: {indices_key}" + assert cache[indices_key].shape == (flat_tokens, bridge.cfg.experts_per_token) diff --git a/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py b/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py index 9775e1114..55168218e 100644 --- a/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py @@ -97,9 +97,6 @@ def test_attn_not_only_and_eager(self, adapter: ArceeArchitectureAdapter) -> Non assert adapter.cfg.attn_only is False assert adapter.cfg.attn_implementation == "eager" - def test_gqa_propagated(self, adapter: ArceeArchitectureAdapter) -> None: - assert adapter.cfg.n_key_value_heads == 4 - class TestArceeAdapterComponentMapping: """Component-mapping structure and HF module names. Key contrasts with Llama: diff --git a/tests/unit/model_bridge/supported_architectures/test_bert_adapter.py b/tests/unit/model_bridge/supported_architectures/test_bert_adapter.py index 948d684b0..28083a3d4 100644 --- a/tests/unit/model_bridge/supported_architectures/test_bert_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_bert_adapter.py @@ -198,7 +198,10 @@ def test_nsp_only_model_uses_hooked_encoder_names(self) -> None: adapter.prepare_model(hf_model) - assert adapter.components["pooler"].name == "bert.pooler.dense" + # The pooler module itself is wrapped so hook_out is the post-tanh + # pooled [CLS]; the projection stays hookable underneath. + assert adapter.components["pooler"].name == "bert.pooler" + assert adapter.components["pooler"].submodules["dense"].name == "dense" assert adapter.components["unembed"].name == "cls.seq_relationship" assert "mlm_head" not in adapter.components assert "ln_final" not in adapter.components @@ -212,7 +215,10 @@ def test_combined_mlm_nsp_model_registers_both_heads(self) -> None: adapter.prepare_model(hf_model) - assert adapter.components["pooler"].name == "bert.pooler.dense" + # The pooler module itself is wrapped so hook_out is the post-tanh + # pooled [CLS]; the projection stays hookable underneath. + assert adapter.components["pooler"].name == "bert.pooler" + assert adapter.components["pooler"].submodules["dense"].name == "dense" assert adapter.components["mlm_head"].name == "cls.predictions.transform.dense" assert adapter.components["nsp_head"].name == "cls.seq_relationship" assert adapter.components["unembed"].name == "cls.predictions.decoder" diff --git a/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py index ede9e019c..319f48567 100644 --- a/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py @@ -123,9 +123,6 @@ def test_not_stateful(self, adapter: FalconH1ArchitectureAdapter) -> None: def test_eps_attr_variance_epsilon(self, adapter: FalconH1ArchitectureAdapter) -> None: assert adapter.cfg.eps_attr == "variance_epsilon" - def test_n_key_value_heads_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None: - assert adapter.cfg.n_key_value_heads == 2 - def test_mamba_intermediate_size_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None: # mamba_d_ssm is the inner SSM width directly. assert getattr(adapter.cfg, "mamba_intermediate_size", None) == 32 diff --git a/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py index 212a6ffeb..054fb3615 100644 --- a/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py @@ -32,6 +32,7 @@ GatedMLPBridge, LinearBridge, MoEBridge, + MoERouterBridge, PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, @@ -222,11 +223,12 @@ def test_ln1_ln2_are_rms_norm_bridges(self, adapter: GPTOSSArchitectureAdapter) assert isinstance(subs["ln1"], RMSNormalizationBridge) assert isinstance(subs["ln2"], RMSNormalizationBridge) - def test_mlp_has_no_submodules(self, adapter: GPTOSSArchitectureAdapter) -> None: - """GPT-OSS exposes no router submodule on the MoE block, the entire MoE module - is wrapped opaquely by MoEBridge.""" + def test_mlp_exposes_a_hookable_router(self, adapter: GPTOSSArchitectureAdapter) -> None: + """The MoE block wraps GPT-OSS's router so the routing observables are hookable.""" mlp = _mapping(adapter)["blocks"].submodules["mlp"] - assert mlp.submodules == {} + assert set(mlp.submodules) == {"router"} + assert isinstance(mlp.submodules["router"], MoERouterBridge) + assert mlp.submodules["router"].name == "router" def test_hf_module_paths(self, adapter: GPTOSSArchitectureAdapter) -> None: mapping = _mapping(adapter) diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py index 41e617cd9..538167664 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py @@ -98,9 +98,6 @@ def test_default_prepend_bos_false(self, adapter: GraniteArchitectureAdapter) -> """Granite models do not prepend BOS by default.""" assert adapter.cfg.default_prepend_bos is False - def test_n_key_value_heads_propagated(self, adapter: GraniteArchitectureAdapter) -> None: - assert adapter.cfg.n_key_value_heads == N_KV_HEADS - # --------------------------------------------------------------------------- # Component mapping tests — dense Granite diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py index 1bd89388e..0a8e01348 100644 --- a/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_granite_moe_adapter.py @@ -119,9 +119,6 @@ def test_uses_rms_norm(self, adapter: GraniteMoeArchitectureAdapter) -> None: def test_default_prepend_bos_false(self, adapter: GraniteMoeArchitectureAdapter) -> None: assert adapter.cfg.default_prepend_bos is False - def test_n_key_value_heads_propagated(self, adapter: GraniteMoeArchitectureAdapter) -> None: - assert adapter.cfg.n_key_value_heads == N_KV_HEADS - # --------------------------------------------------------------------------- # Component mapping tests diff --git a/tests/unit/model_bridge/supported_architectures/test_moe_routing_hooks.py b/tests/unit/model_bridge/supported_architectures/test_moe_routing_hooks.py new file mode 100644 index 000000000..5ea3ab54a --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_moe_routing_hooks.py @@ -0,0 +1,159 @@ +"""MoE routing observables on the bridge (hook_expert_weights / hook_expert_indices). + +Mirrors HookedTransformer's MoE routing hooks: weights are exposed in HT's +``[tokens, num_experts]`` layout, indices as ``[tokens, top_k]``, on both MoE +families (5.13 ``TopKRouter`` blocks and GPT-OSS). +""" + +from __future__ import annotations + +import copy + +import torch +from transformers import Qwen2MoeConfig, Qwen2MoeForCausalLM + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) + +NUM_EXPERTS = 4 +TOP_K = 2 +WEIGHTS = "blocks.0.mlp.gate.hook_expert_weights" +INDICES = "blocks.0.mlp.gate.hook_expert_indices" + + +def _tiny_bridge() -> TransformerBridge: + torch.manual_seed(0) + cfg = Qwen2MoeConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + shared_expert_intermediate_size=96, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOP_K, + max_position_embeddings=64, + decoder_sparse_step=1, + mlp_only_layers=[], + ) + cfg._attn_implementation = "eager" + model = Qwen2MoeForCausalLM(cfg).eval() + bridge = build_bridge_from_module( + model, + "Qwen2MoeForCausalLM", + hf_config=copy.deepcopy(cfg), + tokenizer=None, + device="cpu", + ).eval() + bridge.adapter.setup_component_testing(model, bridge_model=bridge) + return bridge + + +def _tokens() -> torch.Tensor: + return torch.randint(0, 128, (1, 6)) + + +def test_routing_hooks_are_cached_with_hooked_transformer_shapes(): + """Weights arrive at HT's [tokens, num_experts]; indices at [tokens, top_k].""" + bridge, tokens = _tiny_bridge(), _tokens() + _, cache = bridge.run_with_cache(tokens) + + assert cache[WEIGHTS].shape == (tokens.numel(), NUM_EXPERTS) + assert cache[INDICES].shape == (tokens.numel(), TOP_K) + + +def test_routing_hooks_are_observe_only(): + """Firing the hooks does not perturb the forward pass.""" + bridge, tokens = _tiny_bridge(), _tokens() + baseline = bridge(tokens) + cached_logits, _ = bridge.run_with_cache(tokens) + + assert torch.equal(baseline, cached_logits) + + +def test_expanded_weights_carry_only_the_selected_experts(): + """The scatter puts each top-k score at its expert id and zeroes the rest.""" + bridge, tokens = _tiny_bridge(), _tokens() + _, cache = bridge.run_with_cache(tokens) + + weights, indices = cache[WEIGHTS], cache[INDICES] + assert (weights != 0).sum(-1).eq(TOP_K).all() + gathered = weights.gather(-1, indices.long()) + assert torch.equal( + gathered.sort(dim=-1).values, weights.topk(TOP_K, dim=-1).values.sort(dim=-1).values + ) + + +def test_editing_expert_weights_reaches_the_model(): + """An edit in the weights hook changes the output — not a dead hook.""" + bridge, tokens = _tiny_bridge(), _tokens() + baseline = bridge(tokens) + + bridge.add_hook(WEIGHTS, lambda t, hook=None: torch.zeros_like(t)) + edited = bridge(tokens) + bridge.reset_hooks() + + assert not torch.equal(baseline, edited) + + +def test_editing_expert_indices_reroutes(): + """An edit in the indices hook re-routes tokens to different experts.""" + bridge, tokens = _tiny_bridge(), _tokens() + baseline = bridge(tokens) + + bridge.add_hook(INDICES, lambda t, hook=None: torch.zeros_like(t)) + edited = bridge(tokens) + bridge.reset_hooks() + + assert not torch.equal(baseline, edited) + + +def test_hooked_transformer_names_alias_onto_the_router(): + """HT places these on the MoE block itself; migrated code must find them there.""" + bridge, tokens = _tiny_bridge(), _tokens() + _, cache = bridge.run_with_cache(tokens) + + assert torch.equal(cache["blocks.0.mlp.hook_expert_weights"], cache[WEIGHTS]) + assert torch.equal(cache["blocks.0.mlp.hook_expert_indices"], cache[INDICES]) + + +class _StubRouter(torch.nn.Module): + """Router returning a fixed (logits, top-k weights, top-k indices) tuple.""" + + def __init__(self, logits, weights, indices): + super().__init__() + self._out = (logits, weights, indices) + + def forward(self, hidden_states): # noqa: ARG002 - fixed output by design + return self._out + + +def test_rerouting_picks_up_the_weight_at_the_newly_selected_expert(): + """The gather runs after the indices hook, as HookedTransformer does. + + Re-routing a token to a different expert must pick up the weight sitting at + that expert (zero when it was not originally selected), not carry the old + expert's weight over to the new one. + """ + from transformer_lens.model_bridge.generalized_components.moe import MoERouterBridge + + logits = torch.zeros(1, NUM_EXPERTS) + weights = torch.tensor([[0.7, 0.3]]) + indices = torch.tensor([[1, 2]]) + + router = MoERouterBridge(name="router") + router.set_original_component(_StubRouter(logits, weights, indices)) + # Route both slots to expert 0, which held no weight in the original top-k. + router.hook_expert_indices.add_hook(lambda t, hook=None: torch.zeros_like(t)) + + _, out_weights, out_indices = router(torch.zeros(1, 8)) + + assert torch.equal(out_indices, torch.zeros_like(indices)) + assert torch.equal(out_weights, torch.zeros_like(weights)), ( + "weights must be gathered at the edited indices, so a re-route to an " + f"unselected expert yields 0; got {out_weights.tolist()}" + ) diff --git a/tests/unit/model_bridge/test_bridge_weight_properties.py b/tests/unit/model_bridge/test_bridge_weight_properties.py index c81c2ca55..669a954ff 100644 --- a/tests/unit/model_bridge/test_bridge_weight_properties.py +++ b/tests/unit/model_bridge/test_bridge_weight_properties.py @@ -8,6 +8,12 @@ ) from transformer_lens.model_bridge.transformer_bridge import TransformerBridge +_BORROWED_HELPERS = ( + "_enumerate_blocks", + "_resolve_submodule_name", + "_rewrite_submodule_path", +) + class TestReshapeBias: """Tests for AttentionBridge._reshape_bias().""" @@ -67,7 +73,10 @@ def __init__(self, bias): class _FakeBridge: def __init__(self, biases): - self.blocks = [TestStackBlockParams._Block(b) for b in biases] + self.blocks = torch.nn.ModuleList([TestStackBlockParams._Block(b) for b in biases]) + # _stack_block_params walks the registered block lists, so the stub + # exposes the same _modules shape a real bridge does. + self._modules = {"blocks": self.blocks} class Cfg: n_devices = 1 @@ -75,6 +84,12 @@ class Cfg: self.cfg = Cfg() + def __getattr__(self, name): + """Borrow the real block-walking helpers rather than reimplementing them.""" + if name in _BORROWED_HELPERS: + return getattr(TransformerBridge, name).__get__(self) + raise AttributeError(name) + def test_stacks_present_params(self): fake = self._FakeBridge([torch.ones(2, 3), torch.zeros(2, 3)]) stacked = TransformerBridge._stack_block_params(fake, "attn.b_Q") diff --git a/tests/unit/model_bridge/test_encdec_weight_stacking.py b/tests/unit/model_bridge/test_encdec_weight_stacking.py new file mode 100644 index 000000000..995bf989f --- /dev/null +++ b/tests/unit/model_bridge/test_encdec_weight_stacking.py @@ -0,0 +1,136 @@ +"""Stacked weight properties and head labels on encoder-decoder bridges. + +The stacking helpers assumed a single top-level ``blocks``; encoder-decoder +adapters register ``encoder_blocks``/``decoder_blocks`` instead, so every +stacked property raised AttributeError and ``all_head_labels`` silently named +only half the heads. Mirrors ``HookedEncoderDecoder``, which stacks +self-attention over ``chain(encoder, decoder)`` and omits cross-attention. +""" + +from __future__ import annotations + +import copy + +import pytest +import torch +from transformers import T5Config, T5ForConditionalGeneration + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources._bridge_builder import ( + build_bridge_from_module, +) + +N_LAYERS = 2 +N_HEADS = 4 +D_MODEL = 32 +D_FF = 64 + + +@pytest.fixture(scope="module") +def t5_bridge() -> TransformerBridge: + torch.manual_seed(0) + cfg = T5Config( + vocab_size=128, + d_model=D_MODEL, + d_kv=D_MODEL // N_HEADS, + d_ff=D_FF, + num_layers=N_LAYERS, + num_decoder_layers=N_LAYERS, + num_heads=N_HEADS, + ) + model = T5ForConditionalGeneration(cfg).eval() + return build_bridge_from_module( + model, + "T5ForConditionalGeneration", + hf_config=copy.deepcopy(cfg), + tokenizer=None, + device="cpu", + ).eval() + + +def test_attention_weights_stack_over_encoder_then_decoder(t5_bridge): + """W_Q/K/V/O span encoder + decoder layers, not just one stack.""" + total_layers = 2 * N_LAYERS + d_head = D_MODEL // N_HEADS + assert t5_bridge.W_Q.shape == (total_layers, N_HEADS, D_MODEL, d_head) + assert t5_bridge.W_K.shape == (total_layers, N_HEADS, D_MODEL, d_head) + assert t5_bridge.W_V.shape == (total_layers, N_HEADS, D_MODEL, d_head) + assert t5_bridge.W_O.shape == (total_layers, N_HEADS, d_head, D_MODEL) + + +def test_mlp_weights_stack_over_encoder_then_decoder(t5_bridge): + total_layers = 2 * N_LAYERS + assert t5_bridge.W_in.shape == (total_layers, D_MODEL, D_FF) + assert t5_bridge.W_out.shape == (total_layers, D_FF, D_MODEL) + + +def test_factored_circuits_are_available(t5_bridge): + """QK/OV build on the stacked weights, so they follow for free.""" + total_layers = 2 * N_LAYERS + assert t5_bridge.QK.A.shape[0] == total_layers + assert t5_bridge.OV.A.shape[0] == total_layers + + +def test_blocks_with_spans_both_stacks(t5_bridge): + """'attn' means this block's self-attention: encoder attn + decoder self_attn.""" + matching = t5_bridge.blocks_with("attn") + assert [idx for idx, _ in matching] == list(range(2 * N_LAYERS)) + + +def test_stack_params_for_reports_encoder_and_decoder_layers(t5_bridge): + indices, stacked = t5_bridge.stack_params_for("attn", "attn.W_Q") + assert indices == list(range(2 * N_LAYERS)) + assert stacked.shape[0] == 2 * N_LAYERS + + +def test_all_head_labels_uses_the_encoder_decoder_scheme(t5_bridge): + """EL/DL labels name every head; a plain L{l}H{h} list named only half.""" + labels = t5_bridge.all_head_labels + assert len(labels) == 2 * N_LAYERS * N_HEADS + assert labels[0] == "EL0H0" + assert labels[-1] == f"DL{N_LAYERS - 1}H{N_HEADS - 1}" + assert sum(label.startswith("EL") for label in labels) == N_LAYERS * N_HEADS + + +def test_attn_head_labels_cover_both_stacks(t5_bridge): + """Derives from composition_layer_indices, which routes through blocks_with.""" + assert len(t5_bridge.attn_head_labels) == 2 * N_LAYERS * N_HEADS + + +def test_cross_attention_is_excluded_from_stacking(t5_bridge): + """Decoder blocks also carry cross_attn; HookedEncoderDecoder omits it and so do we.""" + decoder_block = t5_bridge.decoder_blocks[0] + assert "cross_attn" in decoder_block._modules, "fixture should have cross-attention" + assert t5_bridge.W_Q.shape[0] == 2 * N_LAYERS + + +def test_labels_follow_actual_block_counts_when_the_stacks_differ(): + """Asymmetric encoder/decoder depths: labels come from the real block lists. + + cfg.n_layers cannot describe both stacks at once, so deriving labels from it + would mislabel every decoder head on a lopsided model. + """ + torch.manual_seed(0) + encoder_layers, decoder_layers = 3, 1 + cfg = T5Config( + vocab_size=128, + d_model=D_MODEL, + d_kv=D_MODEL // N_HEADS, + d_ff=D_FF, + num_layers=encoder_layers, + num_decoder_layers=decoder_layers, + num_heads=N_HEADS, + ) + model = T5ForConditionalGeneration(cfg).eval() + bridge = build_bridge_from_module( + model, + "T5ForConditionalGeneration", + hf_config=copy.deepcopy(cfg), + tokenizer=None, + device="cpu", + ).eval() + + labels = bridge.all_head_labels + assert sum(label.startswith("EL") for label in labels) == encoder_layers * N_HEADS + assert sum(label.startswith("DL") for label in labels) == decoder_layers * N_HEADS + assert bridge.W_Q.shape[0] == encoder_layers + decoder_layers diff --git a/tests/unit/model_bridge/test_gated_hooks.py b/tests/unit/model_bridge/test_gated_hooks.py index 4e212eb0a..2f4eb058c 100644 --- a/tests/unit/model_bridge/test_gated_hooks.py +++ b/tests/unit/model_bridge/test_gated_hooks.py @@ -88,3 +88,65 @@ def test_run_with_cache_warns_on_fully_gated_names_filter(): assert any("gated-off" in str(w.message) for w in caught), ( "Expected a warning naming the gated-off hook, got: " f"{[str(w.message) for w in caught]}" ) + + +def test_add_hook_callable_filter_warns_and_skips_gated_points(): + """A callable filter matching only gated-off points attaches nothing and warns, + rather than leaving dead hooks that never fire.""" + bridge = TransformerBridge.boot_native(_cfg()) + tokens = torch.randint(0, 16, (1, 8)) + + fired = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + bridge.add_hook( + lambda name: name.endswith("hook_mlp_in"), + lambda t, hook=None: fired.append(1) or t, + ) + bridge(tokens, return_type="logits") + + # Not merely "did not fire" — a gated point never fires even when a dead + # hook is attached, so assert nothing was attached in the first place. + assert not bridge.hook_dict["blocks.0.hook_mlp_in"].fwd_hooks + assert not fired, "Gated-off hook fired; the filter should have skipped it" + assert any("gated-off" in str(w.message) for w in caught), ( + "Expected a warning naming the skipped gated-off hook, got: " + f"{[str(w.message) for w in caught]}" + ) + + +def test_add_hook_callable_filter_attaches_once_flag_enabled(): + """The same filter attaches and fires — silently — once the setter enables the flag.""" + bridge = TransformerBridge.boot_native(_cfg()) + bridge.set_use_hook_mlp_in(True) + tokens = torch.randint(0, 16, (1, 8)) + + fired = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + bridge.add_hook( + lambda name: name.endswith("hook_mlp_in"), + lambda t, hook=None: fired.append(1) or t, + ) + bridge(tokens, return_type="logits") + + assert fired, "Hook did not fire after enabling use_hook_mlp_in via the setter" + assert not [w for w in caught if "gated-off" in str(w.message)] + + +def test_add_hook_callable_filter_leaves_ungated_points_alone(): + """Control: a filter over an ungated point attaches and fires without warning.""" + bridge = TransformerBridge.boot_native(_cfg()) + tokens = torch.randint(0, 16, (1, 8)) + + fired = [] + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + bridge.add_hook( + lambda name: name == "blocks.0.hook_resid_post", + lambda t, hook=None: fired.append(1) or t, + ) + bridge(tokens, return_type="logits") + + assert fired, "Ungated hook should still attach and fire" + assert not [w for w in caught if "gated-off" in str(w.message)] diff --git a/tests/unit/model_bridge/test_gpt_oss_moe.py b/tests/unit/model_bridge/test_gpt_oss_moe.py index 81c5bdee1..773c56691 100644 --- a/tests/unit/model_bridge/test_gpt_oss_moe.py +++ b/tests/unit/model_bridge/test_gpt_oss_moe.py @@ -198,3 +198,61 @@ def test_gpt_oss_run_with_cache_with_random_weights(): # GPT-OSS has 32 experts with top-4 routing, so router_scores is (seq_len, 4) router_scores_0 = cache["blocks.0.mlp.hook_router_scores"] assert router_scores_0.shape == (5, 4) # seq_len=5, num_experts_per_tok=4 + + +def test_gpt_oss_routing_hooks_match_hooked_transformer_construction(): + """GPT-OSS softmaxes after top-k, so the scattered weights are tensor-identical + to HookedTransformer's gpt_oss_moe routing_weights, and the hooks are observe-only.""" + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + from transformer_lens.config import TransformerBridgeConfig + from transformer_lens.model_bridge import TransformerBridge + from transformer_lens.model_bridge.sources._hf_format import ( + map_default_transformer_lens_config, + ) + from transformer_lens.model_bridge.supported_architectures.gpt_oss import ( + GPTOSSArchitectureAdapter, + ) + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("openai/gpt-oss-20b", trust_remote_code=True) + config.num_hidden_layers = 2 + config.hidden_size = 128 + config.intermediate_size = 256 + config.num_attention_heads = 8 + config.num_key_value_heads = 2 + model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) + + tl_config = map_default_transformer_lens_config(config) + bridge = TransformerBridge( + model=model, + adapter=GPTOSSArchitectureAdapter( + TransformerBridgeConfig( + d_model=tl_config.d_model, + d_head=tl_config.d_head, + n_layers=tl_config.n_layers, + n_ctx=tl_config.n_ctx, + architecture="GptOssForCausalLM", + ) + ), + tokenizer=AutoTokenizer.from_pretrained("openai/gpt-oss-20b", trust_remote_code=True), + ) + bridge.enable_compatibility_mode(no_processing=True) + + tokens = torch.randint(0, 1000, (1, 5)) + baseline = bridge(tokens) + cached_logits, cache = bridge.run_with_cache(tokens) + + weights = cache["blocks.0.mlp.router.hook_expert_weights"] + indices = cache["blocks.0.mlp.router.hook_expert_indices"] + scores = cache["blocks.0.mlp.hook_router_scores"] + + assert weights.shape == (5, config.num_local_experts) + assert indices.shape == (5, config.num_experts_per_tok) + assert torch.equal(baseline, cached_logits), "routing hooks must be observe-only" + + # HookedTransformer's gpt_oss_moe.py builds the same tensor by scattering the + # post-top-k softmax back over expert ids. + hooked_construction = torch.zeros_like(weights) + hooked_construction.scatter_(-1, indices.long(), scores) + assert torch.equal(weights, hooked_construction) diff --git a/tests/unit/model_bridge/test_stack_block_params_hybrid.py b/tests/unit/model_bridge/test_stack_block_params_hybrid.py index d1a79987f..803d431f1 100644 --- a/tests/unit/model_bridge/test_stack_block_params_hybrid.py +++ b/tests/unit/model_bridge/test_stack_block_params_hybrid.py @@ -35,11 +35,25 @@ def __init__(self, mlp) -> None: self.mlp = mlp # nn.Module assignment registers into _modules +_BORROWED_HELPERS = ( + "_enumerate_blocks", + "_resolve_submodule_name", + "_rewrite_submodule_path", +) + + def _stub(mlps): - return SimpleNamespace( - blocks=[_Block(m) for m in mlps], + blocks = torch.nn.ModuleList([_Block(m) for m in mlps]) + # _stack_block_params walks the registered block lists, so the stub exposes + # the same _modules shape a real bridge does and borrows the real helpers. + stub = SimpleNamespace( + blocks=blocks, + _modules={"blocks": blocks}, cfg=SimpleNamespace(n_devices=1, device=None), ) + for helper in _BORROWED_HELPERS: + setattr(stub, helper, getattr(TransformerBridge, helper).__get__(stub)) + return stub def test_interleaved_model_stacks_dense_layers_only(caplog) -> None: diff --git a/transformer_lens/benchmarks/utils.py b/transformer_lens/benchmarks/utils.py index d3888c4fc..cac7251df 100644 --- a/transformer_lens/benchmarks/utils.py +++ b/transformer_lens/benchmarks/utils.py @@ -68,8 +68,6 @@ def is_tiny_test_model(model_name: str) -> bool: # modules with per-expert hooks (e.g., blocks.0.mlp.experts.3.hook_pre). "mlp.experts.", "mlp.hook_experts", - "mlp.hook_expert_indices", - "mlp.hook_expert_weights", # Parallel attention+MLP architectures (GPT-J, GPT-NeoX): HF has a single # shared layer norm (ln_1), while HT creates a virtual ln2 that shares weights # with ln1. The Bridge only wraps the actual HF ln_1, so ln2 hooks don't exist. diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py index 6233262c8..08ab2fdc7 100644 --- a/transformer_lens/model_bridge/bridge_core.py +++ b/transformer_lens/model_bridge/bridge_core.py @@ -40,6 +40,9 @@ # Block-list container attributes a bridge may expose. _BLOCK_LIST_ATTRS = ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks") +# Encoder blocks name self-attention ``attn``; decoder blocks name it ``self_attn`` +# (``cross_attn`` is a separate submodule, deliberately excluded from stacking). +_SELF_ATTENTION_NAMES = {"attn": ("attn", "self_attn")} def build_alias_to_canonical_map(hook_dict: Any, prefix: str = "") -> dict: @@ -781,18 +784,33 @@ def add_hook( if callable(name) and not isinstance(name, str): hook_dict = self.hook_dict seen_hooks: set = set() + gated_names_skipped: List[str] = [] for hook_name, hook_point in hook_dict.items(): if name(hook_name): hook_id = id(hook_point) if hook_id in seen_hooks: continue seen_hooks.add(hook_id) + # A filter is a sweep, not a targeted request, so a gated-off + # match is skipped rather than raised on — but silently + # attaching here would leave a dead hook that never fires, + # which is the failure this warns about. + if self._gated_hook_reason(hook_name) is not None: + gated_names_skipped.append(hook_name) + continue self._add_fn_to_hook_point(hook_point, hook_name, hook_fn, dir, is_permanent) + if gated_names_skipped: + warnings.warn( + f"add_hook: skipped {len(gated_names_skipped)} gated-off hook name(s) " + f"that would never fire: {gated_names_skipped}. Call the relevant " + "set_use_*(True) setter first to enable them.", + stacklevel=2, + ) return # An explicitly named gated-off hook point is a caller error: the hook # would silently never fire. Raise naming the setter to call (filters - # above add freely — they were not necessarily targeting gated names). + # above skip with a warning — they were not necessarily targeting gated names). reason = self._gated_hook_reason(name) if reason is not None: raise ValueError( diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 5e58a30a3..d4d4476da 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/transformer_lens/model_bridge/generalized_components/__init__.py @@ -74,6 +74,9 @@ MLAAttentionBridge, ) from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge +from transformer_lens.model_bridge.generalized_components.pooler import ( + BertPoolerBridge, +) from transformer_lens.model_bridge.generalized_components.moe import ( MoEBridge, MoERouterBridge, @@ -186,6 +189,7 @@ "GatedRMSNormBridge", "MoEBridge", "MoERouterBridge", + "BertPoolerBridge", "PositionEmbeddingsAttentionBridge", "Qwen3_5VisionBlockBridge", "Qwen3_5VisionEncoderBridge", diff --git a/transformer_lens/model_bridge/generalized_components/moe.py b/transformer_lens/model_bridge/generalized_components/moe.py index 19ba06ae6..695f9b93d 100644 --- a/transformer_lens/model_bridge/generalized_components/moe.py +++ b/transformer_lens/model_bridge/generalized_components/moe.py @@ -91,6 +91,28 @@ def __init__( f"submodules (declared: {sorted(submodules or {})})" ) self._sparse_required = sparse_required + # HookedTransformer exposes the routing observables on the MoE block + # itself; the bridge fires them on the router submodule, whose adapter + # key differs ("gate" on 5.13 SparseMoeBlocks, "router" on GPT-OSS). + # Alias so code migrated from HT finds them under the HT name. + self._router_hook_aliases = self._build_router_hook_aliases(submodules or {}) + self.hook_aliases = {**self.hook_aliases, **self._router_hook_aliases} + + @staticmethod + def _build_router_hook_aliases( + submodules: Mapping[str, GeneralizedComponent], + ) -> Dict[str, str]: + """Map HT's block-level routing hook names onto the router submodule.""" + aliases: Dict[str, str] = {} + for key, component in submodules.items(): + if not isinstance(component, MoERouterBridge): + continue + if component.weights_index is not None: + aliases["hook_expert_weights"] = f"{key}.hook_expert_weights" + if component.indices_index is not None: + aliases["hook_expert_indices"] = f"{key}.hook_expert_indices" + break + return aliases def _binds_dense_projections(self, component: torch.nn.Module) -> bool: """Whether this layer is the dense variant of an interleaved MoE stack. @@ -160,7 +182,7 @@ def set_original_component(self, component: torch.nn.Module) -> 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.hook_aliases = {**type(self).hook_aliases, **self._router_hook_aliases} self.property_aliases = { key: value for key, value in self.property_aliases.items() @@ -329,11 +351,33 @@ class MoERouterBridge(LinearBridge): 5.13 TopKRouters return ``(router_logits, topk_weights, topk_indices)``; hook_out fires on the logits (element ``logits_index`` — JetMoe puts them last) and the tuple is re-packed so HF's unpacking is undisturbed. + + ``hook_expert_weights`` / ``hook_expert_indices`` mirror the HookedTransformer + MoE routing hooks. HF routers hand back top-k-shaped weights + ``[tokens, top_k]``, so the weights are scattered to HT's + ``[tokens, num_experts]`` before firing and gathered back afterwards — an + unedited round trip returns the values bit-for-bit, and an edit reaches HF. """ - def __init__(self, *args: Any, logits_index: int = 0, **kwargs: Any): + def __init__( + self, + *args: Any, + logits_index: int = 0, + weights_index: Optional[int] = 1, + indices_index: Optional[int] = 2, + **kwargs: Any, + ): super().__init__(*args, **kwargs) self.logits_index = logits_index + self.weights_index = weights_index + self.indices_index = indices_index + # None means this router's tuple has no clean [tokens, top_k] pair + # (JetMoe returns a sorted-expert layout); registering the hook anyway + # would advertise an intervention point that can never fire. + if weights_index is not None: + self.hook_expert_weights = HookPoint() + if indices_index is not None: + self.hook_expert_indices = HookPoint() def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any: if self.original_component is None: @@ -344,9 +388,73 @@ def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> Any: output = self.original_component(input, *args, **kwargs) if not isinstance(output, tuple) or len(output) == 0: return self.hook_out(output) - idx = self.logits_index % len(output) - router_logits = self.hook_out(output[idx]) - return output[:idx] + (router_logits,) + output[idx + 1 :] + parts = list(output) + count = len(parts) + logits_at = self.logits_index % count + parts[logits_at] = self.hook_out(parts[logits_at]) + + weights_at = None if self.weights_index is None else self.weights_index % count + indices_at = None if self.indices_index is None else self.indices_index % count + if weights_at is None and indices_at is None: + return tuple(parts) + + indices = None if indices_at is None else parts[indices_at] + expanded = None + if weights_at is not None: + expanded = self._expand_expert_weights(parts[weights_at], indices, parts[logits_at]) + expanded = self.hook_expert_weights(expanded) + if indices_at is not None: + indices = self.hook_expert_indices(indices) + parts[indices_at] = indices + if weights_at is not None: + # Gathered after the indices hook so re-routing picks up the weight + # sitting at the newly selected expert, as HookedTransformer does. + parts[weights_at] = self._collapse_expert_weights(expanded, indices, parts[weights_at]) + return tuple(parts) + + def _expand_expert_weights( + self, + weights: torch.Tensor, + indices: Optional[torch.Tensor], + logits: torch.Tensor, + ) -> torch.Tensor: + """Scatter top-k weights into HT's ``[tokens, num_experts]`` layout.""" + if not self._is_top_k_shaped(weights, indices, logits): + return weights + assert indices is not None + scattered = torch.zeros( + (*weights.shape[:-1], logits.shape[-1]), + dtype=weights.dtype, + device=weights.device, + ) + scattered.scatter_(-1, indices.long(), weights) + return scattered + + def _collapse_expert_weights( + self, + expanded: Optional[torch.Tensor], + indices: Optional[torch.Tensor], + original: torch.Tensor, + ) -> torch.Tensor: + """Gather the expanded weights back to the top-k layout HF expects.""" + if expanded is None or indices is None or expanded.shape == original.shape: + return expanded if expanded is not None else original + return expanded.gather(-1, indices.long()) + + @staticmethod + def _is_top_k_shaped( + weights: torch.Tensor, + indices: Optional[torch.Tensor], + logits: torch.Tensor, + ) -> bool: + """Whether the weights are the top-k slice rather than full expert width.""" + return ( + indices is not None + and isinstance(weights, torch.Tensor) + and isinstance(logits, torch.Tensor) + and weights.shape == indices.shape + and weights.shape[-1] != logits.shape[-1] + ) def set_processed_weights( self, weights: Mapping[str, Optional[torch.Tensor]], verbose: bool = False diff --git a/transformer_lens/model_bridge/generalized_components/pooler.py b/transformer_lens/model_bridge/generalized_components/pooler.py new file mode 100644 index 000000000..7558875df --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/pooler.py @@ -0,0 +1,22 @@ +"""Pooler bridge component. + +This module contains the bridge component for [CLS] pooling heads. +""" + +from __future__ import annotations + +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +class BertPoolerBridge(GeneralizedComponent): + """Bridge component for BERT's [CLS] pooler. + + Wraps the whole pooler, so ``hook_out`` carries the post-tanh pooled + ``[CLS]`` vector rather than the pre-activation projection — the tensor + ``HookedEncoder``'s ``BertPooler`` exposes as ``hook_pooler_out``, which is + aliased here so code migrated from the legacy stack keeps working. + """ + + hook_aliases = {"hook_pooler_out": "hook_out"} diff --git a/transformer_lens/model_bridge/supported_architectures/bert.py b/transformer_lens/model_bridge/supported_architectures/bert.py index f71a237d6..219ab352d 100644 --- a/transformer_lens/model_bridge/supported_architectures/bert.py +++ b/transformer_lens/model_bridge/supported_architectures/bert.py @@ -12,6 +12,7 @@ from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( AttentionBridge, + BertPoolerBridge, BlockBridge, EmbeddingBridge, LinearBridge, @@ -152,7 +153,14 @@ def prepare_model(self, hf_model: Any) -> None: and no MLM-specific LayerNorm. """ if getattr(getattr(hf_model, "bert", None), "pooler", None) is not None: - self.components["pooler"] = LinearBridge(name="bert.pooler.dense") + # Wrap the pooler itself, not its inner Linear: HF applies tanh after + # the projection, so hook_out here is the pooled [CLS] that + # HookedEncoder's BertPooler exposes as hook_pooler_out. The dense + # stays hookable as a submodule for the pre-activation projection. + self.components["pooler"] = BertPoolerBridge( + name="bert.pooler", + submodules={"dense": LinearBridge(name="dense")}, + ) has_predictions = hasattr(getattr(hf_model, "cls", None), "predictions") has_nsp_head = hasattr(getattr(hf_model, "cls", None), "seq_relationship") diff --git a/transformer_lens/model_bridge/supported_architectures/gpt_oss.py b/transformer_lens/model_bridge/supported_architectures/gpt_oss.py index 7195e42bb..1de8ad108 100644 --- a/transformer_lens/model_bridge/supported_architectures/gpt_oss.py +++ b/transformer_lens/model_bridge/supported_architectures/gpt_oss.py @@ -8,6 +8,7 @@ EmbeddingBridge, LinearBridge, MoEBridge, + MoERouterBridge, PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, @@ -67,7 +68,12 @@ def __init__(self, cfg: Any) -> None: ), # GPT-OSS uses batched MoE experts with router scores # MoEBridge handles the (hidden_states, router_scores) tuple returns - "mlp": MoEBridge(name="mlp", config=self.cfg), + "mlp": MoEBridge( + name="mlp", + config=self.cfg, + submodules={"router": MoERouterBridge(name="router")}, + sparse_required=("router",), + ), }, ), "ln_final": RMSNormalizationBridge( diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py index f45e457ea..eef2b5710 100644 --- a/transformer_lens/model_bridge/supported_architectures/granite.py +++ b/transformer_lens/model_bridge/supported_architectures/granite.py @@ -74,7 +74,11 @@ def _build_moe_bridge(self) -> MoEBridge: return MoEBridge( name="block_sparse_moe", config=self.cfg, - submodules={"gate": MoERouterBridge(name="router", logits_index=-1)}, + submodules={ + "gate": MoERouterBridge( + name="router", logits_index=-1, indices_index=0, weights_index=1 + ) + }, ) def _build_component_mapping(self) -> dict: diff --git a/transformer_lens/model_bridge/supported_architectures/jetmoe.py b/transformer_lens/model_bridge/supported_architectures/jetmoe.py index 3f49f5b6f..95e82ff76 100644 --- a/transformer_lens/model_bridge/supported_architectures/jetmoe.py +++ b/transformer_lens/model_bridge/supported_architectures/jetmoe.py @@ -77,7 +77,12 @@ def __init__(self, cfg: Any) -> None: name="experts", submodules={ # JetMoeTopKGating puts logits last in its 5-tuple. - "router": MoERouterBridge(name="router", logits_index=-1), + "router": MoERouterBridge( + name="router", + logits_index=-1, + weights_index=None, + indices_index=None, + ), }, ), }, @@ -87,7 +92,12 @@ def __init__(self, cfg: Any) -> None: name="mlp", config=self.cfg, submodules={ - "gate": MoERouterBridge(name="router", logits_index=-1), + "gate": MoERouterBridge( + name="router", + logits_index=-1, + weights_index=None, + indices_index=None, + ), }, ), }, diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 97ef50940..64c53ce8f 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -37,7 +37,11 @@ from transformer_lens.FactoredMatrix import FactoredMatrix from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter -from transformer_lens.model_bridge.bridge_core import _BLOCK_LIST_ATTRS, BridgeCore +from transformer_lens.model_bridge.bridge_core import ( + _BLOCK_LIST_ATTRS, + _SELF_ATTENTION_NAMES, + BridgeCore, +) from transformer_lens.model_bridge.component_setup import ( refresh_container_state_owners, set_original_components, @@ -804,6 +808,153 @@ def to_tokens( tokens = tokens.to(self.cfg.device) return tokens + def encoder_output( + self, + frames: torch.Tensor, + one_zero_attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run the audio encoder from precomputed frames, skipping feature extraction. + + The audio-path analogue of ``start_at_layer``: ``frames`` is + ``[batch, frames, d_model]``, the tensor the conv front end would have + produced (observable at ``feat_proj.hook_out``), so callers can inject or + reuse frames without re-running the waveform convolutions. Positional + convolution and the encoder layer norm are applied first, exactly as the + full path does, then the blocks run. Mirrors + ``HookedAudioEncoder.encoder_output``. + + Hooks on the bridged components fire as usual, so this composes with + ``add_hook`` / ``get_caching_hooks``. + + Args: + frames: ``[batch, frames, d_model]`` precomputed encoder frames. + one_zero_attention_mask: Optional ``[batch, frames]`` mask, 1 for + real frames and 0 for padding. + + Returns: + The residual stream leaving the final block, ``[batch, frames, d_model]``. + """ + self._require_frame_entry_support() + if frames.ndim != 3: + raise ValueError( + "encoder_output expects precomputed frames [batch, frames, d_model]; " + f"got a {frames.ndim}D tensor. Pass a waveform to forward() instead." + ) + frames = frames.to(self.cfg.device) + + resid = frames + self.conv_pos_embed(frames) + resid = self.embed_ln(resid) + + additive_attention_mask = None + if one_zero_attention_mask is not None: + mask = one_zero_attention_mask.to(self.cfg.device) + additive_attention_mask = torch.where( + mask[:, None, None, :] == 0, + torch.tensor(float("-inf"), dtype=resid.dtype, device=resid.device), + torch.tensor(0.0, dtype=resid.dtype, device=resid.device), + ) + + for block in self.blocks: + output = block(resid, attention_mask=additive_attention_mask) + resid = output[0] if isinstance(output, tuple) else output + return resid + + def _require_frame_entry_support(self) -> None: + """Reject models with no conv-frame stage to re-enter.""" + if not getattr(self.cfg, "is_audio_model", False): + raise NotImplementedError( + "encoder_output is an audio-encoder entry point; this bridge is not an " + "audio model. Use start_at_layer for residual re-entry on text models." + ) + missing = [ + name + for name in ("conv_pos_embed", "embed_ln", "blocks") + if self._modules.get(name) is None + ] + if missing: + raise NotImplementedError( + "encoder_output needs a waveform encoder with a convolutional front end " + f"(missing {missing}). Spectrogram encoders such as AST have no " + "precomputed-frame stage to re-enter, so there is nothing to bypass." + ) + + def to_sentence_pair_tokens( + self, + sentence_a: str, + sentence_b: str, + move_to_device: bool = True, + truncate: bool = True, + ) -> Dict[str, torch.Tensor]: + """Pair-tokenize two sentences as ``[CLS] a [SEP] b [SEP]``. + + Returns ``input_ids``, ``token_type_ids`` and ``attention_mask``. The + segment ids are not decorative: without them a BERT NSP head sees both + sentences as one segment and its logits collapse. Mirrors + ``BertNextSentencePrediction.to_tokens``. + + Args: + sentence_a: First sentence of the pair. + sentence_b: Second sentence of the pair. + move_to_device: Move the returned tensors to ``cfg.device``. + truncate: Truncate to the model's context window. + """ + assert self.tokenizer is not None, "Cannot pair-tokenize without a tokenizer" + encodings = self.tokenizer( + sentence_a, + sentence_b, + return_tensors="pt", + padding=True, + truncation=truncate, + max_length=self.cfg.n_ctx if truncate else None, + ) + if "token_type_ids" not in encodings: + raise ValueError( + f"{type(self.tokenizer).__name__} emits no token_type_ids, so it cannot " + "express a sentence pair. Next-sentence prediction needs a " + "segment-aware tokenizer (e.g. BERT's)." + ) + keys = ("input_ids", "token_type_ids", "attention_mask") + tokens = {key: encodings[key] for key in keys if key in encodings} + if move_to_device: + tokens = {key: value.to(self.cfg.device) for key, value in tokens.items()} + return tokens + + def predict_next_sentence( + self, + sentence_a: str, + sentence_b: str, + return_type: Optional[str] = "predictions", + truncate: bool = True, + ) -> Any: + """Run next-sentence prediction over a sentence pair given as strings. + + Owns the ``token_type_ids`` plumbing that a hand-rolled pair forward has + to remember. Requires a bridge booted onto an NSP head — otherwise the + model has no 2-class output to decode. Mirrors + ``BertNextSentencePrediction.forward``. + + Args: + sentence_a: First sentence of the pair. + sentence_b: Second sentence of the pair. + return_type: ``"predictions"`` for the decoded verdict, or + ``"logits"`` for the raw 2-class scores. + truncate: Truncate to the model's context window. + """ + tokens = self.to_sentence_pair_tokens(sentence_a, sentence_b, truncate=truncate) + forward_kwargs: Dict[str, Any] = { + key: value for key, value in tokens.items() if key != "input_ids" + } + logits = self(tokens["input_ids"], return_type="logits", **forward_kwargs) + if logits.shape[-1] != 2: + raise ValueError( + "predict_next_sentence needs a next-sentence-prediction head, but this " + f"bridge produces {logits.shape[-1]} output classes. Boot it with " + "model_class=BertForNextSentencePrediction." + ) + if return_type == "logits": + return logits + return self._finalize_return(return_type, logits, tokens["input_ids"]) + def to_string( self, tokens: Union[List[int], torch.Tensor, np.ndarray] ) -> Union[str, List[str]]: @@ -962,15 +1113,58 @@ def to_single_str_token(self, int_token: int) -> str: return str(token[0]) raise AssertionError("Expected a single string token.") + def _enumerate_blocks(self) -> List[Tuple[int, Any]]: + """(index, block) over every registered block list, encoder before decoder. + + Decoder-only models register a single ``blocks``, so the indices are the + plain layer indices. Encoder-decoder models register ``encoder_blocks`` + and ``decoder_blocks`` instead; those are concatenated into one index + space, matching ``HookedEncoderDecoder``'s ``chain(encoder, decoder)``. + """ + pairs: List[Tuple[int, Any]] = [] + for list_name in _BLOCK_LIST_ATTRS: + block_list = self._modules.get(list_name) + if not isinstance(block_list, nn.ModuleList): + continue + for block in block_list: + pairs.append((len(pairs), block)) + return pairs + + def _resolve_submodule_name(self, block: Any, submodule: str) -> Optional[str]: + """The block's actual name for ``submodule``. + + Encoder blocks name self-attention ``attn`` while decoder blocks name it + ``self_attn``, so a caller asking for ``attn`` means "this block's + self-attention" on either side. Cross-attention is never resolved here: + ``HookedEncoderDecoder`` omits it from stacked weights, and this mirrors + that so the two stacks line up layer for layer. + """ + for candidate in _SELF_ATTENTION_NAMES.get(submodule, (submodule,)): + if candidate in block._modules: + return candidate + return None + + def _rewrite_submodule_path(self, attr_path: str, submodule: str, actual: Optional[str]) -> str: + """Re-point ``attr_path``'s leading segment at the block's actual submodule.""" + if actual is None or actual == submodule: + return attr_path + if attr_path.split(".")[0] != submodule: + return attr_path + return actual + attr_path[len(submodule) :] + def blocks_with(self, submodule: str) -> List[Tuple[int, "GeneralizedComponent"]]: """Return (index, block) pairs for blocks with the named bridged submodule. Checks _modules (not hasattr) so HF-internal attrs don't match. Use instead of assuming blocks[0] is representative on hybrid models. + On encoder-decoder models the indices span encoder then decoder blocks, + and ``"attn"`` matches the decoder's ``self_attn`` too. """ - if not hasattr(self, "blocks"): - return [] - return [(i, block) for i, block in enumerate(self.blocks) if submodule in block._modules] + return [ + (index, block) + for index, block in self._enumerate_blocks() + if self._resolve_submodule_name(block, submodule) is not None + ] def stack_params_for( self, submodule: str, attr_path: str, reshape_fn: Optional[Callable] = None @@ -988,7 +1182,10 @@ def stack_params_for( indices: List[int] = [] weights: List[torch.Tensor] = [] for idx, block in matching: - w = _resolve_attr_path(block, attr_path) + resolved = self._resolve_submodule_name(block, submodule) + w = _resolve_attr_path( + block, self._rewrite_submodule_path(attr_path, submodule, resolved) + ) if w is None: raise AttributeError( f"blocks[{idx}].{attr_path} is None — this checkpoint has no such " @@ -1011,12 +1208,16 @@ def _stack_block_params( AttributeError killed the accessor for the whole model. """ first_attr = attr_path.split(".")[0] + all_blocks = self._enumerate_blocks() matching_blocks: List[Tuple[int, torch.Tensor]] = [] - for i, block in enumerate(self.blocks): - if first_attr not in block._modules: + for i, block in all_blocks: + resolved = self._resolve_submodule_name(block, first_attr) + if resolved is None: continue try: - weight = _resolve_attr_path(block, attr_path) + weight = _resolve_attr_path( + block, self._rewrite_submodule_path(attr_path, first_attr, resolved) + ) except AttributeError: continue if weight is None: @@ -1036,7 +1237,7 @@ def _stack_block_params( f"Use bridge.blocks_with('{first_attr}') to check availability." ) - if len(matching_blocks) < len(self.blocks): + if len(matching_blocks) < len(all_blocks): indices = [i for i, _ in matching_blocks] logging.warning( "Hybrid model: only %d/%d blocks resolve '%s'. Returning stacked tensor " @@ -1044,7 +1245,7 @@ def _stack_block_params( "indices[i], not layer i. For explicit index mapping, use " "bridge.stack_params_for('%s', '%s').", len(matching_blocks), - len(self.blocks), + len(all_blocks), attr_path, indices, first_attr, @@ -1389,7 +1590,19 @@ def layer_types(self) -> List[str]: @property def all_head_labels(self) -> list[str]: - """Human-readable labels for all attention heads, e.g. ['L0H0', 'L0H1', ...].""" + """Human-readable labels for all attention heads, e.g. ['L0H0', 'L0H1', ...]. + + Encoder-decoder models use ``HookedEncoderDecoder``'s ``EL{l}H{h}`` / + ``DL{l}H{h}`` scheme so encoder and decoder heads stay distinguishable; + a plain ``L{l}H{h}`` list would name only half of them. + """ + encoder_blocks = self._modules.get("encoder_blocks") + decoder_blocks = self._modules.get("decoder_blocks") + if isinstance(encoder_blocks, nn.ModuleList) and isinstance(decoder_blocks, nn.ModuleList): + heads = range(self.cfg.n_heads) + return [f"EL{l}H{h}" for l in range(len(encoder_blocks)) for h in heads] + [ + f"DL{l}H{h}" for l in range(len(decoder_blocks)) for h in heads + ] return [f"L{l}H{h}" for l in range(self.cfg.n_layers) for h in range(self.cfg.n_heads)] @property