* 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>
This commit is contained in:
Daniel Han
GitHub
parent
015fa5772a
commit
a62eb80f7c
5 changed files with 73 additions and 11 deletions
|
|
@ -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" }
|
||||
|
|
|
|||
10
install.sh
10
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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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" <<MOCK
|
||||
#!/bin/sh
|
||||
case "\$1" in
|
||||
-L)
|
||||
echo "GPU 0: NVIDIA GeForce RTX 5090 Laptop GPU (UUID: GPU-fake-uuid)"
|
||||
;;
|
||||
*)
|
||||
cat <<'SMI_OUT'
|
||||
+-----------------------------------------------------------------------------------------+
|
||||
| NVIDIA-SMI 610.47 KMD Version: 610.47 CUDA UMD Version: $1 |
|
||||
+-----------------------------------------------------------------------------------------+
|
||||
SMI_OUT
|
||||
;;
|
||||
esac
|
||||
MOCK
|
||||
chmod +x "$_dir/nvidia-smi"
|
||||
echo "$_dir"
|
||||
}
|
||||
|
||||
# Helper: create a mock amd-smi that prints a given ROCm version string
|
||||
# Supports both "amd-smi version" and "amd-smi list" subcommands so that
|
||||
# the GPU presence check (amd-smi list) also succeeds in tests.
|
||||
|
|
@ -278,6 +301,25 @@ assert_eq "empty mirror env -> 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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue