[ArthurZucker]

What does this PR do?

Output attention and hidden states logic can be reduced to basically attaching hooks to the GradientCheckpointing layer.

Refactored can_return_tuple as it can be much simpler :) only top level will receive False

This makes the code simpler and cleaner!

[ArthurZucker]

Just want to trigger the tests!

[qubvel]

Hey @ArthurZucker, great work 🚀 🚀 🚀

I also started working on removing return_dict from the models in this PR

• Remove return_dict kwarg from all the models #38991

while Lysandre shared your PR, so I'm curious if you going to include all the changes in this PR? I think it might be too huge to apply all at once, WDYT?

[qubvel]

I was also previously looking into removing output_attention and output_hidden_states and had similar idea of utilizing hooks. For intermediate hidden_states, however, it might be harder to apply because it not that consistent across the repo. For example, in some cases we collect the last hidden_state after the norm layer, but there might be no final layer_norm (as in bart or vit) or an additional layer(s) (as in opt). It might be tricky to manage all the cases, and it feels like we should strongly maintain backward compatibility for output_hidden_states because some models are built on these intermediate features. I'm just curious if you have any ideas in mind. Maybe we can have some patterns for the layer to apply hooks to or keep the intermediate hidden_states collection as is and remove just output_attentions. Anyway, I would be happy to help 🤗

[ArthurZucker]

Hey @qubvel the hooks are gonna be attached to the layer class, and we will register explicitly which layers are expected to return what! This way we can deal with any cases!

[ArthurZucker]

It is easier for me to change return dict as well TBH!

[ArthurZucker]

#39120 did it!