[chore] update torch version to 2.9 (#12969)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Fan Yin
2025-11-25 14:47:34 -08:00
committed by GitHub
co-authored by Baizhou Zhang
parent 64a11303ce
commit 36b1bcd242
13 changed files with 44 additions and 87 deletions
+1 -13
View File
@@ -8,7 +8,6 @@ ARG GRACE_BLACKWELL=0
ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2
ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee
ARG TRITON_LANG_COMMIT=4caa0328bf8df64896dd5f6fb9df41b0eb2e750a
ARG BUILD_AND_DOWNLOAD_PARALLEL=8
ARG SGL_KERNEL_VERSION=0.3.17.post2
ARG SGL_VERSION=0.5.5.post3
@@ -134,10 +133,9 @@ RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install --upgrade
fi \
&& if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip uninstall -y torch torchaudio torchvision ; \
python3 -m pip install torch==2.9.0 torchaudio==2.9.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} ; \
else \
echo "No NCCL mapping for CUDA_VERSION=${CUDA_VERSION}" && exit 1 ; \
fi \
@@ -185,16 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/pip cd /sgl-workspace/DeepEP && \
# latest flashinfer_cutedsl. Once 0.4.3 is officially released, remove this
RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install --upgrade --pre "nvidia-cutlass-dsl==4.3.0.dev0" --extra-index-url https://pypi.org/simple/
# For cuda 13, we install triton from source to fix some sm103 issues
# This can be reverted after >3.4.5 is released
# See the conversation in: https://github.com/triton-lang/triton/pull/8536
RUN --mount=type=cache,target=/root/.cache/pip if [ "$CUDA_VERSION" = "13.0.1" ]; then \
wget -q https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/${TRITON_LANG_COMMIT}.zip && \
unzip -q ${TRITON_LANG_COMMIT}.zip && rm ${TRITON_LANG_COMMIT}.zip && mv triton-${TRITON_LANG_COMMIT} triton && \
cd triton && pip install --break-system-packages -r python/requirements.txt && \
MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} pip install --break-system-packages -e .; \
fi
# Python tools
RUN --mount=type=cache,target=/root/.cache/pip python3 -m pip install \
datamodel_code_generator \