use flash-attn 1.0.9 in toch1.14

Azure · Feb 5, 2024 · 17d0186 · 17d0186
1 parent 0665d21
commit 17d0186
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml
@@ -43,6 +43,7 @@ jobs:
         export LD_LIBRARY_PATH="/usr/local/lib:$LD_LIBRARY_PATH"
         python3 -m pip install --upgrade pip
         cd ${{ matrix.dir }}/
+        apt-get update && apt-get install -y python3-mpi4py
         python3 -m pip install .[test]
         make postinstall
     - name: Run code lint

diff --git a/dockerfile/torch1.14-cuda11.8.dockerfile b/dockerfile/torch1.14-cuda11.8.dockerfile
@@ -49,8 +49,9 @@ RUN cd third_party/msccl && \
     -gencode=arch=compute_90,code=sm_90" && \
     make install
 # cache TE build to save time in CI
+ENV MAX_JOBS=1
 RUN python3 -m pip install --upgrade pip && \
-    MAX_JOBS=2 python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    MAX_JOBS=1 python3 -m pip install flash-attn==1.0.9 git+https://github.com/NVIDIA/TransformerEngine.git@stable
 
 ADD . .
 RUN python3 -m pip install . && \