use fa2 258, no build iso (#3224)

mosaicml · Apr 30, 2024 · f0d9093 · f0d9093
1 parent 9da9d81
commit f0d9093
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -283,9 +283,7 @@ RUN if [ -n "$MOFED_VERSION" ] ; then \
 RUN if [ -n "$CUDA_VERSION" ] ; then \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir ninja==1.11.1 && \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir --force-reinstall packaging==22.0 && \
-        git clone --branch v2.5.7 https://github.com/Dao-AILab/flash-attention.git && \
-        cd flash-attention && \
-        MAX_JOBS=1 python${PYTHON_VERSION} setup.py install && \
+        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir --no-build-isolation flash-attn==2.5.8; \
         cd .. ; \
     fi