[Bug Report] [macOS-arm64] Cached eager attention NaNs in transformers v5 — blocks bridge KV-cache generation

[jlarson4]

Summary

TransformerBridge.generate() produces all-NaN logits at the cached step on macOS-arm64 + transformers v5 + PyTorch 2.7.1, even though the bridge already forces eager attention, the cache contents are finite, and the model is on CPU. The same forward computation without past_key_values returns correct logits. The bug appears to live in HuggingFace transformers' cached eager-attention path (and/or PyTorch's CPU matmul on Apple Silicon), not in transformer_lens.

Two tests are skipped on macOS-arm64 to unblock the MPS CI job; they remain active on Linux:

• tests/unit/model_bridge/test_bridge_generate_no_tokenizer.py::test_generate_without_tokenizer_stop_at_eos_false_kv_cache
• tests/unit/model_bridge/test_bridge_generate_no_tokenizer.py::test_generate_return_type_str_without_tokenizer_errors

Reproduction (CI only - have not tried reproduction locally on M1 Mac)

CI matrix where it fires: macos-latest (GitHub Actions managed Apple Silicon) + Python 3.11.9 + PyTorch 2.7.1 + transformers 5.8.1.

import torch
from transformer_lens.model_bridge import TransformerBridge

bridge = TransformerBridge.boot_transformers("distilgpt2", device="cpu")
bridge.tokenizer = None
tokens = torch.tensor([[15496, 11, 314, 1101, 257]], dtype=torch.long)

# NaNs on macOS-arm64 CI; finite locally and on Linux CI:
bridge.generate(
    tokens,
    max_new_tokens=3,
    stop_at_eos=False,
    use_past_kv_cache=True,
    return_type="tokens",
    verbose=False,
)

Direct HF reproduction (no bridge involved):

o0 = bridge.original_model(tokens, use_cache=True)
next_id = o0.logits[:, -1, :].argmax(-1, keepdim=True)

# NaN on macOS-arm64 CI, finite elsewhere:
o1 = bridge.original_model(
    next_id,
    past_key_values=o0.past_key_values,
    use_cache=True,
    attention_mask=torch.ones((1, 6), dtype=torch.long),
    position_ids=torch.tensor([[5]], dtype=torch.long),
)

# Always finite (same effective input):
bridge.original_model(torch.cat([tokens, next_id], dim=1))

Diagnostic evidence

From CI run 26312216246:

[DIAG] step0_logits: nan=False inf=False shape=(1, 5, 50257)            ← step 0 fine
[DIAG] cache_type=DynamicCache seq_len=5 layers=6
[DIAG] cache_layer_0: K_nan=False V_nan=False K_shape=(1, 12, 5, 64)    ← cache contents fine
[DIAG] step1_with_mask_and_pos: nan=True inf=False shape=(1, 1, 50257)  ← fails
[DIAG] step1_full_no_cache: nan=False inf=False shape=(1, 6, 50257)     ← same comp w/o cache: fine

So: cache is valid, kwargs are correct, and the full-recompute path works. Only past_key_values + 1-token input NaNs on the runner.

Environment

	Value
Runner	macos-latest (Apple Silicon, GitHub Actions)
Python	3.11.9
PyTorch	2.7.1
transformers	5.8.1
transformer_lens	dev (commit da370b85 and later)
Model	distilgpt2 (GPT2LMHeadModel, 6 layers, 12 heads, d_head=64)
Attention impl	eager (forced by bridge at bridge.py:210-213)
Cache type	DynamicCache (transformers v5 default)

Investigation already done

1. Confirmed the bridge wrapper and raw bridge.original_model(...) produce bitwise-identical step-0 logits.
2. Confirmed bridge forces attn_implementation="eager" at load via sources/transformers.py:586.
3. Confirmed not missing-kwargs: passing explicit attention_mask=(batch, total_len) and position_ids=(batch, 1) (mirroring HF.generate's contract) does not help.
4. Confirmed not cache corruption: K and V cached after step 0 are NaN-free.
5. Confirmed not a Linux-specific path: all Linux CI configs (Python 3.10 / 3.11 / 3.12) pass.
6. Not reproducible on my M4 Mac (Tahoe-build, both Python 3.11 and 3.12). Strongly suggests runner-image specifics, possibly Apple Silicon CPU generation, macOS version, or Accelerate framework version.
7. Related upstream issues that share the SDPA-NaN pattern but aren't this exact bug (we're already on eager):
   • pytorch/pytorch#163997 — SDPA all-NaN rows in eval mode
   • pytorch/pytorch#163597 — SDPA MPS regression on 2.8.0
   • pytorch/pytorch#127055 — MultiheadAttention NaNs with masks
   • huggingface/transformers#30296 — GPT2 NaN with MPS

Current workaround

Two failing tests gated on macOS-arm64 via:

import platform

_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64"

@pytest.mark.skipif(_MACOS_ARM64, reason="Upstream macOS-arm64 KV-cache NaN; see #<this-issue>.")

