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 collects hidden_states/attentions via hooks on submodules declared in _can_record_outputs. Removes the need for explicit output_attentions/output_hidden_states/return_dict handling.@can_return_tuple: goes on wrapper forwards (ForCausalLM, etc.). Handles return_dict only without hook installation.
How to refactor a model
- Add
_can_record_outputs = {"hidden_states": DecoderLayer, "attentions": Attention} on the PreTrainedModel subclass. Use class references, not strings
- Add
@capture_outputs on the base model forward (the one with the layer loop)
- Add
@can_return_tuple on higher-level forwards
- Drop
output_attentions, output_hidden_states, return_dict from signatures
- Remove parameter resolution lines and manual collection loops
- Decoder layers return
hidden_states directly; 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:
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
_can_record_outputs = {"hidden_states": DecoderLayer, "attentions": Attention}on thePreTrainedModelsubclass. Use class references, not strings@capture_outputson the base model forward (the one with the layer loop)@can_return_tupleon higher-level forwardsoutput_attentions,output_hidden_states,return_dictfrom signatureshidden_statesdirectly; attention always returns(attn_output, attn_weights)You can use
llama,mistral, orqwen2as 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:
barkbeitbitblip_2bloomcaninecodegencolqwen2cpmantctrlcvtdab_detrdacdata2vecdebertadeberta_v2decision_transformerdepth_anythingdepth_prodinatdonutdprefficientnetencodecfalconfalcon_mambafastspeech2_conformerflaubertflavafnetfocalnetfsmtfunnelglpngpt2gpt_neogpt_neox_japanesegptjgranite_speechgrounding_dinohgnet_v2hierahubertibertimagegptkyutai_speech_to_textledlevitliltllama4longformerlongt5lukelxmertmambamamba2mask2formermaskformermegatron_bertmgp_strmimimm_grounding_dinomobilenet_v1mobilenet_v2mobilevitmobilevitv2moshimpnetmptmramt5mvpnystromformeromdet_turbooneformeropenaipatchtsmixerpatchtstperceiverpix2structpoolformerpop2pianoprophetnetpvtpvt_v2ragrecurrent_gemmareformerregnetrembertresnetroformerrt_detrrwkvseamless_m4tseamless_m4t_v2segformerseggptsewsew_dspeech_encoder_decoderspeecht5squeezebertsupergluesuperpointswiftformerswinswin2srswinv2t5table_transformertapastextnettimesformertimm_backbonetrocrtvpudopumt5unispeechunispeech_satunivnetupernetviltvision_encoder_decodervision_text_dual_encodervisual_bertvitdetvitmattevitswav2vec2wav2vec2_bertwav2vec2_conformerwavlmxcodecxlmxlnetyosozambazoedepth