From ea510b5936a9a85c14e0eeec420b2ad066012678 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Mon, 11 May 2026 04:35:20 -0500 Subject: [PATCH] feat(rocm/win): arch-aware wheel selector always picks newest ROCm release MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace HIP-SDK-version-gated wheel selection with GPU arch-based logic. Select-ROCmWheelRelease (PS) and _select_windows_rocm_release (Python) map gcnArchName → minimum ROCm version, then pick the newest available release that satisfies it (currently always rocm-rel-7.2.1 for any supported GPU). Wheels bundle their own ROCm runtime so the installed HIP SDK 7.1 does not prevent using 7.2.1 wheels on gfx1200 (RX 9060 XT) and similar RDNA 4 GPUs. Also installs the bitsandbytes Windows ROCm continuous-release wheel and sets BNB_ROCM_VERSION=72 in worker.py before ML imports so bnb loads the libbitsandbytes_rocm72.dll that ships in that wheel. --- install.ps1 | 118 +++++++++++++++++-------- studio/backend/core/training/worker.py | 6 ++ studio/install_python_stack.py | 106 ++++++++++++++++++---- studio/setup.ps1 | 103 +++++++++++++++------ 4 files changed, 251 insertions(+), 82 deletions(-) diff --git a/install.ps1 b/install.ps1 index dd18b047ea..b3dfed5d65 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1269,6 +1269,7 @@ shell.Run cmd, 0, False $HasROCm = $false $ROCmGpuLabel = $null $ROCmVersion = $null + $ROCmGfxArch = $null if (-not $HasNvidiaSmi) { $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue if ($hipinfoExe) { @@ -1277,7 +1278,8 @@ shell.Run cmd, 0, False if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { $HasROCm = $true if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") { - $ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))" + $ROCmGfxArch = $Matches[1].Trim() + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" } else { $ROCmGpuLabel = "AMD ROCm" } @@ -1373,46 +1375,90 @@ shell.Run cmd, 0, False } $TorchIndexUrl = Get-TorchIndexUrl + # ── GPU arch → newest compatible Windows ROCm wheel release ── + # Wheels bundle their own ROCm runtime; the installed HIP SDK version does + # not constrain which release to use. Always picks the newest release that + # supports the GPU architecture. + function Select-ROCmWheelRelease { + param([string]$GfxArch) + + # Available releases, newest first. + $releases = @( + @{ + Rel = "rocm-rel-7.2.1" + Tag = "rocm7.2" + RocmVer = @(7, 2) + Tarball = "rocm-7.2.1.tar.gz" + Wheels = @( + "rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", + "rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl", + "rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl", + "torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl", + "torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl", + "torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" + ) + }, + @{ + Rel = "rocm-rel-7.1.1" + Tag = "rocm7.1" + RocmVer = @(7, 1) + Tarball = "rocm-0.1.dev0.tar.gz" + Wheels = @( + "rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl", + "rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl", + "torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", + "torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", + "torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl" + ) + } + ) + + # GPU arch → minimum (major, minor) ROCm release needed. + $archMin = @{ + "gfx1201" = @(7,1); "gfx1200" = @(7,1) # RDNA 4 + "gfx1151" = @(7,1); "gfx1150" = @(7,1) # RDNA 3.5 (Strix Halo/Point) + "gfx1103" = @(6,4); "gfx1102" = @(6,4); "gfx1101" = @(6,4); "gfx1100" = @(6,4) # RDNA 3 + "gfx1036" = @(6,4); "gfx1035" = @(6,4); "gfx1034" = @(6,4); "gfx1033" = @(6,4) # RDNA 2 + "gfx1032" = @(6,4); "gfx1031" = @(6,4); "gfx1030" = @(6,4) + "gfx1011" = @(6,4); "gfx1010" = @(6,4) # RDNA 1 + "gfx906" = @(6,4); "gfx908" = @(6,4); "gfx90a" = @(6,4) # Vega/MI + } + $minVer = if ($GfxArch -and $archMin.ContainsKey($GfxArch)) { + $archMin[$GfxArch] + } else { + @(6, 4) # unknown arch: try the latest (7.2.1 supports all modern GPUs) + } + + foreach ($r in $releases) { + $rv = $r.RocmVer + $ok = ($rv[0] -gt $minVer[0]) -or ($rv[0] -eq $minVer[0] -and $rv[1] -ge $minVer[1]) + if ($ok) { return $r } + } + return $null + } + # ── AMD Windows ROCm wheel override ── - # When the HIP SDK is present and Python 3.12, use repo.radeon.com direct wheels. + # Selects the newest wheel release compatible with the GPU arch (HIP SDK + # version is irrelevant; wheels bundle their own ROCm runtime). $ROCmTorchWheelUrl = $null $ROCmTarballUrl = $null + $ROCmWheelTag = $null if ($HasROCm -and -not $SkipTorch) { $pyMajMin = if ($DetectedPython) { ($DetectedPython.Version -split '\.')[0..1] -join '.' } else { "" } if ($pyMajMin -eq "3.12") { $amdWheelBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.radeon.com/rocm/windows" } - if ($ROCmVersion -and $ROCmVersion -match '^7\.2') { - $amdRelBase = "$amdWheelBase/rocm-rel-7.2.1" - $ROCmTarballUrl = "$amdRelBase/rocm-7.2.1.tar.gz" # rocm_sdk namespace - $ROCmAllWheelUrls = @( - "$amdRelBase/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", - "$amdRelBase/rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl", - "$amdRelBase/rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl", - "$amdRelBase/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl", - "$amdRelBase/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl", - "$amdRelBase/torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" - ) - $ROCmTorchWheelUrl = $ROCmAllWheelUrls[3] - $TorchIndexUrl = $null - } elseif ($ROCmVersion -and $ROCmVersion -match '^7\.1') { - $amdRelBase = "$amdWheelBase/rocm-rel-7.1.1" - $ROCmTarballUrl = "$amdRelBase/rocm-0.1.dev0.tar.gz" # rocm_sdk namespace - $ROCmAllWheelUrls = @( - "$amdRelBase/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl", - "$amdRelBase/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl", - "$amdRelBase/torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", - "$amdRelBase/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", - "$amdRelBase/torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl" - ) - $ROCmTorchWheelUrl = $ROCmAllWheelUrls[2] - $TorchIndexUrl = $null - } - if ($ROCmTorchWheelUrl) { - substep "AMD ROCm $ROCmVersion (Python 3.12) -- AMD Windows torch wheel selected" "Cyan" - } elseif ($ROCmVersion) { - substep "No AMD Windows torch wheel for ROCm $ROCmVersion -- falling back to CPU-only PyTorch" "Yellow" + $sel = Select-ROCmWheelRelease -GfxArch $ROCmGfxArch + if ($sel) { + $rb = "$amdWheelBase/$($sel.Rel)" + $ROCmTarballUrl = "$rb/$($sel.Tarball)" + $ROCmAllWheelUrls = $sel.Wheels | ForEach-Object { "$rb/$_" } + $ROCmTorchWheelUrl = ($ROCmAllWheelUrls | Where-Object { $_ -match '/torch-' })[0] + $ROCmWheelTag = $sel.Tag + $TorchIndexUrl = $null + $archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" } + substep "$archLabel -- Windows torch wheel $($sel.Rel) selected" "Cyan" } else { - substep "ROCm version unknown -- falling back to CPU-only PyTorch" "Yellow" + substep "No AMD Windows torch wheel for GPU arch $ROCmGfxArch -- falling back to CPU-only PyTorch" "Yellow" } } else { substep "AMD Windows ROCm wheels require Python 3.12 (detected: $pyMajMin) -- using CPU-only PyTorch" "Yellow" @@ -1421,9 +1467,7 @@ shell.Run cmd, 0, False } $TorchIndexFamily = Get-TauriTorchIndexFamily $( - if ($ROCmTorchWheelUrl) { - if ($ROCmVersion -match '^7\.1') { "rocm7.1" } else { "rocm7.2" } - } else { $TorchIndexUrl } + if ($ROCmTorchWheelUrl) { $ROCmWheelTag } else { $TorchIndexUrl } ) $GpuBranch = Get-TauriGpuBranch $TorchIndexFamily Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version @@ -1512,7 +1556,7 @@ shell.Run cmd, 0, False substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmTorchWheelUrl) { Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" - substep "installing PyTorch (AMD ROCm $ROCmVersion)..." + substep "installing PyTorch ($ROCmWheelTag)..." # rocm_sdk namespace tarball (torch/_rocm_init.py imports it at startup) if ($ROCmTarballUrl) { $tarballExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-deps $ROCmTarballUrl } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index ee0d680f75..ad62ad459a 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1182,6 +1182,12 @@ def run_training_process( except Exception: pass + # ── 1e. Point bitsandbytes at the ROCm 7.2 DLL on Windows ── + # The AMD continuous-release wheel ships libbitsandbytes_rocm72.dll. + # BNB_ROCM_VERSION overrides the version string bnb uses to locate the DLL. + if sys.platform == "win32" and os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": + os.environ.setdefault("BNB_ROCM_VERSION", "72") + # ── 2. Now import ML libraries (fresh in this clean process) ── try: _send_status(event_queue, "Importing Unsloth...") diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 1e71ee8661..a2c15b4542 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -117,6 +117,13 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "download/continuous-release_main/" "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" ), + # Windows ROCm wheel — ships libbitsandbytes_rocm72.dll. + # BNB_ROCM_VERSION=72 must be set in the environment before importing bnb. + "win_amd64": ( + "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" + "download/continuous-release_main/" + "bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl" + ), } _BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1" @@ -228,6 +235,56 @@ def _detect_rocm_version() -> tuple[int, int] | None: return None +# GPU arch → minimum (major, minor) ROCm release that supports it on Windows. +# Wheels bundle their own ROCm runtime, so the installed HIP SDK version does +# not constrain selection — only the GPU's architecture minimum matters. +_GFX_MIN_ROCM_WINDOWS: dict[str, tuple[int, int]] = { + "gfx1201": (7, 1), "gfx1200": (7, 1), # RDNA 4 + "gfx1151": (7, 1), "gfx1150": (7, 1), # RDNA 3.5 (Strix Halo/Point) + "gfx1103": (6, 4), "gfx1102": (6, 4), "gfx1101": (6, 4), "gfx1100": (6, 4), # RDNA 3 + "gfx1036": (6, 4), "gfx1035": (6, 4), "gfx1034": (6, 4), "gfx1033": (6, 4), # RDNA 2 + "gfx1032": (6, 4), "gfx1031": (6, 4), "gfx1030": (6, 4), + "gfx1011": (6, 4), "gfx1010": (6, 4), # RDNA 1 + "gfx906": (6, 4), "gfx908": (6, 4), "gfx90a": (6, 4), # Vega/MI +} + + +def _detect_windows_gfx_arch() -> str | None: + """Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None.""" + import re + + hipinfo = shutil.which("hipinfo") + if not hipinfo: + return None + try: + result = subprocess.run( + [hipinfo], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 10, + ) + if result.returncode != 0: + return None + text = result.stdout.decode(errors = "replace") + m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) + return m.group(1).strip() if m else None + except Exception: + return None + + +def _select_windows_rocm_release(gfx_arch: str | None) -> tuple[str, list[str]] | None: + """Pick the best available Windows ROCm release for the given GPU arch. + + Always selects the newest available release whose ROCm version meets the + GPU's minimum requirement. Returns None when no release qualifies. + """ + min_ver = _GFX_MIN_ROCM_WINDOWS.get(gfx_arch or "", (6, 4)) + for (maj, mn), entry in sorted(_ROCM_WINDOWS_RELEASES.items(), reverse = True): + if (maj, mn) >= min_ver: + return entry + return None + + def _has_rocm_gpu() -> bool: """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" import re @@ -344,8 +401,9 @@ def _ensure_rocm_torch() -> None: return if _has_usable_nvidia_gpu(): return - if not _has_rocm_gpu(): - return + gfx_arch = _detect_windows_gfx_arch() + if not gfx_arch: + return # no AMD GPU visible via hipinfo try: probe = subprocess.run( [ @@ -367,34 +425,44 @@ def _ensure_rocm_torch() -> None: return # already ROCm torch except (OSError, subprocess.TimeoutExpired): pass - ver = _detect_rocm_version() - if ver is None: - print(" ROCm detected but version unreadable -- skipping torch reinstall") - return - entry = next( - ( - v - for (maj, mn), v in sorted(_ROCM_WINDOWS_RELEASES.items(), reverse = True) - if ver >= (maj, mn) - ), - None, - ) + entry = _select_windows_rocm_release(gfx_arch) if entry is None: - print( - f" No AMD Windows torch wheel for ROCm {ver[0]}.{ver[1]} -- skipping" - ) + print(f" No AMD Windows torch wheel for GPU arch {gfx_arch} -- skipping") return rel_tag, wheel_files = entry base = f"{_ROCM_WINDOWS_WHEEL_BASE}/{rel_tag}" wheel_urls = [f"{base}/{fn}" for fn in wheel_files] - print(f" ROCm {ver[0]}.{ver[1]} (Windows) -- installing torch from {base}/") + print(f" {gfx_arch} (Windows) -- installing torch from {base}/") + # Install rocm namespace tarball first (torch/_rocm_init.py imports it) + tarball_url = next((u for u in wheel_urls if u.endswith(".tar.gz")), None) + whl_urls = [u for u in wheel_urls if not u.endswith(".tar.gz")] + if tarball_url: + pip_install( + f"ROCm namespace ({rel_tag})", + "--force-reinstall", + "--no-deps", + tarball_url, + constrain = False, + ) pip_install( f"ROCm torch (Windows, {rel_tag})", "--force-reinstall", "--no-deps", - *wheel_urls, + *whl_urls, constrain = False, ) + # bitsandbytes Windows ROCm wheel (ships libbitsandbytes_rocm72.dll). + # BNB_ROCM_VERSION=72 is set in worker.py before the bnb import. + _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") + if _bnb_win_url is not None: + pip_install_try( + "bitsandbytes (AMD Windows, pre-release main)", + "--force-reinstall", + "--no-cache-dir", + "--no-deps", + _bnb_win_url, + constrain = False, + ) _rocm_windows_torch_installed = True return diff --git a/studio/setup.ps1 b/studio/setup.ps1 index d27e050526..f1c7c10801 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -664,6 +664,7 @@ if (-not $HasNvidiaSmi) { # ── AMD ROCm detection (Windows): probe hipinfo/amd-smi for actual GPU ── $HasROCm = $false $ROCmGpuLabel = $null +$script:ROCmGfxArch = $null if (-not $HasNvidiaSmi) { # hipinfo: present + output contains gcnArchName → real HIP GPU $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue @@ -673,7 +674,8 @@ if (-not $HasNvidiaSmi) { if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { $HasROCm = $true if ($hipOut -match "(?im)^\s*gcnArchName\s*:\s*(\S+)") { - $ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))" + $script:ROCmGfxArch = $Matches[1].Trim() + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" } else { $ROCmGpuLabel = "AMD ROCm" } @@ -1855,38 +1857,87 @@ if ($HasNvidiaSmi) { $CuTag = "cpu" } +# ── GPU arch → newest compatible Windows ROCm wheel release ── +# Wheels bundle their own ROCm runtime; the installed HIP SDK version does +# not constrain which release to use. Always picks the newest release that +# supports the GPU architecture. +function Select-ROCmWheelRelease { + param([string]$GfxArch) + + # Available releases, newest first. + $releases = @( + @{ + Rel = "rocm-rel-7.2.1" + Tag = "rocm7.2" + RocmVer = @(7, 2) + Tarball = "rocm-7.2.1.tar.gz" + Wheels = @( + "rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", + "rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl", + "rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl", + "torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl", + "torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl", + "torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" + ) + }, + @{ + Rel = "rocm-rel-7.1.1" + Tag = "rocm7.1" + RocmVer = @(7, 1) + Tarball = "rocm-0.1.dev0.tar.gz" + Wheels = @( + "rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl", + "rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl", + "torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", + "torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", + "torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl" + ) + } + ) + + # GPU arch → minimum (major, minor) ROCm release needed. + $archMin = @{ + "gfx1201" = @(7,1); "gfx1200" = @(7,1) # RDNA 4 + "gfx1151" = @(7,1); "gfx1150" = @(7,1) # RDNA 3.5 (Strix Halo/Point) + "gfx1103" = @(6,4); "gfx1102" = @(6,4); "gfx1101" = @(6,4); "gfx1100" = @(6,4) # RDNA 3 + "gfx1036" = @(6,4); "gfx1035" = @(6,4); "gfx1034" = @(6,4); "gfx1033" = @(6,4) # RDNA 2 + "gfx1032" = @(6,4); "gfx1031" = @(6,4); "gfx1030" = @(6,4) + "gfx1011" = @(6,4); "gfx1010" = @(6,4) # RDNA 1 + "gfx906" = @(6,4); "gfx908" = @(6,4); "gfx90a" = @(6,4) # Vega/MI + } + $minVer = if ($GfxArch -and $archMin.ContainsKey($GfxArch)) { + $archMin[$GfxArch] + } else { + @(6, 4) # unknown arch: try the latest (7.2.1 supports all modern GPUs) + } + + foreach ($r in $releases) { + $rv = $r.RocmVer + $ok = ($rv[0] -gt $minVer[0]) -or ($rv[0] -eq $minVer[0] -and $rv[1] -ge $minVer[1]) + if ($ok) { return $r } + } + return $null +} + # ── AMD Windows ROCm torch override ────────────────────────────────────────── -# When ROCm HIP SDK is present and Python 3.12 is in use, install AMD's direct -# torch wheels instead of CPU-only PyTorch. +# Selects the newest wheel release compatible with the GPU arch (HIP SDK +# version is irrelevant; wheels bundle their own ROCm runtime). $ROCmVersion = $script:ROCmVersion +$ROCmGfxArch = $script:ROCmGfxArch $ROCmTorchWheelUrls = $null $ROCmTarballUrl = $null +$ROCmWheelTag = $null if ($HasROCm -and $CuTag -eq "cpu") { $pyVer = (& python --version 2>&1 | Out-String) -replace '[^0-9.]','' $pyMajMin = ($pyVer.Trim() -split '\.')[0..1] -join '.' $amdWheelBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.radeon.com/rocm/windows" } - if ($pyMajMin -eq "3.12" -and $ROCmVersion) { - if ($ROCmVersion -match '^7\.2') { - $rb = "$amdWheelBase/rocm-rel-7.2.1" - $ROCmTarballUrl = "$rb/rocm-7.2.1.tar.gz" # rocm_sdk namespace - $ROCmTorchWheelUrls = @( - "$rb/rocm_sdk_core-7.2.1-py3-none-win_amd64.whl", - "$rb/rocm_sdk_devel-7.2.1-py3-none-win_amd64.whl", - "$rb/rocm_sdk_libraries_custom-7.2.1-py3-none-win_amd64.whl", - "$rb/torch-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl", - "$rb/torchvision-0.24.1+rocm7.2.1-cp312-cp312-win_amd64.whl", - "$rb/torchaudio-2.9.1+rocm7.2.1-cp312-cp312-win_amd64.whl" - ) - } elseif ($ROCmVersion -match '^7\.1') { - $rb = "$amdWheelBase/rocm-rel-7.1.1" - $ROCmTarballUrl = "$rb/rocm-0.1.dev0.tar.gz" # rocm_sdk namespace - $ROCmTorchWheelUrls = @( - "$rb/rocm_sdk_core-0.1.dev0-py3-none-win_amd64.whl", - "$rb/rocm_sdk_libraries_custom-0.1.dev0-py3-none-win_amd64.whl", - "$rb/torch-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", - "$rb/torchvision-0.24.0+rocmsdk20251116-cp312-cp312-win_amd64.whl", - "$rb/torchaudio-2.9.0+rocmsdk20251116-cp312-cp312-win_amd64.whl" - ) + if ($pyMajMin -eq "3.12") { + $sel = Select-ROCmWheelRelease -GfxArch $ROCmGfxArch + if ($sel) { + $rb = "$amdWheelBase/$($sel.Rel)" + $ROCmTarballUrl = "$rb/$($sel.Tarball)" + $ROCmTorchWheelUrls = $sel.Wheels | ForEach-Object { "$rb/$_" } + $ROCmWheelTag = $sel.Tag } } } @@ -1894,7 +1945,7 @@ if ($HasROCm -and $CuTag -eq "cpu") { $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } if ($ROCmTorchWheelUrls) { - substep "installing PyTorch (AMD ROCm $ROCmVersion)..." + substep "installing PyTorch ($ROCmWheelTag)..." # Install the rocm namespace tarball first (provides the 'rocm_sdk' Python # package that torch/_rocm_init.py imports at startup). if ($ROCmTarballUrl) {