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

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

scripts/ci/ci_install_deepep.sh

@@ -5,9 +5,6 @@ set -euxo pipefail
 bash scripts/ci/ci_install_dependency.sh
 
 export GDRCOPY_HOME=/usr/src/gdrdrv-2.5.1/
-export NVSHMEM_DIR=/opt/nvshmem/install
-export LD_LIBRARY_PATH="${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH"
-export PATH="${NVSHMEM_DIR}/bin:$PATH"
 export CUDA_HOME=/usr/local/cuda
 
 GRACE_BLACKWELL=${GRACE_BLACKWELL:-0}
@@ -28,7 +25,6 @@ apt install -y curl wget git sudo libibverbs-dev rdma-core infiniband-diags open
 
 # Install GDRCopy
 rm -rf /opt/gdrcopy && mkdir -p /opt/gdrcopy
-rm -rf /opt/nvshmem && mkdir -p /opt/nvshmem
 cd /opt/gdrcopy
 git clone https://github.com/NVIDIA/gdrcopy.git .
 git checkout v2.5.1
@@ -51,26 +47,7 @@ fi
 apt-get update && apt-get install -y libfabric-dev
 
 # Install NVSHMEM
-cd /opt/nvshmem
-wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.4.5/source/nvshmem_src_cuda12-all-all-3.4.5.tar.gz
-tar -xf nvshmem_src_cuda12-all-all-3.4.5.tar.gz
-mv nvshmem_src nvshmem && cd nvshmem
-if [ "$GRACE_BLACKWELL" = "1" ]; then
-    CUDA_ARCH="100;120"
-else
-    CUDA_ARCH="90"
-fi
-NVSHMEM_SHMEM_SUPPORT=0 \
-NVSHMEM_UCX_SUPPORT=0 \
-NVSHMEM_USE_NCCL=0 \
-NVSHMEM_MPI_SUPPORT=0 \
-NVSHMEM_IBGDA_SUPPORT=1 \
-NVSHMEM_PMIX_SUPPORT=0 \
-NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \
-NVSHMEM_USE_GDRCOPY=1 \
-cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=/opt/nvshmem/install -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCH}
-cd build
-make -j$(nproc) install
+pip install nvidia-nvshmem-cu12==3.4.5 --force-reinstall
 
 # Install DeepEP
 DEEPEP_DIR=/root/.cache/deepep
@@ -103,11 +80,7 @@ if [ "$GRACE_BLACKWELL" = "1" ]; then
     if [ "${CUDA_VERSION%%.*}" = "13" ]; then \
         sed -i "/^    include_dirs = \['csrc\/'\]/a\    include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \
     fi
-    NVSHMEM_DIR=/opt/nvshmem/install TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install --no-build-isolation .
+    TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install --no-build-isolation .
 else
     python3 setup.py install
 fi
-
-# Verify configuration
-echo "=== Verify NVSHMEM ==="
-nvshmem-info -a

scripts/ci/ci_install_dependency.sh

@@ -134,6 +134,8 @@ if [ "$IS_BLACKWELL" != "1" ]; then
     git clone --branch v0.5 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git
     $PIP_CMD install -e lmms-eval/ $PIP_INSTALL_SUFFIX
 fi
+# Cudnn with version less than 9.16.0.29 will cause performance regression on Conv3D kernel
+$PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall $PIP_INSTALL_SUFFIX
 $PIP_CMD uninstall xformers || true
 # Show current packages
 $PIP_CMD list