From a62eb80f7c03f627b226795f9475b5483e2fff05 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 27 May 2026 10:37:21 -0700 Subject: [PATCH] Detect CUDA UMD Version from newer nvidia-smi output (fixes #5812) (#5817) * Detect CUDA UMD Version from newer nvidia-smi output (#5812) Newer NVIDIA drivers (e.g. 610.x on Windows) print the driver's max CUDA capability as "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y" header. The installers and Studio setup scripts were only matching the legacy spelling, so on a fresh RTX 5090 laptop with a 13.x driver they failed to detect any CUDA version and fell through to the cu126 wheel default. Accept both spellings everywhere we parse nvidia-smi output: - install.ps1: Get-TorchIndexUrl regex now allows " UMD" - install.sh: two-expression sed (POSIX BRE has no "?"); the two patterns are mutually exclusive per line, head -1 picks the match - studio/setup.ps1: Get-PytorchCudaTag and the $DriverMaxCuda detector both relaxed - studio/install_llama_prebuilt.py: substring scan replaced with a regex search using the same pattern - tests/sh/test_get_torch_index_url.sh: new make_mock_smi_umd helper plus three UMD cases (13.3 -> cu130, 12.8 -> cu128, 11.8 -> cu118); all 30 tests pass locally * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 5 +++- install.sh | 10 +++++-- studio/install_llama_prebuilt.py | 18 ++++++++---- studio/setup.ps1 | 9 ++++-- tests/sh/test_get_torch_index_url.sh | 42 ++++++++++++++++++++++++++++ 5 files changed, 73 insertions(+), 11 deletions(-) diff --git a/install.ps1 b/install.ps1 index 3bb5fb3f8c..c4e8d8d522 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1253,7 +1253,10 @@ shell.Run cmd, 0, False if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = & $NvidiaSmiExe 2>&1 | Out-String - if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') { + # Newer NVIDIA drivers (e.g. 610.x on Windows) print + # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". + # Accept both spellings so we don't fall through to the cu126 default. + if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') { $major = [int]$Matches[1]; $minor = [int]$Matches[2] if ($major -ge 13) { return "$baseUrl/cu130" } if ($major -eq 12 -and $minor -ge 8) { return "$baseUrl/cu128" } diff --git a/install.sh b/install.sh index 0548ee2a57..a3b76a3011 100755 --- a/install.sh +++ b/install.sh @@ -1683,9 +1683,15 @@ get_torch_index_url() { fi echo "$_base/cpu"; return fi - # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) + # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). + # Newer NVIDIA drivers (e.g. 610.x) print "CUDA UMD Version: X.Y" instead + # of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions + # (POSIX sed does not support "?" without -E). The two patterns are + # mutually exclusive per line, so head -1 picks the first emitted match. _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ - | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ + | sed -n \ + -e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ + -e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ | head -1) if [ -z "$_cuda_ver" ]; then echo "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" >&2 diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 394a1c9cd8..b9d63ecee9 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2640,12 +2640,18 @@ def detect_host() -> HostInfo: try: result = run_capture([nvidia_smi], timeout = 20) merged = "\n".join(part for part in (result.stdout, result.stderr) if part) - for line in merged.splitlines(): - if "CUDA Version:" in line: - raw = line.split("CUDA Version:", 1)[1].strip().split()[0] - major, minor = raw.split(".", 1) - driver_cuda_version = (int(major), int(minor)) - break + # Newer NVIDIA drivers (e.g. 610.x on Windows) print + # "CUDA UMD Version: X.Y" instead of the legacy + # "CUDA Version: X.Y"; accept both spellings. + cuda_match = re.search( + r"CUDA(?: UMD)? Version:\s*(\d+)\.(\d+)", + merged, + ) + if cuda_match is not None: + driver_cuda_version = ( + int(cuda_match.group(1)), + int(cuda_match.group(2)), + ) except Exception: pass diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 16df87bbd5..8950627ae8 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -352,7 +352,10 @@ function Get-PytorchCudaTag { # string. Plain 2>$null doesn't fully suppress stderr in PS 5.1 -- # ErrorRecord objects leak into $output and break the -match. $output = & $smiExe 2>&1 | Out-String - if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') { + # Newer NVIDIA drivers (e.g. 610.x on Windows) print + # "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y". + # Accept both spellings so we don't fall through to the cu126 default. + if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') { $major = [int]$Matches[1] $minor = [int]$Matches[2] # PyTorch 2.10 offers: cu124, cu126, cu128, cu130 @@ -842,7 +845,9 @@ if ($HasNvidiaSmi) { $DriverMaxCuda = $null try { $smiOut = & $NvidiaSmiExe 2>&1 | Out-String - if ($smiOut -match "CUDA Version:\s+([\d]+)\.([\d]+)") { + # Newer NVIDIA drivers (e.g. 610.x) report the driver max CUDA as + # "CUDA UMD Version: X.Y" rather than "CUDA Version: X.Y"; accept both. + if ($smiOut -match "CUDA(?: UMD)? Version:\s+([\d]+)\.([\d]+)") { $DriverMaxCuda = "$($Matches[1]).$($Matches[2])" substep "driver supports up to CUDA $DriverMaxCuda" } diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 7235873f53..a9fafa5359 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -59,6 +59,29 @@ MOCK echo "$_dir" } +# Helper: create a mock nvidia-smi that prints the new "CUDA UMD Version" header +# layout used by newer NVIDIA drivers (e.g. 610.x on Windows). See issue #5812. +make_mock_smi_umd() { + _dir=$(mktemp -d) + cat > "$_dir/nvidia-smi" < official/cpu" "https://download.pytorch.org/whl/c _result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl/" run_func "none") assert_eq "trailing slash stripped -> mirror/cpu" "https://mirror.example.com/whl/cpu" "$_result" +# 29) "CUDA UMD Version: 13.3" header (newer NVIDIA driver layout, issue #5812) +# -> cu130, not the cu126 fallback. +_dir=$(make_mock_smi_umd "13.3") +_result=$(run_func "$_dir") +assert_eq "CUDA UMD Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result" +rm -rf "$_dir" + +# 30) "CUDA UMD Version: 12.8" header (newer layout on a 12.x driver) -> cu128 +_dir=$(make_mock_smi_umd "12.8") +_result=$(run_func "$_dir") +assert_eq "CUDA UMD Version 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + +# 31) "CUDA UMD Version: 11.8" header (newer layout on an older driver) -> cu118 +_dir=$(make_mock_smi_umd "11.8") +_result=$(run_func "$_dir") +assert_eq "CUDA UMD Version 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result" +rm -rf "$_dir" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR"