fixed flash-attention build for JP6

remy415 · Apr 13, 2024 · aeae01f · aeae01f
1 parent daec892
commit aeae01f
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/packages/llm/flash-attention/Dockerfile b/packages/llm/flash-attention/Dockerfile
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE}
 
 ARG FLASH_ATTENTION_VERSION="2.5.6"
 
-COPY build.sh install.sh patch.diff /tmp/flash-attention
+COPY build.sh install.sh patch.diff /tmp/flash-attention/
 
 RUN /tmp/flash-attention/install.sh || /tmp/flash-attention/build.sh
 
diff --git a/packages/llm/flash-attention/build.sh b/packages/llm/flash-attention/build.sh
@@ -10,15 +10,17 @@ cd /opt/flash-attention
 git apply /tmp/flash-attention/patch.diff
 git diff
 git status
-    
+
 FLASH_ATTENTION_FORCE_BUILD=1 \
 FLASH_ATTENTION_FORCE_CXX11_ABI=0 \
 FLASH_ATTENTION_SKIP_CUDA_BUILD=0 \
+MAX_JOBS=$(nproc) \
 python3 setup.py --verbose bdist_wheel --dist-dir /opt
 
+ls /opt
 cd /
 
-pip3 install --no-cache-dir --verbose /opt/flash-attn*.whl
-pip3 show flash-attn && python3 -c 'import flash_attn'
+pip3 install --no-cache-dir --verbose /opt/flash_attn*.whl
+#pip3 show flash-attn && python3 -c 'import flash_attn'
 
-twine upload --verbose /opt/flash-attn*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"
+twine upload --verbose /opt/flash_attn*.whl || echo "failed to upload wheel to ${TWINE_REPOSITORY_URL}"
diff --git a/packages/llm/flash-attention/install.sh b/packages/llm/flash-attention/install.sh
@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 set -ex
 
-pip3 install --no-cache-dir --verbose flash-attn=${FLASH_ATTENTION_VERSION}
+pip3 install --no-cache-dir --verbose flash-attn==${FLASH_ATTENTION_VERSION}
 pip3 show flash-attn && python3 -c 'import flash_attn'