Linux coverage of the same bridge code path is unaffected.

Action items

• Reproduce on a known-bad runner image (or pin the exact GitHub Actions image and confirm)
• Determine whether the bug is in transformers' eager_attention_forward cached path or in PyTorch CPU matmul on Apple Silicon
• File upstream (transformers and/or pytorch) once narrowed to a minimal repro
• Remove the skipif markers once the upstream bug is fixed

[jlarson4]

We are primarily waiting for the upstream fix (GitHub CI moving to a more modern M-Series chip), but if someone wants to dig into this looking for a better present-solution, you're more than welcome to

[danra]

Couldn't repro locally on M2 Max, Python 3.11/3.12

[jlarson4]

@danra Your M2 will not reproduce this issue, as it currently only appears on M1 Macs. The CI suite runs on macOS-latest, which primarily runs on M1 chips.

[guptaishaan]

I dug into this on an M1 Pro (macOS 15.0, arm64). I didn't get your NaN — I got a hard
SIGBUS at the same step — but it narrowed to a concrete root cause, and I think it explains
the NaN too.

Short version: this isn't a TransformerLens or transformers bug, and it isn't the KV cache.
It's a PyTorch CPU alignment bug, and it fires at the final vocab projection, not in
attention. Filed upstream as pytorch/pytorch#191238.

What happens

Your reproducer crashes the process (exit 138, deterministic 3/3) rather than returning NaN.
faulthandler puts it at torch/nn/modules/linear.py:134, called from modeling_gpt2.py:699
— which is logits = self.lm_head(hidden_states[:, slice_indices, :]).

Dropping the bridge and using bare transformers reproduces it identically, so TransformerLens
isn't involved.

Root cause

Sweeping lm_head over sequence lengths: 5, 4, 3, 2 all fine — 1 crashes. M == 1 takes a
different BLAS path from M > 1.

The trigger is 4-byte alignment. Dependency-free repro, no transformers, no files:

import torch, torch.nn.functional as F
N, K = 50257, 768
buf = bytearray(N*K*4 + 8)
W = torch.frombuffer(memoryview(buf), dtype=torch.float32, count=N*K, offset=1).view(N, K)
F.linear(torch.randn(2, K), W)   # OK
F.linear(torch.randn(1, K), W)   # SIGBUS

data_ptr() % 4 of 1, 2 or 3 crashes at M=1; 0 is fine; M >= 2 is fine at any offset. Same
for matmul, mv, addmm. A 64×32 matrix is unaffected, so it's the large-matrix path.
Confirmed on torch 2.10.0 and 2.12.0.

Why distilgpt2 specifically hits it — its safetensors header length isn't padded:

header_len            : 8277      (unpadded)
data region starts at : 8285      (8285 % 4 = 1)
=> every fp32 tensor in that checkpoint is offset by one byte

safetensors mmaps the file and hands out views, so lm_head.weight lands misaligned. Cloning
the weights into RAM fixes it, which is what put me onto alignment.

One dead end worth flagging in case it saves anyone time: I first thought this was an over-read
past the end of the mmap. It isn't — I swept the trailing page slack from 0 to 3072 bytes and
none of it crashed. My first "misaligned" control was also invalid, because I offset by 13
floats (52 bytes), which is still 4-byte aligned. Offsetting by bytes is what exposes it.

Why I think this is also your NaN

Whether an unaligned load faults or merely returns garbage depends on which instruction the
kernel selects, which varies with torch version and BLAS build. The same misaligned weight that
hard-faults here could plausibly return garbage — and hence NaN — on your torch 2.7.1 /
transformers 5.8.1 combination instead of killing the process.

Two things in your own report point the same way:

• Your failing tensor is shape (1, 1, 50257) — that's the lm_head output, not an
  attention tensor. And cache_layer_0: K_nan=False V_nan=False shows attention was clean. I
  think the "cached eager attention" framing is one step upstream of where it actually breaks.
• Both fire only on the single-token cached step and never on the full recompute — exactly the
  M == 1 vs M > 1 split. Decode runs the vocab projection at seq_len == 1 every step;
  prefill doesn't.

Being clear about the limit: I could not produce a NaN on this machine at these versions
(torch 2.10.0 / 2.12.0, transformers 5.13.0), so the NaN link is inference from a shared
mechanism, not something I demonstrated.

On the skipif markers

I don't think there's a TransformerLens fix worth making here. The only workaround at this
level would be forcing weights into RAM, which doubles memory for every user to route around an
upstream bug. I'd keep both markers until PyTorch resolves #191238 — though it may be worth
updating the skip reason, since "KV-cache NaN" points at the wrong subsystem and would send the
next person looking at attention.

One thing I can't check myself: if anyone here is on an M1 with torch 2.7.1, running that
5-line repro would directly test the NaN-vs-SIGBUS question and confirm whether these are the
same bug.

[jlarson4]

@guptaishaan Thanks for the additional research on this, great information. I am continuing to keep an eye on that PyTorch issue waiting for a resolution