Fix gpt oss olmo3 parity

[jlarson4]

Description

Fixes the three parity/loading bugs reported for gpt-oss-20b and Olmo-3, plus three deeper correctness bugs that verification against the real checkpoints flushed out along the way. Every fix is gated by an HF-parity test, not just "it loads."

Fixes #1618, #1619, and #1620

#1618

Two independent bugs in the bridge's reconstructed attention (PositionEmbeddingsAttentionBridge):

1. Missing attention sinks. GPT-OSS's eager attention concatenates a learned per-head sink logit as an extra softmax column and drops it after normalization, rescaling every pattern row. The bridge's softmax omitted it entirely. On the real checkpoint this is not a small correction: some trained heads route ~99.9998% of their probability mass to the sink (min pattern row-sum 2e-6), so an implementation without sinks computes a different function, which is why mid-stack cosines went negative.
2. Half-width RoPE misread as partial rotary. GPT-OSS applies full rotation using un-duplicated half-width cos/sin (rotary_dim == head_dim/2). The bridge's partial-rotary branch rotated only the first half of each head and passed the rest through. Since gpt-oss's rotation is exactly llama's rotate-half once the halves are duplicated, the fix reuses that path, gated on rotary_dim * 2 == head_dim and rotary_pct == 1.0 so genuine partial-rotary models (Pythia, Phi) are untouched.

#1619

• Loading by model name now detects quant_method == "mxfp4" in the checkpoint config and passes Mxfp4Config(dequantize=True) to from_pretrained, so the converter always sees plain tensors. A caller-supplied quantization_config takes precedence, and unquantized gpt-oss finetunes are untouched.
• If a pre-loaded quantized model is passed as hf_model, the converter now raises a NotImplementedError that names MXFP4, identifies the triton-kernels wrapper (whose Tensor class name made the original error look like a torch bug), and says how to load dequantized.

Getting past the crash exposed that HookedTransformer's gpt-oss forward had never been numerically correct (logits vs HF: cos 0.85, top-1 agreement 0.00). Three missing pieces, all fixed:

• Attention sinks – new use_attention_sinks config flag; AbstractAttention carries the per-head sinks parameter through the softmax (same semantics as the bridge fix), and the converter maps the weights.
• Sliding-window layers – the config mapping had use_local_attn: False; HT's existing local attention has exactly HF's sliding-window semantics, so attn_types/window_size are now wired from layer_types/sliding_window.
• YaRN rope – the mapping only read rope_theta; yarn parameters are now wired, including a new yarn_truncate flag for gpt-oss's truncate: False variant and HF's default attention factor.

#1620

• The converter's GQA predicate now matches TransformerBlock's (n_key_value_heads is not None), so it emits _W_K/_W_V exactly when the block instantiates GroupedQueryAttention. The previous test that encoded the buggy contract is inverted.
• As suggested in the issue, fill_missing_keys now raises when a missing key is an attention weight matrix (W_Q/W_K/W_V/W_O, _W_K/_W_V) instead of silently zero-filling. Any future converter/component naming mismatch fails loudly. Non-attention weights keep the warn-and-fill behavior.

Verification then caught a further pre-existing bug: Olmo-3 declares per-layer-type rope in transformers 5.x, but the mapping read the defunct rope_scaling attribute, so yarn never engaged. Separately, attn_types without use_local_attn is inert, which left sliding masks unwired. Fixed via a new yarn_global_attn_only flag and a per-layer-type config mapping.

Verification on real checkpoints (RTX PRO 6000, bf16 for gpt-oss / fp32 for Olmo-3)

Test	Issue's broken state	This PR
Bridge, gpt-oss-20b resid_post (L0/12/23)	cos 0.84 / 0.04 / 0.19	cos 1.0000 / 1.0000 / 1.0000
Bridge, gpt-oss-20b layer-0 pattern vs HF output_attentions	—	max diff at bf16 noise
HT, gpt-oss-20b	TypeError crash; post-crash-fix logits top-1 0.00	loads; cos 0.9998, top-1 1.00
HT, Olmo-3-1025-7B	64 missing-key warnings, 32/32 K/V zeroed, cos 0.28–0.84	0 warnings, 0/32 zeroed, cos 1.0000, rel ≤ 1e-6

CI-runnable regression tests use tiny random checkpoints through the production load paths (sinks, sliding windows, and yarn all exist unquantized): bridge-vs-HF parity via boot_transformers, HT-vs-HF parity via convert_hf_model_config. Both at fp32 machine precision (≤ 3e-7).

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)

The only intentionally breaking behavior is narrow: state dicts that are missing attention weight matrices now raise in fill_missing_keys instead of loading with silently zeroed attention. Any code relying on the old behavior was computing wrong numbers.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

@hijohnnylin Thanks for bringing these GPT OSS & OLMo 3 errors to my attention. They should be resolved as of this PR, once CI passes I will bundle this into a release and get that out. Let me know if you have any questions or concerns