[Bug Report] Residual-decomposition APIs raise KeyError on SSM / hybrid bridges (Qwen3.5, Mamba-2, NemotronH, GraniteMoeHybrid)

Aperta
  • Qualcosa non va?
#1,986 1 commento 0 reazioni 1 assegnatario Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

@SBrouck ci sta già lavorando.

Dal 5/11/2026.

  • #1987 di @SBrouck — aperta

Valutazione

Questa issue non è ancora stata valutata.

Descrizione

complexity-moderate enhancement TransformerBridge

Describe the bug

On SSM and hybrid bridges, the ActivationCache residual-decomposition APIs raise KeyError on the first block without attention:

API Error
decompose_resid() KeyError: 'blocks.0.hook_attn_out'
decompose_resid(mode="mlp") KeyError: 'blocks.0.hook_mlp_out' (Mamba-2, NemotronH, GraniteMoeHybrid)
compute_head_results(), stack_head_results() KeyError: 'blocks.0.attn.hook_z'
get_full_resid_decomposition() KeyError: 'blocks.0.attn.hook_z'
logit_attrs(cache.decompose_resid(), ...) KeyError: 'blocks.0.hook_attn_out'

Reproduced on Qwen3.5, Mamba-2, NemotronH and GraniteMoeHybrid. accumulated_resid works.

These methods loop over range(n_layers) and assume every block has attn and mlp. Since #1578, gated-delta-net and Mamba blocks intentionally expose their mixer output (linear_attn.hook_out, mixer.hook_out) instead of hook_attn_out, and Mamba / NemotronH blocks have no MLP. The cache APIs were not updated for that.

The fix looks straightforward because the residual stream still decomposes exactly. On Qwen3.5, NemotronH and Mamba-2, resid_post - resid_pre equals the mixer output (attention or SSM) plus mlp_out when present, within ~1e-8.

Code example

import torch
from transformers import AutoTokenizer
from transformer_lens.model_bridge import TransformerBridge

name = "Qwen/Qwen3.5-0.8B"
bridge = TransformerBridge.boot_transformers(name, device="cpu", dtype=torch.float32)
ids = AutoTokenizer.from_pretrained(name)("The capital of Germany is", return_tensors="pt").input_ids
_, cache = bridge.run_with_cache(ids)

cache.accumulated_resid()      # OK
cache.decompose_resid()        # KeyError: 'blocks.0.hook_attn_out'
cache.stack_head_results()     # KeyError: 'blocks.0.attn.hook_z'

The same happens without Hub access on the tiny random fixtures in tests/integration/model_bridge/test_ssm_effective_attention_dispatch.py.

System Info

  • transformer_lens 4.3.0 from PyPI; the code involved is unchanged on dev (837fb72)
  • macOS (Apple Silicon), CPU, fp32
  • Python 3.12, transformers 5.13 and 5.22

Additional context

I have a fix with tests ready and will open a PR against dev. Blocks without attention contribute their mixer output in the attention slot, only attention layers contribute heads, and only blocks with an MLP contribute mlp_out or neurons. All decompositions still sum to the residual stream.

Checklist
  • I have checked that there is no similar issue in the repo (required)
Lingua principale
Python
Stelle
3.9k
Fork
708
Merge medio
1g 9h
PR unite (30g)
84

Preparare l'ambiente

Apri in Codespaces

Avvia il container di sviluppo del progetto nel browser, con il tuo account GitHub.

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di TransformerLensOrg/TransformerLens

Tutte le issue di TransformerLensOrg/TransformerLens

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.