`output_hidden_states=True` returns `hidden_states=None` on all K2-Horizon repos — missing `@capture_outputs` decorator in `modeling_k2_horizon.py`
Hi — we hit this while running an abliteration pipeline (Heretic's get_residuals) against K2-Horizon-0.9B: with output_hidden_states=True, outputs.hidden_states comes back None, which kills any tool that reads the residual stream per layer (TypeError: 'NoneType' object is not iterable). Generation itself is unaffected.
Environment
- transformers 5.17.0, torch 2.11.0+cu128, Python 3.12
IFM/K2-Horizon-0.9B,trust_remote_code=True(also reproduced by editing the modeling file directly)modeling_k2_horizon.pysha256fb09e010956bd51cfa7d4055b4381cff34c9e06164066b49e3546f38b2e6242f(49,535 bytes)
Root cause
The file is written in the current transformers style: K2HorizonPreTrainedModel._can_record_outputs is present and correct, including the hidden_states and attentions entries. But K2HorizonModel.forward carries only @auto_docstring — the @capture_outputs decorator that reads the output_* flags and lazily installs the capture hooks is missing (stock Qwen3MoeModel.forward has it: @merge_with_config_defaults @capture_outputs @auto_docstring). So the flags are never consumed and the hooks never install.
Minimal fix (2 lines)
--- a/modeling_k2_horizon.py
+++ b/modeling_k2_horizon.py
@@
-from transformers.utils.output_capturing import OutputRecorder
+from transformers.utils.output_capturing import OutputRecorder, capture_outputs
@@ class K2HorizonModel(K2HorizonPreTrainedModel):
@auto_docstring
+ @capture_outputs
def forward(
Evidence (0.9B, greedy decode, 20 tokens)
| call | outputs.hidden_states |
|---|---|
as shipped, output_hidden_states=True |
None |
| as shipped, flag off | None |
decorator added, output_hidden_states=True |
tuple of 29 (28 layers + 1); [0] == embeddings; [-1] == last_hidden_state (standard tie_last_hidden_states semantics) |
| decorator added, flag off | None (no cost) |
Greedy generation is byte-identical with and without the decorator. Verified both by wrapping K2HorizonModel.forward with transformers.utils.output_capturing.capture_outputs at runtime and by editing the file — same result.
Repro:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("IFM/K2-Horizon-0.9B", trust_remote_code=True)
inputs = tok("The capital of France is", return_tensors="pt").to("cuda")
model = AutoModelForCausalLM.from_pretrained(
"IFM/K2-Horizon-0.9B", trust_remote_code=True, dtype=torch.bfloat16
).to("cuda").eval()
out = model(**inputs, output_hidden_states=True)
print(out.hidden_states) # -> None
Scope
modeling_k2_horizon.py is byte-identical (same sha256) and missing the decorator on every full-model repo in the family: 0.9B, 3.7B, 7B, 7B-FP8, 32B, 32B-FP8, 32B-NVFP4, 375B-A23B, 375B-A23B-FP8, 375B-A23B-NVFP4, MoVA-36B-A4B, MoVA-36B-A4B-FP8. Verified empirically on the 0.9B only; the others ship the identical file. The -Uno repos are adapter-only (no custom code) and the GGUF repos load via llama.cpp, so neither is affected.
Optional, separate: dead attentions registry entry on MoVA checkpoints
"attentions": K2HorizonAttention never matches on the 0.9B, whose layers instantiate K2HorizonMoVAAttention — output_attentions=True currently returns an empty tuple rather than per-layer weights. Listing both classes makes it work (verified on the 0.9B with attn_implementation="eager": 28 tensors of shape (1, 32, seq, seq)):
_can_record_outputs = {
"router_logits": OutputRecorder(K2HorizonSparseMoeBlock, index=1),
"hidden_states": K2HorizonDecoderLayer,
- "attentions": K2HorizonAttention,
+ "attentions": [K2HorizonAttention, K2HorizonMoVAAttention],
}
Only produced under eager attention, as with stock models — fine to treat as a separate commit or skip.
Happy to test any of this on other sizes, or send this as a PR if that's easier for you. Thanks!