Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix MusicGen SDPA #31208

Merged
merged 4 commits into from
Jun 14, 2024
Merged

Fix MusicGen SDPA #31208

Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c84cab4
fix sdpa musicgen
ylacombe Jun 3, 2024
cc3bbd6
make style
ylacombe Jun 3, 2024
e559f25
Merge branch 'huggingface:main' into fix-musicgen-sdpa
ylacombe Jun 14, 2024
605df12
remove copied from statement from Musicgen SDPA
ylacombe Jun 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/transformers/models/musicgen/modeling_musicgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,6 +572,24 @@ def forward(
output_attentions=output_attentions,
)

# Ignore copy
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This ignore copy doesn't work, any ideas why @ydshieh ?

Copy link
Collaborator

@ydshieh ydshieh Jun 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi. Copy and ignore copy only works with a named entity (like function, method, class). It could not be used in a place (block) where no declared name is given. Either refactor the code or just remove the copy statement

Copied from transformers.models.bart.modeling_bart.BartSdpaAttention with Bart->Musicgen

if no code re-factorization is possible.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ydshieh, I'll remove the copy statement

if (
attention_mask is not None
and (attention_mask.mean(dim=[1, 2, 3]) <= torch.finfo(attention_mask.dtype).min).any()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the reason for using finfo here and not just 0?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The attention has already been processed and is filled with -inf of the correspondant dtype!

):
logger.warning_once(
'`torch.nn.functional.scaled_dot_product_attention` does not support having an empty attention mask. Falling back to the manual attention implementation. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
"Note that this probably happens because `guidance_scale>1` or because you used `get_unconditional_inputs`. See https://github.com/huggingface/transformers/issues/31189 for more information."
)
return super().forward(
hidden_states,
key_value_states=key_value_states,
past_key_value=past_key_value,
attention_mask=attention_mask,
layer_head_mask=layer_head_mask,
output_attentions=output_attentions,
)
Comment on lines +582 to +589
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than fallback, I would just raise an exception. Otherwise this expensive check and forward pass can easily go unnoticed

Copy link
Contributor Author

@ylacombe ylacombe Jun 4, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Well, sdpa and guidance_scale>1 are used by default, so we'd raise the error almost every time. Also, even if the model uses eager mode for the cross-attention layers (in which the bug happens), it'll still benefit from the speed-up of the self-attention layers.

Should we find a better way of testing the attention mask ? For example, we could raise a warning that it will happens here and here and switch the cross-attention SDPA layers to eager layers by default when it happens?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, I see. I'm a bit concerned about this causing unexpected behaviour, in particular defaulting to eager like this as it's a bit magic. As there's other layers which can still use SDPA, this seems like a pragmatic solution.

Let's leave as-is. If more users raise issues, then we'll have to re-think.


# if key_value_states are provided this layer is used as a cross-attention layer
# for the decoder
is_cross_attention = key_value_states is not None
Expand Down
Loading