Skip to content

Commit d5cd371

Browse files
HuiGao-NVHuiGao-NV
committed
Import proper module
Signed-off-by: Hui Gao <[email protected]>
1 parent 42ae969 commit d5cd371

File tree

3 files changed

+167
-1
lines changed

3 files changed

+167
-1
lines changed

‎tensorrt_llm/_torch/memory_buffer_utils.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import contextlib
12
import math
23
from dataclasses import dataclass
34
from typing import Optional

‎tensorrt_llm/_torch/utils.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import contextlib
21
import threading
32
from dataclasses import dataclass
43
from enum import Enum

‎tensorrt_llm/scripts/install_tensorrt.sh‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
#!/bin/bash
2+
3+
set -ex
4+
5+
TRT_VER="10.13.2.6"
6+
# Align with the pre-installed cuDNN / cuBLAS / NCCL versions from
7+
# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-08.html#rel-25-08
8+
CUDA_VER="13.0" # 13.0.0
9+
# Keep the installation for cuDNN if users want to install PyTorch with source codes.
10+
# PyTorch 2.x can compile with cuDNN v9.
11+
CUDNN_VER="9.12.0.46-1"
12+
# NCCL version 2.26.x used in the NGC PyTorch 25.05 image but has a performance regression issue.
13+
# Use NCCL version 2.27.5 which has the fixes.
14+
NCCL_VER="2.27.7-1+cuda13.0"
15+
# Use cuBLAS version 13.0.0.19 instead.
16+
CUBLAS_VER="13.0.0.19-1"
17+
# Align with the pre-installed CUDA / NVCC / NVRTC versions from
18+
# https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
19+
NVRTC_VER="13.0.48-1"
20+
CUDA_RUNTIME="13.0.48-1"
21+
CUDA_DRIVER_VERSION="580.65.06-1.el8"
22+
23+
for i in "$@"; do
24+
case $i in
25+
--TRT_VER=?*) TRT_VER="${i#*=}";;
26+
--CUDA_VER=?*) CUDA_VER="${i#*=}";;
27+
--CUDNN_VER=?*) CUDNN_VER="${i#*=}";;
28+
--NCCL_VER=?*) NCCL_VER="${i#*=}";;
29+
--CUBLAS_VER=?*) CUBLAS_VER="${i#*=}";;
30+
*) ;;
31+
esac
32+
shift
33+
done
34+
35+
NVCC_VERSION_OUTPUT=$(nvcc --version)
36+
if [[ $(echo $NVCC_VERSION_OUTPUT | grep -oP "\d+\.\d+" | head -n 1) != ${CUDA_VER} ]]; then
37+
echo "The version of pre-installed CUDA is not equal to ${CUDA_VER}."
38+
fi
39+
40+
install_ubuntu_requirements() {
41+
apt-get update && apt-get install -y --no-install-recommends gnupg2 curl ca-certificates
42+
ARCH=$(uname -m)
43+
if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi
44+
if [ "$ARCH" = "aarch64" ];then ARCH="sbsa";fi
45+
46+
curl -fsSLO https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/${ARCH}/cuda-keyring_1.1-1_all.deb
47+
dpkg -i cuda-keyring_1.1-1_all.deb
48+
rm cuda-keyring_1.1-1_all.deb
49+
50+
apt-get update
51+
if [[ $(apt list --installed | grep libcudnn9) ]]; then
52+
apt-get remove --purge -y libcudnn9*
53+
fi
54+
if [[ $(apt list --installed | grep libnccl) ]]; then
55+
apt-get remove --purge -y --allow-change-held-packages libnccl*
56+
fi
57+
if [[ $(apt list --installed | grep libcublas) ]]; then
58+
apt-get remove --purge -y --allow-change-held-packages libcublas*
59+
fi
60+
if [[ $(apt list --installed | grep cuda-nvrtc-dev) ]]; then
61+
apt-get remove --purge -y --allow-change-held-packages cuda-nvrtc-dev*
62+
fi
63+
64+
CUBLAS_CUDA_VERSION=$(echo $CUDA_VER | sed 's/\./-/g')
65+
NVRTC_CUDA_VERSION=$(echo $CUDA_VER | sed 's/\./-/g')
66+
67+
apt-get install -y --no-install-recommends \
68+
libcudnn9-cuda-13=${CUDNN_VER} \
69+
libcudnn9-dev-cuda-13=${CUDNN_VER} \
70+
libcudnn9-headers-cuda-13=${CUDNN_VER} \
71+
libnccl2=${NCCL_VER} \
72+
libnccl-dev=${NCCL_VER} \
73+
libcublas-${CUBLAS_CUDA_VERSION}=${CUBLAS_VER} \
74+
libcublas-dev-${CUBLAS_CUDA_VERSION}=${CUBLAS_VER} \
75+
cuda-nvrtc-dev-${NVRTC_CUDA_VERSION}=${NVRTC_VER}
76+
77+
apt-get clean
78+
rm -rf /var/lib/apt/lists/*
79+
}
80+
81+
install_rockylinux_requirements() {
82+
CUBLAS_CUDA_VERSION=$(echo $CUDA_VER | sed 's/\./-/g')
83+
84+
ARCH=$(uname -m)
85+
if [ "$ARCH" = "x86_64" ];then ARCH1="x86_64" && ARCH2="x64" && ARCH3=$ARCH1;fi
86+
if [ "$ARCH" = "aarch64" ];then ARCH1="aarch64" && ARCH2="aarch64sbsa" && ARCH3="sbsa";fi
87+
88+
# Download and install packages
89+
for pkg in \
90+
"libnccl-${NCCL_VER}.${ARCH1}" \
91+
"libnccl-devel-${NCCL_VER}.${ARCH1}" \
92+
"cuda-compat-${CUBLAS_CUDA_VERSION}-${CUDA_DRIVER_VERSION}.${ARCH1}" \
93+
"cuda-toolkit-${CUBLAS_CUDA_VERSION}-config-common-${CUDA_RUNTIME}.noarch" \
94+
"cuda-toolkit-13-config-common-${CUDA_RUNTIME}.noarch" \
95+
"cuda-toolkit-config-common-${CUDA_RUNTIME}.noarch" \
96+
"libcublas-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER}.${ARCH1}" \
97+
"libcublas-devel-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER}.${ARCH1}"; do
98+
wget --retry-connrefused --timeout=180 --tries=10 --continue "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${ARCH3}/${pkg}.rpm"
99+
done
100+
101+
# Remove old packages
102+
dnf remove -y "libnccl*"
103+
104+
# Install new packages
105+
dnf -y install \
106+
libnccl-${NCCL_VER}.${ARCH1}.rpm \
107+
libnccl-devel-${NCCL_VER}.${ARCH1}.rpm \
108+
cuda-compat-${CUBLAS_CUDA_VERSION}-${CUDA_DRIVER_VERSION}.${ARCH1}.rpm \
109+
cuda-toolkit-${CUBLAS_CUDA_VERSION}-config-common-${CUDA_RUNTIME}.noarch.rpm \
110+
cuda-toolkit-13-config-common-${CUDA_RUNTIME}.noarch.rpm \
111+
cuda-toolkit-config-common-${CUDA_RUNTIME}.noarch.rpm \
112+
libcublas-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER}.${ARCH1}.rpm \
113+
libcublas-devel-${CUBLAS_CUDA_VERSION}-${CUBLAS_VER}.${ARCH1}.rpm
114+
115+
# Clean up
116+
rm -f *.rpm
117+
dnf clean all
118+
nvcc --version
119+
}
120+
121+
install_tensorrt() {
122+
PY_VERSION=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))')
123+
PARSED_PY_VERSION=$(echo "${PY_VERSION//./}")
124+
TRT_CUDA_VERSION=${CUDA_VER}
125+
TRT_VER_SHORT=$(echo $TRT_VER | cut -d. -f1-3)
126+
127+
if [ -z "$RELEASE_URL_TRT" ];then
128+
ARCH=${TRT_TARGETARCH}
129+
if [ -z "$ARCH" ];then ARCH=$(uname -m);fi
130+
if [ "$ARCH" = "arm64" ];then ARCH="aarch64";fi
131+
if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi
132+
RELEASE_URL_TRT="https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/${TRT_VER_SHORT}/tars/TensorRT-${TRT_VER}.Linux.${ARCH}-gnu.cuda-${TRT_CUDA_VERSION}.tar.gz"
133+
fi
134+
135+
wget --retry-connrefused --timeout=180 --tries=10 --continue ${RELEASE_URL_TRT} -O /tmp/TensorRT.tar
136+
tar -xf /tmp/TensorRT.tar -C /usr/local/
137+
mv /usr/local/TensorRT-${TRT_VER} /usr/local/tensorrt
138+
pip3 install --no-cache-dir /usr/local/tensorrt/python/tensorrt-*-cp${PARSED_PY_VERSION}-*.whl
139+
rm -rf /tmp/TensorRT.tar
140+
echo 'export LD_LIBRARY_PATH=/usr/local/tensorrt/lib:$LD_LIBRARY_PATH' >> "${ENV}"
141+
142+
rm -f /usr/local/tensorrt/lib/libnvinfer_vc_plugin_static.a \
143+
/usr/local/tensorrt/lib/libnvinfer_plugin_static.a \
144+
/usr/local/tensorrt/lib/libnvinfer_static.a \
145+
/usr/local/tensorrt/lib/libnvinfer_dispatch_static.a \
146+
/usr/local/tensorrt/lib/libnvinfer_lean_static.a \
147+
/usr/local/tensorrt/lib/libnvonnxparser_static.a \
148+
/usr/local/tensorrt/lib/libnvinfer_builder_resource_win.so.10.10.0
149+
}
150+
151+
# Install base packages depending on the base OS
152+
ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
153+
case "$ID" in
154+
ubuntu)
155+
install_ubuntu_requirements
156+
install_tensorrt
157+
;;
158+
rocky)
159+
install_rockylinux_requirements
160+
install_tensorrt
161+
;;
162+
*)
163+
echo "Unable to determine OS..."
164+
exit 1
165+
;;
166+
esac

0 commit comments

Comments
 (0)