Fix: 'Gemma2Attention' object has no attribute '_flash_attn_uses_top_left_mask'

[jp1924]

What does this PR do?

A PR to fix the 'no attribute '_flash_attn_uses_top_left_mask'' error occurring in each model with the flash_attention module, including gemma2.

Reproduction Code

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer


def main() -> None:
    repo_id = "google/gemma-2-2b-it"
    config = AutoConfig.from_pretrained(repo_id, _attn_implementation="flash_attention_2")
    model = AutoModelForCausalLM.from_pretrained(repo_id, config=config, device_map="cpu")
    tokenizer = AutoTokenizer.from_pretrained(repo_id)

    text = tokenizer.apply_chat_template([{"role": "user", "content": "Hello!"}], tokenize=False)
    input_param = tokenizer(text, return_tensors="pt", return_attention_mask=True)
    input_param["labels"] = input_param["input_ids"].clone()
    output = model(**input_param)


if "__main__" in __name__:
    main()

pip install git+https://github.com/huggingface/transformers.git@5615a393691c81e00251e420c73e4d04c6fe22e5

Env

- `transformers` version: 4.48.0.dev0
- Platform: Linux-5.15.0-124-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.26.2
- Safetensors version: 0.4.5
- Accelerate version: 1.1.1

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@ArthurZucker

[jp1924]

@ArthurZucker
the attention module is being included in the config. Is this working as you intended?

[jp1924]

This pull request is being closed because the issue has been resolved in pull request #35235.

[ArthurZucker]

Thanks ! And sorry for the delay! 🤗