Remove many output_attentions and other traced outputs on 100+ models #43590

vasqu · 2026-03-05T20:18:49Z

This is the new wrapper to make backbones behave like expected: They always output hidden states so we control this here.

output = type(output)(**filtered_output_data) is a bit weird but it allows us to construct our modeling outputs properly as there is no delete function and I don't think we want one

vasqu · 2026-03-05T20:14:52Z

This is needed as encoder oftentimes share the same attention module from the decoder, meaning that if we pass the cache there everything gets messy

Humm, I don't get why we need this suddenly? Even encoder may want the cache with EncoderDecoder cache no?

But the encoder module itself never wants a cache, it's only forwarded once and then "saved" for each subsequent step. Only the decoder needs the cache to properly overwrite states.

If the encoder also gets the cache, then it can update the cache as well which makes generate being broken for certain methods (not sure which anymore but CI was broken for a few tests on bart then).

Ha yes, cause you made sure to propagate kwargs now as well! Makes sense then!

-Original file line number
+Diff line change
@@ Expand Up / @@ -15,6 +15,7 @@ @@
     """Collection of utils to be used by backbones and their components."""
     import enum
+    import functools
     import inspect
     from huggingface_hub import repo_exists
@@ Expand Down Expand Up / @@ -156,6 +157,27 @@ def to_dict(self): @@
             return output
+    def filter_output_hidden_states(forward_function):
+        """
+        Wrapper to filer out `hidden_states` as backbones tend to always use them to get their feature maps, i.e.
+        they also always output `hidden_states`. This controls for user-defined behavior again.
+        NOTE: We assume a `can_return_tuple` decorator to be applied before so that we always expect a dict like
+              object to remove the hidden states.
+        """
+        @functools.wraps(forward_function)
+        def wrapper(self, *args, **kwargs):
+            output_hidden_states = kwargs.get("output_hidden_states", getattr(self.config, "output_hidden_states", False))
+            output = forward_function(self, *args, **kwargs)
+            if not output_hidden_states:
+                filtered_output_data = {k: v for k, v in output.items() if k not in ("hidden_states")}
+                output = type(output)(**filtered_output_data)
+            return output
+        return wrapper
     class BackboneMixin:
         backbone_type: BackboneType | None = None
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -68,6 +68,8 @@ @@
     def _build_checkpoint_conversion_mapping():
         mapping = {
+            "dinov3_convnext": [WeightRenaming(r"(?<!model\.)stages", r"model.stages")],
+            "dinov3_vit": [WeightRenaming(r"(?<!model\.)layer.", r"model.layer.")],
             "timesfm2_5": [
                 WeightRenaming("ff0", "fc1"),
                 WeightRenaming("ff1", "fc2"),
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -765,7 +765,7 @@ def _prepare_encoder_decoder_kwargs_for_generation( @@
                     add_hook_to_module(encoder, AlignDevicesHook(io_same_device=True))
             # 2. Prepare encoder args and encoder kwargs from model kwargs and generation config.
-            irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+            irrelevant_prefix = ["decoder_", "cross_attn", "use_cache", "past_key_values", "cache_params"]
             encoder_kwargs = {
                 argument: value
                 for argument, value in model_kwargs.items()
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -73,10 +73,7 @@ class Aimv2Output(ModelOutput): @@
         vision_model_output: BaseModelOutputWithPooling = None
         def to_tuple(self) -> tuple[Any]:
-            return tuple(
-                self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
-                for k in self.keys()
-            )
+            return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())
     @use_kernel_forward_from_hub("RMSNorm")
@@ Expand Down @@

Navigation Menu

Uh oh!

Remove many output_attentions and other traced outputs on 100+ models #43590

Uh oh!

Diff view

Diff view

There are no files selected for viewing

vasqu Mar 5, 2026

Uh oh!

vasqu Mar 5, 2026

Uh oh!

Cyrilvallez Mar 6, 2026

Uh oh!

vasqu Mar 6, 2026

Uh oh!

Cyrilvallez Mar 6, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Remove many output_attentions and other traced outputs on 100+ models #43590

Uh oh!

Remove many output_attentions and other traced outputs on 100+ models #43590

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

vasqu Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!