support flash-attn in Dockerfile
This commit is contained in:
parent
555ca8d780
commit
c88b1be9f3
|
@ -35,6 +35,11 @@ RUN EXTRA_PACKAGES="metrics"; \
|
||||||
pip install -e .[$EXTRA_PACKAGES] && \
|
pip install -e .[$EXTRA_PACKAGES] && \
|
||||||
pip uninstall -y transformer-engine flash-attn
|
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
|
# Set up volumes
|
||||||
VOLUME [ "/root/.cache/huggingface", "/root/.cache/modelscope", "/app/data", "/app/output" ]
|
VOLUME [ "/root/.cache/huggingface", "/root/.cache/modelscope", "/app/data", "/app/output" ]
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue