Adding support for alibi when using flash attention

[ShashankMosaicML]

ALiBi support has recently been added to FlashAttention2 (PR#540) and is stable with flash-attention=v2.4.2. This PR allows for ALiBi slopes to be passed to flash attention, hence enabling ALiBi for Flash Attention.

Experiments on 125m and 1b models show (nearly) identical training curves for ALiBi with Triton attention (which is the current default attention implementation when using ALiBi) and ALiBi with Flash Attention. Further, using Flash Attention results in higher MFU numbers. In the plots below, 'treat' refers to flash attention, and 'control' refers to triton attention.

Wandb experiment link

Note: Flash Attention seems to use more memory than Triton:

These changes would effectively also enable the use of sliding window attention and memory efficient sequence id masking with ALiBi when using Flash Attention.

---

As an aside, @Skylion007 has been able to successfully update his own version of MosaicBERT with FA2, and has found it to indeed be faster than FA1: w&b link

[Skylion007]

FYI, I added a BERT FA2 PR to the examples repo: mosaicml/examples#440

[irenedea]

Thanks for doing this! Added some questions and comments.

[jacobfulano]

This looks good to me! Could we potentially include an example yaml in scripts/train/pretrain with flash attention 2 as well?

[vchiley]

This looks good to me! Could we potentially include an example yaml in scripts/train/pretrain with flash attention 2 as well?

if you install llm-foundry using pip install -e . [gpu-flash2] (instead of pip install -e . [gpu]) and set `attn_impl: flash, MPT will uses FA2.

[dakinggg]

lgtm, will leave approval for @irenedea