Additional Bloom Cleanup

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"],
+        )

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"],
+        )

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)

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.
 

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.