Optionality of attention_mask argument in Attention classes/functions.

[Godofnothing]

In the current stable version of transformers attention_mask argument is annotated as Optional[torch.Tensor] (see for example modeling_llama.py).

However, in fact it is a required argument.

At the same time, the ancestor LlamaDecoderLayer classes accepts this argument as Optional (see).

Delving deeper, flash_attention_forward annotates attention_mask as Optional[torch.Tensor] and calls inside _flash_attention_forward which takes attention_mask as required torch.Tensor argument, but there is conditional statement checking whether the attention_mask is not None and the function can be called in fact with attention_mask as None.

I suggest correcting the typing by making attention_mask an optional argument with None as its default value.

[Zephyr271828]

Hi! @Godofnothing, I just submiited a PR to fix this issue. Feel free to give me suggestions if there's any problem:)