From d731c5fc6385a69c4c77cc7ea8e1a5f5c38bf1f0 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Thu, 14 May 2026 12:42:04 -0500 Subject: [PATCH] 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 --- install.ps1 | 195 +++-------------- studio/backend/core/training/worker.py | 277 +------------------------ studio/install_python_stack.py | 115 +++------- studio/setup.ps1 | 111 ++-------- 4 files changed, 85 insertions(+), 613 deletions(-) diff --git a/install.ps1 b/install.ps1 index b3dfed5d65..53c0eea4a1 100644 --- a/install.ps1 +++ b/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) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 246f7630c2..f6bebaaa0e 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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...") diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a2c15b4542..bdd80ef4b4 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index f1c7c10801..6519384a1a 100644 --- a/studio/setup.ps1 +++ b/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) {