install flash_attn

allenai · Jul 17, 2024 · eca4972 · eca4972
1 parent e44f991
commit eca4972
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
@@ -45,6 +45,8 @@ runs:
         pip install 'torch${{ inputs.torch-version }}' --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .[all]
         pip install -e hf_olmo
+        pip install --no-cache-dir triton==2.0.0 https://storage.googleapis.com/ai2-python-wheels/flash_attn/flash_attn-0.2.8%2Bcu118torch2.0.0-cp310-cp310-linux_x86_64.whl
+
 
     - if: steps.virtualenv-cache.outputs.cache-hit == 'true'
       shell: bash
@@ -53,6 +55,7 @@ runs:
         . .venv/bin/activate
         pip install --no-deps -e .[all]
         pip install --no-deps -e hf_olmo
+        pip install --no-deps --no-cache-dir triton==2.0.0 https://storage.googleapis.com/ai2-python-wheels/flash_attn/flash_attn-0.2.8%2Bcu118torch2.0.0-cp310-cp310-linux_x86_64.whl
 
     - shell: bash
       run: |