Adding FA2 support for MusicGen #27924
Conversation
|
Hey @staghado, thanks for taking care of this, let us know when it's ready to be reviewed! |
|
I have conducted some tests on an A10 GPU :
|
|
cc @ylacombe could you possibly circle back here when you get the chance! |
ylacombe
left a comment
ylacombe
left a comment
There was a problem hiding this comment.
Hi @staghado, thanks for the update here!
I believe that you have to add _supports_flash_attn_2 to MusicgenForConditionalGeneration and MusicgenForCausalLM, otherwise it won't be flagged as supporting FA2 !
Also, did you make sure to add the flag attn_implementation="flash_attention_2" to from_pretrained ? as indicated here. Let me know.
Besides that, the modeling code looks good except some changes of format that shouldn't happen where you add a coma at the end of the line, all along the modeling file.
You can also add Muscigen to the list of models supported here.
Let me know if you need further help !
staghado
force-pushed
the
musicgen-fa2
branch
from
335aca4 to
7fe8027
Compare
Co-authored-by: Yoach Lacombe <52246514+ylacombe@users.noreply.github.com>
Co-authored-by: Yoach Lacombe <52246514+ylacombe@users.noreply.github.com>
staghado
changed the title
|
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread. Please note that issues that do not follow the contributing guidelines are likely to be ignored. |
ylacombe
left a comment
ylacombe
left a comment
There was a problem hiding this comment.
Hey @staghado , thanks for iterating here ! And sorry for the delay!
There's still some changes that appeared because of added commas, but otherwise looks good to me!
cc @amyeroberts for a review! (btw, there's no need for additional tests right ?)
| for v in [ | ||
| hidden_states, | ||
| next_cache, | ||
| all_hidden_states, | ||
| all_self_attns, | ||
| all_cross_attentions, | ||
| ] |
There was a problem hiding this comment.
Should be a one liner
| for v in [ | |
| hidden_states, | |
| next_cache, | |
| all_hidden_states, | |
| all_self_attns, | |
| all_cross_attentions, | |
| ] | |
| for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_cross_attentions] |
| torch.ones( | ||
| (bsz, num_codebooks, max_length), | ||
| dtype=torch.long, | ||
| device=input_ids.device, | ||
| ) | ||
| * -1 |
| torch.ones((channel_codebooks, max_length), dtype=torch.bool), | ||
| diagonal=max_length - channel_codebooks + 1, |
| input_ids, | ||
| generation_config.pad_token_id, | ||
| generation_config.eos_token_id, |
| self.text_encoder, | ||
| self.decoder._modules[decoder_base_model_prefix], | ||
| self.decoder.base_model_prefix, |
| text_encoder_pretrained_model_name_or_path, | ||
| **kwargs_text_encoder, | ||
| return_unused_kwargs=True, |
There was a problem hiding this comment.
Same here and for all the other changes below
4 participants
What does this PR do?
This PR adds Flash Attention 2 support for MusicGen model. It is based on Bart example and it is a WIP for now.
I could not test the model because FA2 is not supported yet for T4 GPUs.
Fixes #27552
@sanchit-gandhi @ylacombe