[Bug Report] `convert_olmo3_weights` zero-fills K/V weights when `num_key_value_heads == num_attention_heads`

[hijohnnylin]

HookedTransformer.from_pretrained_no_processing("allenai/Olmo-3-1025-7B") loads with 64
Missing key for a weight matrix … filled in with an empty tensor warnings — _W_K and _W_V for each
of the 32 layers. Attention therefore contributes identically zero at every layer, and the model
returns wrong activations and wrong logits with no error.

The trigger is a config that declares MHA through the GQA field (num_key_value_heads == num_attention_heads): convert_olmo3_weights and TransformerBlock use different tests for "is this
GQA?", so the converter writes attn.W_K while the instantiated module expects attn._W_K.

Reproduce

Compares against the HF model the weights came from. The zero-weight and zero-attn_out facts are
prompt-independent; the cosines depend on the prompt.

import torch
from transformers import AutoModelForCausalLM
from transformer_lens import HookedTransformer

HF_ID, LAYERS = "allenai/Olmo-3-1025-7B", (0, 16, 31)
ids = torch.tensor([[100, 200, 300, 400]], device="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()

model = HookedTransformer.from_pretrained_no_processing(HF_ID, device="cuda", dtype=torch.bfloat16)
for name in ("W_Q", "W_K", "W_V"):
    zeroed = sum(1 for b in model.blocks if getattr(b.attn, name).abs().max() == 0)
    print(f"layers with an all-zero {name}: {zeroed}/{len(model.blocks)}")

with torch.no_grad():
    _, cache = model.run_with_cache(ids)
for layer in LAYERS:
    a = ref[layer].flatten()
    b = cache[f"blocks.{layer}.hook_resid_post"][0].float().cpu().flatten()
    print(
        f"layer {layer}: max |attn_out| = {cache[f'blocks.{layer}.hook_attn_out'].abs().max():.1f}, "
        f"resid_post vs HF cos = {torch.dot(a, b) / (a.norm() * b.norm()):.4f}"
    )

Actual

WARNING:root:Missing key for a weight matrix in pretrained, filled in with an empty tensor: blocks.4.attn._W_K
WARNING:root:Missing key for a weight matrix in pretrained, filled in with an empty tensor: blocks.21.attn._W_V
... (64 such warnings: _W_K and _W_V for each of the 32 layers)

layers with an all-zero W_Q: 0/32
layers with an all-zero W_K: 32/32
layers with an all-zero W_V: 32/32
layer 0:  max |attn_out| = 0.0, resid_post vs HF cos = 0.84
layer 16: max |attn_out| = 0.0, resid_post vs HF cos = 0.28
layer 31: max |attn_out| = 0.0, resid_post vs HF cos = 0.30

Per-layer weight norms confirm only K/V are affected — layer 0: |W_Q|=36.25, |W_K|=0, |W_V|=0,
|W_O|=69; layer 31: |W_Q|=69, |W_K|=0, |W_V|=0, |W_O|=84.

Expected

K/V loaded from the checkpoint, and attention output matching the HF forward. Failing that, an
exception: a silently zeroed sublayer is worse than a refusal, because every downstream number still
looks plausible.

Root cause

Two places decide "is this GQA?" with different tests, and they disagree exactly when
num_key_value_heads == num_attention_heads:

• transformer_lens/components/transformer_block.py:77

  attention = Attention if self.cfg.n_key_value_heads is None else GroupedQueryAttention

  n_key_value_heads is set (32), so the block is built as GroupedQueryAttention, whose parameters
  are _W_K / _W_V (components/grouped_query_attention.py).

• transformer_lens/pretrained/weight_conversions/olmo3.py:24

  using_gqa = cfg.n_key_value_heads is not None and cfg.n_key_value_heads < cfg.n_heads
  gqa_uscore = "_" if using_gqa else ""
  ...
  state_dict[f"blocks.{l}.attn.{gqa_uscore}W_K"] = W_K

  32 < 32 is false, so the converter emits blocks.{l}.attn.W_K / W_V.

_W_K / _W_V are then missing from the state dict, fill_missing_keys replaces them with zeros, and
the W_K / W_V entries the converter did write are dropped as unexpected.

Suggested fix

Use the same test as the block, which is what the other converters already do — llama.py:16,
phi3.py:15 and apertus.py:26 are all using_gqa = cfg.n_key_value_heads is not None:

-    using_gqa = cfg.n_key_value_heads is not None and cfg.n_key_value_heads < cfg.n_heads
+    using_gqa = cfg.n_key_value_heads is not None

olmo3.py is the only converter in 3.6.0 carrying the extra < cfg.n_heads clause.

Separately: fill_missing_keys silently zero-filling an attention weight matrix is what turns a name
mismatch into wrong numbers. Raising for W_{Q,K,V,O} (as opposed to, say, a missing bias) would make
any future converter mismatch fail loudly.

Who else is affected

Any checkpoint routed through convert_olmo3_weights whose config sets num_key_value_heads equal to
num_attention_heads. allenai/Olmo-3-1025-7B (32 == 32) hits it; allenai/Olmo-3-1125-32B
(num_key_value_heads=8, 40 heads) is genuinely GQA and loads correctly.

Environment

transformer_lens	3.6.0, and main at 69e98ab — same result on both
transformers	5.14.1
torch	2.13.0
python	3.11.10
GPU	NVIDIA A100 80GB PCIe, driver 550.127.05
checkpoint	allenai/Olmo-3-1025-7B (Olmo3ForCausalLM, 32 layers, 32 heads, num_key_value_heads=32)