d8a05744d7
The non-root user approach had multiple issues with RunPod compatibility, sudo PATH handling, and tmux in exec sessions. Restoring root as the default user for now.
58 lines
2.1 KiB
Plaintext
58 lines
2.1 KiB
Plaintext
ARG CUDA_VERSION="12.6.3"
|
|
ARG CUDNN_VERSION=""
|
|
ARG UBUNTU_VERSION="22.04"
|
|
ARG MAX_JOBS=4
|
|
ARG TARGETARCH
|
|
|
|
FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION AS base-builder
|
|
|
|
ARG TARGETARCH
|
|
ARG PYTHON_VERSION="3.11"
|
|
ARG PYTORCH_VERSION="2.6.0"
|
|
ARG CUDA="126"
|
|
ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
|
|
|
|
ENV PYTHON_VERSION=$PYTHON_VERSION
|
|
ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
|
|
ENV UV_TORCH_BACKEND="cu${CUDA}"
|
|
|
|
RUN apt-get update \
|
|
&& apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev pkg-config curl && rm -rf /var/lib/apt/lists/* \
|
|
&& git lfs install --skip-repo \
|
|
&& curl -LsSf https://astral.sh/uv/install.sh | sh
|
|
|
|
ENV PATH="/root/.local/bin:${PATH}"
|
|
|
|
RUN uv python install ${PYTHON_VERSION}
|
|
|
|
WORKDIR /workspace
|
|
|
|
RUN uv venv --no-project --relocatable axolotl-venv
|
|
|
|
ENV PATH="/workspace/axolotl-venv/bin:${PATH}"
|
|
|
|
RUN uv pip install packaging setuptools wheel psutil \
|
|
&& uv pip install torch==${PYTORCH_VERSION} torchvision \
|
|
&& uv pip install awscli pydantic
|
|
|
|
RUN if [ "$TARGETARCH" = "amd64" ]; then \
|
|
uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main"; \
|
|
uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"; \
|
|
fi
|
|
|
|
# Map Python version (e.g., 3.12 -> cp312)
|
|
RUN PYTHON_CP="cp$(echo $PYTHON_VERSION | tr -d '.')" && \
|
|
# Map PyTorch version (e.g., 2.9.1 -> torch2.9, 2.10.0 -> torch2.10)
|
|
TORCH_TAG="torch$(echo $PYTORCH_VERSION | grep -oP '^\d+\.\d+')" && \
|
|
# Map architecture
|
|
case "$TARGETARCH" in \
|
|
amd64) ARCH_TAG="x86_64" ;; \
|
|
arm64) ARCH_TAG="aarch64" ;; \
|
|
*) echo "Unsupported architecture: $TARGETARCH"; exit 1 ;; \
|
|
esac && \
|
|
WHL_VERSION="v0.7.16" && \
|
|
WHL_FILE="flash_attn-2.8.3+cu${CUDA}${TORCH_TAG}-${PYTHON_CP}-${PYTHON_CP}-linux_${ARCH_TAG}.whl" && \
|
|
wget -nv "https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/${WHL_VERSION}/${WHL_FILE}" && \
|
|
uv pip install --no-cache-dir "${WHL_FILE}" && \
|
|
rm "${WHL_FILE}"
|