update flash-atto

dptech-corp · Jul 17, 2024 · 5dcd0a6 · 5dcd0a6
1 parent ba5962e
commit 5dcd0a6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docker/rdma/Dockerfile b/docker/rdma/Dockerfile
@@ -66,7 +66,7 @@ RUN cd /tmp && \
 
 RUN pip3 install --no-cache-dir tokenizers lmdb biopython ml-collections timeout-decorator urllib3 tree dm-tree && rm -rf ~/.cache/pip 
 
-RUN MAX_JOBS=1 pip3 install -U 'flash-attn<2.5.0' --no-build-isolation --no-cache-dir
+RUN MAX_JOBS=1 pip3 install -U 'flash-attn==2.6.1' --no-build-isolation --no-cache-dir
 
 RUN ldconfig && \
  apt-get clean && \