Update Flash Attention to 2.5.6

mosaicml · Mar 13, 2024 · e8e5467 · e8e5467
1 parent 63905d6
commit e8e5467
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -276,7 +276,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 && \
-        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir flash-attn==2.5.0; \
+        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir flash-attn==2.5.5; \
     fi
 
 ###############