issue with training with flash attention on

[ashunaveed]

Describe the bug

while training a llama 3.1 model with flash attention on and text file, rank:1024, microbatch size:1 or 2 or 4, got following error:
/Documents/text-gen/installer_files/env/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 460, in forward
attn_output = _flash_attention_forward(
^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: _flash_attention_forward() got an unexpected keyword argument 'num_items_in_batch'

22:42:13-041312 INFO Training complete, saving
22:42:15-181788 INFO Training complete!

Is there an existing issue for this?

• I have searched the existing issues

Reproduction

Training a text file with flsh attention on.

Screenshot

No response

Logs

/Documents/text-gen/installer_files/env/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py", line 460, in forward
    attn_output = _flash_attention_forward(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: _flash_attention_forward() got an unexpected keyword argument 'num_items_in_batch'

22:42:13-041312 INFO     Training complete, saving                                                                                                                                      
22:42:15-181788 INFO     Training complete!

System Info

2 nvidia 4090 gpu