423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-11-04 18:02:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

support flash-attn in Dockerfile

Browse Source 
Former-commit-id: 0dba000aa178f915cea7d75bf0c9d47e671a21d2
This commit is contained in:
hzhaoy 2024-06-25 15:13:07 +08:00
parent 98f382fda3
commit daea86e047
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
docker/docker-cuda/Dockerfile
View File

@ -35,6 +35,11 @@ RUN EXTRA_PACKAGES="metrics"; \
pip install -e .[$EXTRA_PACKAGES] && \
pip uninstall -y transformer-engine flash-attn
# Rebuild flash-attn
RUN ninja --version || \
(pip uninstall -y ninja && pip install ninja) && \
MAX_JOBS=4 pip install --no-cache-dir flash-attn --no-build-isolation
# Set up volumes
VOLUME [ "/root/.cache/huggingface", "/root/.cache/modelscope", "/app/data", "/app/output" ]