From 5dcd0a652fb1fc0dc1bd7cd3951e571c85ff3651 Mon Sep 17 00:00:00 2001
From: Guolin Ke <guolin.ke@outlook.com>
Date: Wed, 17 Jul 2024 11:30:42 +0800
Subject: [PATCH] update flash-atto

---
 docker/rdma/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/rdma/Dockerfile b/docker/rdma/Dockerfile
index 0cf248b..d17029f 100644
--- 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 && \