[Dependency] Upgrade to Torch 2.11.0 (#21247)
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com> Co-authored-by: b8zhong <b8zhong@users.noreply.github.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Brayden Zhong
co-authored by
Kangyan Zhou
Claude Opus 4.7
Baizhou Zhang
b8zhong
Mick
Kangyan Zhou
Claude Opus 4.7
Baizhou Zhang
b8zhong
Mick
parent
24a6b3084d
commit
88bb5dffe4
@@ -1,44 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Cache and pre-install nvidia wheels that torch pins.
|
||||
#
|
||||
# pypi.nvidia.com returns Cache-Control: no-store, so pip re-downloads
|
||||
# ~2 GB of NVIDIA wheels on every CI run. This script:
|
||||
# 1. Caches cudnn + nvshmem wheels locally and pre-installs them
|
||||
# 2. Points pip at a local wheel directory via PIP_FIND_LINKS so that
|
||||
# all NVIDIA torch dependencies (cublas, cufft, nvrtc, etc.) are
|
||||
# installed from local files instead of re-downloading.
|
||||
#
|
||||
# Pre-cache the wheels on the host at /opt/ci-cache/nvidia-pip-wheels/
|
||||
# (mounted as /root/.cache/nvidia-pip-wheels inside containers).
|
||||
# See the 5090 ops guide post-reboot checklist for how to populate this.
|
||||
#
|
||||
# Integrity: uses `unzip -t` to detect partial/corrupt downloads.
|
||||
#
|
||||
# Usage: source scripts/ci/cuda/cache_nvidia_wheels.sh
|
||||
|
||||
NVIDIA_WHEEL_CACHE="/root/.cache/nvidia-wheels"
|
||||
NVIDIA_PIP_WHEELS="/root/.cache/nvidia-pip-wheels"
|
||||
mkdir -p "$NVIDIA_WHEEL_CACHE"
|
||||
|
||||
for url in \
|
||||
"https://pypi.nvidia.com/nvidia-cudnn-cu13/nvidia_cudnn_cu13-9.16.0.29-py3-none-manylinux_2_27_x86_64.whl" \
|
||||
"https://pypi.nvidia.com/nvidia-nvshmem-cu13/nvidia_nvshmem_cu13-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; do
|
||||
whl="$NVIDIA_WHEEL_CACHE/$(basename "$url")"
|
||||
[ -f "$whl" ] && unzip -tq "$whl" &>/dev/null || curl -fL -o "$whl" "$url"
|
||||
done
|
||||
|
||||
# Caller (ci_install_dependency.sh) sets $PIP_CMD/$PIP_INSTALL_SUFFIX to route
|
||||
# installs into the active environment (venv or system). The `:-pip` fallback
|
||||
# keeps the file runnable ad-hoc for debugging; in CI the caller always sets
|
||||
# these. Silent failure here is deliberate — the pinned cudnn/nvshmem installs
|
||||
# later in ci_install_dependency.sh are the source of truth; this is only a
|
||||
# download optimization.
|
||||
${PIP_CMD:-pip} install --no-deps "$NVIDIA_WHEEL_CACHE"/nvidia_cudnn_cu13-*.whl \
|
||||
"$NVIDIA_WHEEL_CACHE"/nvidia_nvshmem_cu13-*.whl ${PIP_INSTALL_SUFFIX:-} 2>/dev/null || true
|
||||
|
||||
# If pre-cached NVIDIA pip wheels exist, tell pip to check there first.
|
||||
# This avoids re-downloading ~2 GB of cublas/cufft/nvrtc/etc. every run
|
||||
# (pypi.nvidia.com sends Cache-Control: no-store).
|
||||
if [ -d "$NVIDIA_PIP_WHEELS" ] && ls "$NVIDIA_PIP_WHEELS"/*.whl &>/dev/null; then
|
||||
export PIP_FIND_LINKS="${PIP_FIND_LINKS:+$PIP_FIND_LINKS }$NVIDIA_PIP_WHEELS"
|
||||
fi
|
||||
@@ -31,7 +31,10 @@ if [ "$ARCH" != "x86_64" ] && [ "$ARCH" != "aarch64" ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if python3 -c "import deep_ep" >/dev/null 2>&1; then
|
||||
if [ "${FORCE_REBUILD_DEEPEP:-0}" = "1" ]; then
|
||||
echo "FORCE_REBUILD_DEEPEP=1; uninstalling any cached deep_ep before rebuild."
|
||||
${PIP_UNINSTALL_CMD:-pip uninstall -y} deep_ep ${PIP_UNINSTALL_SUFFIX:-} || true
|
||||
elif python3 -c "import deep_ep" >/dev/null 2>&1; then
|
||||
echo "deep_ep is already installed or importable. Skipping installation."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -34,9 +34,6 @@ configure_environment() {
|
||||
CU_STRIP="${CU_VERSION#cu}"
|
||||
CU_MAJOR="${CU_STRIP:0:2}"
|
||||
|
||||
# Nvidia package versions we pin (torch ships older versions).
|
||||
NVIDIA_CUDNN_VERSION="9.16.0.29"
|
||||
NVIDIA_NVSHMEM_VERSION="3.4.5"
|
||||
OPTIONAL_DEPS="${1:-}"
|
||||
|
||||
# Whether to create a uv venv (set USE_VENV=1). Default: 0.
|
||||
@@ -288,19 +285,11 @@ install_sglang_kernel() {
|
||||
$PIP_CMD install "torch==${TORCH_VER}" "torchaudio==${TORCHAUDIO_VER}" "torchvision==${TORCHVISION_VER}" --index-url "https://download.pytorch.org/whl/${CU_VERSION}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# Reinstall sglang-kernel with matching CUDA version if needed
|
||||
SGL_KERNEL_FULL_VER=$(pip show sglang-kernel 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
SGL_KERNEL_CUDA_VER=$(printf '%s' "$SGL_KERNEL_FULL_VER" | sed -n 's/.*+//p')
|
||||
echo "Detected sglang-kernel version: ${SGL_KERNEL_FULL_VER} (CUDA tag: ${SGL_KERNEL_CUDA_VER:-none})"
|
||||
if [ -n "$SGL_KERNEL_CUDA_VER" ] && [ "$SGL_KERNEL_CUDA_VER" != "$CU_VERSION" ]; then
|
||||
SGL_KERNEL_VER="${SGL_KERNEL_FULL_VER%+*}"
|
||||
echo "Reinstalling sglang-kernel==${SGL_KERNEL_VER} from ${CU_VERSION} index to match torch..."
|
||||
if [ "$CU_MAJOR" = "13" ]; then
|
||||
$PIP_CMD install "sglang-kernel==${SGL_KERNEL_VER}" --index-url "https://docs.sglang.ai/whl/${CU_VERSION}/" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
else
|
||||
$PIP_CMD install "sglang-kernel==${SGL_KERNEL_VER}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
fi
|
||||
# install_sglang above pulls sglang-kernel from PyPI, whose default wheel
|
||||
# tracks one CUDA version (currently cu130). Force-reinstall from the
|
||||
# CU_VERSION-matched sglang wheel index so runners on a different CUDA
|
||||
# (e.g. h20 / cu129) get a wheel linked against the right libnvrtc.
|
||||
$PIP_CMD install "sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT}" --index-url "https://docs.sglang.ai/whl/${CU_VERSION}/" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
|
||||
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
@@ -407,34 +396,6 @@ install_extra_deps() {
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
|
||||
fix_nvidia_deps() {
|
||||
if [ "$CU_MAJOR" = "13" ]; then
|
||||
NVSHMEM_PKG="nvidia-nvshmem-cu13"
|
||||
CUDNN_PKG="nvidia-cudnn-cu13"
|
||||
else
|
||||
NVSHMEM_PKG="nvidia-nvshmem-cu12"
|
||||
CUDNN_PKG="nvidia-cudnn-cu12"
|
||||
fi
|
||||
|
||||
# DeepEP depends on nvshmem 3.4.5
|
||||
INSTALLED_NVSHMEM=$(pip show ${NVSHMEM_PKG} 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_NVSHMEM" = "$NVIDIA_NVSHMEM_VERSION" ]; then
|
||||
echo "${NVSHMEM_PKG}==${NVIDIA_NVSHMEM_VERSION} already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install ${NVSHMEM_PKG}==${NVIDIA_NVSHMEM_VERSION} $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
# cudnn < 9.16.0.29 causes Conv3D performance regression
|
||||
INSTALLED_CUDNN=$(pip show ${CUDNN_PKG} 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
|
||||
if [ "$INSTALLED_CUDNN" = "$NVIDIA_CUDNN_VERSION" ]; then
|
||||
echo "${CUDNN_PKG}==${NVIDIA_CUDNN_VERSION} already installed, skipping reinstall"
|
||||
else
|
||||
$PIP_CMD install ${CUDNN_PKG}==${NVIDIA_CUDNN_VERSION} $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
|
||||
install_test_tools() {
|
||||
# Download kernels from kernels community
|
||||
kernels download python || true
|
||||
@@ -506,7 +467,6 @@ main() {
|
||||
download_flashinfer_cache
|
||||
stabilize_flashinfer_jit_paths
|
||||
install_extra_deps
|
||||
fix_nvidia_deps
|
||||
install_test_tools
|
||||
prepare_runner
|
||||
setup_ld_library_path
|
||||
|
||||
Reference in New Issue
Block a user