Add Flash Attention 2 support to Musicgen and Musicgen Melody #29939
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. |
sanchit-gandhi
left a comment
sanchit-gandhi
left a comment
There was a problem hiding this comment.
Thanks for adding this!
| return self.audio_encoder.sampling_rate | ||
|
|
||
| @property | ||
| def _attn_implementation(self): |
There was a problem hiding this comment.
This method is one-to-one the same as in the PreTrainedConfig class:
Can we remove it from here?
There was a problem hiding this comment.
Not if we want to keep the setter part!
|
|
||
| MUSICGEN_ATTENTION_CLASSES = { | ||
| "eager": MusicgenAttention, | ||
| "flash_attention_2": MusicgenFlashAttention2, |
There was a problem hiding this comment.
Worth adding sdpa in one go as well? Would enable you to showcase attention implementation through sdpa on free tier Colab T4 GPU (where FA2 is not available)
| return self.audio_encoder.sampling_rate | ||
|
|
||
| @property | ||
| def _attn_implementation(self): |
| else outputs_fa.decoder_hidden_states[-1] | ||
| ) | ||
|
|
||
| assert torch.allclose(logits_fa[1:], logits[1:], atol=4e-2, rtol=4e-2) |
There was a problem hiding this comment.
Good enough for a generative audio model with FA2
There was a problem hiding this comment.
I've copied out the same tolerance threshold than any other models (regardless of modality) btw
|
I've also added SDPA! cc @amyeroberts or @ArthurZucker could you review when you have time? |
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
LGTM! Tests are ... huge, would be nice if you can use copied from, would help the review 😅
| self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" | ||
| self._use_sdpa = config._attn_implementation == "sdpa" |
There was a problem hiding this comment.
let's only save self._attn_implementation please
|
|
||
| return attn_output, None, past_key_value | ||
|
|
||
|
|
There was a problem hiding this comment.
copied from can be used here as well!
ArthurZucker
left a comment
ArthurZucker
left a comment
There was a problem hiding this comment.
Ouf! Thanks for the big PR and adding those tests!
4 participants
What does this PR do?
Supersedes #27924
The attention tests all pass but there are no integration equivalence between the original attention models and the FA ones. I don't hear any difference in quality despite not being the same song, though.
cc @sanchit-gandhi and @amyeroberts, could you review please?