refactor(rocm/win): switch to repo.amd.com arch-aware index, remove stubs
AMD recommends repo.amd.com/rocm/whl/{arch}/ as the Windows ROCm wheel
source. These wheels bundle their own ROCm runtime, support all Python
versions (not just cp312), and include the full torch._C extension set
(including _distributed_c10d) that the old repo.radeon.com wheel omitted.
Changes:
- install.ps1: remove Select-ROCmWheelRelease + hardcoded cp312 wheel
URLs; remove Python 3.12 forced-preference logic; install via
--index-url repo.amd.com/rocm/whl/{arch-family}/
- studio/setup.ps1: same -- remove Select-ROCmWheelRelease, switch to
repo.amd.com arch-aware index URL
- studio/install_python_stack.py: replace _ROCM_WINDOWS_RELEASES /
_select_windows_rocm_release with _windows_rocm_index_url() using the
_GFX_TO_AMD_INDEX_ARCH map; drop Python 3.12 restriction
- studio/backend/core/training/worker.py: remove all stub machinery
(_make_mod_stub, _StubSubpackageFinder, _StubSubpackageLoader,
_StubClassMeta, torchao/fsdp/dtensor stubs, _c10d_functional ops
stubs, BNB DLL detection) -- no longer needed with new wheel source
This commit is contained in:
LeoBorcherding
parent
32643198b2
commit
d731c5fc63
4 changed files with 85 additions and 613 deletions
195
install.ps1
195
install.ps1
|
|
@ -950,71 +950,13 @@ shell.Run cmd, 0, False
|
|||
return $null
|
||||
}
|
||||
|
||||
# ── Quick AMD GPU probe (before Python selection) ──
|
||||
# Checks hipinfo and WMI now so we can pick Python 3.12 upfront if AMD is
|
||||
# present (ROCm Windows wheels are cp312-only). The full GPU detection with
|
||||
# version strings and display labels runs after venv creation below.
|
||||
$_EarlyAmdDetected = $false
|
||||
try {
|
||||
$hipinfoEarly = Get-Command hipinfo -ErrorAction SilentlyContinue
|
||||
if ($hipinfoEarly) {
|
||||
$hipEarlyOut = & $hipinfoEarly.Source 2>&1 | Out-String
|
||||
if ($LASTEXITCODE -eq 0 -and $hipEarlyOut -match "(?i)gcnArchName") {
|
||||
$_EarlyAmdDetected = $true
|
||||
}
|
||||
}
|
||||
} catch {}
|
||||
if (-not $_EarlyAmdDetected) {
|
||||
try {
|
||||
$wmiGpuEarly = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue |
|
||||
Where-Object { $_.Name -match "AMD|Radeon" } | Select-Object -First 1
|
||||
if ($wmiGpuEarly) { $_EarlyAmdDetected = $true }
|
||||
} catch {}
|
||||
}
|
||||
|
||||
# ── Install Python if no compatible version (3.11-3.13) found ──
|
||||
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
|
||||
Write-TauriLog "STEP" "Installing Python"
|
||||
$DetectedPython = Find-CompatiblePython
|
||||
|
||||
# If AMD GPU is present and we didn't land on Python 3.12, find 3.12 now
|
||||
# before creating the venv -- avoids creating the environment twice.
|
||||
if ($_EarlyAmdDetected -and $DetectedPython -and (($DetectedPython.Version -split '\.')[0..1] -join '.') -ne "3.12") {
|
||||
$py312Pre = $null
|
||||
$pyLauncherPre = Get-Command py -CommandType Application -ErrorAction SilentlyContinue
|
||||
if ($pyLauncherPre -and $pyLauncherPre.Source -notmatch $script:CondaSkipPattern) {
|
||||
try {
|
||||
$out312 = & $pyLauncherPre.Source "-3.12" --version 2>&1 | Out-String
|
||||
if ($out312 -match "Python 3\.12\.\d+") {
|
||||
$resolvedExe312 = (& $pyLauncherPre.Source "-3.12" -c "import sys; print(sys.executable)" 2>$null | Out-String).Trim()
|
||||
if ($resolvedExe312 -and (Test-Path $resolvedExe312) -and -not (Test-IsCondaPython $resolvedExe312)) {
|
||||
$py312Pre = @{ Version = "3.12"; Path = $resolvedExe312 }
|
||||
}
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
if (-not $py312Pre) {
|
||||
foreach ($name312 in @("python3.12", "python3", "python")) {
|
||||
foreach ($cmd312 in @(Get-Command $name312 -All -ErrorAction SilentlyContinue)) {
|
||||
if (-not $cmd312.Source -or $cmd312.Source -like "*\WindowsApps\*") { continue }
|
||||
if (Test-IsCondaPython $cmd312.Source) { continue }
|
||||
try {
|
||||
$out312 = & $cmd312.Source --version 2>&1 | Out-String
|
||||
if ($out312 -match "Python 3\.12\.\d+") { $py312Pre = @{ Version = "3.12"; Path = $cmd312.Source }; break }
|
||||
} catch {}
|
||||
}
|
||||
if ($py312Pre) { break }
|
||||
}
|
||||
}
|
||||
if ($py312Pre) { $DetectedPython = $py312Pre }
|
||||
}
|
||||
|
||||
if ($DetectedPython) {
|
||||
$pyStepLabel = "Python $($DetectedPython.Version) already installed"
|
||||
if ($_EarlyAmdDetected -and (($DetectedPython.Version -split '\.')[0..1] -join '.') -eq "3.12") {
|
||||
$pyStepLabel = "Python 3.12 selected (ROCm wheels are cp312-only)"
|
||||
}
|
||||
step "python" $pyStepLabel
|
||||
step "python" "Python $($DetectedPython.Version) already installed"
|
||||
}
|
||||
if (-not $DetectedPython) {
|
||||
substep "installing Python ${PythonVersion}..."
|
||||
|
|
@ -1345,14 +1287,6 @@ shell.Run cmd, 0, False
|
|||
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow"
|
||||
}
|
||||
|
||||
# Warn if AMD GPU is present but the venv still isn't Python 3.12.
|
||||
# The early probe above covers the normal case; this fires only when the
|
||||
# full GPU detection reveals AMD that the early probe missed (e.g. hipinfo
|
||||
# not yet on PATH) and 3.12 still wasn't found.
|
||||
if (($HasROCm -or $ROCmGpuLabel) -and $DetectedPython -and (($DetectedPython.Version -split '\.')[0..1] -join '.') -ne "3.12") {
|
||||
substep "AMD GPU requires Python 3.12 for ROCm wheels -- install it from python.org and re-run." "Yellow"
|
||||
}
|
||||
|
||||
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
|
||||
# Mirrors Get-PytorchCudaTag in setup.ps1.
|
||||
function Get-TorchIndexUrl {
|
||||
|
|
@ -1379,101 +1313,41 @@ shell.Run cmd, 0, False
|
|||
# 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
|
||||
# ── AMD Windows ROCm: arch-aware pip index (repo.amd.com) ──
|
||||
# Wheels bundle their own ROCm runtime and support all Python versions.
|
||||
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
|
||||
$ROCmIndexUrl = $null
|
||||
if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
|
||||
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
|
||||
$archFamilyMap = @{
|
||||
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
|
||||
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
|
||||
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
|
||||
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
|
||||
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
|
||||
}
|
||||
$minVer = if ($GfxArch -and $archMin.ContainsKey($GfxArch)) {
|
||||
$archMin[$GfxArch]
|
||||
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
|
||||
if ($archFamily) {
|
||||
$ROCmIndexUrl = "$amdIndexBase/$archFamily/"
|
||||
$archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" }
|
||||
substep "$archLabel -- AMD repo.amd.com index selected" "Cyan"
|
||||
} elseif ($ROCmGfxArch) {
|
||||
substep "AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
|
||||
} 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 ──
|
||||
# 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" }
|
||||
$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 "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"
|
||||
substep "To enable ROCm training, reinstall with Python 3.12." "Yellow"
|
||||
substep "AMD GPU detected but arch unknown -- falling back to CPU-only PyTorch" "Yellow"
|
||||
}
|
||||
}
|
||||
|
||||
$TorchIndexFamily = Get-TauriTorchIndexFamily $(
|
||||
if ($ROCmTorchWheelUrl) { $ROCmWheelTag } else { $TorchIndexUrl }
|
||||
)
|
||||
if ($ROCmIndexUrl) {
|
||||
$TorchIndexFamily = "rocm"
|
||||
} else {
|
||||
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
|
||||
}
|
||||
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
|
||||
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
if (-not $SkipTorch -and -not $ROCmTorchWheelUrl -and $TorchIndexUrl -like "*/cpu") {
|
||||
if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") {
|
||||
Write-Host ""
|
||||
if ($HasROCm -or $ROCmGpuLabel) {
|
||||
substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow"
|
||||
|
|
@ -1551,20 +1425,13 @@ shell.Run cmd, 0, False
|
|||
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
|
||||
}
|
||||
}
|
||||
} elseif ($TorchIndexUrl -or $ROCmTorchWheelUrl) {
|
||||
} elseif ($TorchIndexUrl -or $ROCmIndexUrl) {
|
||||
if ($SkipTorch) {
|
||||
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
|
||||
} elseif ($ROCmTorchWheelUrl) {
|
||||
} elseif ($ROCmIndexUrl) {
|
||||
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
|
||||
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 }
|
||||
if ($tarballExit -ne 0) {
|
||||
Write-Host "[WARN] ROCm namespace tarball install failed (exit $tarballExit) -- continuing" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --no-deps @ROCmAllWheelUrls }
|
||||
substep "installing PyTorch from $ROCmIndexUrl..."
|
||||
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl torch torchvision torchaudio }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
|
|
|
|||
|
|
@ -1084,10 +1084,10 @@ def run_training_process(
|
|||
'Install for better performance: pip install "triton-windows<3.7"'
|
||||
)
|
||||
|
||||
# ── 1d. Ensure torch.distributed is importable before ML libs load ──
|
||||
# Windows ROCm wheel lacks torch._C._distributed_c10d. Pre-stub it to handle
|
||||
# both ImportError and lazy-load crashes from trl/transformers. The
|
||||
# `not in sys.modules` guard preserves a real NVIDIA implementation.
|
||||
# ── 1d. Ensure torch.distributed helper attrs are present ──
|
||||
# Single-GPU training never initialises the process group, so these helpers
|
||||
# are never called — but transformers/trl import them unconditionally at the
|
||||
# module level and crash when they're missing.
|
||||
import types as _types
|
||||
|
||||
_td_stubs = {
|
||||
|
|
@ -1099,290 +1099,23 @@ def run_training_process(
|
|||
"barrier": lambda: None,
|
||||
}
|
||||
|
||||
# Helper: build a ModuleType stub whose __getattr__ auto-creates child stubs.
|
||||
# __path__ is set to [] so Python treats the stub as a package — without it,
|
||||
# any attempt to import a submodule (e.g. "torch.distributed.tensor._foo")
|
||||
# raises "is not a package" because Python checks __path__ before looking
|
||||
# in sys.modules for the child.
|
||||
import importlib.machinery as _ilm
|
||||
import importlib.abc as _ilabc
|
||||
|
||||
_STUB_SENTINEL = object() # identity tag placed on every stub module
|
||||
|
||||
def _make_mod_stub(mod_name):
|
||||
m = _types.ModuleType(mod_name)
|
||||
m.__path__ = [] # marks this as a package to the import system
|
||||
m.__package__ = mod_name
|
||||
# _STUB_SENTINEL survives __spec__ being replaced by the import
|
||||
# machinery (which overwrites __spec__ with the spec returned by
|
||||
# find_spec, breaking the loader=None sentinel we used before).
|
||||
m._unsloth_stub = _STUB_SENTINEL
|
||||
# importlib.util.find_spec() raises ValueError when __spec__ is None.
|
||||
m.__spec__ = _ilm.ModuleSpec(mod_name, loader=None, is_package=True)
|
||||
def _ga(attr, _m=m, _n=mod_name):
|
||||
if attr.startswith("__"):
|
||||
raise AttributeError(attr)
|
||||
child_name = f"{_n}.{attr}"
|
||||
child = _make_mod_stub(child_name)
|
||||
sys.modules.setdefault(child_name, child)
|
||||
setattr(_m, attr, child)
|
||||
return child
|
||||
m.__getattr__ = _ga
|
||||
return m
|
||||
|
||||
class _StubSubpackageLoader(_ilabc.Loader):
|
||||
"""Creates a stub module for any subpackage of a stub package."""
|
||||
def __init__(self, mod_name):
|
||||
self._mod_name = mod_name
|
||||
def create_module(self, spec):
|
||||
return _make_mod_stub(self._mod_name)
|
||||
def exec_module(self, module):
|
||||
pass # stub is fully initialised in create_module
|
||||
|
||||
class _StubSubpackageFinder(_ilabc.MetaPathFinder):
|
||||
"""Auto-stubs any subpackage import whose parent is one of our stubs.
|
||||
|
||||
Uses _unsloth_stub sentinel on the module object — NOT __spec__.loader,
|
||||
which the import machinery overwrites with the loader from find_spec,
|
||||
breaking the loader=None check for second-level subpackages.
|
||||
"""
|
||||
def find_spec(self, fullname, path, target=None):
|
||||
if "." not in fullname:
|
||||
return None
|
||||
parent_name = fullname.rsplit(".", 1)[0]
|
||||
parent = sys.modules.get(parent_name)
|
||||
if parent is None:
|
||||
return None
|
||||
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
|
||||
return None # real installed module — don't intercept
|
||||
loader = _StubSubpackageLoader(fullname)
|
||||
return _ilm.ModuleSpec(fullname, loader, is_package=True)
|
||||
|
||||
sys.meta_path.append(_StubSubpackageFinder())
|
||||
|
||||
# Metaclass for stub *classes* so class-level attribute access works too.
|
||||
# e.g. torchao / distributed_c10d does ProcessGroup.BackendType.NCCL —
|
||||
# plain type() has no __getattr__ on the metaclass, so we need this to
|
||||
# avoid AttributeError on arbitrary class-level attribute access.
|
||||
#
|
||||
# We intentionally do NOT use a real enum.Enum here: the C++ BackendType
|
||||
# enum gains new members across PyTorch versions (XCCL was added in 2.6+)
|
||||
# and hard-coding the list means every new member causes another crash.
|
||||
# Instead _StubClassMeta auto-creates child stubs for any attr access, and
|
||||
# the __members__ safety net satisfies Enum-duck-typing checks in torchao.
|
||||
class _StubClassMeta(type):
|
||||
def __getattr__(cls, attr):
|
||||
if attr == "__members__":
|
||||
# torchao checks ProcessGroup.BackendType.__members__ (Enum
|
||||
# interface). Return an empty dict — we have no real members
|
||||
# to enumerate and the caller just iterates / checks membership.
|
||||
return {}
|
||||
if attr.startswith("__"):
|
||||
raise AttributeError(attr)
|
||||
# Auto-create a child stub for any member access (BackendType.NCCL,
|
||||
# BackendType.XCCL, BackendType.UNDEFINED, …). We cache it on the
|
||||
# class so repeated accesses return the same object (identity
|
||||
# comparisons stay consistent).
|
||||
child = _StubClassMeta(attr, (), {"__init__": lambda self, *a, **kw: None})
|
||||
setattr(cls, attr, child)
|
||||
return child
|
||||
|
||||
def _make_stub_class(name):
|
||||
return _StubClassMeta(name, (), {"__init__": lambda self, *a, **kw: None})
|
||||
|
||||
if sys.platform == "win32":
|
||||
# torchao is not supported on ROCm Windows and its import chain
|
||||
# transitively pulls in torch._C._distributed_c10d (absent from the
|
||||
# ROCm Windows wheel), causing cascading AttributeErrors. We don't
|
||||
# use torchao quantization (unsloth uses bitsandbytes), so stub the
|
||||
# entire package up-front. transformers falls back gracefully when
|
||||
# torchao is importable but empty.
|
||||
for _tao_name in (
|
||||
"torchao",
|
||||
"torchao.quantization",
|
||||
"torchao.dtypes",
|
||||
"torchao.float8",
|
||||
"torchao.utils",
|
||||
):
|
||||
if _tao_name not in sys.modules:
|
||||
sys.modules[_tao_name] = _make_mod_stub(_tao_name)
|
||||
|
||||
_c10d_key = "torch._C._distributed_c10d"
|
||||
if _c10d_key not in sys.modules: # guard: never overwrite real NVIDIA impl
|
||||
_c10d_stub = _types.ModuleType(_c10d_key)
|
||||
|
||||
# ROCm Windows wheels omit this C extension; auto-stub every
|
||||
# missing symbol so torch._dynamo's fsdp imports don't crash.
|
||||
def _c10d_stub_getattr(_attr):
|
||||
if _attr.startswith("__"):
|
||||
raise AttributeError(_attr)
|
||||
_cls = _make_stub_class(_attr)
|
||||
setattr(_c10d_stub, _attr, _cls)
|
||||
return _cls
|
||||
|
||||
_c10d_stub.__getattr__ = _c10d_stub_getattr
|
||||
sys.modules[_c10d_key] = _c10d_stub
|
||||
try:
|
||||
import torch._C as _torch_C_mod # C ext — always importable
|
||||
if not hasattr(_torch_C_mod, "_distributed_c10d"):
|
||||
_torch_C_mod._distributed_c10d = _c10d_stub
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Pre-register torch.distributed.fsdp submodules as stubs so
|
||||
# torch._dynamo's module-level fsdp import short-circuits before
|
||||
# the real package loads (it has a circular import on ROCm Windows).
|
||||
for _fsdp_name in (
|
||||
"torch.distributed.fsdp",
|
||||
"torch.distributed.fsdp._flat_param",
|
||||
"torch.distributed.fsdp._fully_shard",
|
||||
"torch.distributed.fsdp._fsdp_param_group",
|
||||
"torch.distributed.fsdp._common_utils",
|
||||
):
|
||||
if _fsdp_name not in sys.modules:
|
||||
sys.modules[_fsdp_name] = _make_mod_stub(_fsdp_name)
|
||||
|
||||
# torch._dynamo.trace_rules.get_torch_obj_rule_map() eagerly loads
|
||||
# torch.distributed.tensor, which in turn imports
|
||||
# torch.distributed._functional_collectives. That module registers
|
||||
# Meta kernels for ops in the _c10d_functional C++ namespace, but
|
||||
# that namespace only exists when torch._C._distributed_c10d (the
|
||||
# C extension absent from ROCm Windows wheels) has been loaded.
|
||||
# Without it the impl() call raises "operator does not exist".
|
||||
# Pre-stubbing these modules short-circuits the real import so
|
||||
# torch._dynamo gets empty stub objects instead of crashing.
|
||||
for _dist_name in (
|
||||
"torch.distributed._functional_collectives",
|
||||
"torch.distributed._functional_collectives_impl",
|
||||
"torch.distributed.tensor",
|
||||
"torch.distributed.tensor._ops",
|
||||
"torch.distributed.tensor._ops._conv_ops",
|
||||
"torch.distributed.tensor._dtensor_spec",
|
||||
"torch.distributed.tensor.placement_types",
|
||||
# torch.distributed._tensor is the canonical private package;
|
||||
# its __init__.py tries to re-export submodules from
|
||||
# torch.distributed.tensor (which we stubbed above), causing
|
||||
# "is not a package" errors. Stubbing _tensor directly
|
||||
# short-circuits that __init__ so torchao's
|
||||
# `from torch.distributed._tensor import DTensor` gets a stub.
|
||||
"torch.distributed._tensor",
|
||||
"torch.distributed._tensor.placement_types",
|
||||
"torch.distributed._tensor.api",
|
||||
):
|
||||
if _dist_name not in sys.modules:
|
||||
sys.modules[_dist_name] = _make_mod_stub(_dist_name)
|
||||
|
||||
try:
|
||||
import torch.distributed as _td
|
||||
|
||||
for _name, _stub in _td_stubs.items():
|
||||
if not hasattr(_td, _name):
|
||||
setattr(_td, _name, _stub)
|
||||
# Stub C-extension-backed class attrs (Store, ProcessGroup, …) that
|
||||
# the ROCm Windows wheel omits. __getattr__ checks sys.modules first
|
||||
# so it never intercepts real subpackage lookups as plain classes.
|
||||
if not hasattr(_td, "__getattr__"):
|
||||
def _td_getattr(_attr):
|
||||
if _attr.startswith("__"):
|
||||
raise AttributeError(_attr)
|
||||
_full = f"torch.distributed.{_attr}"
|
||||
if _full in sys.modules:
|
||||
_mod = sys.modules[_full]
|
||||
setattr(_td, _attr, _mod)
|
||||
return _mod
|
||||
_cls = _make_stub_class(_attr)
|
||||
setattr(_td, _attr, _cls)
|
||||
return _cls
|
||||
_td.__getattr__ = _td_getattr
|
||||
except Exception:
|
||||
_td_mock = _make_mod_stub("torch.distributed")
|
||||
_td_mock = _types.ModuleType("torch.distributed")
|
||||
for _name, _stub in _td_stubs.items():
|
||||
setattr(_td_mock, _name, _stub)
|
||||
sys.modules["torch.distributed"] = _td_mock
|
||||
if "torch._C._distributed_c10d" not in sys.modules:
|
||||
sys.modules["torch._C._distributed_c10d"] = _make_mod_stub("torch._C._distributed_c10d")
|
||||
try:
|
||||
import torch as _torch
|
||||
_torch.distributed = _td_mock
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── 1e. Stub torch.ops._c10d_functional (and c10d.scatter_) ──
|
||||
# torchao.dtypes.nf4tensor accesses these at *import time* as dict keys:
|
||||
# NF4_OPS_TABLE = {
|
||||
# torch.ops._c10d_functional.all_gather_into_tensor.default: ...,
|
||||
# torch.ops._c10d_functional.wait_tensor.default: ..., (decorator)
|
||||
# torch.ops.c10d.scatter_.default: ...,
|
||||
# }
|
||||
# The real _c10d_functional ops are registered by torch._C._distributed_c10d
|
||||
# (absent on ROCm Windows). We replace the whole namespace with a stub
|
||||
# whose op objects are hashable so dict-key usage doesn't crash.
|
||||
try:
|
||||
import torch as _torch_ops
|
||||
|
||||
class _C10dFunctionalOpDefault:
|
||||
"""Hashable stub for op.default — used as dict keys."""
|
||||
__slots__ = ("_name",)
|
||||
def __init__(self, name):
|
||||
self._name = name
|
||||
def __hash__(self):
|
||||
return hash(("_c10d_functional_stub", self._name))
|
||||
def __eq__(self, other):
|
||||
return (type(other) is _C10dFunctionalOpDefault
|
||||
and self._name == other._name)
|
||||
def __call__(self, *a, **kw):
|
||||
return a[0] if a else None
|
||||
def __repr__(self):
|
||||
return f"torch.ops._c10d_functional.{self._name}.default"
|
||||
|
||||
class _C10dFunctionalOp:
|
||||
"""Stub for a single _c10d_functional op (has a .default attr)."""
|
||||
__slots__ = ("_name", "default")
|
||||
def __init__(self, name):
|
||||
self._name = name
|
||||
self.default = _C10dFunctionalOpDefault(name)
|
||||
def __call__(self, *a, **kw):
|
||||
return self.default(*a, **kw)
|
||||
def __repr__(self):
|
||||
return f"torch.ops._c10d_functional.{self._name}"
|
||||
|
||||
class _C10dFunctionalNamespace:
|
||||
"""Drop-in for torch.ops._c10d_functional; auto-stubs every op."""
|
||||
def __getattr__(self, name):
|
||||
if name.startswith("_"):
|
||||
raise AttributeError(name)
|
||||
op = _C10dFunctionalOp(name)
|
||||
object.__setattr__(self, name, op)
|
||||
return op
|
||||
|
||||
_torch_ops.ops._c10d_functional = _C10dFunctionalNamespace()
|
||||
|
||||
# Also stub torch.ops.c10d.scatter_ if it's missing (same root cause).
|
||||
try:
|
||||
_ = _torch_ops.ops.c10d.scatter_.default
|
||||
except AttributeError:
|
||||
# c10d namespace exists but scatter_ op isn't registered; inject stub.
|
||||
_c10d_scatter_stub = _C10dFunctionalOp("scatter_")
|
||||
try:
|
||||
setattr(_torch_ops.ops.c10d, "scatter_", _c10d_scatter_stub)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── 1g. Point bitsandbytes at the ROCm 7.2 DLL on Windows ──
|
||||
# The AMD continuous-release wheel ships libbitsandbytes_rocm72.dll.
|
||||
# Only set BNB_ROCM_VERSION when that DLL is actually present — setting it
|
||||
# when the DLL is absent makes bnb fail harder than the default detection.
|
||||
if sys.platform == "win32" and os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
|
||||
import importlib.util as _ilu
|
||||
_bnb_spec = _ilu.find_spec("bitsandbytes")
|
||||
if _bnb_spec and _bnb_spec.origin:
|
||||
import pathlib as _pl
|
||||
_bnb_dll = _pl.Path(_bnb_spec.origin).parent / "libbitsandbytes_rocm72.dll"
|
||||
if _bnb_dll.exists():
|
||||
os.environ.setdefault("BNB_ROCM_VERSION", "72")
|
||||
|
||||
# ── 2. Now import ML libraries (fresh in this clean process) ──
|
||||
try:
|
||||
_send_status(event_queue, "Importing Unsloth...")
|
||||
|
|
|
|||
|
|
@ -69,38 +69,22 @@ _PYTORCH_WHL_BASE = (
|
|||
os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl"
|
||||
).rstrip("/")
|
||||
|
||||
# AMD Windows ROCm wheels — repo.radeon.com (cp312 only)
|
||||
_ROCM_WINDOWS_WHEEL_BASE = (
|
||||
# AMD Windows ROCm wheels — repo.amd.com (arch-specific pip index)
|
||||
# Format: https://repo.amd.com/rocm/whl/{arch_family}/
|
||||
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
|
||||
_ROCM_WINDOWS_INDEX_BASE = (
|
||||
os.environ.get("UNSLOTH_ROCM_WINDOWS_MIRROR")
|
||||
or "https://repo.radeon.com/rocm/windows"
|
||||
or "https://repo.amd.com/rocm/whl"
|
||||
).rstrip("/")
|
||||
# Maps (major, minor) → (release_folder, [wheel_filename, ...])
|
||||
_ROCM_WINDOWS_RELEASES: dict[tuple[int, int], tuple[str, list[str]]] = {
|
||||
(7, 2): (
|
||||
"rocm-rel-7.2.1",
|
||||
[
|
||||
# rocm tarball provides the 'rocm_sdk' Python namespace package
|
||||
"rocm-7.2.1.tar.gz",
|
||||
"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",
|
||||
],
|
||||
),
|
||||
(7, 1): (
|
||||
"rocm-rel-7.1.1",
|
||||
[
|
||||
# rocm tarball provides the 'rocm_sdk' Python namespace package
|
||||
"rocm-0.1.dev0.tar.gz",
|
||||
"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",
|
||||
],
|
||||
),
|
||||
|
||||
# Maps gfx arch → AMD index arch-family suffix.
|
||||
# Each family is a separate pip index on repo.amd.com.
|
||||
_GFX_TO_AMD_INDEX_ARCH: dict[str, str] = {
|
||||
"gfx1201": "gfx120X-all", "gfx1200": "gfx120X-all", # RDNA 4
|
||||
"gfx1151": "gfx1151", "gfx1150": "gfx1150", # RDNA 3.5 (Strix Halo/Point)
|
||||
"gfx1103": "gfx110X-all", "gfx1102": "gfx110X-all", # RDNA 3
|
||||
"gfx1101": "gfx110X-all", "gfx1100": "gfx110X-all",
|
||||
"gfx90a": "gfx90a", "gfx908": "gfx908", # MI200/MI100
|
||||
}
|
||||
|
||||
# bitsandbytes continuous-release_main wheels with the ROCm 4-bit GEMV fix
|
||||
|
|
@ -235,20 +219,6 @@ 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
|
||||
|
|
@ -272,17 +242,12 @@ def _detect_windows_gfx_arch() -> str | None:
|
|||
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 _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
|
||||
"""Return the AMD pip index URL for the given GPU arch, or None if unsupported."""
|
||||
arch_family = _GFX_TO_AMD_INDEX_ARCH.get(gfx_arch or "")
|
||||
if arch_family is None:
|
||||
return None
|
||||
return f"{_ROCM_WINDOWS_INDEX_BASE}/{arch_family}/"
|
||||
|
||||
|
||||
def _has_rocm_gpu() -> bool:
|
||||
|
|
@ -378,7 +343,7 @@ def _ensure_rocm_torch() -> None:
|
|||
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
|
||||
|
||||
On Linux x86_64: uses pytorch.org ROCm wheel index tags.
|
||||
On Windows (cp312 only): uses AMD's repo.radeon.com direct wheel releases.
|
||||
On Windows: uses AMD's repo.amd.com arch-specific pip index.
|
||||
No-op on macOS, non-x86_64 Linux, NVIDIA-primary hosts, or when torch
|
||||
already links against HIP.
|
||||
Uses pip_install() to respect uv, constraints, and --python targeting.
|
||||
|
|
@ -392,13 +357,6 @@ def _ensure_rocm_torch() -> None:
|
|||
return
|
||||
|
||||
if IS_WINDOWS:
|
||||
# AMD only publishes Windows ROCm wheels for Python 3.12 (cp312)
|
||||
if sys.version_info[:2] != (3, 12):
|
||||
print(
|
||||
f" ROCm torch on Windows requires Python 3.12 "
|
||||
f"(current: {sys.version_info[0]}.{sys.version_info[1]}) -- skipping"
|
||||
)
|
||||
return
|
||||
if _has_usable_nvidia_gpu():
|
||||
return
|
||||
gfx_arch = _detect_windows_gfx_arch()
|
||||
|
|
@ -425,34 +383,19 @@ def _ensure_rocm_torch() -> None:
|
|||
return # already ROCm torch
|
||||
except (OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
entry = _select_windows_rocm_release(gfx_arch)
|
||||
if entry is None:
|
||||
print(f" No AMD Windows torch wheel for GPU arch {gfx_arch} -- skipping")
|
||||
index_url = _windows_rocm_index_url(gfx_arch)
|
||||
if index_url is None:
|
||||
print(f" No AMD Windows torch index 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" {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,
|
||||
)
|
||||
print(f" {gfx_arch} (Windows) -- installing torch from {index_url}")
|
||||
pip_install(
|
||||
f"ROCm torch (Windows, {rel_tag})",
|
||||
f"ROCm torch (Windows, {gfx_arch})",
|
||||
"--force-reinstall",
|
||||
"--no-deps",
|
||||
*whl_urls,
|
||||
"--index-url", index_url,
|
||||
"torch", "torchvision", "torchaudio",
|
||||
constrain = False,
|
||||
)
|
||||
# bitsandbytes Windows ROCm wheel (ships libbitsandbytes_rocm72.dll).
|
||||
# BNB_ROCM_VERSION=72 is set in worker.py before the bnb import.
|
||||
# bitsandbytes Windows ROCm wheel.
|
||||
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
|
||||
if _bnb_win_url is not None:
|
||||
pip_install_try(
|
||||
|
|
|
|||
111
studio/setup.ps1
111
studio/setup.ps1
|
|
@ -1861,114 +1861,43 @@ if ($HasNvidiaSmi) {
|
|||
# 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 ──────────────────────────────────────────
|
||||
# Selects the newest wheel release compatible with the GPU arch (HIP SDK
|
||||
# version is irrelevant; wheels bundle their own ROCm runtime).
|
||||
$ROCmVersion = $script:ROCmVersion
|
||||
# Uses AMD's arch-specific pip index (repo.amd.com/rocm/whl/{arch}/).
|
||||
# Wheels bundle their own ROCm runtime; HIP SDK version is irrelevant.
|
||||
$ROCmGfxArch = $script:ROCmGfxArch
|
||||
$ROCmTorchWheelUrls = $null
|
||||
$ROCmTarballUrl = $null
|
||||
$ROCmWheelTag = $null
|
||||
$ROCmIndexUrl = $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") {
|
||||
$sel = Select-ROCmWheelRelease -GfxArch $ROCmGfxArch
|
||||
if ($sel) {
|
||||
$rb = "$amdWheelBase/$($sel.Rel)"
|
||||
$ROCmTarballUrl = "$rb/$($sel.Tarball)"
|
||||
$ROCmTorchWheelUrls = $sel.Wheels | ForEach-Object { "$rb/$_" }
|
||||
$ROCmWheelTag = $sel.Tag
|
||||
}
|
||||
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
|
||||
$archFamilyMap = @{
|
||||
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
|
||||
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
|
||||
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
|
||||
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
|
||||
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
|
||||
}
|
||||
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
|
||||
if ($archFamily) {
|
||||
$ROCmIndexUrl = "$amdIndexBase/$archFamily/"
|
||||
}
|
||||
}
|
||||
|
||||
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
|
||||
|
||||
if ($ROCmTorchWheelUrls) {
|
||||
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) {
|
||||
$tarballOut = Fast-Install --force-reinstall --no-deps $ROCmTarballUrl | Out-String
|
||||
if ($LASTEXITCODE -ne 0) {
|
||||
Write-Host "[WARN] ROCm namespace tarball install failed -- continuing" -ForegroundColor Yellow
|
||||
Write-Host $tarballOut -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
# Install remaining SDK + torch wheels using array splatting.
|
||||
$output = Fast-Install --force-reinstall --no-deps @ROCmTorchWheelUrls | Out-String
|
||||
if ($ROCmIndexUrl) {
|
||||
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
|
||||
$output = Fast-Install --force-reinstall --index-url $ROCmIndexUrl torch torchvision torchaudio | Out-String
|
||||
$torchInstallExit = $LASTEXITCODE
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow
|
||||
Write-Host $output -ForegroundColor Yellow
|
||||
$ROCmTorchWheelUrls = $null
|
||||
$ROCmIndexUrl = $null
|
||||
} else {
|
||||
# Tell install_python_stack.py to skip probe + suppress manual-install warning.
|
||||
$env:UNSLOTH_ROCM_TORCH_INSTALLED = "1"
|
||||
}
|
||||
}
|
||||
|
||||
if (-not $ROCmTorchWheelUrls -and $CuTag -eq "cpu") {
|
||||
if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
|
||||
substep "installing PyTorch (CPU-only)..."
|
||||
if ($script:UnslothVerbose) {
|
||||
Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/cpu"
|
||||
|
|
@ -1983,7 +1912,7 @@ if (-not $ROCmTorchWheelUrls -and $CuTag -eq "cpu") {
|
|||
Write-Host $output -ForegroundColor Red
|
||||
exit 1
|
||||
}
|
||||
} elseif (-not $ROCmTorchWheelUrls) {
|
||||
} elseif (-not $ROCmIndexUrl) {
|
||||
substep "installing PyTorch with CUDA support ($CuTag)..."
|
||||
substep "(This download is ~2.8 GB -- may take a few minutes)"
|
||||
if ($script:UnslothVerbose) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue