From 9325c1fb8afb4af4317df189a4ad05f09ac158f6 Mon Sep 17 00:00:00 2001 From: jlarson4 Date: Tue, 11 Aug 2026 09:41:58 -0500 Subject: [PATCH] Additional Bloom cleanup --- .../model_bridge/test_bloom_hook_semantics.py | 96 ++++++++++++++++++- .../test_residual_decomposition_identities.py | 49 ++++++++++ .../test_bloom_adapter.py | 8 ++ .../generalized_components/bloom_mlp.py | 9 ++ .../supported_architectures/bloom.py | 3 + 5 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 tests/integration/model_bridge/test_residual_decomposition_identities.py diff --git a/tests/integration/model_bridge/test_bloom_hook_semantics.py b/tests/integration/model_bridge/test_bloom_hook_semantics.py index f2aead104..2c0747854 100644 --- a/tests/integration/model_bridge/test_bloom_hook_semantics.py +++ b/tests/integration/model_bridge/test_bloom_hook_semantics.py @@ -30,14 +30,106 @@ def test_residual_branch_hooks_decompose_stream(bloom_bridge: TransformerBridge) ) -def test_residual_branch_hooks_are_writable(bloom_bridge: TransformerBridge) -> None: +def test_attn_out_ablation_collapses_residual_step(bloom_bridge: TransformerBridge) -> None: + """Zeroing hook_attn_out must yield resid_mid == resid_pre — pins that writes + land on the additive contribution, not the residual-added module output.""" tokens = bloom_bridge.to_tokens("The capital of France is Paris.") + captured = {} + + def grab(key: str): + def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor: + captured[key] = tensor.detach().clone() + return tensor + + return hook_fn + + with torch.no_grad(): + baseline = bloom_bridge(tokens) + ablated = bloom_bridge.run_with_hooks( + tokens, + fwd_hooks=[ + ("blocks.0.hook_attn_out", lambda tensor, hook: torch.zeros_like(tensor)), + ("blocks.0.hook_resid_pre", grab("resid_pre")), + ("blocks.0.hook_resid_mid", grab("resid_mid")), + ], + ) + + assert not torch.equal(ablated, baseline) + torch.testing.assert_close(captured["resid_mid"], captured["resid_pre"]) + + +def test_mlp_out_ablation_collapses_residual_step(bloom_bridge: TransformerBridge) -> None: + """Zeroing hook_mlp_out must yield resid_post == resid_mid.""" + tokens = bloom_bridge.to_tokens("The capital of France is Paris.") + captured = {} + + def grab(key: str): + def hook_fn(tensor: torch.Tensor, hook) -> torch.Tensor: + captured[key] = tensor.detach().clone() + return tensor + + return hook_fn with torch.no_grad(): baseline = bloom_bridge(tokens) ablated = bloom_bridge.run_with_hooks( tokens, - fwd_hooks=[("blocks.0.hook_attn_out", lambda tensor, hook: torch.zeros_like(tensor))], + fwd_hooks=[ + ("blocks.0.hook_mlp_out", lambda tensor, hook: torch.zeros_like(tensor)), + ("blocks.0.hook_resid_mid", grab("resid_mid")), + ("blocks.0.hook_resid_post", grab("resid_post")), + ], ) assert not torch.equal(ablated, baseline) + torch.testing.assert_close(captured["resid_post"], captured["resid_mid"]) + + +def test_slow_but_exact_checkpoint_still_fires_mlp_out() -> None: + """Checkpoints shipping pretraining_tp>1 + slow_but_exact=True (e.g. this + bigscience testing model) make HF's BloomMLP bypass the dense_4h_to_h module + call, so hook_mlp_out would silently vanish from the cache. The bridge must + force the module-call path its hooks attach to.""" + bridge = TransformerBridge.boot_transformers( + "bigscience/bigscience-small-testing", device="cpu", dtype=torch.float32 + ) + tokens = bridge.to_tokens("The capital of France is Paris.") + + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) + + for layer in range(bridge.cfg.n_layers): + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_resid_post"], + cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"], + ) + + +def test_compatibility_mode_preserves_residual_semantics() -> None: + """Compat-mode weight processing must not crash on BLOOM and must keep the + residual identities and the model function (log_softmax; raw logits shift + under center_unembed). Regression: the tiny fixture's stray + num_key_value_heads sent fold_value_biases down a GQA branch.""" + bridge = TransformerBridge.boot_transformers(MODEL, device="cpu", dtype=torch.float32) + tokens = bridge.to_tokens("The capital of France is Paris.") + + with torch.no_grad(): + base_log_probs = torch.log_softmax(bridge(tokens), dim=-1) + + bridge.enable_compatibility_mode() + + with torch.no_grad(): + logits, cache = bridge.run_with_cache(tokens) + + torch.testing.assert_close( + torch.log_softmax(logits, dim=-1), base_log_probs, atol=1e-4, rtol=1e-5 + ) + for layer in range(bridge.cfg.n_layers): + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_resid_mid"], + cache[f"blocks.{layer}.hook_resid_pre"] + cache[f"blocks.{layer}.hook_attn_out"], + ) + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_resid_post"], + cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"], + ) diff --git a/tests/integration/model_bridge/test_residual_decomposition_identities.py b/tests/integration/model_bridge/test_residual_decomposition_identities.py new file mode 100644 index 000000000..0eac7339c --- /dev/null +++ b/tests/integration/model_bridge/test_residual_decomposition_identities.py @@ -0,0 +1,49 @@ +"""Residual-stream decomposition identities across sequential-residual architectures. + +Guards the HookedTransformer contract that ``hook_attn_out`` / ``hook_mlp_out`` +are additive contributions (#1639: BLOOM's HF modules add the residual +internally, so the default block aliases exposed accumulated states): + + resid_pre + attn_out == resid_mid + resid_mid + mlp_out == resid_post + +The fixture list is curated, not exhaustive — one tiny checkpoint per +residual-wiring pattern the bridge handles: + +- gpt2: block-level residual adds, the HookedTransformer reference wiring +- mistral: Llama-style pre-RMSNorm block-level adds +- bloom: residual added *inside* the HF attention/MLP modules (the #1639 case) + +Parallel-residual architectures (Falcon, GPT-J, NeoX, Cohere) are out of scope: +they have no ``hook_resid_mid``. +""" + +import pytest +import torch + +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.mark.parametrize("model_name", SEQUENTIAL_RESIDUAL_MODELS) +def test_residual_decomposition_identities(model_name: str) -> None: + bridge = TransformerBridge.boot_transformers(model_name, device="cpu", dtype=torch.float32) + tokens = bridge.to_tokens("The capital of France is Paris.") + + with torch.no_grad(): + _, cache = bridge.run_with_cache(tokens) + + for layer in range(bridge.cfg.n_layers): + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_resid_mid"], + cache[f"blocks.{layer}.hook_resid_pre"] + cache[f"blocks.{layer}.hook_attn_out"], + ) + torch.testing.assert_close( + cache[f"blocks.{layer}.hook_resid_post"], + cache[f"blocks.{layer}.hook_resid_mid"] + cache[f"blocks.{layer}.hook_mlp_out"], + ) diff --git a/tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py b/tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py index a77c22183..ae7c04e26 100644 --- a/tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py @@ -108,6 +108,14 @@ def test_residual_branch_hook_aliases(self, adapter: BloomArchitectureAdapter) - assert blocks.hook_aliases["hook_attn_out"] == "attn.o.hook_out" assert blocks.hook_aliases["hook_mlp_out"] == "mlp.out.hook_out" + def test_stray_kv_head_count_is_discarded(self) -> None: + """A stray num_key_value_heads in a checkpoint config must not put the + adapter into GQA mode — HF BloomAttention ignores the field entirely.""" + cfg = _make_cfg() + cfg.n_key_value_heads = 4 + adapter = BloomArchitectureAdapter(cfg) + assert adapter.cfg.n_key_value_heads is None + def test_ln_final_type_and_name(self, adapter: BloomArchitectureAdapter) -> None: mapping = self._mapping(adapter) assert isinstance(mapping["ln_final"], NormalizationBridge) diff --git a/transformer_lens/model_bridge/generalized_components/bloom_mlp.py b/transformer_lens/model_bridge/generalized_components/bloom_mlp.py index 69111c24d..63af2d9d7 100644 --- a/transformer_lens/model_bridge/generalized_components/bloom_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/bloom_mlp.py @@ -38,6 +38,15 @@ def __init__( """ super().__init__(name, config, submodules or {}) + def set_original_component(self, original_component: torch.nn.Module) -> None: + super().set_original_component(original_component) + # The Megatron-TP replay path (pretraining_tp>1 + slow_but_exact) computes + # dense_4h_to_h via F.linear on weight slices, bypassing the module call the + # out-projection hooks attach to. Force the module path; the only difference + # is fp summation order. + if getattr(original_component, "slow_but_exact", False): + setattr(original_component, "slow_but_exact", False) + def forward(self, *args: Any, **kwargs: Any) -> Any: """Forward pass through BLOOM MLP with hooks. diff --git a/transformer_lens/model_bridge/supported_architectures/bloom.py b/transformer_lens/model_bridge/supported_architectures/bloom.py index 2744d6d2e..c82a365de 100644 --- a/transformer_lens/model_bridge/supported_architectures/bloom.py +++ b/transformer_lens/model_bridge/supported_architectures/bloom.py @@ -33,6 +33,9 @@ def __init__(self, cfg: Any) -> None: self.cfg.final_rms = False self.cfg.gated_mlp = False self.cfg.attn_only = False + # HF BloomAttention is always full MHA and ignores num_key_value_heads; + # a stray config field would otherwise send weight processing down GQA paths. + self.cfg.n_key_value_heads = None self.cfg.default_prepend_bos = False # After split_qkv_matrix, Q/K/V are individual [n_heads*d_head, d_model] weights.