Fixes the inconsistency of the optionality of attention_mask #37153
Conversation
|
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the |
Zephyr271828
changed the title
| @@ -280,7 +280,7 @@ def _flash_attention_forward( | |||
| query_states: torch.Tensor, | |||
| key_states: torch.Tensor, | |||
| value_states: torch.Tensor, | |||
| attention_mask: torch.Tensor, | |||
| attention_mask: Optional[torch.Tensor], | |||
There was a problem hiding this comment.
You should then set it to None by default
There was a problem hiding this comment.
Hi! @Godofnothing Yes! I noticed this is different from your proposed solution in issue #37046. My point is that attention_mask will be passed from flash_attention_forward to _flash_attention_forward, though the type can be either torch.Tensor or NoneType, so I think it's not necessary to explicit set the default value as None here.
A similar case is in modeling_llama.py, I noticed LlamaModel uses Option[torch.Tensor] = None for attention_mask where as LlamaAttention simply uses Optional[torch.Tensor], because LlamaModel is responsible to pass the value of attention_mask to LlamaAttention. Technically I could modify the types of attention_mask in all involved functions to Optional[torch.Tensor] = None, but I did not do that for the sake of simplicity.
Please correct me if I'm wrong.
There was a problem hiding this comment.
No, that's fine, thanks for your efforts!
…face#37153) * debugging issue 36758 * debugging issue 36758 * debugging issue 36758 * updated attn_mask type specification in _flash_attention_forward * removed pdb * added a blank line * removed indentation
…face#37153) * debugging issue 36758 * debugging issue 36758 * debugging issue 36758 * updated attn_mask type specification in _flash_attention_forward * removed pdb * added a blank line * removed indentation
…face#37153) * debugging issue 36758 * debugging issue 36758 * debugging issue 36758 * updated attn_mask type specification in _flash_attention_forward * removed pdb * added a blank line * removed indentation
What does this PR do?
This PR fixes Issue #37046.
This issue discusses the inconsistency of the optionality of the parameter "attention_mask" in different functions. Specifically, the type of attention_mask in
LlamaForCausalLM,LlamaModel, andLlamaDecoderLayerareOptional[torch.Tensor] = None. InLlamaAttention(a class which wraps the attention interface), the type isOptional[torch.Tensor]becauseattention_maskwill be passed to it fromLlamaModel, whether the type ofattention_maskistorch.TensororNoneType. Therefore,attention_maskwill be passed to the attention_interface, no matter it's a tensor or NoneType.The key problem lies in the type specification in flash_attention. In
flash_attention_forward, the type ofattention_maskis stillOptional[torch.Tensor]whereas in_flash_attention_forward, a function called inflash_attention_forward, the type istorch.Tensor, which is unreasonable because:flash_attention_forward._flash_attention_forwardis usable even ifattention_maskis None.Therefore, this PR fixes the specification in
_flash_attention_forwardas well as the docstring to address the aforementioned issue.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.