Skip to content

Ignore non-causal mask in more cases with SDPA#30138

Merged
fxmarty merged 7 commits intohuggingface:mainfrom
fxmarty:sdpa-non-causal-mask-fix
Jun 3, 2024
Merged

Ignore non-causal mask in more cases with SDPA#30138
fxmarty merged 7 commits intohuggingface:mainfrom
fxmarty:sdpa-non-causal-mask-fix

Conversation

@fxmarty
Copy link
Contributor

@fxmarty fxmarty commented Apr 9, 2024

Fixes #30095. In the non-causal case, we can ignore the mask even for key_value_length != tgt_len as we will never have any fully masked row and will never hit the SDPA's mem-efficient attention backend issue.

Depends on #28802

@fxmarty fxmarty requested a review from ArthurZucker April 9, 2024 08:28
@HuggingFaceDocBuilderDev

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.

The target length or query length the created mask shall have.
"""
batch_size, key_value_length = mask.shape
_, key_value_length = mask.shape
Copy link
Contributor

@minostauros minostauros Apr 9, 2024

Choose a reason for hiding this comment

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

Can we change the input arg mask: torch.Tensor to mask: Optional[torch.Tensor] and return None immediately if mask is None? The docstring is not compliant with the actual input. (mask ("torch.Tensor" or "None" ):)
Will it break the is_tracing check?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@minostauros Yes indeed ideally we would want to do that. In practice, the calls to these functions in modeling files are always guarded by:

if attention_mask is not None:

but we should IMO indeed accept Optional[torch.Tensor]. I'll leave that to an other PR.

@fxmarty fxmarty requested review from ArthurZucker, LysandreJik and amyeroberts and removed request for ArthurZucker April 15, 2024 08:10
Copy link
Collaborator

@ArthurZucker ArthurZucker left a comment

Choose a reason for hiding this comment

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

Good cleanup IMO. Let's make sure the workflow is triggered for this

@github-actions
Copy link
Contributor

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.

@ArthurZucker
Copy link
Collaborator

This is probably fixed on main

Comment on lines +435 to +437
is_causal = (
True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False
)
Copy link
Contributor Author

Choose a reason for hiding this comment

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

FYI @hackyon

@fxmarty fxmarty merged commit 221aaec into huggingface:main Jun 3, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

_prepare_4d_attention_mask_for_sdpa is not for causal attention but claims...

4 participants