add fa3 compilation

dptech-corp · Jul 12, 2024 · 17848b4 · 17848b4
1 parent ba5962e
commit 17848b4
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/docker/rdma/Dockerfile b/docker/rdma/Dockerfile
@@ -66,7 +66,13 @@ 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 cd /tmp && \
+    git clone https://github.com/Dao-AILab/flash-attention.git && \
+    cd flash-attention && \
+    python setup.py install && \
+    cd hopper && \
+    python setup.py install && \
+    rm -rf  /tmp/* && rm -rf ~/.cache/pip
 
 RUN ldconfig && \
     apt-get clean && \