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

未关闭
  • 哪里不对？
#1,986 1 条评论 0 个 reaction 已指派 1 人 在 GitHub 查看

维护者通常 1 天内回复

@SBrouck 已经在做这个了。

开始于 2026年11月5日。

评估

这个 Issue 还没有评估数据。

描述

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)
主要语言
Python
星标
3.9k
派生
708
平均合并
1 天 4 小时
30 天内合并 PR
87

环境准备

在 Codespaces 中打开

在浏览器里用你自己的 GitHub 账号启动这个项目的开发容器。

  • 没有 Dockerfile 或 Docker Compose 文件
  • 有 Pull Request 模板
  • 没有贡献指南

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

TransformerLensOrg/TransformerLens 的其他 Issue

查看 TransformerLensOrg/TransformerLens 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。