feat(rocm/win): arch-aware wheel selector always picks newest ROCm release
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.
This commit is contained in:
LeoBorcherding
parent
73ae40c34b
commit
ea510b5936
4 changed files with 251 additions and 82 deletions
118
install.ps1
118
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 }
|
||||
|
|
|
|||
|
|
@ -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...")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
103
studio/setup.ps1
103
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) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue