fix

allenai · Jul 17, 2024 · 74cc618 · 74cc618
1 parent eca4972
commit 74cc618
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
@@ -45,7 +45,6 @@ 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'
@@ -55,7 +54,6 @@ 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: |

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -128,6 +128,10 @@ jobs:
         run: |
           echo "COMMIT_SHA=$GITHUB_SHA" >> $GITHUB_ENV
 
+      - name: Install flash attn
+        - run: |
+            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
+
       - name: GPU Tests
         uses: allenai/[email protected]
         if: env.BEAKER_TOKEN != ''