Skip to content

Fix BLOOM residual branch hook semantics - #1640

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:fix/bloom-hook-output-semantics
Aug 11, 2026
Merged

Fix BLOOM residual branch hook semantics#1640
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
koriyoshi2041:fix/bloom-hook-output-semantics

Conversation

@koriyoshi2041

Copy link
Copy Markdown
Contributor

Description

Fixes #1639.

BLOOM's Hugging Face attention and MLP modules return states after adding the residual stream, so the default block aliases made hook_attn_out and hook_mlp_out point at accumulated residual states. This breaks the HookedTransformer-compatible decomposition identities and makes interventions target the wrong tensors.

This change overrides those two aliases to use BLOOM's existing output-projection hooks, which sit immediately before each residual addition and remain writable. The integration coverage uses the tiny BLOOM fixture to verify both residual identities on every layer and confirms that an attention-output ablation changes the model output.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Validation

  • uv run pytest tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py tests/integration/model_bridge/test_bloom_hook_semantics.py — 31 passed
  • uv run pytest tests/unit/model_bridge/test_hook_alias_resolution.py tests/unit/model_bridge/test_component_hooks_fire.py tests/integration/model_bridge/test_pretrain_adapter.py — 212 passed, 3 skipped, 10 xfailed
  • make format
  • uv run mypy transformer_lens/model_bridge/supported_architectures/bloom.py tests/unit/model_bridge/supported_architectures/test_bloom_adapter.py tests/integration/model_bridge/test_bloom_hook_semantics.py

Checklist

  • I have commented my code, particularly in hard-to-understand areas (no new comments were needed for the alias mapping)
  • I have made corresponding changes to the documentation (not applicable; behavior is covered by tests)
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective
  • New and existing targeted tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4
jlarson4 merged commit 7db6018 into TransformerLensOrg:dev Aug 11, 2026
25 checks passed
@jlarson4 jlarson4 mentioned this pull request Aug 12, 2026
7 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants