Combining StableLM and Flash Attention 2
First, make sure to install the latest version of Flash Attention v2.

pip install -U flash-attn --no-build-isolation
Also make sure that your hardware is compatible with Flash-Attention 2. Read more about it in the official documentation of the flash-attn repository. Note: you must load your model in half-precision (e.g. torch.bfloat16).
Now, to run the model with Flash Attention 2, refer to the snippet below:
thon