[Deps] Upgrade CUDA PyTorch stack to 2.13 (#28836)
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
This commit is contained in:
co-authored by
Brayden Zhong
parent
4ad990ba7d
commit
434e646282
@@ -352,6 +352,22 @@ uninstall_stale_flashinfer() {
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
|
||||
install_pytorch_stack() {
|
||||
PYTORCH_SPECS=()
|
||||
for package in torch torchaudio torchvision torchao torchcodec; do
|
||||
spec=$(grep -Po -m1 "\"${package}([<>=!~ ;][^\"]*)?\"" python/pyproject.toml | tr -d '"' || true)
|
||||
if [ -n "$spec" ]; then
|
||||
PYTORCH_SPECS+=("$spec")
|
||||
fi
|
||||
done
|
||||
|
||||
$PIP_CMD install \
|
||||
"${PYTORCH_SPECS[@]}" \
|
||||
--index-url "https://download.pytorch.org/whl/${CU_VERSION}"
|
||||
|
||||
mark_step_done "${FUNCNAME[0]}"
|
||||
}
|
||||
|
||||
require_prebuilt_rust_exts() {
|
||||
# Stages whose download succeeded set this to none. Runs before
|
||||
# setup_pip_toolchain uninstalls sglang, so clearing it here still reaches
|
||||
@@ -470,43 +486,6 @@ install_sglang_kernel() {
|
||||
fi
|
||||
fi
|
||||
|
||||
# Reinstall torch with matching CUDA version if needed
|
||||
# TODO: Remove after torch 2.11 where cu13 is enabled by default
|
||||
REINSTALL_TORCH=false
|
||||
if TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')" 2>&1); then
|
||||
echo "Detected torch CUDA version: ${TORCH_CUDA_VER}"
|
||||
else
|
||||
TORCH_IMPORT_ERROR="${TORCH_CUDA_VER}"
|
||||
TORCH_CUDA_VER=""
|
||||
echo "WARNING: importing torch failed while probing CUDA version; force-reinstalling torch packages."
|
||||
printf '%s\n' "${TORCH_IMPORT_ERROR}"
|
||||
REINSTALL_TORCH=true
|
||||
fi
|
||||
TORCHAUDIO_CUDA_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true)
|
||||
TORCHVISION_CUDA_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+\(cu[0-9][0-9]*\)$/\1/p' || true)
|
||||
if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then
|
||||
REINSTALL_TORCH=true
|
||||
else
|
||||
for cuda_ver in "${TORCHAUDIO_CUDA_VER}" "${TORCHVISION_CUDA_VER}"; do
|
||||
if [ -n "${cuda_ver}" ] && [ "${cuda_ver}" != "${CU_VERSION}" ]; then
|
||||
REINSTALL_TORCH=true
|
||||
break
|
||||
fi
|
||||
done
|
||||
fi
|
||||
if [ "${REINSTALL_TORCH}" = true ]; then
|
||||
TORCH_VER=$(pip show torch 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
|
||||
if [ -z "${TORCH_VER}" ] || [ -z "${TORCHAUDIO_VER}" ] || [ -z "${TORCHVISION_VER}" ]; then
|
||||
echo "ERROR: could not determine installed torch package versions before reinstall."
|
||||
pip show torch torchaudio torchvision || true
|
||||
exit 1
|
||||
fi
|
||||
echo "Reinstalling torch==${TORCH_VER} torchaudio==${TORCHAUDIO_VER} torchvision==${TORCHVISION_VER} from ${CU_VERSION} index to match torch..."
|
||||
$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
|
||||
|
||||
if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" != "true" ]; then
|
||||
# The PyPI default wheel tracks one CUDA version (currently cu130); other
|
||||
# runners (e.g. h20 / cu129) need the +${CU_VERSION}-tagged wheel from the
|
||||
@@ -779,6 +758,7 @@ main() {
|
||||
setup_pip_toolchain
|
||||
remove_stale_cuda12_nvidia_wheels
|
||||
uninstall_stale_flashinfer
|
||||
install_pytorch_stack
|
||||
install_sglang
|
||||
# Diffusion B200 CI imports torch inside install_sglang_kernel after removing
|
||||
# stale CUDA 12 NVIDIA wheels, so opt into one early LD_LIBRARY_PATH refresh.
|
||||
|
||||
Reference in New Issue
Block a user