Repository navigation
Call to contributions: refactor output tracing in transformers #43979
Description
Following #43590 that updates 112 models, we want to finish migrating all models to the standardized output collection interface. So, this is a call to contributors to open PRs to link to this meta-issue and learn about the codebase!
Two decorators replace the old manual boilerplate:
@capture_outputs: goes on the base model forward. Automatically collectshidden_states/attentionsvia hooks on submodules declared in_can_record_outputs. Removes the need for explicitoutput_attentions/output_hidden_states/return_dicthandling.@can_return_tuple: goes on wrapper forwards (ForCausalLM, etc.). Handlesreturn_dictonly without hook installation.
How to refactor a model
- Add
_can_record_outputs = {"hidden_states": DecoderLayer, "attentions": Attention}on thePreTrainedModelsubclass. Use class references, not strings - Add
@capture_outputson the base model forward (the one with the layer loop) - Add
@can_return_tupleon higher-level forwards - Drop
output_attentions,output_hidden_states,return_dictfrom signatures - Remove parameter resolution lines and manual collection loops
- Decoder layers return
hidden_statesdirectly; attention always returns(attn_output, attn_weights)
You can use llama, mistral, or qwen2 as reference implementations. Modern agents are of some help in this issue, but correctly placing the decorators can be tricky.
Models to update
Please claim one or several models by commenting before starting. One PR per model (or small group of related models) could be ideal, for instance!
As of today here's the remaining list of models:
-
bark -
beit -
bit -
blip_2 -
bloom -
canine -
codegen -
colqwen2 -
cpmant -
ctrl -
cvt -
dab_detr -
dac -
data2vec -
deberta -
deberta_v2 -
decision_transformer -
depth_anything -
depth_pro -
dinat -
donut -
dpr -
efficientnet -
encodec -
falcon -
falcon_mamba -
fastspeech2_conformer -
flaubert -
flava -
fnet -
focalnet -
fsmt -
funnel -
glpn -
gpt2 -
gpt_neo -
gpt_neox_japanese -
gptj -
granite_speech -
grounding_dino -
hgnet_v2 -
hiera -
hubert -
ibert -
imagegpt -
kyutai_speech_to_text -
led -
levit -
lilt -
llama4 -
longformer -
longt5 -
luke -
lxmert -
mamba -
mamba2 -
mask2former -
maskformer -
megatron_bert -
mgp_str -
mimi -
mm_grounding_dino -
mobilenet_v1 -
mobilenet_v2 -
mobilevit -
mobilevitv2 -
moshi -
mpnet -
mpt -
mra -
mt5 -
mvp -
nystromformer -
omdet_turbo -
oneformer -
openai -
patchtsmixer -
patchtst -
perceiver -
pix2struct -
poolformer -
pop2piano -
prophetnet -
pvt -
pvt_v2 -
rag -
recurrent_gemma -
reformer -
regnet -
rembert -
resnet -
roformer -
rt_detr -
rwkv -
seamless_m4t -
seamless_m4t_v2 -
segformer -
seggpt -
sew -
sew_d -
speech_encoder_decoder -
speecht5 -
squeezebert -
superglue -
superpoint -
swiftformer -
swin -
swin2sr -
swinv2 -
t5 -
table_transformer -
tapas -
textnet -
timesformer -
timm_backbone -
trocr -
tvp -
udop -
umt5 -
unispeech -
unispeech_sat -
univnet -
upernet -
vilt -
vision_encoder_decoder -
vision_text_dual_encoder -
visual_bert -
vitdet -
vitmatte -
vits -
wav2vec2 -
wav2vec2_bert -
wav2vec2_conformer -
wavlm -
xcodec -
xlm -
xlnet -
yoso -
zamba -
zoedepth
Reacted by Lysandre Debut, Raushan Turganbay, Omkar Kabde, Amin Mahjoub, Marc Sun, Huy X. Dang, Paarth Sharma and Daniel Tiourine
Activity
Aki-07 commentedon Feb 13, 2026
I woud like to claim gpt2
11happy commentedon Feb 13, 2026
I will work on llama4.
Thank you : )
akeemlh commentedon Feb 13, 2026
I'll work on falcon
beelapranay commentedon Feb 14, 2026
i would like to work on fnet and cvt
Aatman09 commentedon Feb 14, 2026
I would like to work on vits
nexiouscaliver commentedon Feb 14, 2026
I'd like to work on segformer.
nexiouscaliver commentedon Feb 14, 2026
I'll work on segformer. Refactoring to use the new @capture_outputs and @can_return_tuple decorators for standardized output collection.
preetam1407 commentedon Feb 14, 2026
I'd like to pick squeezebert, T5, speecht5
feat: import output capturing decorators for segformer
feat: add _can_record_outputs property to SegformerPreTrainedModel
refactor: use @capture_outputs and @can_return_tuple decorators
feat: add @capture_outputs and @can_return_tuple to SegformerForImage…
refactor: add decorators to SegformerForImageClassification.forward
83 remaining items
RohanPatil2 commentedon Feb 17, 2026
Hi @molbap, I would like to claim the rag and grounding_dino model. I see no active PRs for this specific model in the comments. I will submit a PR shortly. Thanks!
23atharvaS commentedon Feb 18, 2026
Completed and opened PR #44114 for this task (
wav2vec2/wav2vec2_conformer/wav2vec2_bertoutput-collection decorator migration); CI checks are passing.TaherTadpatri commentedon Feb 18, 2026
Hi @molbap I will claim visual_bert.
aman-coder03 commentedon Feb 18, 2026
@TaherTadpatri i have already opened the visual_bert PR, kindly check before claiming
cogniera commentedon Feb 18, 2026
hi @ArthurZucker I was working on the longt5 model , could you please reopen the issue so I can contribute
Aatman09 commentedon Feb 19, 2026
@ArthurZucker could you open the issue , i was working on vit and was going to raise pr today
mmahjoub5 commentedon Feb 26, 2026
Are these PRs going to be reviewed? @molbap
kulkarni-rohan commentedon Mar 14, 2026
I would like to claim colqwen2 model.
[ColQwen2] Refactor output tracing (issue huggingface#43979)
raunakb-stack commentedon Mar 15, 2026
Hi! I'd like to work on refactoring output tracing for the albert model.
Is this model still available to claim?
chandan11248 commentedon Mar 15, 2026
Hi, I'd like to claim
gptjfor this refactor. Will open a PR once ready.Refactor gptj output tracing to use @capture_outputs and @can_return_…
codybarrett108-pixel commentedon Mar 17, 2026
AbdulAhad1189 commentedon Oct 29, 2026
Hi! I'd like to work on refactoring output tracing for gpt_neo using @capture_outputs and @can_return_tuple. I couldn't find an active claim/PR for it. I'd like to take this one. Thanks!