Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[Bug Report] get_caching_hooks slices the wrong axis when remove_batch_dim=True is combined with pos_slice

Chiusa
  • Qualcosa non va?
#1,856 1 commento 0 reazioni 1 assegnatario Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

@JoeyTan21 ci sta già lavorando.

Dal 6/10/2026.

Valutazione

Questa issue non è ancora stata valutata.

Descrizione

Describe the bug

TransformerBridge.get_caching_hooks(remove_batch_dim=True, pos_slice=...) (and add_caching_hooks, which delegates to it) returns activations sliced on the wrong axis. save_hook drops the batch dimension first and then applies pos_slice on _pos_slice_dim(name), which is the positive axis 1 for every activation except attention maps. Once the batch dim is gone, axis 1 is d_model for the residual stream and n_heads for head-split tensors, so e.g. blocks.0.hook_out is cached as [pos, 1] instead of [1, d_model]. Attention patterns are unaffected because their axis is -2.

run_with_cache with the same arguments slices first and removes the batch dim afterwards, so the two caching APIs silently disagree for the same inputs.

Code example

import torch
from transformer_lens.config import TransformerBridgeConfig
from transformer_lens.model_bridge import TransformerBridge

cfg = TransformerBridgeConfig(
    d_model=32, d_head=16, n_heads=2, n_layers=1, n_ctx=8, d_vocab=16,
    d_mlp=64, act_fn="gelu", normalization_type="LN", seed=0,
)
bridge = TransformerBridge.boot_native(cfg)
tokens = torch.randint(0, cfg.d_vocab, (1, 6))
names = ["blocks.0.hook_out", "blocks.0.attn.o.hook_in", "blocks.0.attn.hook_pattern"]

with torch.no_grad():
    _, ref = bridge.run_with_cache(tokens, names_filter=names, remove_batch_dim=True, pos_slice=-1)

cache, fwd_hooks, _ = bridge.get_caching_hooks(names_filter=names, remove_batch_dim=True, pos_slice=-1)
with torch.no_grad(), bridge.hooks(fwd_hooks=fwd_hooks):
    bridge.forward(tokens)

for name in names:
    print(name, tuple(ref[name].shape), tuple(cache[name].shape))

Output on dev (1012730f):

blocks.0.hook_out (1, 32) (6, 1)
blocks.0.attn.o.hook_in (1, 2, 16) (6, 1, 16)
blocks.0.attn.hook_pattern (2, 1, 6) (2, 1, 6)

Expected: the get_caching_hooks shapes (and values) match run_with_cache, i.e. (1, 32), (1, 2, 16), (2, 1, 6).

Without remove_batch_dim=True both APIs agree, so the bug is specifically the order of the two operations in get_caching_hooks.save_hook (transformer_lens/model_bridge/bridge_core.py). Swapping them — slice on the batched tensor, then drop the batch dim, as run_with_cache._store already does — fixes the reproduced cases; I have a patch with a unit test ready and can open a PR against dev.

System Info

  • Source checkout installed with uv sync, dev at 1012730f; same code on main.
  • macOS (Apple Silicon), CPU only; Python 3.12.13; torch 2.11.0; transformers 5.x.
  • No model downloads needed: reproduced with a random-init TransformerBridge.boot_native model.

Additional context

The legacy HookedRootModule.get_caching_hooks used negative position dims (-2 / -3) and so was immune to this ordering. Searched issues/PRs for pos_slice, remove_batch_dim, get_caching_hooks; #574 / #578 (2024) concern the HookedTransformer-era pos_slice=None crash, not this.

Checklist
  • I have checked that there is no similar issue in the repo (required)

Disclosure: found, reproduced and the draft fix tested locally with the help of an AI coding agent (Claude Code); reviewed before filing.

Lingua principale
Python
Stelle
3.9k
Fork
708
Merge medio
2g 1h
PR unite (30g)
78

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.