This commit is contained in:
OlivierDehaene 2023-11-23 11:28:21 +01:00
parent 703302ae19
commit 289eee535e
3 changed files with 9 additions and 12 deletions

View File

@ -37,15 +37,12 @@ RUN cargo build --release
# Python builder # Python builder
# Adapted from: https://github.com/pytorch/pytorch/blob/master/Dockerfile # Adapted from: https://github.com/pytorch/pytorch/blob/master/Dockerfile
FROM debian:bullseye-slim as pytorch-install FROM nvidia/cuda:12.1.0-devel-ubuntu20.04 as pytorch-install
ARG PYTORCH_VERSION=2.1.0 ARG PYTORCH_VERSION=2.1.0
ARG PYTHON_VERSION=3.10 ARG PYTHON_VERSION=3.10
# Keep in sync with `server/pyproject.toml # Keep in sync with `server/pyproject.toml
ARG CUDA_VERSION=11.8
ARG MAMBA_VERSION=23.3.1-1 ARG MAMBA_VERSION=23.3.1-1
ARG CUDA_CHANNEL=nvidia
ARG INSTALL_CHANNEL=pytorch
# Automatically set by buildx # Automatically set by buildx
ARG TARGETPLATFORM ARG TARGETPLATFORM
@ -75,21 +72,21 @@ RUN chmod +x ~/mambaforge.sh && \
RUN case ${TARGETPLATFORM} in \ RUN case ${TARGETPLATFORM} in \
"linux/arm64") exit 1 ;; \ "linux/arm64") exit 1 ;; \
*) /opt/conda/bin/conda update -y conda && \ *) /opt/conda/bin/conda update -y conda && \
/opt/conda/bin/conda config --remove channels conda-forge && \ /opt/conda/bin/conda install -y "python=${PYTHON_VERSION}" ;; \
/opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y "python=${PYTHON_VERSION}" pytorch==$PYTORCH_VERSION "pytorch-cuda=$(echo $CUDA_VERSION | cut -d'.' -f 1-2)" ;; \
esac && \ esac && \
/opt/conda/bin/conda clean -ya /opt/conda/bin/conda clean -ya
RUN pip install torch==${PYTORCH_VERSION} --extra-index-url https://download.pytorch.org/whl/cu121 --no-cache-dir
# CUDA kernels builder image # CUDA kernels builder image
FROM pytorch-install as kernel-builder FROM pytorch-install as kernel-builder
ARG MAX_JOBS=8
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
ninja-build \ ninja-build \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
RUN /opt/conda/bin/conda install -c "nvidia/label/cuda-11.8.0" cuda==11.8 && \
/opt/conda/bin/conda clean -ya
# Build Flash Attention CUDA kernels # Build Flash Attention CUDA kernels
FROM kernel-builder as flash-att-builder FROM kernel-builder as flash-att-builder
@ -149,7 +146,7 @@ COPY server/Makefile-vllm Makefile
RUN make build-vllm RUN make build-vllm
# Text Generation Inference base image # Text Generation Inference base image
FROM nvidia/cuda:11.8.0-base-ubuntu20.04 as base FROM nvidia/cuda:12.1.0-base-ubuntu20.04 as base
# Conda env # Conda env
ENV PATH=/opt/conda/bin:$PATH \ ENV PATH=/opt/conda/bin:$PATH \

View File

@ -18,7 +18,7 @@ gen-server:
install-torch: install-torch:
# Install specific version of torch # Install specific version of torch
pip install torch==2.1.0 --extra-index-url https://download.pytorch.org/whl/cu118 --no-cache-dir pip install torch==2.1.0 --extra-index-url https://download.pytorch.org/whl/cu121 --no-cache-dir
install: gen-server install-torch install: gen-server install-torch
pip install pip --upgrade pip install pip --upgrade

View File

@ -49,7 +49,7 @@ pytest = "^7.3.0"
[[tool.poetry.source]] [[tool.poetry.source]]
name = "pytorch-gpu-src" name = "pytorch-gpu-src"
url = "https://download.pytorch.org/whl/cu118" url = "https://download.pytorch.org/whl/cu121"
priority = "explicit" priority = "explicit"
[tool.pytest.ini_options] [tool.pytest.ini_options]