switch to build from source

mosaicml · Apr 25, 2024 · 8f0851e · 8f0851e
1 parent e86d915
commit 8f0851e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -283,7 +283,10 @@ 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 && \
-        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir flash-attn==2.5.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 && \
+        cd .. ; \
     fi
 
 ###############