[Bug Report] `TransformerBridge` activations on gpt-oss-20b disagree with the HF model it wraps

[hijohnnylin]

TransformerBridge.boot_transformers("openai/gpt-oss-20b") loads without warnings, but its
hook_resid_post disagrees with the underlying HF forward from layer 0 onwards and becomes
anti-correlated by mid-stack: cosine 0.84 at layer 0, 0.04 at layer 12, and a final-layer residual
whose norm is ~200x the reference. Since the bridge wraps the same HF module it is being compared
against, both sides should be the same tensor up to bf16 noise.

Reproduce

Records HF hidden_states, frees the model, then boots the bridge on the same checkpoint and compares
hook_resid_post for the same token ids. (hidden_states[i + 1] is block i's output, which is what
blocks.{i}.hook_resid_post reports.)

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

HF_ID, LAYERS = "openai/gpt-oss-20b", (0, 12, 23)
ids = torch.tensor([[15496, 11, 616, 1438, 318, 1757, 13, 314, 1101, 257]]).cuda()

hf = AutoModelForCausalLM.from_pretrained(HF_ID, dtype=torch.bfloat16, device_map="cuda")
with torch.no_grad():
    out = hf(ids, output_hidden_states=True)
ref = {layer: out.hidden_states[layer + 1][0].float().cpu() for layer in LAYERS}
del hf, out
torch.cuda.empty_cache()

bridge = TransformerBridge.boot_transformers(HF_ID, device="cuda", dtype=torch.bfloat16)
wanted = {f"blocks.{layer}.hook_resid_post" for layer in LAYERS}
with torch.no_grad():
    _, cache = bridge.run_with_cache(ids, names_filter=lambda n: n in wanted)

for layer in LAYERS:
    a = ref[layer].flatten()
    b = cache[f"blocks.{layer}.hook_resid_post"][0].float().cpu().flatten()
    cos = torch.dot(a, b) / (a.norm() * b.norm())
    print(f"resid_post.{layer}: cos={cos:.4f} rel={(a - b).norm() / a.norm():.4f}")

Actual

resid_post.0:  cos=0.8449 rel=0.5946
resid_post.12: cos=0.0379 rel=1.0023
resid_post.23: cos=0.1892 rel=210.6264

Forward hooks on the blocks' attn / mlp submodules — raw module outputs on both sides, so no
post-norm convention is involved — go negative in the middle of the stack: attention output cosine
−0.60 and MLP output cosine −0.79 at layer 12.

Expected

cos ≈ 1.0 and a relative error at the bf16 noise floor (~1e-3), as the bridge produces on other
checkpoints. A cosine near zero, or negative, is a different quantity rather than a precision
difference.

Root cause

Not diagnosed. Two observations narrow it: the divergence is already 0.59 relative at layer 0, so it
is not accumulation, and resid_post comes from the bridge's own cache rather than from a user hook, so
it is on the bridge's forward path.

Candidates specific to this checkpoint, in the order we would check them:

1. MXFP4 experts. experts.gate_up_proj here is not a torch.Tensor but a
   triton_kernels.tensor.Tensor holding the packed 4-bit payload plus scales. Any code that reads
   those weights, rather than calling the module, has to unwrap or dequantize first. (The legacy
   HookedTransformer conversion path trips over the same object — filed separately.)
2. Attention sinks. gpt-oss adds a learned per-head sink logit to the attention softmax; a wrapper
   that reimplements or reorders the softmax has to carry it.
3. MoE routing. Top-k router weights and expert combination happen inside the block; a bridge that
   re-derives them can select different experts.

Environment

transformer_lens	3.6.0, and main at 69e98ab — identical cosines to four decimals
transformers	5.14.1
torch	2.13.0
python	3.11.10
GPU	NVIDIA A100 80GB PCIe, driver 550.127.05
checkpoint	openai/gpt-oss-20b (GptOssForCausalLM, 24 layers, MXFP4 experts, attention sinks)

[jlarson4]

From my testing, it appears intermediate hidden_states[i + 1] entries are bit-exact to the raw block outputs. The "~200× norm at layer 23" was a measurement artifact specific to the final entry: output_hidden_states[-1] is post-model.norm, so it never matches the last block's resid_post. The norm ratio is residual-norm vs unit-RMS, huge on the trained 20b. The parity tests added in #1621 hook the decoder layers directly instead.