Skip to content

Commit

Permalink
unfreeze torch version in dockerfile (InternLM#2906)
Browse files Browse the repository at this point in the history
  • Loading branch information
RunningLeon authored Dec 18, 2024
1 parent 8afb84c commit bafa3d2
Showing 1 changed file with 1 addition and 5 deletions.
6 changes: 1 addition & 5 deletions docker/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,6 @@ FROM ${CUDA_VERSION} AS final

ARG PYTHON_VERSION=3.10

ARG TORCH_VERSION=2.3.0
ARG TORCHVISION_VERSION=0.18.0

RUN apt-get update -y && apt-get install -y software-properties-common wget vim git curl openssh-server ssh sudo &&\
curl https://sh.rustup.rs -sSf | sh -s -- -y &&\
add-apt-repository ppa:deadsnakes/ppa -y && apt-get update -y && apt-get install -y --no-install-recommends \
Expand Down Expand Up @@ -43,7 +40,6 @@ ENV LD_LIBRARY_PATH=/usr/local/nccl/lib:$LD_LIBRARY_PATH


RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install --upgrade pip setuptools==69.5.1 &&\
python3 -m pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/${CUDA_VERSION_SHORT} &&\
python3 -m pip install cmake packaging wheel

ENV NCCL_LAUNCH_MODE=GROUP
Expand All @@ -54,7 +50,7 @@ COPY . /opt/lmdeploy
WORKDIR /opt/lmdeploy

RUN --mount=type=cache,target=/root/.cache/pip cd /opt/lmdeploy &&\
python3 -m pip install -r requirements.txt &&\
python3 -m pip install -r requirements_cuda.txt --extra-index-url https://download.pytorch.org/whl/${CUDA_VERSION_SHORT} &&\
mkdir -p build && cd build &&\
sh ../generate.sh &&\
ninja -j$(nproc) && ninja install &&\
Expand Down

0 comments on commit bafa3d2

Please sign in to comment.