Incorrectness in Flash Attention

[mayank31398]

We completely ignore attention_mask here: https://github.com/pacman100/DHS-LLM-Workshop/blob/53672e1b774da7798fb10a50ef8ca5b2750c5608/personal_copilot/training/starcoder_flash_attn_monkey_patch.py#L60

If the input had padding, then this is incorrect (not by a major amount I think but that might depend on how much padding the input has).
We need to maintain cu_seqlens and use the packed version of flash here.
But the current implementementation is easier to implement and maintain I guess?

Can we add a note regarding the incorrect behaviour?

[pacman100]

Hello, I use it for continued pretraining with packing wherein there is no padding and attention mask involved. As such, it works as intended, I have mentioned this caveat here: huggingface/accelerate#1864 (comment)

Note:

Flash V2 support that I have implemented above ignores padding/attention_mask/custom_mask. It is meant for continued pre-training with packing inputs to consume the entire sequence lengths.

I will raise a PR to raise this warning. Thank you!

[pacman100]

Added the warning:
https://github.com/pacman100/DHS-LLM-Workshop/blob/main/personal_copilot/training/train.py#L381-L384

[mayank31398]

Hey, great
Sorry, didn't know you were using it with dense packing.
Closing this issue.