tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

update docker args to minimums used and use MAX_JOBS already set as arg

Browse Source
This commit is contained in:
Wing Lian
2025-05-17 08:12:25 -04:00
parent 79daf5b934
commit e1b74d710b
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
docker/Dockerfile-base
View File

@@ -1,5 +1,5 @@
ARG CUDA_VERSION="11.8.0" ARG CUDA_VERSION="12.4.1"
ARG CUDNN_VERSION="8" ARG CUDNN_VERSION=""
ARG UBUNTU_VERSION="22.04" ARG UBUNTU_VERSION="22.04"
ARG MAX_JOBS=4 ARG MAX_JOBS=4
@@ -7,9 +7,9 @@ FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION A
ENV PATH="/root/miniconda3/bin:${PATH}" ENV PATH="/root/miniconda3/bin:${PATH}"
ARG PYTHON_VERSION="3.10" ARG PYTHON_VERSION="3.11"
ARG PYTORCH_VERSION="2.1.2" ARG PYTORCH_VERSION="2.5.1"
ARG CUDA="118" ARG CUDA="124"
ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX" ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
ENV PYTHON_VERSION=$PYTHON_VERSION ENV PYTHON_VERSION=$PYTHON_VERSION
@@ -43,7 +43,7 @@ RUN if [ "$TORCH_CUDA_ARCH_LIST" = "9.0+PTX" ] ; then \
cd flash-attention; \ cd flash-attention; \
git checkout v2.7.4.post1; \ git checkout v2.7.4.post1; \
cd hopper; \ cd hopper; \
FLASH_ATTENTION_DISABLE_SM80=TRUE MAX_JOBS=16 python setup.py install; \ FLASH_ATTENTION_DISABLE_SM80=TRUE MAX_JOBS=${MAX_JOBS} python setup.py install; \
cd ../..; \ cd ../..; \
elif [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \ elif [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \
pip3 install flash-attn==2.7.4.post1; \ pip3 install flash-attn==2.7.4.post1; \