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"