Repository navigation
Fix warning for output_attentions=True - #40597
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
qubvel
changed the title
|
cc @ydshieh to merge, CI failures are unrelated |
|
|
||
| # Check attention implementation is properly set for capturing attention outputs | ||
| if recordable_keys.get("output_attentions", False): | ||
| supported_attn = ["eager", "eager_paged", "flex_attention"] |
There was a problem hiding this comment.
Flex attention returns the lse (like flash attention) - not really the same as what you would get with eager
There was a problem hiding this comment.
yeah, I noticed that, but that would be captured as attentions and returned in ModelOutput, if I'm no missing anything. Do you have an idea of a better warning?
There was a problem hiding this comment.
Don't have an immediate solution, but 2 suggestions:
- Properly change what
output_attentionsdoes in the docs and enable outputting the lse (would include flash attn then for example) - Another separate kwarg/attr
output_lsewhich works for flash, flex etc and not on eager for example.
Currently, we have a weird situation with flex 😅 doesnt need immediate change but maybe put it up for discussion internally?
What does this PR do?
With the
@check_model_inputsdecorator, we no longer check and fallback to theeagerattention in the Attention module's code, let's notify users to set proper attention implementation. Otherwise, we are silently gettingNoneattentions even if we specifyoutput_attentions=True(because the default attention issdpa)