docker: add CUDA13 support in dockerfile and update GDRCopy/NVSHMEM for blackwell support (#11517)
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
co-authored by
Baizhou Zhang
parent
813bd6f85c
commit
285a8e6986
@@ -4,7 +4,7 @@ set -euxo pipefail
|
||||
|
||||
bash scripts/ci/ci_install_dependency.sh
|
||||
|
||||
export GDRCOPY_HOME=/usr/src/gdrdrv-2.4.4/
|
||||
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"
|
||||
@@ -27,9 +27,9 @@ 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.4.4
|
||||
git checkout v2.5.1
|
||||
apt update
|
||||
apt install -y nvidia-dkms-535
|
||||
apt install -y nvidia-dkms-580
|
||||
apt install -y build-essential devscripts debhelper fakeroot pkg-config dkms
|
||||
apt install -y check libsubunit0 libsubunit-dev python3-venv
|
||||
cd packages
|
||||
@@ -46,8 +46,8 @@ apt-get update && apt-get install -y libfabric-dev
|
||||
|
||||
# Install NVSHMEM
|
||||
cd /opt/nvshmem
|
||||
wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.3.9/source/nvshmem_src_cuda12-all-all-3.3.9.tar.gz
|
||||
tar -xf nvshmem_src_cuda12-all-all-3.3.9.tar.gz
|
||||
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
|
||||
NVSHMEM_SHMEM_SUPPORT=0 \
|
||||
NVSHMEM_UCX_SUPPORT=0 \
|
||||
@@ -57,7 +57,7 @@ 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=90
|
||||
cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=/opt/nvshmem/install -DCMAKE_CUDA_ARCHITECTURES="90;100;103;121"
|
||||
cd build
|
||||
make -j$(nproc) install
|
||||
|
||||
|
||||
@@ -55,7 +55,7 @@ else
|
||||
$PIP_CMD install flashinfer-python==0.4.1 --prerelease=allow $PIP_INSTALL_SUFFIX
|
||||
|
||||
# Install the main package
|
||||
$PIP_CMD install -e "python[dev]" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX --upgrade
|
||||
$PIP_CMD install -e "python" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX --upgrade
|
||||
fi
|
||||
|
||||
# Install router for pd-disagg test
|
||||
@@ -68,7 +68,7 @@ echo "SGL_KERNEL_VERSION_FROM_KERNEL=${SGL_KERNEL_VERSION_FROM_KERNEL} SGL_KERNE
|
||||
|
||||
if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ]; then
|
||||
ls -alh sgl-kernel/dist
|
||||
$PIP_CMD install sgl-kernel/dist/sgl_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}-cp310-abi3-manylinux2014_x86_64.whl --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
$PIP_CMD install sgl-kernel/dist/sgl_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}+${CU_VERSION}-cp310-abi3-manylinux2014_x86_64.whl --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
else
|
||||
$PIP_CMD install sgl-kernel==${SGL_KERNEL_VERSION_FROM_SRT} --force-reinstall $PIP_INSTALL_SUFFIX
|
||||
fi
|
||||
|
||||
Reference in New Issue
Block a user