Repository files navigation

Code to reproduce results for the paper "Faster Causal Attention Over Large Sequences Through Sparse Flash Attention"

To install the required python dependencies, first run:

pip install -r ./requirements.txt

Then install Triton:

git clone https://github.com/openai/triton.gitcd triton git checkout b2a757d00028fe844a93904036a18e8670bfe92fcd python pip install cmakepip install -e.

In the command above we set the Triton library to the commit used in our experiments. Feel free to experiment with later Triton versions.

GPU requirements: Preferably, you need at least one A100. Some of our experiments use data-parallelism with up to 3 A100s. You should have no problem running those experiments on any GPU supportingbfloat16, you might have to change the model parameters to adapt to the memory available.

Go in theopenwebtext2-experiments folder and run thescript/train-LMs.sh command.

GPU requirements: We used one A100.

For the Hash-sparse and QK-sparse results, go in theruntime-experiments folder and check thetimeperf-hash-and-qk-sparse.ipynb notebook.

Coming soon