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. |
| @@ -239,3 +239,20 @@ def from_sub_models_config( | |||
| # This is a property because you might want to change the codec model on the fly | |||
| def sampling_rate(self): | |||
| 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)
| @@ -254,3 +252,20 @@ def from_sub_models_config( | |||
| # This is a property because you might want to change the codec model on the fly | |||
| def sampling_rate(self): | |||
| 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? |
| 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!
* add FA2 to o.g Musicgen * make style * add FA2 support to Musicgen Melody * add generation FA2 tests to o.g Musicgen * make style and fix copies * add Musicgen to FA2 docs + deprecate list * add sdpa supports to Musicgen's * make style and fix copies * refactor attention implementation arguments * add Copied from to sdpa tests * add copied form in sdpa tests melody * add copied for FA2 generation tests * add FA2 inference copied from * make style
* add FA2 to o.g Musicgen * make style * add FA2 support to Musicgen Melody * add generation FA2 tests to o.g Musicgen * make style and fix copies * add Musicgen to FA2 docs + deprecate list * add sdpa supports to Musicgen's * make style and fix copies * refactor attention implementation arguments * add Copied from to sdpa tests * add copied form in sdpa tests melody * add copied for FA2 generation tests * add FA2 inference copied from * make style
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?