🔴🔴🔴 [Attention] Refactor Attention Interface for Bart-based Models
#38108
Conversation
|
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the |
|
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. |
Attention] Refactor Attention Interface and Enable Flex Attention Attention] Refactor Attention Interface for Bart-based Modelsand Enable Flex Attention
Attention] Refactor Attention Interface for Bart-based Modelsand Enable Flex Attention Attention] Refactor Attention Interface for Bart-based Models and Enable Flex Attention
| attn_mask=causal_mask, | ||
| dropout_p=self.dropout if self.training else 0.0, | ||
| is_causal=is_causal, | ||
| attn_output, attn_weights = ALL_ATTENTION_FUNCTIONS( |
There was a problem hiding this comment.
no here we want one explicit thing, that the default attention is eager_attention_forward!
The core philosophy when we abstract is we keep the inderections to a minimal, and thus the core logic (eager attention) should be explicit!
But very nice to put the rest (warning and etc) inside each function, but not in ALL_ATTENTION_FUNCTIONS's call!
Thus each sdpa, flex or flash have their own warning, we should not abstract!
| attention_mask=attention_mask, | ||
| training=self.training, | ||
| dropout=self.dropout, | ||
| attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] |
There was a problem hiding this comment.
| attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] | |
| attention_interface = eager_attn_forward | |
| if self.config._attn_implementation != "eager": | |
| attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] |
let's make it explicit what is the default!
There was a problem hiding this comment.
You're too fast :D will change in a sec, wouldnt work either way with attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] since eager is not registered in the interface
Attention] Refactor Attention Interface for Bart-based Models and Enable Flex Attention Attention] Refactor Attention Interface for Bart-based Models and Enable Flex Attention
tests/test_modeling_common.py
Outdated
| config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() | ||
| config._attn_implementation = "flex_attention" | ||
| model = model_class(config).to(device=torch_device, dtype=torch.float16) | ||
| # Flex Attention can not use dropout | ||
| if hasattr(config, "attention_droput"): |
There was a problem hiding this comment.
Need to fix this, typo (discovered during whisper)
There was a problem hiding this comment.
Rerunning flex tests to see which models will fail.
There was a problem hiding this comment.
Seems like a lot of models fail on flex attention... Not sure if I should submit a torch issue. Disabled flex on them for now - I don't think it's a high priority atm.
| "Falling back to eager attention because `flash_attention_2` does not support" | ||
| " `output_attentions=True` or `head_mask`." | ||
| ) | ||
| return eager_fallback( |
There was a problem hiding this comment.
if eager fallback is none we are failing hard! Let's maybe just emit the warning no? And return None output attentions
| "Falling back to eager attention because `flex_attention` does not support" | ||
| " `output_attentions=True`, `head_mask`, or `dropout`." | ||
| ) | ||
| return eager_fallback( |
There was a problem hiding this comment.
same here! And head mask is something we kinda deprecated so let's just return None prob
Attention] Refactor Attention Interface for Bart-based Models and Enable Flex Attention Attention] Refactor Attention Interface for Bart-based Models
|
Another day, another merge conflict |
|
Checked offline again with Arthur, merging |
…uggingface#38108) * starting attn refactor for encoder decoder models via bart (eager + sdpa) * flash attention works, remove unnecessary code * flex attention support for bart!, gotta check if the renaming is not too aggressive * some comments * skip flex grad test for standalone as done with the other test * revert flex attn rename (for now), sdpa simplify, and todos * more todos * refactor mask creation for reuse * modular attempt at biogpt * first batch of other models * fix attn dropout * fix autoformer copies * hubert * another batch of models * copies/style + last round of bart models --> whisper next? * remove unnecessary _reshape function and remove copy to whisper * add skip for decoder-only models out of enc-dec (same as in bart) * bring back licences * remove comment, added to pr read instead * mostly docs * disable sew flex attn as it's unclear attn mask for now * oops * test fixes for enc-dec * torch fx fixes + try at flex attn * skip on mbart * some more fixes * musicgen skip / delete old attn class logic + sdpa compose compile skip * disable flex attn for musicgen, not worth the effort * more fixes and style * flex attention test for dropout and encoder decoder that dont have main input names * informer fixes * the weirdest thing I've encountered yet... * style * remove empty tensor attempt, found core root in previous commits * disable time series due to tests being very text centric on inputs * add speech to text to be ignoring the other attns, also due to tests * update docs * remaining issues resolved ? * update docs for current state --> nllb moe and pegasus x sdpa is questionable :D * some models have not set the is_causal flag... * change dtype in softmax tol old behaviour + some modular fixes * I hate it but it is what it is * fixes from main for bart * forgot this one * some model fixes * style * current status * marian works now * fixing some copies * some copy fixes + time series x informer * last models possibly and fixes on style/copies * some post merge fixes * more fixes * make attention interface callable and move warnings there * style lol * add comment to "unsupported" * remove callable interface and change interface warnings + some copies * fix * ternary is ugly af, make it simpler * how did that happen * fix flex attn test * failing the test * no more fallback! fixing copies next * style + attn fixed * fixing copies and mask creation * wrong copy * fixup tests and disable flex attn for now * fixup last tests?
This PR is gonna tackle two things in general:
Affected models (will be updated when I have enough time) - probably not 100% accurate
Possibly doable in this PR:
Worth a discussion (?):
Worth a discussion TL;DR:
Future PRs will address: