[Bug Report] `get_caching_hooks` fires for alias names in `names_filter` but caches only under the canonical key

Đã đóng
  • Có gì không ổn?
#1,857 2 bình luận 0 reaction 1 người được giao Xem trên GitHub

Maintainer thường phản hồi trong vòng 1 ngày

@Mudassiruddin7 đang làm issue này rồi.

Từ ngày 6/10/2026.

Đánh giá

Issue này chưa được đánh giá.

Mô tả

bug complexity-moderate TransformerBridge

Describe the bug

TransformerBridge.get_caching_hooks(names_filter=<alias name>) matches the alias entry in hook_dict and the hook fires, but save_hook keys the cache by hook.name, which is the canonical name. The returned cache therefore has no entry under the name the caller asked for: names_filter="blocks.0.hook_resid_pre" produces a cache whose only key is blocks.0.hook_in, and cache["blocks.0.hook_resid_pre"] raises KeyError. Because the hook did fire and the cache is non-empty, nothing warns.

run_with_cache with the same filter returns both the canonical key and the alias key (it copies alias entries into the cache after the forward), so the two caching APIs disagree for the same input.

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))
alias = "blocks.0.hook_resid_pre"

with torch.no_grad():
    _, ref = bridge.run_with_cache(tokens, names_filter=alias)
print(sorted(ref.cache_dict))

cache, fwd_hooks, _ = bridge.get_caching_hooks(names_filter=alias)
with torch.no_grad(), bridge.hooks(fwd_hooks=fwd_hooks):
    bridge.forward(tokens)
print(sorted(cache))
cache[alias]  # KeyError

Output on dev (1012730f):

['blocks.0.hook_in', 'blocks.0.hook_resid_pre']
['blocks.0.hook_in']
KeyError: 'blocks.0.hook_resid_pre'

System Info

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

Additional context

Found while verifying #1856. Note the fix cannot be a post-forward alias pass like run_with_cache's, because get_caching_hooks hands back the live cache dict before the model runs, it needs alias awareness inside save_hook (e.g. also store under the hook point's alias names, or under the filter name that matched).

Checklist
  • I have checked that there is no similar issue in the repo (required)
Ngôn ngữ chính
Python
Star
3.9k
Fork
708
Merge trung bình
1 ngày 14 giờ
Pull request đã merge (30 ngày)
74

Chuẩn bị môi trường

Mở trong Codespaces

Khởi chạy dev container của dự án ngay trên trình duyệt, bằng tài khoản GitHub của bạn.

  • Không có Dockerfile hay tệp Docker Compose
  • Có mẫu pull request
  • Không có hướng dẫn đóng góp

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của TransformerLensOrg/TransformerLens

Tất cả issue của TransformerLensOrg/TransformerLens

Issue tương tự

Thêm issue về Python

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.