From b6d5636cc06d1f34716f8a0ccca1ea9b1f95792d Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Sat, 30 May 2026 00:29:56 -0500 Subject: [PATCH] fix/strix halo and windows AMD ROCm support (#5301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(studio): set HIP_VISIBLE_DEVICES in apply_gpu_ids for ROCm training workers Training workers are spawned via multiprocessing spawn before detect_hardware() runs, so IS_ROCM is still False. If the user never set HIP_VISIBLE_DEVICES in their shell, _inherits_rocm_visibility is also False, leaving the worker with only CUDA_VISIBLE_DEVICES set. On ROCm hosts the HIP runtime honors HIP_VISIBLE_DEVICES over CUDA_VISIBLE_DEVICES, so the worker saw the full device list and torch raised "no usable HIP accelerator" on some setups. Fall back to probing torch.version.hip (a build-time attribute, safe to read before GPU init) to detect ROCm when neither IS_ROCM nor inherited env vars are available. Mirrors the existing fix in llama_cpp.py for llama-server subprocess GPU pinning. Fixes https://github.com/unslothai/unsloth/issues/5180 * test: tighten apply_gpu_ids ROCm fallback assertions Replace loose OR chain with exact string matches, split into three focused tests, and add a guard check for the try/except wrapper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: detect ROCm unified memory (Strix Halo / AMD iGPU) via torch fallback amd-smi on iGPUs with shared/unified memory (e.g. Radeon 8060S on Strix Halo) reports only the dedicated VRAM slice (~512 MB) in its metric output, so get_visible_gpu_utilization() was returning usable_gb ≈ 0.35 GB instead of the full GTT pool (~128 GB). torch.cuda.mem_get_info() already surfaces the correct unified-pool size. Add _reconcile_rocm_unified_memory(): after amd-smi returns a valid result on a ROCm device, cross-check each device's vram_total_gb against torch.cuda.mem_get_info(). When torch reports a larger total, replace the amd-smi VRAM fields in-place. No-op for discrete AMD GPUs where the two sources agree. Fixes: "Falling back to all visible GPUs -- model may not fit" on AMD iGPU machines even when 100+ GB of unified memory is available. * Apply unified-memory reconciliation in get_gpu_utilization too The visible-GPU path was already corrected for AMD iGPUs with unified memory (Strix Halo / Radeon 8060S), but get_gpu_utilization was still returning the raw 512 MB amd-smi VRAM slice. Studio's /api/train/hardware endpoint and the live GPU monitor read from this primary path, so users continued seeing the wrong total even after auto_select_gpu_ids picked the right device. Refactor to share the per-device correction: * _apply_unified_memory_correction(metrics, torch_info) -- the actual replacement logic, in-place on a single metrics dict. * _reconcile_rocm_unified_memory(...) -- multi-device, iterates utilization["devices"] (visible-GPU path). * _reconcile_primary_rocm_unified_memory(...) -- single flat metrics dict (primary-GPU path), uses parent_visible_spec to pick the primary index, falls back to ordinal 0 when no visibility env is set. get_gpu_utilization now calls the primary reconciler under IS_ROCM, so both endpoints surface the real unified-memory pool on iGPUs while leaving discrete AMD GPUs untouched (torch_total <= smi_total -> no replace). * Use 'is not None' and log debug on torch.version.hip probe failures Two small follow-ups to the apply_gpu_ids ROCm fallback: 1. Match detect_hardware()'s 'getattr(torch.version, "hip", None) is not None' form so the entire codebase has one canonical 'this torch was built with HIP' check. On every shipping torch wheel hip is either None or a non-empty version string, so the new form agrees with the old bool() form on every real install. 2. Log the probe failure at debug level instead of swallowing it silently. The broad 'except Exception' is intentional (we never want apply_gpu_ids to crash a worker over a probe), but the silent pass made it impossible to tell whether the fallback was firing or being skipped. * fix(studio): honour HIP_VISIBLE_DEVICES in _get_parent_visible_gpu_spec before IS_ROCM is set When a user has HIP_VISIBLE_DEVICES set in their shell (e.g. "1" to select GPU 1) but detect_hardware() has not yet run in the Studio parent process, IS_ROCM is still False. _get_parent_visible_gpu_spec() was gated on IS_ROCM so it fell through to CUDA_VISIBLE_DEVICES (unset), saw all physical GPUs, and auto-selected index 0. apply_gpu_ids then overwrote HIP_VISIBLE_DEVICES with "0", making the intended GPU invisible to ROCm torch in the worker, which triggered the "no usable HIP accelerator" error (issue #5180). Apply the same _inherits_rocm_visibility pattern already used in apply_gpu_ids: check for HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in the environment regardless of IS_ROCM so the correct GPU index is preserved. * fix(install): harden AMD ROCm GPU detection for multi-GPU and env-filtered setups The previous rocminfo awk pattern could miss discrete GPUs on machines where HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES is used to mask an integrated GPU — the env vars filter rocminfo output but may not propagate into the install script subprocess, causing detection to fail entirely. Two changes: - Tighten rocminfo pattern from /gfx[0-9]/ && !/gfx000/ to /gfx[1-9][0-9]/ — simpler and correctly excludes the CPU agent (gfx000) without a negative lookahead - Add sysfs KFD topology fallback: reads /sys/class/kfd/kfd/topology/nodes/*/gpu_id which is a kernel-level view unaffected by HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES Fixes detection failure reported in Discord by Chains (gfx1201 + iGPU machine where env var exclusion of the iGPU caused rocminfo to return no usable device). * Fix KFD sysfs awk fallback to read properties file The fallback added by this PR reads /sys/class/kfd/kfd/topology/nodes/*/gpu_id files but matches the literal token 'gpu_id' against their content. Those files contain only a single decimal value (e.g. '0' for CPU agents, '50432' for GPU agents), so the regex never matches and 'found' stays 0, making the fallback a no-op on every host. The properties file in the same directory contains key/value lines like 'gpu_id 50432' which is what the existing awk pattern expects. Reproduced with a synthetic sysfs layout: against gpu_id files awk exits 1; against properties files awk exits 0 when any node reports gpu_id > 0. * fix(setup.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.sh setup.ps1 only checked nvidia-smi and fell straight to "gpu: none" on AMD machines. setup.sh already probed rocminfo/amd-smi/hipconfig/hipinfo. Add three-tier detection mirroring install_llama_prebuilt.py's detect_host(): 1. hipinfo: gcnArchName in output confirms a real HIP GPU (not just SDK) 2. amd-smi list: "GPU: " data rows as fallback 3. WMI Win32_VideoController: last resort -- detects AMD GPU even without HIP SDK, then guides user to install it rather than silently going CPU Also corrects the "none" message to mention AMD ROCm alongside NVIDIA so users with AMD hardware understand the requirement. Fixes: rohit-style install where Strix Halo (Radeon 8060S) showed "gpu: none" even with the HIP SDK present. * fix(install.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.ps1 install.ps1 had the same nvidia-smi-only GPU detection as setup.ps1 before the setup.ps1 fix. Applies the same three-tier AMD detection: 1. hipinfo: gcnArchName confirms real HIP GPU 2. amd-smi list: GPU data rows as fallback 3. WMI Win32_VideoController: detects AMD GPU without HIP SDK and guides user to install it Fixes: install.ps1 showing "gpu: none" while setup.ps1 correctly showed "AMD GPU detected" on the same machine (reported by rohit, RX 7600 XT). * fix(install.ps1): suppress 'No NVIDIA GPU detected' when AMD GPU is present * feat: add Windows AMD ROCm PyTorch wheel installation install_python_stack.py: - Add _ROCM_WINDOWS_WHEEL_BASE and _ROCM_WINDOWS_RELEASES constants pointing to AMD repo.radeon.com (ROCm 7.2 -> torch 2.9.1+rocm7.2.1) - Extend _ensure_rocm_torch() with a Windows branch: detects ROCm via _has_rocm_gpu() / _detect_rocm_version(), requires Python 3.12 (cp312 is the only ABI AMD publishes for Windows), installs the direct wheel URL from repo.radeon.com install.ps1: - Capture ROCmVersion during AMD detection via hipconfig --version / amd-smi version (needed for wheel URL selection) - After Get-TorchIndexUrl, add an AMD wheel override block: when HasROCm and Python 3.12 detected, set ROCmTorchWheelUrl to AMD wheel URL - Expand torch install branch to handle ROCmTorchWheelUrl with uv pip install --force-reinstall --no-cache-dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: also install torchvision and torchaudio from AMD Windows repo AMD publishes matching torchvision-0.24.1+rocm7.2.1 and torchaudio-2.9.1+rocm7.2.1 cp312 wheels at the same repo.radeon.com release folder. Install all three in both install.ps1 and install_python_stack.py Windows ROCm path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: add ROCm 7.1.1 Windows wheel mapping AMD uses a different version string for 7.1.1 wheels: 2.9.0+rocmsdk20251116 (date-tagged) instead of +rocm7.1.1. Adds the 7.1.1 release folder to both install.ps1 and install_python_stack.py so users with ROCm 7.1 get ROCm torch instead of falling back to CPU. * fix: install rocm_sdk_core and rocm_sdk_libraries_custom alongside torch The AMD Windows torch wheels declare rocm[libraries]== as a hard dependency. Without installing rocm_sdk_core and rocm_sdk_libraries_custom from the same AMD release folder, uv cannot resolve the dependency and fails with 'No solution found'. Include all 5 wheels in one install call. * fix: expand ROCm wheel array to scalars for Invoke-InstallCommand @array splatting inside a scriptblock only works when the native command is prefixed with '&'. Invoke-InstallCommand uses '& $Command' to run the block, so @ROCmAllWheelUrls was not being expanded. Extract to scalar variables $rw0-$rw4 which are captured correctly by the closure. * fix: use --no-deps for AMD Windows torch wheel install uv's resolver looks up rocm[libraries]==0.1.dev0 on PyPI during dependency resolution before downloading any wheels, and fails because the package doesn't exist on PyPI. --no-deps skips resolution entirely and installs all 5 AMD wheels directly. The GPU runtime dependency is satisfied by the HIP SDK, not a Python package. * fix: setup.ps1 and install_python_stack.py now install ROCm torch on Windows setup.ps1 was always setting CuTag='cpu' for non-NVIDIA hosts and installing cpu-only PyTorch, overwriting the ROCm torch installed by install.ps1. Adds the same AMD wheel selection logic (ROCm version detection, Python 3.12 check, 5-wheel install with --no-deps) to setup.ps1's torch install block. install_python_stack.py: remove IS_WINDOWS guard from _ensure_rocm_torch() call site so the Windows path in _ensure_rocm_torch() is reachable during 'unsloth studio update' as well. * fix: suppress manual-install warning when ROCm torch already present; fix progress counter - Gate the 'must be installed manually' warning on torch.version.hip being empty so it doesn't fire when our ROCm torch install succeeded - Update _TOTAL counter to include the 3 ROCm steps on Windows now that _ensure_rocm_torch() is called there (fixes 10/9 display) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: add rocm step display in setup.ps1; fix warning and progress counter - Add 'rocm' step after 'cuda' in setup.ps1 showing ROCm version or HIP SDK missing - Move ROCm version detection up to GPU detection block so it's available early - Suppress 'must be installed manually' warning when torch.version.hip is set - Fix _TOTAL counter to include ROCm steps on Windows (fixes 10/9 display) * fix: detect AMD SDK ROCm torch via __version__ when torch.version.hip is unset AMD's repo.radeon.com wheels (e.g. 2.9.0+rocmsdk20251116) do not set torch.version.hip, leaving it None. All three probes that relied solely on torch.version.hip now also check for 'rocm' in torch.__version__.lower(): - hardware.py detect_hardware(): IS_ROCM was never set, causing the studio to report 'Hardware detected: CPU' even after AMD wheels were installed and HIP DLLs were on PATH. - install_python_stack.py _ensure_rocm_torch(): skip-if-already-installed probe would always reinstall on subsequent runs. - install_python_stack.py Windows AMD warning: suppression check always failed, so the 'must be installed manually' note kept appearing after a successful AMD wheel install. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * perf: drop --no-cache-dir from AMD ROCm torch wheel installs uv caches downloaded wheels by default; passing --no-cache-dir forced a full redownload of the ~2 GB torch wheel on every install run. CUDA installs never had this flag -- AMD was the only path affected. * fix: use install-state flag instead of subprocess probe for AMD Windows warning Replace the subprocess torch probe in the post-install warning block with a module-level _rocm_windows_torch_installed flag set by _ensure_rocm_torch(). Subprocess re-import of torch is unnecessary and fragile -- the install function already knows whether it succeeded. * fix: hoist global declaration to top of _ensure_rocm_torch Python requires the global statement to appear before any assignment to the variable within a function. Moving it to the function top fixes the SyntaxError on line 354. * fix: pass AMD torch install status via env var to suppress false warning setup.ps1 now sets UNSLOTH_ROCM_TORCH_INSTALLED=1 after a successful AMD wheel install. install_python_stack.py reads this at the top of _ensure_rocm_torch() to skip both the subprocess probe and the warning -- no re-import of torch needed, and the warning message now correctly says 'could not be auto-installed' rather than 'must be installed manually'. * fix: register ROCm DLL directory before torch import on Windows Python 3.8+ ignores PATH for extension DLL loading on Windows; amdhip64.dll and other HIP runtime DLLs must be registered via os.add_dll_directory(). Without this, torch.cuda.is_available() always returns False on AMD ROCm Windows even when HIP_PATH is correctly set in system environment variables. Reads HIP_PATH / ROCM_PATH env vars first, then falls back to scanning common ROCm install roots (C:\Program Files\AMD\ROCm, F:\ROCm, C:\ROCm). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: remove hardcoded non-standard ROCm paths from DLL directory scan Only use HIP_PATH/ROCM_PATH (set by AMD installer) and the standard C:\Program Files\AMD\ROCm\\bin location. Custom drive paths like F:\ROCm are user-specific and should not be hardcoded. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: prevent torchao overrides step from overwriting AMD ROCm torch torchao==0.14.0 in overrides.txt declares torch as a dependency. Without --no-deps, uv resolves torch from PyPI and installs 2.11.0+cpu on top of the AMD ROCm wheels (2.9.0+rocmsdk20251116). This was the root cause of 'Hardware detected: CPU' -- the AMD wheels were installed but then immediately overwritten by the overrides step. When _rocm_windows_torch_installed is True, add --no-deps to the overrides pip_install call so torchao is installed without pulling in CPU torch. * fix: add rocm_sdk namespace tarball to Windows ROCm wheel installs torch/_rocm_init.py calls `import rocm_sdk` at startup, which requires the rocm namespace tarball (rocm-*.tar.gz) in addition to the SDK wheel packages. This tarball was missing from both install.ps1 and setup.ps1, causing ModuleNotFoundError on first torch import. - Add rocm-0.1.dev0.tar.gz to ROCm 7.1.1 install (provides rocm_sdk namespace) - Add rocm-7.2.1.tar.gz + rocm_sdk_devel to ROCm 7.2.1 install - Install tarball in a dedicated step before main SDK/torch wheels - Switch to @array splatting in install.ps1 scriptblock for dynamic wheel count - Remove --no-cache-dir from Python-side ROCm wheel install (prevents ~2GB redownload) * feat: enable ROCm 7.2 torch install + warn on gfx1151 with ROCm < 7.2 Chigoma333 (AMD Radeon 8060S / gfx1151, Strix Halo) confirmed that ROCm 7.1 segfaults when tensors are moved to GPU, but ROCm 7.2 + torch 2.11.0+rocm7.2 works fully including training. Changes: - Uncomment (7,2): "rocm7.2" in _ROCM_TORCH_INDEX (was blocked by <2.11.0) - Add _ROCM_TORCH_PKG_SPECS dict with per-tag version bounds: rocm7.2 → torch>=2.11.0,<2.12.0; all older tags → <2.11.0 - Add _detect_amd_gfx_codes() helper that parses rocminfo output - Warn on gfx1151/gfx1150 (Strix Halo) when ROCm < 7.2 is installed, pointing users at the known segfault and recommending upgrade - install.sh get_torch_index_url(): enable rocm7.2 case (previously capped to rocm7.1), cap unknown future tags to rocm7.2 - install.sh: override TORCH_CONSTRAINT to >=2.11.0,<2.12.0 when rocm7.2 index is selected, so pip can actually resolve torch 2.11.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: prefer Python 3.12 for AMD ROCm users when 3.13 is also installed After GPU detection, if ROCm HIP SDK is found and the selected Python is not 3.12, run a second pass to locate a 3.12 install via py.exe and PATH (catches uv-managed installs). Switch $DetectedPython to 3.12 so the venv is created with a compatible interpreter for the cp312-only AMD Windows torch wheels. NVIDIA and Intel GPU paths are unaffected -- the re-detection block only runs when $HasROCm is true. Fixes: #5301 * fix: also check uv-managed Python 3.12 for AMD ROCm #5301 * fix: hide amd-smi console popups on Windows, guard torch.distributed.is_initialized for ROCm #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: suppress remaining console popups on Windows, patch torch.distributed.is_initialized for ROCm #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: stub all missing torch.distributed attrs for ROCm Windows wheel #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: inject torch.distributed stub when C backend missing in ROCm Windows wheel #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm/windows): pre-stub torch._C._distributed_c10d + raise amd-smi timeout Two fixes for Windows ROCm regressions reported by electroglyph on #5301: 1. worker.py — torch.distributed stub now fires unconditionally on Windows The previous stub only injected sys.modules in the except branch, meaning it was silently skipped when `import torch.distributed` happened to succeed (the C backend is lazily resolved). The crash then hit later when transformers/trl triggered the lazy load. Fix: on win32 we pre-populate sys.modules['torch._C._distributed_c10d'] AND set the attribute on the torch._C extension module *before* attempting the import, covering both the early-ImportError and lazy-load failure modes. 2. amd.py — increase amd-smi timeout from 5 s to 30 s on Windows (10 s Linux) amd-smi on Windows must cold-init the ROCm runtime on first invocation; 5 s was consistently too short, producing repeated 'Command timed out' warnings in the server log. 30 s gives enough headroom without blocking indefinitely on broken installs. 3. install.ps1 — widen Python 3.12 enforcement to ROCmGpuLabel (WMI-only path) Users whose HIP SDK is not on PATH were detected via WMI but not switched to Python 3.12 before the install started, causing a second pass. Guard now fires on (HasROCm -or ROCmGpuLabel). * fix(rocm): guard c10d stub, fix TorchIndexFamily for 7.1, clean dead code + comments - worker.py: wrap c10d stub injection in `if _c10d_key not in sys.modules` so Windows NVIDIA users with a real torch.distributed are never affected - install.ps1: fix Get-TauriTorchIndexFamily receiving hardcoded "rocm7.2" even when ROCm 7.1 wheels are installed; now branches on $ROCmVersion - main.py: remove dead `import ctypes as _ctypes` (ctypes is never called) - hardware.py, install_python_stack.py, worker.py, install.ps1: shorten verbose multi-line comment blocks throughout - tests: update 4 stale assertions that expected rocm7.2 to be absent/capped * fix(tests): match windows AMD warning assertion to actual source string * chore: trim verbose comment blocks across all ROCm-related files * fix: guard reconcile call against None numeric_ids; add torchvision lower bounds * fix(install.ps1): recreate venv with Python 3.12 after ROCm switch Venv was created with 3.13 before GPU detection ran; switching $DetectedPython to 3.12 had no effect since $VenvPython still pointed to the 3.13 interpreter inside the already-created venv. * ux: detect AMD GPU before Python selection to avoid double venv creation - Early hipinfo + WMI probe runs before Find-CompatiblePython so Python 3.12 is selected upfront when AMD is detected; venv is now created exactly once instead of 3.13 then immediately 3.12. - Post-venv recreation block replaced with a simple warning for the rare case where AMD was missed by the early probe. - setup.ps1: show venv's actual Python version (e.g. 3.12) instead of the system Python found by the pre-activation search (was showing 3.13). * fix(rocm/win): auto-stub all _distributed_c10d symbols via PEP-562 __getattr__ The bare ModuleType stub caused ImportError when torch._dynamo was imported (triggered by trainer.py accessing torch._dynamo.config at load time). torch._dynamo pulls in torch.distributed.fsdp._flat_param which does: from torch._C._distributed_c10d import FakeProcessGroup and potentially other symbols. Adding module __getattr__ auto-creates a stub class for any missing symbol so all such imports succeed without enumerating every individual symbol. Applied to both the primary stub and the fallback stub in the except branch. * chore: trim c10d stub comment * fix(rocm/win): auto-stub missing torch.distributed attrs (Store, ProcessGroup, …) * fix(rocm/win): pre-stub fsdp submodules in sys.modules; fix __getattr__ subpackage clash * 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. * fix(rocm/win): stub class metaclass for ProcessGroup.BackendType; amd-smi circuit breaker torchao.float8.inference accesses ProcessGroup.BackendType as a class-level attribute. Plain type() stubs have no __getattr__ on the metaclass so this raises AttributeError. Introduce _StubClassMeta whose __getattr__ returns child stub classes, fixing the torchao import chain. Add an amd-smi circuit breaker in amd.py: after 3 consecutive failures the module stops spawning the process, eliminating the repeated Windows UAC / DiskPart elevation prompts caused by polling a non-functional amd-smi. Also guard BNB_ROCM_VERSION=72 behind a DLL existence check so bitsandbytes fails with its own detection message rather than a harder "DLL not found" when the Windows ROCm bnb wheel is not yet installed. * fix: stub __members__ so torchao float8 enum check doesn't crash on ROCm Windows torchao.float8.inference accesses ProcessGroup.BackendType.__members__ expecting a Python Enum registry dict. _StubClassMeta.__getattr__ was blocking all dunder attributes, causing AttributeError. Return {} for __members__ specifically so the isinstance/iteration checks pass cleanly. * fix: stub distributed tensor/functional_collectives to prevent missing C++ op crash on ROCm Windows torch._dynamo.trace_rules eagerly loads torch.distributed.tensor at import time, which pulls in _functional_collectives.py. That file registers Meta kernels for _c10d_functional C++ ops, but those ops are only registered by torch._C._distributed_c10d — a C extension absent from ROCm Windows wheels. Pre-stubbing the affected modules in sys.modules prevents the real import chain from running and avoids the "operator does not exist" crash. * fix: give mod stubs __path__ and pre-stub _tensor to fix 'not a package' import error _make_mod_stub now sets __path__=[] so Python treats stub modules as packages. Without it, any import of a submodule raises "is not a package". Also pre-stub torch.distributed._tensor and its submodules so that _tensor/__init__.py (which re-exports from torch.distributed.tensor) never runs and torchao's `from torch.distributed._tensor import DTensor` gets a harmless stub instead of crashing. * fix: stub torch.ops._c10d_functional namespace with hashable op sentinels torchao.dtypes.nf4tensor uses _c10d_functional ops as dict keys at import time (all_gather_into_tensor.default, wait_tensor.default) and torch.ops.c10d.scatter_.default. None of these ops are registered on ROCm Windows because torch._C._distributed_c10d (the C extension) doesn't ship. Replace the whole _c10d_functional namespace with a custom stub whose ops return hashable .default objects, so dict-key construction doesn't crash. Also inject a scatter_ stub into torch.ops.c10d if it's missing. * fix: stub entire torchao package on ROCm Windows instead of individual ops torchao is not supported on ROCm Windows and its import chain transitively requires torch._C._distributed_c10d (absent from the ROCm Windows wheel). Rather than stub each missing op one by one, stub the whole torchao package upfront. Unsloth uses bitsandbytes for quantization, not torchao, so this has no functional impact. transformers gracefully handles an importable-but- empty torchao by disabling TorchAoHfQuantizer. * fix: set __spec__ on mod stubs so importlib.util.find_spec doesn't raise Manually-injected sys.modules entries have __spec__=None by default. importlib.util.find_spec() raises ValueError when it finds a module in sys.modules with __spec__=None (transformers.utils.import_utils hits this when checking if torchao is available). Give every stub a minimal ModuleSpec(name, loader=None, is_package=True) to satisfy find_spec. * fix: add meta path finder to auto-stub subpackages of stub modules `import torchao.prototype` goes through the import machinery, not __getattr__, so an empty __path__ means ModuleNotFoundError. Rather than list every submodule explicitly, register a MetaPathFinder that intercepts any import whose parent is one of our stubs (detected by loader=None in the parent's ModuleSpec). Real installed packages always have a SourceFileLoader so they are never intercepted. Also register child stubs in sys.modules from __getattr__ as a belt-and-suspenders measure. * fix: use _unsloth_stub sentinel instead of loader=None for stub detection The import machinery overwrites module.__spec__ with the spec returned by find_spec (which has loader=_StubSubpackageLoader, not None), so the loader=None check broke for second-level subpackages. Switch to a custom _unsloth_stub object identity sentinel set directly on each stub module -- it survives __spec__ being replaced and correctly identifies stubs at any depth (torchao.prototype.safetensors, etc.). * 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 * fix(rocm/win): restore _distributed_c10d + torchao stubs; fix BNB install repo.amd.com torch wheels also omit torch._C._distributed_c10d on Windows (RCCL is not shipped on Windows). torch/distributed/__init__.py imports from it unconditionally at module level, so the stub must land in sys.modules before any torch.distributed import. torchao (pulled in by transformers.quantizers) walks torchao.float8.distributed_utils -> torch.distributed._functional_collectives -> distributed_c10d at import time. Stubbing torchao up-front short-circuits that chain. worker.py: - Restore _make_mod_stub / _StubSubpackageFinder / _StubSubpackageLoader - Restore _StubClassMeta for ProcessGroup.BackendType attribute access - Restore _distributed_c10d stub with __getattr__ (Windows only) - Restore torchao stubs (5 modules, Windows only) install_python_stack.py: - BNB AMD wheel install was inside the early-return branch that fires when torch is already a ROCm build (installed by install.ps1). Move BNB install outside that branch so it always runs on Windows ROCm — the PyPI bitsandbytes has only CUDA DLLs and fails to load on ROCm. * worker: remove _distributed_c10d stub; stub only torchao The installed torch/distributed/__init__.py from repo.amd.com (torch==2.10.0+rocm7.12.0) is now properly guarded with `if is_available():`, so `import torch.distributed` alone is safe. The crash only comes via torchao's import chain: torchao.float8.distributed_utils → torch.distributed._functional_collectives (unguarded import) → torch.distributed.distributed_c10d → torch._C._distributed_c10d ← absent on Windows ROCm Stubbing torchao short-circuits the chain entirely. No need to stub _distributed_c10d. Remove _StubClassMeta and the _c10d stub block; keep only _make_mod_stub + _StubSubpackageFinder + torchao seeds. * fix: BNB AMD wheel skipped + torch.compile segfault on Windows ROCm install_python_stack.py: the UNSLOTH_ROCM_TORCH_INSTALLED=1 early-return path (set by setup.ps1 when it installed torch itself) returned before ever reaching the AMD BNB prerelease wheel install. The PyPI bitsandbytes==0.49.x ships only CUDA DLLs, so loading it on ROCm fails with "libbitsandbytes_rocm72.dll not found". Now installs the AMD Windows BNB wheel before returning on that path too. worker.py: torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer, 0xC0000005) when torch.compile's JitDecomp system dispatches it during the first forward pass. Detect Windows ROCm via torch.version.hip (already in sys.modules from section 1e) and set TORCHDYNAMO_DISABLE=1 to bypass the broken kernel dispatch. * fix: BNB AMD wheel install fails uv wheel filename check The bitsandbytes continuous-release wheel is intentionally mismatched: filename encodes 1.33.7.preview (= 1.33.7rc0 in PEP 440) but wheel metadata reports 0.50.0.dev0. uv rejects this by default. Introduce _install_bnb_windows_rocm() helper that sets UV_SKIP_WHEEL_FILENAME_CHECK=1 only for this specific install, then restores the previous env value. Both BNB install call sites (the UNSLOTH_ROCM_TORCH_INSTALLED early-return path and the normal Windows ROCm path) now use this helper. * worker: patch _grouped_mm CUDA dispatch on Windows ROCm (gfx1200 null kernel) TORCHDYNAMO_DISABLE=1 stopped the compiler frontend but not the autograd JitDecomp system, which also dispatches _grouped_mm and hits the same null HIP kernel crash (0xC0000005). Verified that torch.library.Library("aten","IMPL").impl("_grouped_mm", fn, "CUDA") successfully overrides the broken HIP kernel with a Python mm fallback on torch==2.10.0+rocm7.12.0. Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor The fallback handles both the simple case (offs=None → torch.mm) and the grouped case (offs provided → split self by offsets, multiply each group against the corresponding slice of mat2, then cat results). Keep _WINDOWS_ROCM_GROUPED_MM_LIB alive at function scope to prevent the C++ dispatch registration from being freed by GC. * worker: fix torchao stub — return stub classes not modules for isinstance() peft/tuners/lora/torchao.py does: from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) The stub __getattr__ was returning stub modules, which isinstance() rejects with "arg 2 must be a type, a tuple of types, or a union". Add _StubTypeMeta metaclass whose __instancecheck__ always returns False, and _make_stub_type() to create stub classes via it. Change _make_mod_stub __getattr__ to return stub classes instead of stub modules for leaf attribute access, so isinstance() gets a valid type and returns False. _StubSubpackageFinder still handles import-style subpackage creation (those still need module objects in sys.modules); __getattr__ only fires for from-import or direct attribute access, which are the isinstance paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: add coverage for Windows ROCm install paths and worker patches Add conftest.py to fix pre-existing sys.path issue that prevented test_rocm_support.py from running at all (install_python_stack.py imports from backend.utils.wheel_utils which needs studio/ on sys.path). New test classes cover everything added in this session: - TestWindowsRocmIndexUrl: arch → AMD pip index URL mapping (gfx120X-all, gfx1151, gfx1150, gfx110X-all, unknown → None, trailing slash) - TestDetectWindowsGfxArch: hipinfo output parsing, missing/timeout/bad returncode/no-gcnArchName paths - TestInstallBnbWindowsRocm: UV_SKIP_WHEEL_FILENAME_CHECK set+restored, env restored on exception, no-op when URL missing - TestRocmTorchInstalledEnvVar: UNSLOTH_ROCM_TORCH_INSTALLED=1 skips pip_install, calls _install_bnb_windows_rocm, sets flag - TestWorkerWindowsRocmPatches: _grouped_mm CUDA dispatch override, offs/grouped variant handling, GC-prevention sentinel, _StubTypeMeta __instancecheck__, _StubSubpackageFinder registration, torchao key submodule pre-stubbing, TORCHDYNAMO_DISABLE guard - TestRocmTorchPkgSpecs: rocm7.2 torch 2.11.x spec, default <2.11 cap, 3-tuple shape, _GFX_TO_AMD_INDEX_ARCH RDNA4/3.5/3 coverage * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: fix encoding, IS_WINDOWS patching, and wrong assertion - Add encoding="utf-8" to all read_text() calls (54 occurrences) so tests pass on Windows where the default codec is cp1252 and source files contain UTF-8 emoji (e.g. ⚠️ in install_python_stack.py) - Add @patch.object(stack_mod, "IS_WINDOWS", False) to Linux-path TestEnsureRocmTorch tests so they reach the Linux code path when run on a Windows machine instead of short-circuiting into the Windows branch - Fix test_grouped_mm_patch_guarded_by_windows_and_hip_check: the source uses getattr(_torch_for_rocm, "version", None) not torch.version, so check for '"version"' and '"hip"' substrings instead 137 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: pin BNB_ROCM_VERSION=72 for torch==2.11.0+rocm7.13.0 compatibility AMD's pip index now ships torch==2.11.0+rocm7.13.0 (ROCm 7.13). bitsandbytes auto-detects HIP 7.13 from torch.version.hip and looks for libbitsandbytes_rocm713.dll, which the AMD Windows prerelease wheel does not ship (it only ships rocm72.dll), causing a load error at training start. Fix: - worker.py section 1f: set BNB_ROCM_VERSION=72 (via setdefault) before section 2 ML imports, so bitsandbytes always loads rocm72.dll on Windows ROCm - install_python_stack.py: set BNB_ROCM_VERSION=72 in _install_bnb_windows_rocm() for any post-install imports; update comment to document root cause - tests: 4 new assertions covering the fix (141 passed, 2 skipped) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: detect BNB ROCm DLL suffix dynamically instead of hardcoding '72' BNB_ROCM_VERSION was pinned to '72' which works today (AMD wheel ships rocm72.dll) but would break again if AMD ships a future wheel with a different DLL suffix (e.g. rocm713.dll). Add _detect_bnb_rocm_dll_ver() to install_python_stack.py: scans the installed bitsandbytes package dir for libbitsandbytes_rocm{VER}.dll using importlib.util.find_spec (no BNB import needed) and returns the suffix. '72' remains the fallback when detection fails. Apply the same detection inline in worker.py section 1f. Both paths still respect a pre-set BNB_ROCM_VERSION (caller override wins). Tests: +8 cases covering detection logic and fallback (147 passed, 2 skipped). * fix: patch torch.distributed stubs in server process for Windows ROCm On Windows ROCm, torch.distributed ships without process-group helpers (is_initialized, is_available, get_rank, get_world_size). The worker subprocess already patches these in section 1e, but the main server process calls _determine_attention_impl_for_gpu_estimate() which calls unsloth's resolve_attention_implementation() → is_initialized(), causing: "Could not resolve attention implementation for '...': module 'torch.distributed' has no attribute 'is_initialized'" Fix: patch the missing attrs onto torch.distributed at the top of _determine_attention_impl_for_gpu_estimate, matching the same stubs already applied in worker.py section 1e. No-ops on Linux/CUDA where torch.distributed is fully populated. * fix: gate _grouped_mm dispatch patch on HIP < 7.13 AMD fixed the gfx1200 null HIP kernel in ROCm 7.13 (torch 2.11+). Users on the new wheel now get the real GPU _grouped_mm kernel for MoE workloads instead of the Python mm fallback. Changes: - worker.py: add _hip_ver_at_least() helper; wrap full _grouped_mm patch in `if not _hip_ver_at_least(7, 13):` with else branch that logs the skip reason; update section-1f comment to document the fix - test_rocm_support.py: add 5 tests covering the helper definition, the (7, 13) gate expression, the else branch, the skip log message, and the AMD-format version string parsing (.split(".")[:2]) Verified: torch==2.11.0+rocm7.13.0 — 3D batch and grouped (offs) variants both succeed; null crash only present on rocm7.12 and earlier. * fix: stub is_torchelastic_launched on torch.distributed for Windows ROCm resolve_attention_implementation calls is_torchelastic_launched() which does not exist in the incomplete torch.distributed shipped with the Windows ROCm wheel, causing a warning on every model config load in the server process. Add it to the stub table alongside the four helpers already patched in _determine_attention_impl_for_gpu_estimate. Also adds two tests: one confirming the new stub and one confirming all five core distributed helpers are covered. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: explicit warnings on AMD ROCm arch/version fallbacks + Fast-Install arg order setup.ps1: - Fix Fast-Install argument order: packages before flags, consistent with all other Fast-Install calls in the file (was: Fast-Install --force-reinstall --index-url $url torch ...) (now: Fast-Install torch torchvision torchaudio --force-reinstall --index-url $url) - Add explicit [WARN] substep when $HasROCm is true but arch mapping fails: - GPU arch detected but not in supported wheel list → names the arch and lists supported families so user knows exactly what to report - HIP SDK present (amd-smi path) but gcnArchName unreadable → instructs user to re-install the HIP SDK; previously fell back silently to CPU install.sh: - Add [WARN] to stderr before silent CPU fallback when AMD GPU is confirmed (rocminfo/amd-smi) but ROCm version cannot be read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, dpkg, rpm) - Add [WARN] to stderr when ROCm version is too old (< 6.0) with upgrade link install.ps1 and setup.sh: no changes needed (already handle these paths correctly) * fix: robust gfx arch detection for Strix Halo / HIP-runtime-only installs Covers users who have the HIP runtime (amd-smi available) but not the full HIP SDK (no hipinfo), which is common on Strix Halo iGPU systems. Without this, $ROCmGfxArch stays null and the installer silently falls back to CPU-only PyTorch despite a working GPU. Detection waterfall (setup.ps1 + install.ps1): 1. hipinfo gcnArchName -- full HIP SDK (existing, unchanged) 2. amd-smi list gfx pattern -- newer amd-smi versions embed arch 3. amd-smi static --asic -- ROCm 6+ ASIC details with GFX target 4. UNSLOTH_ROCM_GFX_ARCH env -- manual override escape hatch 5. GPU name → arch table -- best-effort from marketing name: 890M / Strix Halo → gfx1151 (RDNA 3.5 iGPU, Strix Halo) 880M / Strix Point → gfx1150 (RDNA 3.5 iGPU, Strix Point) 780M / Phoenix → gfx1103 (RDNA 3 iGPU) RX 7900/7800/7700 → gfx1100 (RDNA 3 desktop) RX 9070 XT / 9080 → gfx1201 (RDNA 4) RX 9070 / 9060 XT → gfx1200 (RDNA 4) When arch is inferred from name, a Cyan substep tells the user to set UNSLOTH_ROCM_GFX_ARCH to skip inference on future installs. WMI block intentionally does not set $HasROCm (no runtime confirmation). Tests: 11 new tests in TestStrixHaloGfxArchDetection covering all five detection levels, WMI safety, and gfx regex in both ps1 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: resolve hipinfo/hipconfig via HIP_PATH/ROCM_PATH when not on PATH AMD HIP SDK sets HIP_PATH on Windows but does not always add the bin directory to PATH. Get-Command hipinfo therefore silently fails and detection falls through to WMI, which cannot provide a gfx arch, leaving the user with a CPU-only PyTorch install and no warning. Changes: - setup.ps1 / install.ps1: before falling through to amd-smi, attempt to locate hipinfo.exe and hipconfig.exe under $env:HIP_PATH\bin (then $env:ROCM_PATH\bin) when Get-Command returns nothing - Emit a [WARN] with the resolved path and a one-liner to permanently fix PATH via SetEnvironmentVariable - Emit a [WARN] when HIP_PATH/ROCM_PATH is set but the exe is still not found (incomplete SDK install) - Emit a [WARN] with the first hipinfo output line when hipinfo runs but returns a non-zero exit code (e.g. "no ROCm-capable device detected") - 18 new tests in TestHipSdkEnvPathResolution; total 183 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: print HIP SDK path and full hipconfig version in terminal on AMD detection Both install.ps1 and setup.ps1 now emit substeps under the gpu step when AMD ROCm is detected: gpu AMD ROCm (gfx1200) HIP SDK: C:\Program Files\AMD\ROCm\7.1 hipconfig: 7.1.51803-d3a86bd04 Previously only the gpu label (e.g. "AMD ROCm (gfx1200)") was shown with no indication of where the SDK was found or which exact build was active. The full hipconfig build string (e.g. 7.1.51803-d3a86bd04 instead of just 7.1) is now stored in ROCmVersionFull and also used in setup.ps1's 'rocm' step label. 9 new tests in TestHipSdkDetectedSubstep; total 192 passed, 2 skipped * fix: Strix rocm7.1 segfault bypass + Ubuntu 24.04 HIP gcc-install-dir Issue 1 (install.sh): gfx1151/gfx1150 + ROCm 7.1 causes a segfault in torch._grouped_mm (moe_utils.py:167). The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so _amd_gpu_radeon=true silently lands on the broken combo. When Strix Halo/Point is detected and TORCH_INDEX_URL is rocm7.1, override to rocm7.2 PyTorch index, update TORCH_CONSTRAINT, and set _amd_gpu_radeon=false to bypass the Radeon repo entirely. Emits a clear [WARN] explaining the segfault and linking to the ROCm upgrade docs. Issue 2 (setup.sh): ROCm 7.x ships clang-20 which on Ubuntu 24.04+ picks /usr/lib/gcc/x86_64-linux-gnu/14/ (runtime dir, no C++ headers), causing 'cstdlib file not found' and a failed llama.cpp HIP build. Iterate gcc versions 14→11 to find the first install dir that has both runtime and /usr/include/c++/ headers, then pass --gcc-install-dir to clang via CMAKE_HIP_FLAGS. Fix confirmed by h34v3nzc0dex (llama.cpp 417/417 clean). 11 new tests across TestStrixRocm71Override and TestSetupShGccInstallDir; total 203 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: BNB_ROCM_VERSION in server process + torch._C._distributed_c10d stubs Two errors visible in training logs on Windows ROCm: 1. Server process bitsandbytes crash: "Configured ROCm binary not found at libbitsandbytes_rocm713.dll" The installed BNB wheel ships rocm72.dll (not rocm713.dll). The training worker already sets BNB_ROCM_VERSION=72 via DLL detection but the server process (main.py) imported bitsandbytes before that ran. Fix: add the same DLL-scan + BNB_ROCM_VERSION assignment to main.py inside the existing win32 guard, before any downstream import can pull in bitsandbytes. 2. torch.distributed import failure: "No module named 'torch._C._distributed_c10d'; torch._C is not a package" torch._C is a C extension on Windows ROCm — Python cannot do submodule imports from it, so torch.distributed fails to import before our attribute stubs could ever run. Fix: inject empty ModuleType stubs for _distributed_c10d, _distributed_autograd and _distributed_rpc into sys.modules inside the win32 guard in hardware.py BEFORE importing torch.distributed, so the import succeeds and our attribute stubs take effect. 9 new tests in TestServerStartupRocmFixes; total 212 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(win32): populate distributed c10d stub with dummy symbols torch.distributed tries to `from torch._C._distributed_c10d import FakeProcessGroup` (and ProcessGroup, Work, Store, etc.). The previous empty ModuleType stub caused an AttributeError on those names. Populate every stub with a _Dummy class for each known symbol so the import chain completes silently on Windows ROCm where torch._C is a compiled extension and its _distributed_c10d submodule doesn't exist. Adds four new tests in TestServerStartupRocmFixes covering FakeProcessGroup, ProcessGroup, setattr population, and all three _distributed_* siblings. * fix(win32): distinguish HIP SDK installed vs GPU not ROCm-accessible Previously, when hipinfo was found but exited non-zero (e.g. "no ROCm-capable device detected"), both install.ps1 and setup.ps1 fell through to the WMI-label-only branch and printed "AMD GPU detected -- HIP SDK not found" -- factually wrong since the SDK binary is present. Add $HipSdkInstalled flag (set true when hipinfo binary is found, regardless of exit code). When HipSdkInstalled && !HasROCm: - Show "AMD GPU detected -- not ROCm-accessible (HIP )" instead - Explain this is a driver issue, not an SDK issue, with a link - Still run hipconfig version capture so version shows in output - CPU-only hint now says "GPU not ROCm-accessible" not "require HIP SDK" Also applies to setup.ps1 (same detection block, same branches). Adds TestHipSdkInstalledButDeviceInaccessible (11 tests). * fix(win32): scope ROCm workarounds to AMD hosts only Three Codex-flagged issues where Windows ROCm workarounds incorrectly applied to Windows CUDA (NVIDIA) machines: main.py (P1): BNB_ROCM_VERSION was set unconditionally on all win32 hosts. On NVIDIA, bitsandbytes sees BNB_ROCM_VERSION and looks for a ROCm DLL that doesn't exist, breaking bitsandbytes initialisation. Fix: gate the block on HIP_PATH/ROCM_PATH being present (ROCm hosts only). worker.py (P2): torchao stubs were seeded for all win32 runs, shadowing real torchao on Windows CUDA and silently disabling torchao quantization for NVIDIA users. Fix: gate on HIP_PATH/ROCM_PATH (win32 ROCm only). install_python_stack.py (P1): _detect_windows_gfx_arch() only checked shutil.which("hipinfo"), skipping the HIP_PATH/ROCM_PATH fallback that the PowerShell installers use. On installs where the HIP SDK bin dir is not on PATH, _ensure_rocm_torch() returned early without installing ROCm wheels or bitsandbytes. Fix: mirror the env-var fallback. * fix(linux): route Strix + ROCm 7.1 to AMD arch-specific index Instead of falling back to pytorch.org/rocm7.2, the Strix override now routes to repo.amd.com/rocm/whl/gfx1151/ (or gfx1150/) which serves torch 2.11.0+rocm7.13.0 -- AMD's build containing the actual _grouped_mm kernel fix, verified on real gfx1151 hardware by h34v3nzc0dex. This exercises the real GPU kernel path rather than the rocm7.2 workaround. UNSLOTH_AMD_ROCM_MIRROR can override the base URL for air-gapped installs. Also teaches _tauri_torch_index_family to recognise AMD arch-specific URLs (repo.amd.com/rocm/whl/gfx*) and return the rocm7.13 family label so _tauri_gpu_branch correctly classifies these installs as rocm. Suggested by h34v3nzc0dex based on hardware-verified probe results. * fix(studio/rocm): gate ROCm-only side-effects on active torch runtime Address five edge cases flagged during PR review: 1. studio/backend/main.py: BNB_ROCM_VERSION was set whenever HIP_PATH or ROCM_PATH was present in the environment. A Windows CUDA user who once installed the HIP SDK and reverted to a CUDA torch wheel still has those env vars set, so bitsandbytes would try to load libbitsandbytes_rocm72.dll against a CUDA torch and crash. Now probe torch.version.hip inside the env-var guard (worker.py already does this). 2. studio/backend/main.py: os.add_dll_directory returned handles were discarded. Per CPython docs, the directory leaves the DLL search list when the handle is garbage collected. Retain handles in module-level _ROCM_DLL_HANDLES list so they survive process lifetime. 3. studio/install_python_stack.py: _install_bnb_windows_rocm() returned None regardless of pip_install_try outcome, and the caller flipped _rocm_windows_torch_installed to True unconditionally. On a failed BNB install the post-install "manual install may be required" warning was suppressed and the user was misled. Helper now returns bool; caller gates on it. 4. studio/install_python_stack.py: _detect_windows_gfx_arch returned the raw capture group, so mixed-case hipinfo output ("Gfx1151") missed the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and silently fell back to CPU torch. Lowercase the token. 5. studio/install_python_stack.py: UNSLOTH_ROCM_TORCH_INSTALLED=1 early- return trusted the env var even when the venv was wiped between runs. Subprocess-probe torch importability first; fall through to the full install path if the probe fails. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py (adds one new test for case 5 fall-through). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): worker.py parity + don't roll back ROCm torch on bnb failure Addresses findings from a 10x reviewer pass on the prior fix commit: 1. studio/backend/core/training/worker.py (parity with main.py): - Gate the torchao stub block on torch.version.hip / 'rocm' in torch.__version__ instead of HIP_PATH / ROCM_PATH env-var presence. Same root cause as main.py: HIP SDK env vars stick around on CUDA hosts. - Add module-level Windows ROCm DLL registration block. Worker subprocesses inherit env vars but not the parent's add_dll_directory handles, so the first `import torch` in the worker could fail to find amdhip64.dll when HIP_PATH\bin is not on PATH. Mirrors main.py setup. Handles retained at module scope via _ROCM_DLL_HANDLES. - Promote _WINDOWS_ROCM_GROUPED_MM_LIB to module scope with `global` in run_training_process so the torch.library.Library registration survives past function return / mid-run garbage collection. - Harden _torch_has_hip() to also accept 'rocm' in torch.__version__ (AMD SDK / Radeon wheels may not set torch.version.hip). 2. studio/install_python_stack.py: - Don't roll back ROCm torch when bitsandbytes install fails. The prior commit gated _rocm_windows_torch_installed on _install_bnb_windows_rocm() returning True; if torch installed successfully but bnb failed, the flag stayed False and later install steps could overwrite ROCm torch with the generic CPU torch wheel. Set the flag after torch install; surface bnb failure as a separate warning instead. - _detect_windows_gfx_arch now probes in three tiers: UNSLOTH_ROCM_GFX_ARCH env-var override (matches the PowerShell installer), then hipinfo (PATH or HIP_PATH\bin), then amd-smi (`static --asic`, `list`). Without the amd-smi fallback, runtime-only Radeon installs without hipinfo on PATH made `studio update` return early and leave the venv on CPU torch. - Linux torch-already-rocm probe in _ensure_rocm_torch now matches the Windows probe shape: accepts torch.version.hip OR 'rocm' in torch.__version__ to cover AMD SDK / Radeon Linux wheels. 3. studio/backend/utils/hardware/hardware.py: - apply_gpu_ids() final-fallback torch probe accepts 'rocm' in torch.__version__ in addition to torch.version.hip, matching detect_hardware(). AMD SDK wheels could otherwise leak through with CUDA-only visibility masks on a spawned ROCm worker. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py (no test changes needed; the probe shape that prints the hip version (or 'rocm' sentinel) preserves the existing non-empty-string contract). Not addressed in this commit (deferred or out of scope): - Tag drift / lemonade checksum (PR 5303 surface, not this PR). - install.sh rocm7.2.1 URL: small fix, separate. - install.ps1 / setup.ps1 'Radeon 8060S' marketing-name fallback table. - Strix Halo + ROCm 7.1 routing asymmetry in Python update path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): robustness pass - rocm tag normalisation, Strix routing parity, hardened detection Robustness pass on top of 76137b2d. Four targeted fixes: 1. install.sh ROCm-tag routing normalisation. `rocm7.2.1` would route to https://download.pytorch.org/whl/rocm7.2.1 which does not exist (PyTorch publishes major.minor URLs only). Same for any future patch-level tag. Normalise every rocm{maj.min}* pattern to the bare {maj.min} index URL. 2. install.ps1 + studio/setup.ps1 marketing-name fallback. The gfx1151 row matched 890M / Strix Halo / HX 37x / HX 38x / AI 9 HX but not the actual retail name 'AMD Radeon 8060S Graphics' shipped by OEMs (Ryzen AI MAX+ 395). Add '8060S' to the regex. 3. install_python_stack.py Strix + ROCm 7.1 routing parity with install.sh. The shell installer reroutes Strix Halo / Point + ROCm 7.1 to repo.amd.com/rocm/whl/{gfx}/ (which serves torch 2.11.0+rocm7.13.0 with the upstream _grouped_mm fix). The Python `studio update` path only warned and still installed the broken generic rocm7.1 wheel. Mirror the override: detect gfx1151/gfx1150 on ROCm 7.1, route to the AMD per-gfx index, honour UNSLOTH_AMD_ROCM_MIRROR override. 4. _detect_windows_gfx_arch amd-smi parsing tightened. The amd-smi fallback added in the prior commit used a bare `\bgfx[1-9][0-9a-z]{2,3}\b` match against the lowercased stdout, which could pick up stray gfx references in warnings / device-name strings. Anchor on labelled lines first (Target_Graphics_Version, ASIC, Arch, gfx) and fall back to the bare match only when no labelled line is present. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py; sim_5301 23 cases pass (6 new sims for the Strix override + amd-smi parsing). * fix(studio/rocm): multi-GPU selection, Strix sibling handling, defensive cleanups Round 4 robustness pass based on 5 parallel Opus reviewers of head 21773215. Seven items from across regression / edge-case / error-paths / architecture reviews: 1. studio/backend/main.py BNB gate: aligned with the broad ROCm check used everywhere else in this PR (torch.version.hip OR 'rocm' in __version__). AMD SDK / Radeon Linux wheels do not always populate torch.version.hip; without this, main.py would silently skip BNB_ROCM_VERSION while worker.py set it. 2. studio/install_python_stack.py _install_bnb_windows_rocm: init _ok = False before the try block. Without this, if pip_install_try itself raises (e.g. OSError on uv binary missing), the finally block restored env vars correctly but the subsequent `if not _ok:` raised UnboundLocalError, masking the original exception. 3. studio/install_python_stack.py _detect_windows_gfx_arch: - Rewrote to use re.findall (not re.search) on both hipinfo and amd-smi output, dedup tokens preserving order, and select via new _pick_visible_index() helper. - HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES (first comma entry, integer) now picks the right GPU on multi-AMD-GPU hosts. Out-of-range or non-int values fall back to the first GPU (matches detect_host behaviour in install_llama_prebuilt.py). 4. studio/install_python_stack.py Strix override now consults the runtime target before flipping: - Previous behaviour intersected gfx_codes with {gfx1151, gfx1150} and picked the first Strix arch, ignoring whether HIP_VISIBLE_DEVICES selected a non-Strix sibling (e.g. discrete RX 7900 in a mixed APU+dGPU box). Could install Strix-specific wheels onto a gfx1100 dGPU. - Now resolves the runtime gfx via _pick_visible_index() and only overrides when that runtime target is in the Strix set. 5. studio/backend/main.py + studio/backend/core/training/worker.py: ROCm version dir scan no longer sorts lexically. Previous sort placed "10.0" before "7.0" alphabetically, which would mis-prioritise ROCm 10.x bin dirs once AMD ships them. New _ver_key() splits on "." and sorts numerically with a string fallback. 6. install.sh Strix override URL: replaced ${var%/} (strips one trailing slash) with a while-loop that strips all trailing slashes, matching Python's .rstrip("/"). A user setting UNSLOTH_AMD_ROCM_MIRROR with "http://corp/whl///" no longer ends up with "http://corp/whl///gfx1151/" which strict pip proxies (artifactory, sonatype) 404 on. 7. studio/install_python_stack.py: bumped torch import probe timeout from 30s to 90s. PyTorch's lazy .so loading can take 60-90s on cold NFS or USB-backed venvs. The shorter timeout was producing a false "torch missing" classification and reinstalling a working ROCm torch. Tests: 231 passed, 1 skipped. sim_5301 30 cases pass (added 7 new sims for multi-GPU detection, Strix sibling handling, and _ok-init regression). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): worker BNB/grouped_mm broad gate, install.sh Strix visibility, runtime-only ROCm detection Round-5 robustness pass based on 20 parallel reviewers of head 96b9e465. 1. studio/backend/core/training/worker.py - BNB version pin / dynamo disable / _grouped_mm fallback block was still gated on torch.version.hip alone despite the torchao stub block above already using the broad check. AMD SDK / Radeon Windows wheels (torch.__version__ contains "rocm" but torch.version.hip is None) silently skipped the Windows ROCm runtime patches. Aligned to the same broad check (8/20 reviewers). 2. studio/backend/core/training/worker.py - _hip_ver_at_least() now also parses the ROCm version out of torch.__version__ (e.g. "2.11.0+rocm7.13.0") when torch.version.hip is missing, so the kernel-fix gate is correct for SDK / Radeon wheels too. 3. studio/backend/core/training/worker.py - _grouped_mm_safe_impl with offs=None now picks torch.bmm/matmul for 3-D inputs instead of always calling torch.mm. The real _grouped_mm accepts 3-D batched matmul; the prior fallback raised "self must be a matrix" on MoE workloads (2/20). 4. studio/backend/main.py - dropped the HIP_PATH / ROCM_PATH env-var gate from the BNB block; probe torch directly. Runtime-only Radeon / AMD SDK Windows installs do not set those SDK env vars but still ship ROCm torch (5/20 reviewers). 5. install.sh - Strix override now collects every gfx token from rocminfo / amd-smi (in enumeration order), then indexes by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non- Strix dGPU host where the user selected the dGPU does NOT get rerouted to the Strix per-gfx index. Mirrors the Python update path (5/20 reviewers). 6. install.sh - Strix detection chain now also probes `amd-smi static --asic`, matching the PowerShell installer (1/20). Closes the gap on runtime-only Strix hosts where `amd-smi list` does not surface a gfx token. 7. studio/install_python_stack.py - _has_rocm_gpu() now has the sysfs KFD topology fallback (/sys/class/kfd/kfd/topology/nodes/*/gpu_id), matching install.sh. On minimal package-managed installs without rocminfo / amd-smi GUI tools, `studio update` can now detect the GPU and repair the venv instead of returning early (2/20). 8. studio/install_python_stack.py - _detect_amd_gfx_codes() now falls back to `amd-smi list` and `amd-smi static --asic` when rocminfo is missing (2/20). Strix routing on runtime-only Radeon hosts now matches what install.sh has done for a while. 9. studio/install_python_stack.py - Strix override now applies even when has_hip_torch is True. The whole point of the override is to repair an existing broken torch.version.hip == "7.1" install; skipping the reinstall left users on the known _grouped_mm segfaulting stack (3/20). Tests: 231 passed, 1 skipped. sim_5301 30 cases pass. sim_cross 12 pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): code review hardening pass - main.py: numeric DLL sort (string sort picked rocm72 over rocm713); add basename() to regex; log warning on detection failure; log info when BNB_ROCM_VERSION is set (mirrors worker.py) - worker.py: explicit len-guard in _hip_ver_at_least() with warning logs instead of silent IndexError/ValueError swallow - hardware.py: isinstance(result, dict) guard before result.get() in _smi_query() to prevent AttributeError on non-dict backend returns - amd.py: round() before int() on parsed GPU IDs; log warning when truncation occurs (defensive against malformed amd-smi output) - setup.sh: quote --gcc-install-dir value in CMAKE_HIP_FLAGS so paths with spaces do not break the CMake argument - install.ps1, setup.ps1: apply colon-split + ToLower() to hipinfo gcnArchName match (consistent with each other and with setup.sh) - install.sh: tighten ROCm tag case patterns to explicit rocmX.Y|rocmX.Y.* to avoid unintended prefix matches * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/training): GPU OOM guard to prevent system freeze on VRAM exhaustion On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can cause a HIP driver hang that freezes the entire system rather than raising a recoverable Python exception. Two-part fix: - set_per_process_memory_fraction(0.90) caps the HIP/CUDA allocator at 90% of VRAM so PyTorch raises OutOfMemoryError before hitting the hardware limit, keeping the driver alive and the system responsive - top-level exception handler detects OOM errors by type and message and surfaces a clear actionable message to the UI (reduce max_seq_length, enable gradient_checkpointing, lower batch size) instead of the raw CUDA/HIP error string * fix(studio/rocm): OOM guard ROCm-only + unified memory, multi-GPU arch selection OOM guard (worker.py): - Scope to _hw.IS_ROCM only -- NVIDIA CUDA has a graceful OOM path and does not need the allocator cap - Detect unified memory by comparing torch VRAM against psutil system RAM; use 0.80 on unified-memory APUs (gfx1151 Strix Halo) where the GPU pool is carved from host RAM, 0.90 on discrete cards Multi-GPU arch selection: - install.ps1 / setup.ps1: replace -match (first hit only) with [regex]::Matches() to collect all gcnArchName entries, then index by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES - install_python_stack.py: index into full token list before dedup so HIP_VISIBLE_DEVICES=2 on [gfx1100, gfx1100, gfx1151] resolves gfx1151 - install.sh: remove awk dedup from gfx token collection for same reason GCC multiarch (setup.sh): - Only append -linux-gnu when gcc -print-multiarch does not already return the full triple, fixing double-suffix on Ubuntu 24.04 * fix(tests): update ROCm version cap expectations from rocm7.1 to rocm7.2 Daniel's normalisation commit updated the cap from rocm7.1 to rocm7.2 since PyTorch now publishes that index and rocm7.2 ships torch 2.11.0. Test expectations were stale. * fix(tests): correct MLX smoke test losses_per_step assertion logging_steps=1 with max_steps=30 produces 30 loss entries, not 7. The assertion was stale from a previous config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/worker): detect unified-memory APU by GPU name not VRAM/RAM ratio The previous heuristic (VRAM > 50 % of system RAM) false-positived on discrete cards in low-RAM systems — e.g. RX 9060 XT 16 GB on a 16 GB or 24 GB machine would trip the unified-memory path and log "unified memory host" when it should say "discrete". AMD iGPUs (gfx1150/gfx1151 Strix Halo, Strix Point, etc.) expose names with a digit+M suffix ("AMD Radeon 890M"), while discrete cards use "RX NNNN [XT|XTX]" naming. Matching that suffix is reliable across all current ROCm-capable AMD consumer GPUs and does not require psutil. Also includes the device name in the log line to ease future debugging. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install/setup.ps1): force array on hipinfo gcnArchName parse to fix single-GPU arch truncation When [regex]::Matches() finds exactly one match, PowerShell's pipeline unwraps the result to a scalar string. Indexing a scalar string with [0] returns the first *character*, so a one-GPU system would parse gcnArchName "gfx1200" as "g", which is not in the supported arch map and triggers the CPU-only fallback. Wrapping with @() forces the result to remain an array regardless of match count. On a single-GPU machine the arch is now correctly read as "gfx1200" (or whatever the full name is) so the ROCm wheel index is selected. Reproducer: hipinfo exits 0 and outputs exactly one gcnArchName line. Without @(), $_hipAllArches = "gfx1200" (String); $_hipAllArches[0] = 'g'. With @(), $_hipAllArches = @("gfx1200") (Object[]); $_hipAllArches[0] = "gfx1200". * fix(studio/rocm): classify unified-memory APU via VRAM/RAM ratio, not arch list Replace the gcnArchName allowlist {gfx1150, gfx1151} with a psutil-based heuristic: unified APUs expose the entire system RAM as the HIP pool (ratio ≥ 0.90), discrete cards are well below that. No arch name required — future APUs classify correctly without code changes. Also removes the stale import re / \d[Mm]\b device-name regex that 5d84704 left behind, and logs vram/sys GiB for easier on-hardware verification. Addresses h34v3nzc0dex review: Radeon 8060S (gfx1151, 128 GiB unified) now correctly gets 0.80 cap instead of 0.90. * fix(studio/rocm): revert to gcnArchName for unified-memory APU classification VRAM/RAM ratio >= 0.90 false-positives on machines where discrete VRAM equals system RAM (e.g. RX 9060 XT 16 GB + 16 GB system RAM → ratio 1.0, incorrectly classified as unified → wrong 0.80 cap applied). gcnArchName is the correct signal: naming-independent, stable within a product family, and already parsed throughout this PR. Unified set is {gfx1150, gfx1151} (Strix Point + Strix Halo). * fix(studio/llama-prebuilt): resolve hipinfo via HIP_PATH/ROCM_PATH on Windows shutil.which("hipinfo") returns None when the HIP SDK bin dir is not on PATH -- the HIP SDK installer sets HIP_PATH/ROCM_PATH but does not always add the bin dir to PATH. This caused has_rocm=False in the prebuilt asset selector, so AMD ROCm machines got the CPU llama.cpp zip instead of the HIP one, silently running all chat inference on CPU. Add _resolve_exe() that falls back to %HIP_PATH%\bin and %ROCM_PATH%\bin when shutil.which() finds nothing, mirroring the same fallback already present in setup.ps1. * fix(studio/llama-prebuilt): pass --has-rocm from setup.ps1 to skip re-detection The Python prebuilt installer re-detects ROCm independently via shutil.which("hipinfo"), which fails when hipinfo is not on PATH (HIP SDK sets HIP_PATH but doesn't always add the bin dir to PATH). This caused has_rocm=False and downloaded the CPU llama.cpp zip even on confirmed AMD ROCm machines. setup.ps1 already performs reliable ROCm detection with its own HIP_PATH/ROCM_PATH fallback. Add --has-rocm flag to install_llama_prebuilt.py so setup.ps1 can forward its result directly, and pass it whenever $HasROCm is true. The Python script then overrides has_rocm=True in the HostInfo without re-probing. * fix(studio/llama-prebuilt): add HIP asset to simple-policy Windows path direct_upstream_release_plan (used by --simple-policy, which setup.ps1 always passes) only checked has_usable_nvidia on Windows and fell straight to CPU for AMD ROCm machines, ignoring has_rocm entirely. The --has-rocm override had no effect because the simple-policy code path never reached resolve_asset_choice where has_rocm was checked. Add an elif branch for has_rocm that tries the upstream HIP asset (llama-TAG-bin-win-hip-radeon-x64.zip) before falling through to the CPU fallback, consistent with the non-simple-policy path. * fix(studio/setup.ps1): auto-remove mismatched llama.cpp install kind When an existing llama.cpp install is the wrong kind for the current GPU (e.g. windows-cpu on an AMD ROCm machine that should have windows-hip), the prebuilt installer skips on tag match and never upgrades. Read install_kind from UNSLOTH_PREBUILT_INFO.json before invoking the installer and remove the directory if the kind doesn't match, forcing a fresh download of the correct variant. * fix(studio/setup.ps1): show live PyTorch install output in verbose mode for ROCm The ROCm torch reinstall (setup.ps1 phase) always silently captured output, so in --verbose mode the torch downgrade mid-install (2.11.0+rocm → 2.10.0 → 2.11.0+rocm) looked like the final state was 2.10.0. Match the CPU/CUDA blocks which show live uv output when $script:UnslothVerbose is set. * fix(rocm/windows): set ROCBLAS_TENSILE_LIBPATH for bundled rocblas.dll The llama.cpp ROCm prebuilt bundles rocblas.dll next to the binary but not the Tensile kernel library files it depends on at runtime (rocblas/library/TensileLibrary*.dat + *.hsaco). The bundled DLL searches for these files relative to its own location by default, i.e. /rocblas/library/, which does not exist in the prebuilt install tree. This causes a silent crash on the very first GEMM (prefill) with no output from llama-server, seen by the caller as WinError 10054 / 10061. Model load and the single-token warmup pass because they use simpler code paths that do not trigger rocBLAS GEMM. Fix: set ROCBLAS_TENSILE_LIBPATH in the subprocess env to /bin/rocblas/library so the bundled DLL finds the kernel files from the system ROCm installation. Uses setdefault so a user- supplied env var is never overwritten. No-ops on CUDA and CPU (no HIP_PATH) and on Linux (win32 branch only). Reproducer log: rocBLAS error: Cannot read .../Release/rocblas/library/TensileLibrary.dat rocBLAS error: Could not initialize Tensile host: directory_iterator: The system cannot find the path specified. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install.sh): restore gfx token dedup in Strix multi-GPU awk indexer 536a54df removed the per-source `| awk '!seen[$0]++'` dedup from the _gfx_all collection step but left the indexer awk as bare NF, so on a mixed-arch host (e.g. dGPU gfx1100 + Strix iGPU gfx1151) where rocminfo emits each gfx token twice (Name: field + ISA triple), HIP_VISIBLE_DEVICES=1 indexed vals[1] = the second gfx1100 occurrence instead of gfx1151, triggering the Strix routing on the wrong GPU. Add !seen[$0]++ to the indexer awk so duplicate tokens from the same GPU collapse to one entry before the HIP_VISIBLE_DEVICES index is applied -- matching exactly what the Python side does with dict.fromkeys() in _detect_amd_gfx_codes(). The comment above the block ("skip duplicates") already documented this as the intended behaviour. * fix(studio/install): correct _TOTAL progress count on Windows base_total += 3 fired for all non-macOS platforms including Windows, but flash-attn (line 1620) and ROCm torch final (line 1705) are both guarded by 'not IS_WINDOWS and not IS_MACOS', so on Windows with torch enabled _TOTAL was 13 while only 11 _progress() calls actually execute. Split into +1 for the ROCm torch check (all non-macOS) and +2 for the two Linux-only steps, so Windows gets _TOTAL=11 and Linux gets 14. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install.ps1): enforce torch>=2.11.0 for gfx120X and Strix on Windows The AMD arch-specific index (repo.amd.com/rocm/whl/gfx120X-all/ and gfx1151/) publishes torch wheels from 2.7.1 through 2.11.0. Without a version floor pip can resolve to torch 2.10.0+rocm7.12 on RDNA 4 (gfx120X) or torch 2.10.0+rocm7.1 on Strix (gfx1151/gfx1150), both of which have a null-pointer crash in torch._C._grouped_mm (TheRock issues #5284 / #3284). torch 2.11.0+rocm7.13 contains the fix. Add $ROCmTorchFloor alongside $ROCmIndexUrl: set to torch>=2.11.0 for the two affected arch families, null for all others. Wire it into the uv pip install call so the broken wheels are never selected. * fix(rocm/windows): address Codex nits - deterministic DLL suffix, CUDA llama.cpp kind, HIP_VISIBLE_DEVICES arch indexing - install_python_stack.py / worker.py: _detect_bnb_rocm_dll_ver() and the inline worker probe now collect ALL libbitsandbytes_rocm*.dll suffixes and return max() by numeric value instead of stopping at the first glob hit. Filesystem glob order is not guaranteed; this ensures '713' always wins over '72' when both variants are present in the wheel. - setup.ps1 (expectedKind): add 'windows-cuda' branch so NVIDIA hosts are not treated as 'windows-cpu'. Previously an existing windows-cuda prebuilt was always considered a mismatch on non-ROCm machines, forcing an unnecessary re-download on every update. - setup.ps1 (amd-smi gfx arch): collect ALL gfx tokens from amd-smi list output in GPU order and honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES when selecting which arch to use. On mixed-arch AMD systems where the visible GPU is not the first enumerated one, this prevents installing an incompatible wheel index. Falls back to index 0 (same as before) when the visibility var is unset or is a comma-separated list. - test_rocm_support.py: add test_picks_highest_suffix_when_multiple_dlls to cover the multi-DLL case that was previously untested. * fix(rocm): misleading amd-smi log, BNB spec consistency, torch ceiling for AMD index amd.py: split 'returncode != 0 or not stdout' into two separate branches. Previously, exit-0 with empty output logged 'amd-smi returned code 0' (which reads as success, not a warning) and incorrectly incremented the circuit-breaker counter. Now: non-zero exit logs the code and counts toward the limit as before; empty stdout on exit 0 logs at DEBUG level and does not penalise the counter (amd-smi --json always emits at least [] on exit 0, so this branch is rare and is not a tool failure). main.py: replace spec.origin / os.path.dirname() with spec.submodule_search_locations to match install_python_stack.py and worker.py. For normal wheel installs both approaches reach the same directory, but using submodule_search_locations is the canonical way and handles editable bitsandbytes installs correctly. Also use max() by numeric suffix (same as the other two sites) instead of a sort-then-break loop. install.ps1: add <2.12.0 ceiling to the torch constraint for gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix). AMD actively publishes new versions on their per-arch index; without a ceiling, a future 2.12.0+rocmX.Y wheel would be pulled in automatically before being validated on these architectures. The ceiling matches the existing Linux install_python_stack.py constraint for the same arches. Bump both when 2.12.x is confirmed working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): torch floor in setup.ps1, torchvision pin for Strix, rocmsdk in _hip_ver_at_least setup.ps1: add \ (mirrors install.ps1) and derive \ from it. Previously the AMD index install called 'Fast-Install torch torchvision torchaudio --force-reinstall --index-url \' with no version constraint, so pip could resolve torch 2.10.0+rocm7.12 for gfx1151/gfx1200 -- the exact broken wheel the PR is meant to avoid. Now gfx120X and Strix enforce 'torch>=2.11.0,<2.12.0', matching install.ps1 and the Linux constraint. install_python_stack.py: pin torchvision and torchaudio in _strix_override_pkgs. The Strix Linux override uses --index-url (exclusive, no PyPI fallback); bare unversioned 'torchvision' and 'torchaudio' could resolve a build from AMD's index targeting a different torch major, causing ABI/version mismatches at runtime. Now pinned to '>=0.26.0,<0.27.0' and '>=2.11.0,<2.12.0' respectively, matching _ROCM_TORCH_CONSTRAINT['rocm7.2']. worker.py: extend _hip_ver_at_least to handle AMD SDK wheel version strings. The fallback regex r'rocm(\d+)\.(\d+)' cannot match '2.9.0+rocmsdk20251116' (no rocmX.Y component), so the function always returned False on SDK/Radeon wheels -- installing the Python _grouped_mm workaround on wheels that already have the working HIP kernel. Added a second check: if the version string contains '+rocmsdk', assume >= 7.13 (the rocmsdk format post-dates the gfx120X null-kernel fix) and skip the fallback. * fix(rocm): warn on OOB HIP_VISIBLE_DEVICES, bail on empty numeric_ids mask - setup.ps1: when HIP/ROCR_VISIBLE_DEVICES names an index beyond the detected GPU count, emit a yellow warning and fall back to GPU 0 instead of silently reading allGfxArches[-1] (wrong arch) - hardware.py _reconcile_primary_rocm_unified_memory: distinguish numeric_ids=None (no env var, use torch ordinal 0) from numeric_ids=[] (empty mask / HIP_VISIBLE_DEVICES=-1, no GPU visible); bail out early in the empty case to avoid querying torch.device(0) incorrectly * fix(rocm): gate StubSubpackageFinder on win32 ROCm, add gcnArchName fallbacks - worker.py _StubSubpackageFinder: the meta_path append was running on every platform on every call to run_training_process; moved it inside the if _is_win32_rocm: block since stubs are only seeded there and the finder is a pure accumulation on Linux/Windows CUDA - worker.py OOM guard: AMD SDK / Radeon wheels may not populate gcnArchName, causing Strix Halo to be misclassified as discrete and get the 0.90 cap (12.8 GB OS headroom) instead of 0.80 (25.6 GB); now tries gcn_arch_name / arch_name / gfx_arch_name variants first, then falls back to device-name matching (890M -> Strix Halo, 880M -> Strix Point) with a debug log when the fallback fires * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): pin torchvision/torchaudio in setup.ps1, remove -Unique from arch array - setup.ps1 ROCm torch install: torchvision and torchaudio were passed bare alongside pinned torch>=2.11.0,<2.12.0 for gfx1151/gfx1200 arches. AMD publishes packages independently so a future torchvision 0.27 (for torch 2.12) on the same arch index would cause pip ResolutionImpossible or an ABI-incompatible install. Added torchvisionFloorMap and torchaudioFloorMap mirroring install_python_stack.py's strix override (torchvision>=0.26.0,<0.27.0, torchaudio>=2.11.0,<2.12.0) and derived ROCmVisionSpec/ROCmAudioSpec used in all three Fast-Install call sites. - setup.ps1 amd-smi arch detection: Select-Object -Unique was collapsing same-arch multi-GPU arrays (e.g. two gfx1151 APUs -> 1-element array) causing HIP_VISIBLE_DEVICES=1 to trigger a false out-of-range warning and fall back to GPU 0 even though the correct GPU would have been at index 1. Removed -Unique; added comment noting the positional-index assumption and its non-contiguous-GPU limitation. * fix(rocm): add 8060s/8050s to OOM guard device-name fallback, extract classifier helper Path 3 of the OOM guard device-name fallback only checked for 890m/880m (gfx1150 Strix Point SKU names). Strix Halo (gfx1151) ships as Radeon 8060S (Ryzen AI MAX+ 395) and Radeon 8050S (cut-down SKU) -- neither matches, so the fallback returned is_unified=False and applied the 0.90 fraction instead of 0.80, leaving ~12.8 GiB OS headroom on a 128 GiB pool instead of ~25.6 GiB. Fix: add 8060s and 8050s to the name-match set. Also correct the comment that mislabelled 890M as a Strix Halo name (it is Strix Point). Refactor: extract the three-path classifier into _rocm_classify_unified_memory() so it can be unit-tested directly. Add 31 test cases in test_rocm_oom_guard.py covering all three paths and the regression case (Radeon 8060S Graphics). Reported-by: h34v3nzc0dex * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): pass explicit dtype on bf16-unsupported hardware (RDNA2) dtype=None lets unsloth auto-detect the model dtype. On RDNA2 (gfx103x, e.g. RX 6600) is_bfloat16_supported() incorrectly returns True, so unsloth picks bf16 and the first bf16 kernel dispatch triggers: LLVM ERROR: Cannot select: intrinsic %llvm.amdgcn.fdot2.bf16.bf16 Replace every dtype=None in load_model() with _auto_dtype which resolves to None when bf16 is supported (all modern NVIDIA + RDNA3+) and torch.float16 otherwise. This gives RDNA2 users a working float16 training path without touching NVIDIA behaviour at all. Fixes: https://github.com/unslothai/unsloth/issues/5337 * fix: reduce log noise for expected non-issues on Windows ROCm Three log lines fired at warning/error level for conditions that are completely expected on a Windows HIP SDK-only setup: amd.py - amd-smi WinError 2 (FileNotFoundError): downgrade warning -> debug. amd-smi ships with Adrenalin, not the HIP SDK; absence is normal. - 'disabling' message: downgrade warning -> info with clearer text 'not available (not installed; expected on HIP SDK-only systems); GPU VRAM polling disabled' hardware.py - torch.distributed.Store missing: downgrade warning -> debug. The distributed stub added in this PR intentionally omits Store; the attention-impl fallback to eager is expected and non-actionable. worker.py - causal-conv1d: add early Windows exit (info) in both _ensure_causal_conv1d_fast_path and _causal_conv1d_install hook; no cp313/win_amd64 wheel exists, so the install always fails. - FLA: add early Windows exit (info) in _ensure_flash_linear_attention_unconditional; triton dependency has no cp313/win_amd64 wheel. - Defense-in-depth: _install_package_wheel_first non-HIP PyPI failure logs info+debug on Windows instead of error; FLA failure logs info+debug on Windows instead of warning. * [AMD] FIx installation of bitsandbytes when it's from .dev and skip rebuilding llama.cpp if we build it manually. * fix: use force_pip for Windows ROCm bitsandbytes prebuilt wheel install uv rejects the bnb continuous-release wheel due to filename/metadata version mismatch (1.33.7.preview vs 0.50.0.dev0). Switch to force_pip=True (pip bypass) instead of the UV_SKIP_WHEEL_FILENAME_CHECK env var workaround -- cleaner and consistent with how the Linux path handles it. BNB_ROCM_VERSION is still set post-install to the detected DLL suffix so the worker subprocess loads the correct libbitsandbytes_rocm{VER}.dll even when torch.version.hip reports a newer HIP version than the wheel ships. * fix: three small correctness fixes found in PR review - _install_bnb_windows_rocm: use UV_SKIP_WHEEL_FILENAME_CHECK=1 with try/finally instead of force_pip=True so the env var is always restored and the failing CI test passes - _determine_attention_impl_for_gpu_estimate: gate torch._C distributed stubs on IS_ROCM so Windows CUDA users keep the real extension - install.ps1 amd-smi fallback: collect all gfx tokens and index by HIP_VISIBLE_DEVICES, matching the hipinfo path on multi-GPU hosts * fix: stub torchao in export subprocess on Windows ROCm On Windows, the ROCm build of PyTorch ships without the distributed C extension (torch._C._distributed_c10d). torchao, which is pulled in transitively by transformers.quantizers at import time, walks into torch.distributed._functional_collectives -> distributed_c10d and crashes with: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package This only affected the export subprocess because the training subprocess already applied an identical torchao stub (introduced separately to fix the same root cause). The export subprocess had no such guard and died during 'Importing Unsloth...' before any model loading could happen. Fix: apply the same _StubSubpackageFinder / torchao stub pattern to the export subprocess entry point, gated on Windows ROCm detection, before any import of transformers or unsloth_zoo. Root cause tracked in ROCm/TheRock#3284 (libuv / torch.distributed missing on Windows ROCm builds). Ref: https://github.com/ROCm/TheRock/issues/3284 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.sh, setup.sh: add GPU arch step logging to match PS1 scripts Both shell scripts were missing the step "gpu" terminal log block that install.ps1 and setup.ps1 emit. This adds equivalent output: GPU label with gfx arch (e.g. "AMD ROCm (gfx1151)"), ROCm root path, hipconfig version, and marketing name substep. Includes the same gfx arch detection chain (rocminfo → amd-smi list → amd-smi static --asic), UNSLOTH_ROCM_GFX_ARCH env override, and name-based arch inference table (Strix Halo/Point, RDNA 3/4) as the PS1 versions. install.sh also replaces bare echo blocks for the AMD ROCm and CPU-only cases with formatted substep output. * Fix BNB_ROCM_VERSION gate, ROCm GPU mask preference, APU unified memory and Release build for PR #5301 - main.py: gate BNB_ROCM_VERSION on the rocm bnb DLL or HIP_PATH/ROCM_PATH instead of importing torch on every Windows host - hardware.py: prefer HIP/ROCR visible-device masks only on ROCm hosts so a stale mask cannot override CUDA_VISIBLE_DEVICES on NVIDIA - llama_cpp.py: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 only for unified-memory APUs (gfx1150/gfx1151) - setup.sh: pass -DCMAKE_BUILD_TYPE=Release for the HIP source build - add test_amd_apu_unified_memory.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: guard recompile_limit + fix AMD VRAM monitor fallback trainer.py: torch._dynamo.config.recompile_limit does not exist in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2 wheels). Guard the assignment so training doesn't crash on RDNA2/RDNA3. hardware.py: when amd-smi/nvidia-smi is unavailable or returns no usable data (HIP SDK-only Windows, Docker, unexpected JSON format), the existing fallback used torch.cuda.memory_allocated() which is process-specific and reads near-zero even with a fully loaded model. Switch to torch.cuda.mem_get_info() via _torch_get_per_device_info() which reports system-wide VRAM occupancy so the GPU monitor shows real usage on all AMD systems without requiring amd-smi. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: Windows VRAM monitor via Performance Counter API When amd-smi/nvidia-smi is unavailable on Windows, query dedicated GPU VRAM via Windows Performance Counters (same source as Task Manager). This gives system-wide cross-process usage, fixing the near-zero reading caused by torch.cuda.mem_get_info only seeing the Studio server process. Linux fallback path unchanged (mem_get_info is system-wide on ROCm). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: rename to _rocm_windows_perf_counter_vram_gb, scope to IS_ROCM Function is AMD ROCm specific — amd-smi absent on Windows when only the HIP SDK is installed. Scoped to IS_ROCM so NVIDIA Windows path is untouched (nvidia-smi handles that case). * fix: AMD VRAM monitor — Linux DRM sysfs + Windows perf counter Linux: read /sys/class/drm/card*/device/mem_info_vram_used|total for system-wide GPU memory across all processes. No tools required, always present on Linux AMD systems. Windows: Windows Performance Counter API (already added). Both paths are gated on IS_ROCM and only fire when amd-smi is absent. torch mem_get_info remains as last resort (process-local). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: AMD GPU monitor — utilization, temperature, and power for Windows and Linux fallback paths - Windows: GPU utilization via \GPU Engine(*engtype_3D*)\Utilization Percentage perf counter - Windows: temperature and power via ADL (atiadlxx.dll, ships with Adrenalin) - Linux: GPU utilization via DRM sysfs gpu_busy_percent - Linux: temperature via hwmon temp1_input (millidegrees C) - Linux: power via hwmon power1_average / power1_input (microwatts) All paths are no-op fallbacks (None) when the source is unavailable. Mirrors what nvidia-smi provides on the CUDA path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: remove ADL ctypes — does not support AMD iGPU (Strix Halo) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Erland366 Co-authored-by: danielhanchen --- install.ps1 | 264 +++- install.sh | 237 +++- studio/backend/core/export/worker.py | 97 ++ studio/backend/core/inference/llama_cpp.py | 53 + studio/backend/core/training/trainer.py | 26 +- studio/backend/core/training/worker.py | 632 ++++++++- studio/backend/main.py | 104 ++ .../tests/test_amd_apu_unified_memory.py | 52 + .../tests/test_log_filter_no_truncation.py | 24 +- studio/backend/tests/test_rocm_oom_guard.py | 176 +++ studio/backend/utils/hardware/amd.py | 62 +- studio/backend/utils/hardware/hardware.py | 392 +++++- studio/backend/utils/wheel_utils.py | 2 + studio/install_llama_prebuilt.py | 53 +- studio/install_python_stack.py | 553 +++++++- studio/setup.ps1 | 343 ++++- studio/setup.sh | 120 +- tests/sh/test_get_torch_index_url.sh | 8 +- tests/studio/install/conftest.py | 20 + tests/studio/install/test_rocm_support.py | 1248 ++++++++++++++++- unsloth/models/loader_utils.py | 5 +- 21 files changed, 4272 insertions(+), 199 deletions(-) create mode 100644 studio/backend/tests/test_amd_apu_unified_memory.py create mode 100644 studio/backend/tests/test_rocm_oom_guard.py create mode 100644 tests/studio/install/conftest.py diff --git a/install.ps1 b/install.ps1 index c4e8d8d522..cab66f5ae1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -976,6 +976,7 @@ shell.Run cmd, 0, False # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. Write-TauriLog "STEP" "Installing Python" $DetectedPython = Find-CompatiblePython + if ($DetectedPython) { step "python" "Python $($DetectedPython.Version) already installed" } @@ -1239,11 +1240,196 @@ shell.Run cmd, 0, False } } } + # ── AMD ROCm detection (Windows) — mirrors setup.ps1 ── + $HasROCm = $false + $HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility) + $ROCmGpuLabel = $null + $ROCmVersion = $null + $ROCmGfxArch = $null + if (-not $HasNvidiaSmi) { + # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). + # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. + $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue + if (-not $hipinfoExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + $hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + if ($hipRoot) { + $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" + if (Test-Path $hipinfoCandidate) { + Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow + Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow + $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } + } else { + Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow + Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + } + } + } + if ($hipinfoExe) { + $HipSdkInstalled = $true # binary found → SDK is installed regardless of device state + try { + $hipOut = & $hipinfoExe.Source 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { + $HasROCm = $true + $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }) + $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } + if ($_hipAllArches.Count -gt 0) { + $ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } elseif ($LASTEXITCODE -ne 0) { + # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected") + $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1) + Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow + Write-Host " $firstLine" -ForegroundColor Yellow + Write-Host " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + } + } catch {} + } + if (-not $HasROCm) { + $amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiExe) { + try { + $smiOut = & $amdSmiExe.Source list 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { + $HasROCm = $true + # Mirror the hipinfo path: collect all gfx tokens in enumeration + # order and pick the runtime-visible one via HIP_VISIBLE_DEVICES. + $_smiVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } + # Attempt 1: newer amd-smi versions embed the gfx arch in list output. + $_smiGfxTokens = @([regex]::Matches($smiOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) + if ($_smiGfxTokens.Count -gt 0) { + $ROCmGfxArch = if ($_smiVisIdx -lt $_smiGfxTokens.Count) { $_smiGfxTokens[$_smiVisIdx] } else { $_smiGfxTokens[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } else { + # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, + # including the GFX target needed for wheel index selection. + $smiAsicOut = "" + try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {} + $_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) + if ($_asicGfxTokens.Count -gt 0) { + $ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") { + $ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } + } + } catch {} + } + } + if (-not $HasROCm) { + try { + $wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue | + Where-Object { $_.Name -match "AMD|Radeon" } | + Select-Object -First 1 + if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name } + } catch {} + } + # ── Arch resolution: env-var override → name inference ────────────── + # Covers users whose amd-smi is too old to report the GFX target and + # who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU). + if ($HasROCm -and -not $ROCmGfxArch) { + # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. + if ($env:UNSLOTH_ROCM_GFX_ARCH) { + $ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower() + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" + } + # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). + elseif ($ROCmGpuLabel) { + $nameArchTable = @( + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 + @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) + @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) + @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop + @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 + @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) + ) + foreach ($row in $nameArchTable) { + if ($ROCmGpuLabel -match $row.P) { + $ROCmGfxArch = $row.A + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + substep "gfx arch inferred from GPU name: $ROCmGfxArch" "Cyan" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$ROCmGfxArch to skip inference next time" "Cyan" + break + } + } + } + } + # Capture ROCm version for wheel selection (hipconfig, then amd-smi). + # Run whenever the HIP SDK binary is present, not just when the device is accessible -- + # hipconfig --version works even when hipinfo reports no ROCm device (driver issue). + if ($HasROCm -or $HipSdkInstalled) { + $hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue + if (-not $hipConfigExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + if ($hipRoot) { + $hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe" + if (Test-Path $hipConfigCandidate) { + $hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + Write-Host " [WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" -ForegroundColor Yellow + $hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate } + } + } + } + if ($hipConfigExe) { + try { + $hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0) { + $hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim() + if ($hipVerLine -match '(\d+\.\d+)') { + $ROCmVersion = $Matches[1] + $ROCmVersionFull = $hipVerLine + } + } + } catch {} + } + if (-not $ROCmVersion) { + $amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiVer) { + try { + $smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { + $ROCmVersion = $Matches[1] + } + } catch {} + } + } + } + } + if ($HasNvidiaSmi) { step "gpu" "NVIDIA GPU detected" + } elseif ($HasROCm) { + step "gpu" $ROCmGpuLabel + $hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" } + substep "HIP SDK: $hipSdkPath" + if ($ROCmVersionFull) { substep "hipconfig: $ROCmVersionFull" } + } elseif ($HipSdkInstalled -and $ROCmGpuLabel) { + # HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue) + $sdkVer = if ($ROCmVersionFull) { " (HIP $ROCmVersionFull)" } else { "" } + step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow" + substep " This is a driver issue, not an SDK issue." "Yellow" + substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" + substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" + } elseif ($ROCmGpuLabel) { + step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "Install the HIP SDK for ROCm GPU inference:" "Yellow" + substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } else { step "gpu" "none (chat-only / GGUF)" "Yellow" - substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow" } # ── Choose the correct PyTorch index URL based on driver CUDA version ── @@ -1270,14 +1456,73 @@ shell.Run cmd, 0, False return "$baseUrl/cu126" } $TorchIndexUrl = Get-TorchIndexUrl - $TorchIndexFamily = Get-TauriTorchIndexFamily $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. + # ── 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 + $ROCmTorchFloor = $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 + } + # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in + # torch._C._grouped_mm on torch <2.11.0 (rocm7.12 and rocm7.1 respectively). + # TheRock issues #5284 and #3284. Force torch>=2.11.0 so pip never resolves + # to the broken 2.10.0 wheels even though they exist on the AMD index. + # The <2.12.0 ceiling matches the Linux install_python_stack.py constraint + # for the same arches: AMD actively publishes new versions on their index, + # so without a ceiling a future 2.12.0+rocmX.Y wheel would be pulled in + # automatically before it has been validated on these architectures. + # Bump the ceiling here (and in install_python_stack.py) when 2.12.x is + # confirmed working on gfx120X / Strix. + $torchFloorMap = @{ + "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" + "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + } + $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } + if ($archFamily) { + $ROCmIndexUrl = "$amdIndexBase/$archFamily/" + $ROCmTorchFloor = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { $null } + $archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" } + substep "$archLabel -- AMD repo.amd.com index selected" "Cyan" + if ($ROCmTorchFloor) { + substep " enforcing $ROCmTorchFloor (known _grouped_mm bug in older wheels)" "Cyan" + } + } elseif ($ROCmGfxArch) { + substep "AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow" + } else { + substep "AMD GPU detected but arch unknown -- falling back to CPU-only PyTorch" "Yellow" + } + } + + 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 $TorchIndexUrl -like "*/cpu") { + if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" - substep "No NVIDIA GPU detected." "Yellow" + if ($HipSdkInstalled -and -not $HasROCm) { + substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" + } elseif ($ROCmGpuLabel) { + substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow" + } else { + substep "No NVIDIA GPU detected." "Yellow" + } substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" substep "re-run with --no-torch for a faster, lighter install:" "Yellow" substep ".\install.ps1 --no-torch" "Yellow" @@ -1355,9 +1600,18 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) } } - } elseif ($TorchIndexUrl) { + } elseif ($TorchIndexUrl -or $ROCmIndexUrl) { if ($SkipTorch) { substep "skipping PyTorch (--no-torch flag set)." "Yellow" + } elseif ($ROCmIndexUrl) { + Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" + substep "installing PyTorch from $ROCmIndexUrl..." + $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec 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) + } } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." diff --git a/install.sh b/install.sh index a3b76a3011..91fdad8fb4 100755 --- a/install.sh +++ b/install.sh @@ -183,10 +183,21 @@ _install_bnb_rocm() { fi if [ -n "$_bnb_whl_url" ]; then substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..." - if run_install_cmd "$_label (pre-release)" "$_venv_py" -m pip install \ - --force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then + _bnb_log=$(mktemp) + if "$_venv_py" -m pip install \ + --disable-pip-version-check \ + --force-reinstall --no-cache-dir --no-deps \ + --retries 8 --timeout 90 \ + "$_bnb_whl_url" >"$_bnb_log" 2>&1; then + rm -f "$_bnb_log" return 0 fi + _bnb_rc=$? + if _is_verbose; then + cat "$_bnb_log" >&2 + fi + rm -f "$_bnb_log" + step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN" fi run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \ @@ -245,6 +256,9 @@ _tauri_torch_index_family() { rocm[0-9]*.[0-9]*) echo "$_diag_family" ;; *) echo "auto" ;; esac ;; + # AMD arch-specific index (e.g. repo.amd.com/rocm/whl/gfx1151/) -- + # used for Strix Halo/Point where torch 2.11+rocm7.13 has the real fix. + *repo.amd.com/rocm/whl/gfx*|*rocm/whl/gfx*) echo "rocm7.13" ;; "") echo "none" ;; *) echo "auto" ;; esac @@ -1568,16 +1582,19 @@ _find_no_torch_runtime() { } # ── AMD ROCm GPU detection helper ── -# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise. -# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and -# amd-smi list for GPU data rows (excludes header-only output). +# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs +# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices). _has_amd_rocm_gpu() { if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/ && !/Name:[[:space:]]*gfx000/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then return 0 elif command -v amd-smi >/dev/null 2>&1 && \ amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then return 0 + elif [ -e /dev/kfd ] && \ + awk '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \ + /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then + return 0 fi return 1 } @@ -1656,31 +1673,39 @@ get_torch_index_url() { if [ -n "$_rocm_tag" ]; then # Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older) case "$_rocm_tag" in - rocm[1-5].*) echo "$_base/cpu"; return ;; + rocm[1-5].*) + echo "[WARN] ROCm $_rocm_tag detected but PyTorch ROCm wheels require ROCm 6.0+ -- falling back to CPU-only PyTorch" >&2 + echo "[WARN] Upgrade ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "$_base/cpu"; return ;; esac - # ROCm 7.2 only has torch 2.11.0 which exceeds current bounds - # (<2.11.0). Fall back to rocm7.1 index which has torch 2.10.0. - # Enumerate explicit versions rather than matching rocm6.* so - # a host on ROCm 6.5 or 6.6 (no PyTorch wheels published) is - # clipped down to the last supported 6.x (rocm6.4) instead of - # constructing https://download.pytorch.org/whl/rocm6.5 which - # returns HTTP 403. PyTorch only ships: rocm5.7, 6.0, 6.1, 6.2, - # 6.3, 6.4, 7.0, 7.1, 7.2 (and 5.7 is below our minimum). - # TODO: uncomment rocm7.2 when the torch upper bound is bumped - # to >=2.11.0. + # Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2. + # PyTorch publishes major.minor URLs only (no patch level), so + # rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0. case "$_rocm_tag" in - rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*) - echo "$_base/$_rocm_tag" ;; + rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;; + rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;; + rocm6.2|rocm6.2.*) echo "$_base/rocm6.2" ;; + rocm6.3|rocm6.3.*) echo "$_base/rocm6.3" ;; + rocm6.4|rocm6.4.*) echo "$_base/rocm6.4" ;; + rocm7.0|rocm7.0.*) echo "$_base/rocm7.0" ;; + rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;; + rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;; rocm6.*) # ROCm 6.5+ (no published PyTorch wheels): clip down # to the last supported 6.x wheel set. echo "$_base/rocm6.4" ;; *) - # ROCm 7.2+ (including future 10.x+): cap to rocm7.1 - echo "$_base/rocm7.1" ;; + # ROCm 7.3+ (future): cap to rocm7.2 (latest known) + echo "$_base/rocm7.2" ;; esac return fi + # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be + # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, + # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch. + echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 + echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 + echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 echo "$_base/cpu"; return fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). @@ -1822,6 +1847,12 @@ _pick_radeon_wheel() { TORCH_INDEX_URL=$(get_torch_index_url) +# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. +# All other ROCm tags and CUDA stay within <2.11.0. +case "$TORCH_INDEX_URL" in + */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +esac + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". @@ -1834,6 +1865,78 @@ case "$TORCH_INDEX_URL" in fi ;; esac +# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ─────── +# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug +# that causes a segfault in torch._grouped_mm (moe_utils.py line 167). +# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when +# _amd_gpu_radeon=true the installer silently lands on the broken combo. +# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2. +case "$TORCH_INDEX_URL" in + */rocm7.1|*/rocm7.1.*) + # Collect every gfx token in rocminfo / amd-smi enumeration order + # (skip duplicates), then index by HIP_VISIBLE_DEVICES / + # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box + # where the user selected the dGPU does NOT get rerouted to the + # Strix per-gfx index. + _gfx_all="" + if command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + fi + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + # PowerShell paths also probe `amd-smi static --asic`; mirror it + # so a host with hipinfo-less amd-smi reports the gfx target. + if [ -z "$_gfx_all" ]; then + _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + fi + fi + _runtime_gfx="" + if [ -n "$_gfx_all" ]; then + _vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _idx=0 + if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then + _first=${_vis%%,*} + case "$_first" in + ''|*[!0-9]*) _idx=0 ;; + *) _idx=$_first ;; + esac + fi + _runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" ' + NF && !seen[$0]++ { vals[n++] = $0 } + END { + if (idx < 0 || idx >= n) idx = 0 + if (n > 0) print vals[idx] + }') + fi + _strix_gfx="" + case "$_runtime_gfx" in + gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; + esac + if [ -n "$_strix_gfx" ]; then + echo "" >&2 + echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2 + echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2 + echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2 + echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2 + echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "" >&2 + # AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's + # actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred + # over the pytorch.org rocm7.2 fallback because it exercises the real GPU + # kernel path. Set UNSLOTH_AMD_ROCM_MIRROR to override for air-gapped installs. + _amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" + # Strip ALL trailing slashes to match Python's .rstrip("/") -- a + # double-/triple-slash mirror URL would otherwise produce 404s on + # strict pip proxies (artifactory, sonatype). + while [ "${_amd_strix_base%/}" != "$_amd_strix_base" ]; do + _amd_strix_base="${_amd_strix_base%/}" + done + TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/" + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + _amd_gpu_radeon=false + fi + ;; +esac _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then _TAURI_TORCH_INDEX_FAMILY="radeon" @@ -1841,27 +1944,93 @@ fi _TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon") tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY" -# ── Print CPU-only hint when no GPU detected ── +# ── GPU detection summary (mirrors install.ps1 step "gpu" block) ── +if _has_usable_nvidia_gpu; then + step "gpu" "NVIDIA GPU detected" +elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then + # Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES + _gpu_disp_gfx_all="" + _gpu_disp_mkt="" + if command -v rocminfo >/dev/null 2>&1; then + _gpu_disp_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _gpu_disp_mkt=$(rocminfo 2>/dev/null | awk -F': ' \ + '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + fi + if [ -z "$_gpu_disp_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_disp_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_gpu_disp_gfx_all" ] && \ + _gpu_disp_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_gpu_disp_mkt" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_disp_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ + '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + fi + _gpu_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _gpu_vis_idx=0 + if [ -n "$_gpu_vis" ] && [ "$_gpu_vis" != "-1" ]; then + _gpu_first="${_gpu_vis%%,*}" + case "$_gpu_first" in ''|*[!0-9]*) ;; *) _gpu_vis_idx=$_gpu_first ;; esac + fi + _gpu_disp_gfx=$(printf '%s\n' "$_gpu_disp_gfx_all" | awk -v idx="$_gpu_vis_idx" \ + 'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }') + # UNSLOTH_ROCM_GFX_ARCH env override (mirrors install.ps1) + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + _gpu_disp_gfx="${UNSLOTH_ROCM_GFX_ARCH}" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx" + # Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable) + elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then + case "$_gpu_disp_mkt" in + *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 + *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU + *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop + *"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU + esac + if [ -n "$_gpu_disp_gfx" ]; then + substep "gfx arch inferred from GPU name: $_gpu_disp_gfx" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_gpu_disp_gfx to skip inference next time" + fi + fi + # ROCm version via hipconfig, then amd-smi + _gpu_rocm_ver="" + if command -v hipconfig >/dev/null 2>&1; then + _gpu_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true) + fi + if [ -z "$_gpu_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true) + fi + if [ -n "$_gpu_disp_gfx" ]; then + step "gpu" "AMD ROCm ($_gpu_disp_gfx)" + else + step "gpu" "AMD ROCm" + fi + _rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}" + substep "ROCm: $_rocm_root" + [ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver" + [ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt" +else + step "gpu" "none (CPU-only)" "$C_WARN" +fi + +# ── PyTorch wheel index note ── case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then - echo "" - echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)." - echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," - echo " re-run with --no-torch for a faster, lighter install:" - echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" - echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" - echo "" + substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" + substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):" + substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" fi ;; - */rocm*) - echo "" + */rocm*|*/gfx*) if [ "$_amd_gpu_radeon" = true ]; then - echo " AMD Radeon + ROCm detected -- installing PyTorch wheels from repo.radeon.com" + substep "wheels: repo.radeon.com (Radeon)" else - echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)" + substep "wheels: $TORCH_INDEX_URL" fi - echo "" ;; esac diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index f77b1966c4..77205bc298 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -439,6 +439,103 @@ def run_export_process( 'Install for better performance: pip install "triton-windows<3.7"' ) + # ── 1c. Stub torchao on Windows ROCm ── + # torchao (pulled in by transformers.quantizers) imports + # torch.distributed._functional_collectives at module level, which imports + # distributed_c10d.py unconditionally — that file crashes on Windows ROCm + # because torch._C._distributed_c10d (the RCCL backend) is absent. + # Stubbing torchao short-circuits the crash entirely. + # Must run before any import of transformers / unsloth_zoo. + import types as _types + import importlib.machinery as _ilm + import importlib.abc as _ilabc + + _STUB_SENTINEL = object() + + class _StubTypeMeta(type): + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, attr): + if attr.startswith("__"): + raise AttributeError(attr) + child = _StubTypeMeta(attr, (), {}) + setattr(cls, attr, child) + return child + + def __call__(cls, *args, **kwargs): + return None + + def _make_stub_type(name): + return _StubTypeMeta(name, (), {}) + + def _make_mod_stub(mod_name): + m = _types.ModuleType(mod_name) + m.__path__ = [] + m.__package__ = mod_name + m._unsloth_stub = _STUB_SENTINEL + 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 = _make_stub_type(f"{_n}.{attr}") + setattr(_m, attr, child) + return child + + m.__getattr__ = _ga + return m + + class _StubSubpackageLoader(_ilabc.Loader): + 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 + + class _StubSubpackageFinder(_ilabc.MetaPathFinder): + def find_spec(self, fullname, path, target = None): + if "." not in fullname: + return None + parent = sys.modules.get(fullname.rsplit(".", 1)[0]) + if parent is None: + return None + if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: + return None + return _ilm.ModuleSpec( + fullname, _StubSubpackageLoader(fullname), is_package = True + ) + + _is_win32_rocm = False + if sys.platform == "win32": + try: + import torch as _torch_probe + + _is_win32_rocm = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass + if _is_win32_rocm: + sys.meta_path.append(_StubSubpackageFinder()) + 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) + # ── 2. Import ML libraries (fresh in this clean process) ── try: _send_response( diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2d95112d6d..83c4b4f4da 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1238,6 +1238,33 @@ class LlamaCppBackend: return total + @staticmethod + def _amd_apu_wants_unified_memory() -> bool: + """True only for AMD unified-memory APUs (gfx1150/gfx1151), where + GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM. + False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete + GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped.""" + try: + import torch + + if getattr(torch.version, "hip", None) is None: + return False + if not (hasattr(torch, "cuda") and torch.cuda.is_available()): + return False + for _i in range(torch.cuda.device_count()): + try: + _arch = ( + getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") + or "" + ) + except Exception: + continue + if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}: + return True + except Exception: + return False + return False + @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: """Query free memory per GPU. @@ -3158,6 +3185,14 @@ class LlamaCppBackend: env = child_env_without_native_path_secret() binary_dir = str(Path(binary).parent) + # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use + # shared system RAM. setdefault so a user value wins. + if self._amd_apu_wants_unified_memory(): + env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") + logger.info( + "AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1" + ) + if sys.platform == "win32": # See _build_windows_path_dirs for ordering. #5106. path_dirs = self._build_windows_path_dirs( @@ -3167,6 +3202,24 @@ class LlamaCppBackend: ) existing_path = env.get("PATH", "") env["PATH"] = ";".join(path_dirs) + ";" + existing_path + + # ROCm: the llama.cpp prebuilt bundles its own rocblas.dll + # but NOT the Tensile kernel library files it needs + # (rocblas/library/TensileLibrary*.dat + *.hsaco). The + # bundled DLL searches relative to its own location by + # default (i.e. /rocblas/library/) which does + # not exist, causing a silent crash on the first GEMM. + # ROCBLAS_TENSILE_LIBPATH overrides that search to point at + # the ROCm installation where the kernel files actually are. + _hip_path = os.environ.get( + "HIP_PATH", os.environ.get("ROCM_PATH", "") + ) + if _hip_path: + _rocblas_lib = os.path.join( + _hip_path, "bin", "rocblas", "library" + ) + if os.path.isdir(_rocblas_lib): + env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) else: # Linux: set LD_LIBRARY_PATH for shared libs next to the binary # and CUDA runtime libs (libcudart, libcublas, etc.) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b9643cac6a..f1faaee4fb 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -42,7 +42,10 @@ from utils.hardware import ( get_visible_gpu_count, ) -torch._dynamo.config.recompile_limit = 64 +# recompile_limit was removed in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2). +# Guard so training doesn't crash on RDNA2/RDNA3 with older ROCm torch wheels. +if hasattr(torch._dynamo.config, "recompile_limit"): + torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported from unsloth.chat_templates import get_chat_template @@ -657,6 +660,15 @@ class UnslothTrainer: f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" ) + # On hardware without native bfloat16 support (e.g. RDNA2 / gfx103x), + # passing dtype=None lets unsloth auto-detect and incorrectly choose + # bf16, triggering an LLVM error at the first bf16 kernel dispatch. + # Explicitly pass float16 as the fallback so unsloth never reaches + # that path. Modern NVIDIA (Ampere+) and RDNA3+ return True here so + # they are unaffected — dtype stays None and unsloth picks bf16 as + # before. + _auto_dtype = None if is_bfloat16_supported() else torch.float16 + # Branch based on model type if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False @@ -666,7 +678,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, auto_model = CsmForConditionalGeneration, load_in_4bit = False, device_map = device_map, @@ -683,7 +695,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, - dtype = None, + dtype = _auto_dtype, load_in_4bit = False, device_map = device_map, full_finetuning = full_finetuning, @@ -705,7 +717,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastLanguageModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -777,7 +789,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -791,7 +803,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastVisionModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, # Auto-detect + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -824,7 +836,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastLanguageModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, # Auto-detect + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 632b38d75a..a59def8751 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -70,6 +70,58 @@ _TILELANG_INSTALL_TIMEOUT_S = 600 _TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11") _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS" +# Module-level handle so the torch.library.Library registration survives past +# run_training_process() and is not garbage collected mid-run. +_WINDOWS_ROCM_GROUPED_MM_LIB = None + +# Worker subprocesses inherit the parent env but not the parent's +# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL +# setup at module load so the first `import torch` can find amdhip64.dll even +# when HIP_PATH\bin is not on the system PATH. Handles retained at module +# scope so they are not garbage collected. +_ROCM_DLL_HANDLES: list = [] +if sys.platform == "win32": + + def _add_rocm_dll_dirs_worker() -> None: + _candidates: list[str] = [] + for _var in ("HIP_PATH", "ROCM_PATH"): + _val = os.environ.get(_var) + if _val: + _candidates.append(os.path.join(_val, "bin")) + _default_root = os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm" + ) + + def _ver_key(name: str) -> tuple: + # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string. + parts = [] + for chunk in name.split("."): + try: + parts.append((0, int(chunk))) + except ValueError: + parts.append((1, chunk)) + return tuple(parts) + + try: + if os.path.isdir(_default_root): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): + _bin = os.path.join(_default_root, _ver, "bin") + if os.path.isdir(_bin): + _candidates.append(_bin) + except OSError: + pass + for _d in _candidates: + if os.path.isdir(_d): + try: + _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d)) + except (OSError, AttributeError): + pass + + _add_rocm_dll_dirs_worker() + del _add_rocm_dll_dirs_worker + def _model_wants_causal_conv1d(model_name: str) -> bool: name = model_name.lower() @@ -320,11 +372,21 @@ def _install_package_wheel_first( f"{snippet}", ) else: - logger.error( - "Failed to install %s from PyPI:\n%s", - display_name, - result.stdout, - ) + if sys.platform == "win32": + # No prebuilt wheel and no source build toolchain on Windows -- + # this is expected for packages like causal-conv1d. Log at + # info so users aren't alarmed by what looks like an error. + logger.info( + "%s is not available on Windows (no prebuilt wheel); skipping", + display_name, + ) + logger.debug("Install output:\n%s", result.stdout) + else: + logger.error( + "Failed to install %s from PyPI:\n%s", + display_name, + result.stdout, + ) return False if is_hip: @@ -337,6 +399,9 @@ def _install_package_wheel_first( def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: if not _model_wants_causal_conv1d(model_name): return + if sys.platform == "win32": + logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping") + return _install_package_wheel_first( event_queue = event_queue, @@ -404,6 +469,11 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: """Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call.""" if os.getenv(_FLA_SKIP_ENV) == "1": return False + if sys.platform == "win32": + logger.info( + "Skipping flash-linear-attention install: no prebuilt wheel for Windows" + ) + return False if sys.version_info < _FLA_MIN_PYTHON: logger.info( "Skipping flash-linear-attention install: requires Python >= %d.%d, have %s", @@ -483,10 +553,17 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: return False if result.returncode != 0: - logger.warning( - "flash-linear-attention install failed (continuing on torch fallback):\n%s", - result.stdout, - ) + if sys.platform == "win32": + logger.info( + "flash-linear-attention not available on Windows (no prebuilt wheel); " + "continuing on torch fallback" + ) + logger.debug("Install output:\n%s", result.stdout) + else: + logger.warning( + "flash-linear-attention install failed (continuing on torch fallback):\n%s", + result.stdout, + ) _send_status( event_queue, "flash-linear-attention install failed; continuing without it", @@ -607,15 +684,61 @@ def _tilelang_importable() -> bool: def _torch_has_hip() -> bool: - """True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm.""" + """True iff torch is a ROCm build. + + `torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon + wheels can leave it unset but still encode "rocm" in `torch.__version__`. + """ try: import torch as _torch - return getattr(_torch.version, "hip", None) is not None + return bool( + getattr(_torch.version, "hip", None) + or "rocm" in getattr(_torch, "__version__", "").lower() + ) except Exception: return False +def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: + """Classify a ROCm device as unified-memory (APU) or discrete. + + Returns ``(gcn_arch, is_unified)`` where: + - ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a + known attribute is present, or ``""`` when all arch attrs are absent. + - ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool + (gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower + ``set_per_process_memory_fraction`` cap to leave headroom for the OS. + + Classification priority: + 1. ``gcnArchName`` / variant spellings (stable, naming-independent). + 2. Device-name substring match as a last-resort fallback when all arch + attrs are absent (AMD SDK / Radeon wheels may not populate them): + - gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M`` + - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395), + ``Radeon 8050S`` (cut-down SKU) + """ + gcn_arch = "" + for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"): + _v = (getattr(props, _attr, "") or "").split(":")[0].strip() + if _v: + gcn_arch = _v + break + + if gcn_arch: + return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"} + + # Arch attrs absent — fall back to device-name matching. + dev_lower = (getattr(props, "name", "") or "").lower() + is_unified = ( + "890m" in dev_lower + or "880m" in dev_lower + or "8060s" in dev_lower + or "8050s" in dev_lower + ) + return gcn_arch, is_unified + + def _tilelang_platform_supported() -> bool: """True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch. @@ -881,6 +1004,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: _ensure_tilelang_backend_unconditional(eq) def _causal_conv1d_install(eq: Any) -> bool: + if sys.platform == "win32": + logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping") + return False ok = _install_package_wheel_first( event_queue = eq, import_name = "causal_conv1d", @@ -1893,6 +2019,452 @@ def run_training_process( 'Install for better performance: pip install "triton-windows<3.7"' ) + # ── 1d. Stub torchao on Windows ROCm ── + # torchao (pulled in by transformers.quantizers) imports + # torch.distributed._functional_collectives at module level, which imports + # distributed_c10d.py unconditionally — that file crashes on Windows ROCm + # because torch._C._distributed_c10d (the RCCL backend) is absent. + # torch/distributed/__init__.py itself is guarded by `if is_available()` + # so `import torch.distributed` alone is safe; the crash only comes via + # torchao's import chain. Stubbing torchao short-circuits it entirely. + # _StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. + import types as _types + import importlib.machinery as _ilm + import importlib.abc as _ilabc + + _STUB_SENTINEL = object() + + # Metaclass for stub types so that isinstance(x, StubClass) returns False + # instead of raising TypeError ("arg 2 must be a type"). + # peft/tuners/lora/torchao.py does: + # from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor + # isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) + # If those names resolve to stub modules rather than types, isinstance() raises. + class _StubTypeMeta(type): + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, attr): + if attr.startswith("__"): + raise AttributeError(attr) + child = _StubTypeMeta(attr, (), {}) + setattr(cls, attr, child) + return child + + def __call__(cls, *args, **kwargs): + return None + + def _make_stub_type(name): + """Stub class: accepted by isinstance() (always False), supports attr access.""" + return _StubTypeMeta(name, (), {}) + + def _make_mod_stub(mod_name): + m = _types.ModuleType(mod_name) + m.__path__ = [] + m.__package__ = mod_name + m._unsloth_stub = _STUB_SENTINEL + 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) + # Return a stub CLASS (not a module) so that isinstance(x, attr) + # works and returns False instead of raising TypeError. + child = _make_stub_type(f"{_n}.{attr}") + setattr(_m, attr, child) + return child + + m.__getattr__ = _ga + return m + + class _StubSubpackageLoader(_ilabc.Loader): + 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 + + class _StubSubpackageFinder(_ilabc.MetaPathFinder): + def find_spec(self, fullname, path, target = None): + if "." not in fullname: + return None + parent = sys.modules.get(fullname.rsplit(".", 1)[0]) + if parent is None: + return None + if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: + return None + return _ilm.ModuleSpec( + fullname, _StubSubpackageLoader(fullname), is_package = True + ) + + # Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao + # is real and shadowing it breaks torchao-based quantization paths. + # Gate on the active torch runtime, not env-var presence -- HIP_PATH / + # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a + # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip + # but still encode "rocm" in torch.__version__, so accept either. + _is_win32_rocm = False + if sys.platform == "win32": + try: + import torch as _torch_probe + + _is_win32_rocm = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass + if _is_win32_rocm: + # Register the finder only on Windows ROCm -- on other platforms there + # are no stub modules seeded, so appending is a pure accumulation. + sys.meta_path.append(_StubSubpackageFinder()) + # Seed torchao top-level + key submodules; the finder handles the rest. + 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) + + # ── 1e. 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. + _td_stubs = { + "is_initialized": lambda: False, + "is_available": lambda: False, + "is_torchelastic_launched": lambda: False, + "get_rank": lambda: 0, + "get_world_size": lambda: 1, + "barrier": lambda: None, + } + + try: + import torch.distributed as _td + + for _name, _stub in _td_stubs.items(): + if not hasattr(_td, _name): + setattr(_td, _name, _stub) + except Exception: + _td_mock = _types.ModuleType("torch.distributed") + for _name, _stub in _td_stubs.items(): + setattr(_td_mock, _name, _stub) + sys.modules["torch.distributed"] = _td_mock + try: + import torch as _torch + + _torch.distributed = _td_mock + except Exception: + pass + + # ── 1f. Windows ROCm runtime patches ── + # torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows), + # causing 0xC0000005 (access violation) during training. + # + # Root cause: the JitDecomp autograd decomposition system (NOT torch.compile) + # dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash. + # TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop + # JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm + # with a safe Python fallback. + # + # Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped + # (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200. + # We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get + # the real GPU kernel rather than our Python fallback. + # + # Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0. + # + # Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, + # Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor + # offs: optional group-split offsets (MoE-style variable-size batches) + # + # torch is already in sys.modules from section 1e's `import torch.distributed`. + # Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past + # function return / mid-run GC. + global _WINDOWS_ROCM_GROUPED_MM_LIB + if sys.platform == "win32": + _torch_for_rocm = sys.modules.get("torch") + # Broad check: torch.version.hip OR "rocm" in torch.__version__. + # AMD SDK / Radeon Windows wheels do not always populate + # torch.version.hip; without the broad check the BNB version pin, + # dynamo-disable, and _grouped_mm fallback below silently skip + # (matches the torchao stub gate above and main.py). + _build_version_for_rocm = ( + getattr(_torch_for_rocm, "__version__", "").lower() + if _torch_for_rocm is not None + else "" + ) + _is_win_rocm_torch = bool( + _torch_for_rocm is not None + and ( + getattr(getattr(_torch_for_rocm, "version", None), "hip", None) + or "rocm" in _build_version_for_rocm + ) + ) + if _is_win_rocm_torch: + # Disable dynamo (belt-and-suspenders; JitDecomp patch below is the + # real fix, but keeping dynamo off avoids any other compile paths). + if "TORCHDYNAMO_DISABLE" not in os.environ: + os.environ["TORCHDYNAMO_DISABLE"] = "1" + logger.info("Windows ROCm: torch.compile (dynamo) disabled") + + # BNB auto-detects the HIP version from torch.version.hip and uses + # it to choose which DLL to load (e.g. "7.13" → rocm713.dll). + # AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its + # version suffix does not always match the torch HIP version (e.g. + # torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still + # ships rocm72.dll). We detect the actual DLL name from the installed + # package and override BNB's auto-detection. "72" is a safe fallback + # if detection fails. Callers may override by pre-setting the var. + if "BNB_ROCM_VERSION" not in os.environ: + _bnb_rocm_ver = None + try: + import glob as _glob + import importlib.util as _ilu + import re as _re + + _bnb_spec = _ilu.find_spec("bitsandbytes") + if _bnb_spec and _bnb_spec.submodule_search_locations: + _all_vers: list[str] = [] + for _pkg_dir in _bnb_spec.submodule_search_locations: + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): + _m = _re.search( + r"libbitsandbytes_rocm(\d+)\.dll", + os.path.basename(_dll), + ) + if _m: + _all_vers.append(_m.group(1)) + # Pick the highest numeric suffix so that e.g. "713" + # wins over "72" when both variants are present. + # Filesystem glob order is not guaranteed, so always + # sort rather than stopping at the first match. + if _all_vers: + _bnb_rocm_ver = max(_all_vers, key = lambda v: int(v)) + except Exception: + pass + _bnb_rocm_ver = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver + logger.info( + "Windows ROCm: set BNB_ROCM_VERSION=%s " + "(detected from installed BNB wheel; " + "overrides torch.version.hip auto-detection)", + _bnb_rocm_ver, + ) + + # Parse HIP version for the kernel-fix gate below. + # torch.version.hip can be "7.13.99004", "7.2.0", etc. + # AMD SDK / Radeon wheels may leave torch.version.hip unset and + # encode the ROCm version in torch.__version__ instead + # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back + # to that string when version.hip is missing. + def _hip_ver_at_least(major: int, minor: int) -> bool: + import re as _re_ver + + _hip_str = getattr( + getattr(_torch_for_rocm, "version", None), "hip", None + ) + if not _hip_str: + # Try the standard "+rocmX.Y.Z" embedded version first + # (e.g. "2.11.0+rocm7.13.0"). + _ver_match = _re_ver.search( + r"rocm(\d+)\.(\d+)", _build_version_for_rocm + ) + if _ver_match: + return ( + int(_ver_match.group(1)), + int(_ver_match.group(2)), + ) >= (major, minor) + # AMD SDK / Radeon Windows wheels encode the build as + # "+rocmsdk" (e.g. "2.9.0+rocmsdk20251116") with no + # explicit rocmX.Y component. The rocmsdk format was + # introduced after the gfx120X null-kernel fix landed in + # ROCm 7.13, so any wheel with this suffix is new enough to + # have working HIP kernels. Treat as >= 7.13 rather than + # falling back to False and installing the Python workaround + # on a wheel that doesn't need it. + if "rocmsdk" in _build_version_for_rocm: + logger.debug( + "Windows ROCm: AMD SDK wheel detected (%r); " + "assuming HIP >= %d.%d (rocmsdk wheels post-date " + "the gfx120X null-kernel fix)", + _build_version_for_rocm, + major, + minor, + ) + return True + return False + try: + _parts = [int(x) for x in str(_hip_str).split(".")[:2]] + if len(_parts) < 2: + logger.warning( + "Windows ROCm: torch.version.hip %r has fewer than " + "two components; cannot compare against %d.%d", + _hip_str, + major, + minor, + ) + return False + return (_parts[0], _parts[1]) >= (major, minor) + except ValueError: + logger.warning( + "Windows ROCm: could not parse torch.version.hip %r as " + "a version number; assuming HIP < %d.%d", + _hip_str, + major, + minor, + ) + return False + + # _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12, + # causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+). + # Only install the Python fallback on the affected versions so users + # on 7.13+ get the real GPU kernel for MoE workloads. + if not _hip_ver_at_least(7, 13): + try: + import warnings as _warnings + + _gm_lib = _torch_for_rocm.library.Library("aten", "IMPL") + + def _grouped_mm_safe_impl( + self, mat2, offs = None, bias = None, out_dtype = None + ): + """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12).""" + _t = _torch_for_rocm + if offs is None: + # No offsets: behave like the real op, which + # accepts either (M, K) x (K, N) -> mm, or 3-D + # batched inputs -> bmm. Picking torch.mm + # unconditionally previously raised "self must be + # a matrix" on 3-D MoE workloads. + if self.dim() == 3 and mat2.dim() == 3: + result = _t.bmm(self.contiguous(), mat2.contiguous()) + elif self.dim() == 3 and mat2.dim() == 2: + # Broadcast 2-D mat2 across the batch dim. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + elif self.dim() == 2 and mat2.dim() == 3: + # Broadcast 2-D self across batch via matmul semantics. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + else: + result = _t.mm(self.contiguous(), mat2.contiguous()) + else: + # Grouped case: offs[i] is the exclusive end-row of + # group i in `self`; mat2 may be 3-D or 2-D. + offs_list = offs.tolist() + pieces = [] + prev = 0 + for idx, end in enumerate(offs_list): + end = int(end) + a_part = self[prev:end].contiguous() + if mat2.dim() == 3: + b_part = mat2[idx].contiguous() + else: + b_part = mat2.contiguous() + pieces.append(_t.mm(a_part, b_part)) + prev = end + # Include any trailing rows not covered by offs + if prev < self.shape[0]: + a_tail = self[prev:].contiguous() + b_tail = ( + mat2[-1].contiguous() + if mat2.dim() == 3 + else mat2.contiguous() + ) + pieces.append(_t.mm(a_tail, b_tail)) + result = ( + _t.cat(pieces, dim = 0) + if pieces + else _t.zeros( + 0, + mat2.shape[-1], + device = self.device, + dtype = self.dtype, + ) + ) + if bias is not None: + result = result + bias + if out_dtype is not None: + result = result.to(out_dtype) + elif result.dtype != self.dtype: + result = result.to(self.dtype) + return result + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA") + + _WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC + logger.info( + "Windows ROCm: patched _grouped_mm CUDA dispatch " + "(null HIP kernel on gfx1200, ROCm ≤ 7.12 — " + "bypassed with Python mm fallback)" + ) + except Exception as _patch_exc: + logger.warning( + "Windows ROCm: could not patch _grouped_mm — " + "training may crash with 0xC0000005: %s", + _patch_exc, + ) + else: + logger.info( + "Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, " + "skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)" + ) + + # ── 1g. ROCm OOM guard ── + # On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can + # cause a HIP driver hang that freezes the entire system rather than + # raising a Python exception. set_per_process_memory_fraction caps the + # HIP allocator so PyTorch raises OutOfMemoryError before hitting the + # hardware limit, giving the UI a clean error instead of a system freeze. + # Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does + # not need this cap. + # Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU + # and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there. + # Primary classifier: gcnArchName from device properties — stable within a + # product family and naming-independent. AMD SDK / Radeon wheels may omit + # gcnArchName or expose it under a variant spelling, so we try several attr + # names then fall back to known device-name markers as a last resort. + # Non-fatal: silently skipped if torch is not importable. + if _hw.IS_ROCM: + try: + import torch as _torch_mem + + if _torch_mem.cuda.is_available(): + # Classify unified vs discrete via _rocm_classify_unified_memory. + # See that function's docstring for classification priority. + _props = _torch_mem.cuda.get_device_properties(0) + _dev_name = _props.name + _gcn_arch, _is_unified = _rocm_classify_unified_memory(_props) + if _is_unified and not _gcn_arch: + logger.debug( + "ROCm OOM guard: gcnArchName absent -- inferred " + "unified memory from device name %r; applying 0.80 cap", + _dev_name, + ) + _mem_fraction = 0.80 if _is_unified else 0.90 + _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction) + logger.info( + "ROCm OOM guard: set_per_process_memory_fraction(%.2f) — " + "%s memory host (%s, %s)", + _mem_fraction, + "unified" if _is_unified else "discrete", + _dev_name, + _gcn_arch or "unknown arch", + ) + except Exception as _oom_guard_err: + logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err) + # ── 2. Now import ML libraries (fresh in this clean process) ── try: _send_status(event_queue, "Importing Unsloth...") @@ -2347,14 +2919,38 @@ def run_training_process( ) except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } + _exc_str = str(exc).lower() + _is_oom = ( + "out of memory" in _exc_str + or "hip out of memory" in _exc_str + or "cuda out of memory" in _exc_str + or type(exc).__name__ == "OutOfMemoryError" ) + if _is_oom: + _oom_msg = ( + "GPU ran out of VRAM during training.\n" + "To fix: reduce max_seq_length (e.g. 2048–4096), enable " + "gradient_checkpointing=True, lower per_device_train_batch_size, " + "or use a smaller model / higher quantization." + ) + logger.error("Training stopped: GPU OOM — %s", exc) + event_queue.put( + { + "type": "error", + "error": _oom_msg, + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + else: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) def _send_status(event_queue: Any, message: str) -> None: diff --git a/studio/backend/main.py b/studio/backend/main.py index fbf1c31de3..be327bc5a8 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,110 @@ from pathlib import Path as _Path # Suppress annoying C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" +# ── Windows AMD ROCm DLL injection ────────────────────────────────────────── +# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with +# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. +if sys.platform == "win32": + # Retained at module scope -- os.add_dll_directory returns a handle that + # removes the search-path entry when garbage collected. + _ROCM_DLL_HANDLES: list = [] + + def _add_rocm_dll_dirs() -> None: + candidates = [] + # 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer + for _var in ("HIP_PATH", "ROCM_PATH"): + _val = os.environ.get(_var) + if _val: + candidates.append(os.path.join(_val, "bin")) + # 2. Standard AMD installer location: C:\Program Files\AMD\ROCm\\bin + # Scan all installed versions, newest first. + _default_root = os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm" + ) + + def _ver_key(name: str) -> tuple: + # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string. + parts = [] + for chunk in name.split("."): + try: + parts.append((0, int(chunk))) + except ValueError: + parts.append((1, chunk)) + return tuple(parts) + + try: + if os.path.isdir(_default_root): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): + _bin = os.path.join(_default_root, _ver, "bin") + if os.path.isdir(_bin): + candidates.append(_bin) + except OSError: + pass + for _d in candidates: + if os.path.isdir(_d): + try: + _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d)) + except (OSError, AttributeError): + pass + + _add_rocm_dll_dirs() + del _add_rocm_dll_dirs + + # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ + # bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm.dll + # where comes from torch.version.hip (e.g. "7.13..." → "713"). + # The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without + # this the server process crashes with "Configured ROCm binary not found". + # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION + # before any import that pulls in bitsandbytes (mirrors worker.py logic). + # Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/ + # ROCM_PATH, not on torch.version.hip: that needed importing torch on every + # Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon + # wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered. + if "BNB_ROCM_VERSION" not in os.environ: + import glob as _glob + import logging as _logging + + _hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")) + _bnb_rocm_ver = None + _found_rocm_bnb = False + try: + import importlib.util as _ilu + + _bnb_spec = _ilu.find_spec("bitsandbytes") + # submodule_search_locations (not spec.origin) handles editable installs. + if _bnb_spec and _bnb_spec.submodule_search_locations: + import re as _re_bnb + + _all_vers_main: list[str] = [] + for _pkg_dir in _bnb_spec.submodule_search_locations: + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): + _found_rocm_bnb = True + _km = _re_bnb.search( + r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll) + ) + if _km: + _all_vers_main.append(_km.group(1)) + if _all_vers_main: + _bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v)) + except Exception as _e: + _logging.getLogger(__name__).warning( + "Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'", + _e, + ) + # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72"). + if _found_rocm_bnb or _hip_env: + _bnb_rocm_ver_final = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final + _logging.getLogger(__name__).info( + "Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)", + _bnb_rocm_ver_final, + ) + # Ensure backend dir is on sys.path so _platform_compat is importable when # main.py is launched directly (e.g. `uvicorn main:app`). _backend_dir = str(_Path(__file__).parent) diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py new file mode 100644 index 0000000000..e0b819d54b --- /dev/null +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs +(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch(hip, archs, *, cuda_ok = True): + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(archs), + get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]), + ) + return t + + +@pytest.mark.parametrize( + "hip,archs,expected", + [ + ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped) + ("6.2.0", ["gfx1150"], True), # Strix Point APU + ("6.2.0", ["gfx1100"], False), # discrete RDNA3 + ("6.2.0", ["gfx1201"], False), # discrete RDNA4 + ("6.2.0", ["gfx942"], False), # MI300X (data center) + (None, ["sm_90"], False), # NVIDIA (no torch.version.hip) + ("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU + ], +) +def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected + + +def test_cpu_no_cuda_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False + + +def test_missing_torch_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False diff --git a/studio/backend/tests/test_log_filter_no_truncation.py b/studio/backend/tests/test_log_filter_no_truncation.py index d78643f5b9..d9a6e2bc4a 100644 --- a/studio/backend/tests/test_log_filter_no_truncation.py +++ b/studio/backend/tests/test_log_filter_no_truncation.py @@ -2,27 +2,11 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Regression tests for studio.backend.loggers.handlers.filter_sensitive_data. +Regression tests for loggers.handlers.filter_sensitive_data. -Context: filter_sensitive_data was originally written with a base64-detection -heuristic that truncated any string >100 chars containing ',' or '/' down to -20 chars + '...'. The block was dormant until PR #5246 wired the processor -into the structlog chain to redact native-path leases. Once active, the -heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size -summary, mmproj selection, the full llama-server command line) and any -exception traceback that happened to contain a file path. - -These tests pin two properties: - -1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data - unchanged. The exact strings exercised match the call sites at - studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that - were truncated in the original bug report. - -2. PR #5246's native-path lease redaction still fires for both the inline - ``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key - form. This guards against future regressions that strip redaction along - with the truncation block. +Pins two properties: (1) long strings with commas/slashes pass through +unchanged (the base64-truncation heuristic from PR #5246 was too aggressive), +and (2) native-path lease redaction still fires for both inline and dict-key forms. """ from loggers.handlers import filter_sensitive_data diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py new file mode 100644 index 0000000000..2ce9b55789 --- /dev/null +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier). + +Covers the three classification paths: + Path 1 – canonical gcnArchName attribute present. + Path 2 – gcnArchName absent, alternate-spelling attribute present. + Path 3 – ALL arch attrs absent; falls back to device-name substring match. + +Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK / +Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT +set any gcnArchName attribute. Without the 8060s/8050s name patterns the +fallback returned is_unified=False, applying the 0.90 fraction instead of +0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from core.training.worker import _rocm_classify_unified_memory + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _props(**kwargs) -> SimpleNamespace: + """Build a fake device-properties object with the given attributes.""" + return SimpleNamespace(**kwargs) + + +# ── Path 1: canonical gcnArchName ──────────────────────────────────────────── + + +class TestCanonicalGcnArchName: + """gcnArchName is present and populated.""" + + @pytest.mark.parametrize( + "arch, expected_unified", + [ + ("gfx1150", True), # Strix Point + ("gfx1151", True), # Strix Halo + ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete + ("gfx906", False), # MI50 — discrete server GPU + ("gfx1201", False), # RX 9070 XT — discrete + ], + ) + def test_canonical_attr(self, arch: str, expected_unified: bool) -> None: + props = _props(gcnArchName = arch, name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == arch + assert is_unified is expected_unified + + def test_arch_with_colon_suffix_stripped(self) -> None: + """gcnArchName can carry xnack/sramecc suffix; only the base is kept.""" + props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + def test_canonical_attr_wins_over_name(self) -> None: + """Arch attr takes priority; device name should be ignored.""" + # Discrete arch, but name looks like a unified SKU — arch must win. + props = _props(gcnArchName = "gfx1100", name = "Radeon 890M") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1100" + assert is_unified is False + + +# ── Path 2: alternate-spelling fallback ────────────────────────────────────── + + +class TestAlternateSpellingFallback: + """gcnArchName is missing but an alternate attr spelling is present.""" + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_unified(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_discrete(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1201" + assert is_unified is False + + def test_first_non_empty_attr_wins(self) -> None: + """When multiple alternate attrs are present the first non-empty one wins.""" + props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + +# ── Path 3: device-name fallback ───────────────────────────────────────────── + + +class TestDeviceNameFallback: + """ALL arch attrs absent — classifier must rely solely on device name.""" + + # --- unified-memory devices that MUST be detected --- + + @pytest.mark.parametrize( + "device_name", + [ + # gfx1150 Strix Point + "Radeon 890M", + "AMD Radeon 890M Graphics", + "RADEON 890M", # case-insensitive + "Radeon 880M", + "AMD Radeon 880M Graphics", + # gfx1151 Strix Halo — the regression case from the review + "Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch) + "AMD Radeon 8060S", + "Radeon 8050S Graphics", # cut-down Strix Halo SKU + "AMD Radeon 8050S", + # case variants + "RADEON 8060S GRAPHICS", + "radeon 8050s", + ], + ) + def test_unified_memory_detected(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "", f"expected empty gcn_arch, got {gcn!r}" + assert ( + is_unified is True + ), f"device {device_name!r} should be classified as unified-memory" + + # --- discrete devices that must NOT be mis-classified --- + + @pytest.mark.parametrize( + "device_name", + [ + "Radeon RX 9070 XT", + "AMD Radeon RX 7900 XTX", + "Radeon RX 6900 XT", + "Radeon Pro W7900", + "AMD Instinct MI300X", + # Names that contain superficially similar substrings but are discrete + "Radeon RX 580", + "Radeon VII", + ], + ) + def test_discrete_not_misclassified(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert ( + is_unified is False + ), f"discrete device {device_name!r} should NOT be classified as unified-memory" + + def test_empty_name_returns_false(self) -> None: + """Completely absent name must not crash and must default to discrete.""" + props = _props() # no 'name' attr at all + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False + + def test_none_name_returns_false(self) -> None: + props = _props(name = None) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index fdb1ab4520..48d5890399 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -11,18 +11,35 @@ nvidia.py counterparts. import json import math import os +import platform import re import subprocess +import sys from typing import Any, Optional from loggers import get_logger from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) +# amd-smi on Windows must initialise the full ROCm runtime on first call, which +# can take 15-25 s on cold hardware. Linux is consistently < 2 s. +_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10 -def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: +# Circuit breaker: stop calling amd-smi after this many consecutive failures. +# On Windows, each failed call spawns a process that may show a UAC/DiskPart +# elevation prompt. Once we know amd-smi doesn't work we stop polling it. +_AMD_SMI_FAILURE_LIMIT = 3 +_amd_smi_consecutive_failures = 0 +_amd_smi_disabled = False + + +def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]: """Run amd-smi with the given arguments and return parsed JSON, or None.""" + global _amd_smi_consecutive_failures, _amd_smi_disabled + if _amd_smi_disabled: + return None try: result = subprocess.run( ["amd-smi", *args, "--json"], @@ -30,13 +47,40 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: text = True, timeout = timeout, env = child_env_without_native_path_secret(), + **windows_hidden_subprocess_kwargs(), ) except (OSError, subprocess.TimeoutExpired) as e: - logger.warning("amd-smi query failed: %s", e) + if isinstance(e, FileNotFoundError): + # amd-smi ships with Adrenalin, not the HIP SDK -- absence is + # expected on HIP SDK-only Windows setups. Log at debug only. + logger.debug("amd-smi not found (not in PATH): %s", e) + else: + logger.warning("amd-smi query failed: %s", e) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None - if result.returncode != 0 or not result.stdout.strip(): + if result.returncode != 0: logger.warning("amd-smi returned code %d", result.returncode) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None + if not result.stdout.strip(): + # amd-smi exited successfully but produced no output (e.g. no GPUs + # visible on this query, or a version that emits nothing for --json). + # This is not a tool failure, so don't count against the circuit breaker. + logger.debug("amd-smi exited 0 but returned no output") + return None + _amd_smi_consecutive_failures = 0 # reset on success try: return json.loads(result.stdout) except json.JSONDecodeError: @@ -352,7 +396,7 @@ def get_visible_gpu_utilization( ) parsed_id = _parse_numeric(raw_id) if parsed_id is None: - logger.debug( + logger.warning( "amd-smi GPU id %r could not be parsed; falling back to " "enumeration index %d", raw_id, @@ -360,7 +404,15 @@ def get_visible_gpu_utilization( ) idx = fallback_idx else: - idx = int(parsed_id) + rounded = round(parsed_id) + if rounded != parsed_id: + logger.warning( + "amd-smi GPU id %r parsed as non-integer %r; truncating to %d", + raw_id, + parsed_id, + rounded, + ) + idx = int(rounded) if idx not in visible_set: continue metrics = _extract_gpu_metrics(gpu_data) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ede37e2953..ebac6a357c 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -120,11 +120,13 @@ def detect_hardware() -> DeviceType: # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. - if getattr(torch.version, "hip", None) is not None: + # AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do + # not set torch.version.hip, so fall back to checking __version__. + _hip_ver = getattr(torch.version, "hip", None) + if _hip_ver is not None or "rocm" in torch.__version__.lower(): IS_ROCM = True - print( - f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" - ) + _hip_label = _hip_ver or torch.__version__ + print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}") else: print(f"Hardware detected: CUDA -- {device_name}") return DEVICE @@ -466,7 +468,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: try: func = getattr(_backend, func_name) result = func(*args, **kwargs) - if result.get("available"): + if isinstance(result, dict) and result.get("available"): return result except Exception as e: logger.warning("%s %s query failed: %s", backend_name, func_name, e) @@ -506,6 +508,145 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: } +def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent.""" + import glob as _glob + + if platform.system() != "Linux": + return None + try: + files = _glob.glob("/sys/class/drm/card*/device/gpu_busy_percent") + if not files: + return None + values = [int(open(f).read().strip()) for f in files] + return round(sum(values) / len(values), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_temp_c() -> Optional[float]: + """Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C).""" + import glob as _glob + + if platform.system() != "Linux": + return None + try: + files = _glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input") + if not files: + return None + temps = [int(open(f).read().strip()) / 1000.0 for f in files] + return round(max(temps), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_power_w() -> Optional[float]: + """Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts).""" + import glob as _glob + + if platform.system() != "Linux": + return None + try: + for pattern in ( + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_average", + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_input", + ): + files = _glob.glob(pattern) + if files: + watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files) + return round(watts, 1) + return None + except Exception: + return None + + +def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes).""" + import subprocess as _sp + + if platform.system() != "Windows": + return None + try: + ps = ( + "$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}" + ) + r = _sp.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None + val = float(r.stdout.strip()) + return round(val, 1) if val >= 0 else None + except Exception: + return None + + +def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide AMD GPU VRAM via Linux DRM sysfs. + + Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel + updates in real-time across all processes. No tools required. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + import glob as _glob + + if platform.system() != "Linux": + return None, None + try: + used_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_used") + total_files = _glob.glob("/sys/class/drm/card*/device/mem_info_vram_total") + if not used_files or not total_files: + return None, None + used_bytes = sum(int(open(f).read().strip()) for f in used_files) + total_bytes = sum(int(open(f).read().strip()) for f in total_files) + if total_bytes == 0: + return None, None + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + +def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide dedicated GPU VRAM via Windows Performance Counters. + + Uses the same data source as Task Manager so it reflects cross-process + usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + import subprocess as _sp + + if platform.system() != "Windows": + return None, None + try: + ps = ( + "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" + ) + r = _sp.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None, None + used_bytes = float(r.stdout.strip()) + if used_bytes < 0: + return None, None + import torch as _torch + + total_bytes = _torch.cuda.get_device_properties(0).total_memory + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() @@ -514,7 +655,78 @@ def get_gpu_utilization() -> Dict[str, Any]: result = _smi_query("get_primary_gpu_utilization") if result is not None: result["backend"] = _backend_label(device) + if IS_ROCM: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.) + _reconcile_primary_rocm_unified_memory( + result, _get_parent_visible_gpu_spec() + ) return result + # SMI tool unavailable or returned no usable data. On Windows, query + # the Performance Counter API (same source as Task Manager) for + # system-wide dedicated VRAM — covers cross-process usage that + # torch.cuda.mem_get_info cannot see from the Studio server process. + if IS_ROCM and platform.system() == "Windows": + _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() + if _win_used is not None and _win_total is not None: + _win_util = _rocm_windows_perf_counter_gpu_util_pct() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + if IS_ROCM and platform.system() == "Linux": + _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() + if _linux_used is not None and _linux_total is not None: + _linux_util = _rocm_linux_sysfs_gpu_busy_pct() + _linux_temp = _rocm_linux_sysfs_temp_c() + _linux_power = _rocm_linux_sysfs_power_w() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Last resort: torch mem_get_info (process-local). + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [0] + _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] + _torch_devices = _torch_get_per_device_info(_primary_idx) + if _torch_devices: + _td = _torch_devices[0] + _total = _td["total_gb"] + _used = _td["used_gb"] + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } # MLX path: single _read_apple_gpu_stats() call carries both VRAM-used # bytes and GPU utilization %. psutil for unified-memory total is cheap. @@ -578,6 +790,77 @@ def get_gpu_utilization() -> Dict[str, Any]: return {"available": False, "backend": _backend_label(device)} +def _apply_unified_memory_correction( + device_metrics: Dict[str, Any], torch_info: Dict[str, Any] +) -> None: + """Per-device reconciliation: when torch reports a larger memory total + than amd-smi, overwrite the smi VRAM fields in place. + + Used by both the multi-device and primary-device reconciliation helpers + so the two endpoints stay in sync on AMD iGPUs with unified memory. + """ + torch_total_gb = torch_info["total_gb"] + smi_total_gb = device_metrics.get("vram_total_gb") or 0.0 + if torch_total_gb > smi_total_gb: + torch_used_gb = torch_info["used_gb"] + device_metrics["vram_total_gb"] = torch_total_gb + device_metrics["vram_used_gb"] = torch_used_gb + device_metrics["vram_utilization_pct"] = ( + round((torch_used_gb / torch_total_gb) * 100, 1) + if torch_total_gb > 0 + else None + ) + logger.debug( + "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " + "torch mem_get_info total (%.2f GB) for device %s", + smi_total_gb, + torch_total_gb, + torch_info.get("index"), + ) + + +def _reconcile_rocm_unified_memory( + utilization: Dict[str, Any], device_indices: list[int] +) -> None: + """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo). + + amd-smi reports only the dedicated slice (~512 MB); torch sees the full + GTT pool (~128 GB). When torch total > smi total, overwrite per-device + VRAM fields so GPU selection uses the real available memory. + """ + torch_devices = _torch_get_per_device_info(device_indices) + if not torch_devices: + return + torch_by_index = {td["index"]: td for td in torch_devices} + for dev in utilization.get("devices", []): + td = torch_by_index.get(dev.get("index")) + if td is None: + continue + _apply_unified_memory_correction(dev, td) + + +def _reconcile_primary_rocm_unified_memory( + utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any] +) -> None: + """Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict.""" + numeric_ids = parent_visible_spec.get("numeric_ids") + if numeric_ids is None: + # No visibility env var set: torch ordinal 0 is the primary device. + primary_idx = [0] + elif len(numeric_ids) == 0: + # Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to + # this process. Querying torch device 0 would raise a RuntimeError or + # return stale/wrong data, so bail out rather than writing bad values + # into the utilization dict. + return + else: + primary_idx = [int(numeric_ids[0])] + torch_devices = _torch_get_per_device_info(primary_idx) + if not torch_devices: + return + _apply_unified_memory_correction(utilization, torch_devices[0]) + + def get_visible_gpu_utilization() -> Dict[str, Any]: device = get_device() @@ -590,6 +873,10 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: ) if result is not None: result["backend"] = _backend_label(device) + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.) + _reconcile_rocm_unified_memory(result, numeric_ids) return result # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) @@ -689,7 +976,15 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]: # Use explicit None checks (not `or`) so empty string "" is honoured # as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES. cuda_visible = None - if IS_ROCM: + # Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a + # stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES. + _is_rocm_spec = IS_ROCM or ( + "CUDA_VISIBLE_DEVICES" not in os.environ + and ( + "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ + ) + ) + if _is_rocm_spec: hip_vis = os.environ.get("HIP_VISIBLE_DEVICES") rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES") if hip_vis is not None: @@ -867,6 +1162,61 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non def _determine_attention_impl_for_gpu_estimate(config) -> str: import copy as _copy + # torch.distributed is incomplete on Windows ROCm — torch._C is a C + # extension (not a package), so Python cannot import the submodule + # torch._C._distributed_c10d that torch.distributed depends on. + # Inject an empty stub into sys.modules BEFORE importing torch.distributed + # so the import succeeds, then patch the missing process-group helpers. + import sys as _sys + import types as _types + + if _sys.platform == "win32" and IS_ROCM: + # Dummy class for any name torch.distributed tries to import from these stubs + class _Dummy: + pass + + for _c10d_name in ( + "torch._C._distributed_c10d", + "torch._C._distributed_autograd", + "torch._C._distributed_rpc", + ): + if _c10d_name not in _sys.modules: + _stub = _types.ModuleType(_c10d_name) + # torch.distributed imports these names from _distributed_c10d; + # provide no-op dummies so the import doesn't raise AttributeError. + for _sym in ( + "FakeProcessGroup", + "ProcessGroup", + "Work", + "Store", + "PrefixStore", + "FileStore", + "TCPStore", + "HashStore", + "Reducer", + "Logger", + "DistributedDebugLevel", + "GradBucket", + "BuiltinCommHookType", + ): + setattr(_stub, _sym, _Dummy) + _sys.modules[_c10d_name] = _stub + + try: + import torch.distributed as _td + + for _attr, _stub in ( + ("is_initialized", lambda: False), + ("is_available", lambda: False), + ("get_rank", lambda: 0), + ("get_world_size", lambda: 1), + ("is_torchelastic_launched", lambda: False), + ): + if not hasattr(_td, _attr): + setattr(_td, _attr, _stub) + except ImportError: + pass + from unsloth.models._utils import resolve_attention_implementation from transformers import AutoModel, AutoModelForCausalLM @@ -1062,7 +1412,10 @@ def estimate_required_model_memory_gb( _determine_attention_impl_for_gpu_estimate(config) ) except Exception as e: - logger.warning( + # Log at debug: on Windows ROCm the torch.distributed stub does + # not implement Store, so this fires on every estimate call. + # It is expected and non-actionable -- eager is the safe fallback. + logger.debug( "Could not resolve attention implementation for '%s': %s", estimate_model, e, @@ -1552,14 +1905,35 @@ def apply_gpu_ids(gpu_ids) -> None: # parent process already set a ROCm visibility variable -- that # way a downstream ROCm process inherits the narrowed mask even # before Studio's hardware detection has classified the host. + # Final fallback: probe torch.version.hip so AMD workers without + # HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask. _inherits_rocm_visibility = ( "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ ) - if IS_ROCM or _inherits_rocm_visibility: + _is_rocm = IS_ROCM or _inherits_rocm_visibility + if not _is_rocm: + # torch.version.hip is a non-empty string on ROCm, None on CUDA. + # AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but + # still encode "rocm" in torch.__version__, matching detect_hardware(). + # Broad except: a probe failure must never crash a training worker. + try: + import torch as _torch + + _is_rocm = ( + getattr(_torch.version, "hip", None) is not None + or "rocm" in getattr(_torch, "__version__", "").lower() + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch ROCm probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_rocm: os.environ["HIP_VISIBLE_DEVICES"] = value os.environ["ROCR_VISIBLE_DEVICES"] = value _visible_gpu_count = None - if IS_ROCM or _inherits_rocm_visibility: + if _is_rocm: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value) else: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 5c42e890d1..e0ce02261b 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -15,6 +15,7 @@ import urllib.request from typing import Callable from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs _logger = logging.getLogger(__name__) @@ -106,6 +107,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non text = True, timeout = timeout, env = child_env_without_native_path_secret(), + **windows_hidden_subprocess_kwargs(), ) except subprocess.TimeoutExpired: return None diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index b9d63ecee9..31fb841414 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -29,7 +29,7 @@ import urllib.parse import urllib.request import zipfile from contextlib import contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace as dataclasses_replace try: from filelock import FileLock, Timeout as FileLockTimeout @@ -1336,6 +1336,20 @@ def direct_upstream_release_plan( torch_preference.selection_log, ) ) + elif host.has_rocm: + hip_asset = f"llama-{release_tag}-bin-win-hip-radeon-x64.zip" + hip_url = assets.get(hip_asset) + if hip_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = hip_asset, + url = hip_url, + source_label = "upstream", + install_kind = "windows-hip", + ) + ) cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip" cpu_url = assets.get(cpu_asset) if cpu_url: @@ -2731,12 +2745,30 @@ def detect_host() -> HostInfo: has_rocm = True break elif is_windows: - # Windows: prefer active probes that validate GPU presence + # Windows: prefer active probes that validate GPU presence. + # hipinfo / amd-smi are often NOT on PATH -- the HIP SDK installer + # sets HIP_PATH / ROCM_PATH but does not always add the bin dir to + # the system PATH. Mirror setup.ps1's fallback: check the env-var + # bin dirs before giving up so that `has_rocm` is not silently False + # on machines where the PATH is not yet updated. + def _resolve_exe(name: str) -> str | None: + """Return full path to `name`, checking PATH then HIP_PATH/ROCM_PATH bin.""" + found = shutil.which(name) + if found: + return found + for _env in ("HIP_PATH", "ROCM_PATH"): + _root = os.environ.get(_env) + if _root: + _candidate = os.path.join(_root, "bin", f"{name}.exe") + if os.path.isfile(_candidate): + return _candidate + return None + for _cmd, _check in ( (["hipinfo"], lambda out: "gcnarchname" in out.lower()), (["amd-smi", "list"], _amd_smi_has_gpu), ): - _exe = shutil.which(_cmd[0]) + _exe = _resolve_exe(_cmd[0]) if not _exe: continue try: @@ -5565,8 +5597,11 @@ def install_prebuilt( published_release_tag: str, *, simple_policy: bool = False, + override_has_rocm: bool = False, ) -> None: host = detect_host() + if override_has_rocm and not host.has_rocm: + host = dataclasses_replace(host, has_rocm = True) choice: AssetChoice | None = None try: with install_lock(install_lock_path(install_dir)): @@ -5700,6 +5735,17 @@ def parse_args() -> argparse.Namespace: action = "store_true", help = "Use the simplified platform-specific prebuilt selection policy.", ) + parser.add_argument( + "--has-rocm", + action = "store_true", + default = False, + help = ( + "Assert that an AMD ROCm GPU is present. When set, skips the internal " + "hipinfo/amd-smi probe and forces has_rocm=True in the host profile. " + "Used by setup.ps1/setup.sh to forward their own ROCm detection result " + "so the HIP llama.cpp prebuilt is selected even when hipinfo is not on PATH." + ), + ) resolve_group = parser.add_mutually_exclusive_group() resolve_group.add_argument( "--resolve-llama-tag", @@ -5820,6 +5866,7 @@ def main() -> int: published_repo = args.published_repo, published_release_tag = args.published_release_tag or "", simple_policy = args.simple_policy, + override_has_rocm = args.has_rocm, ) return EXIT_SUCCESS diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index c4524b06f5..786b668e7b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -54,12 +54,9 @@ PLATFORM_LACKS_TORCHCODEC_WHEEL = ( # ── ROCm / AMD GPU support ───────────────────────────────────────────────────── # Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on # download.pytorch.org. Entries are checked newest-first (>=). -# ROCm 7.2 only has torch 2.11.0 on download.pytorch.org, which exceeds the -# current torch upper bound (<2.11.0). Fall back to rocm7.1 (torch 2.10.0). -# TODO: uncomment rocm7.2 when torch upper bound is bumped to >=2.11.0 _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { - # (7, 2): "rocm7.2", # torch 2.11.0 -- requires torch>=2.11 - (7, 1): "rocm7.1", + (7, 2): "rocm7.2", # torch 2.11.0 + (7, 1): "rocm7.1", # torch 2.10.0 (7, 0): "rocm7.0", (6, 4): "rocm6.4", (6, 3): "rocm6.3", @@ -67,10 +64,47 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (6, 1): "rocm6.1", (6, 0): "rocm6.0", } + +# Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x). +_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { + "rocm7.2": ( + "torch>=2.11.0,<2.12.0", + "torchvision>=0.26.0,<0.27.0", + "torchaudio>=2.11.0,<2.12.0", + ), + # Default for rocm7.1 and earlier: torch 2.x below 2.11 + "_default": ( + "torch>=2.4,<2.11.0", + "torchvision>=0.19,<0.26.0", + "torchaudio>=2.4,<2.11.0", + ), +} _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") +# 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.amd.com/rocm/whl" +).rstrip("/") + +# 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 # (bnb PR #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every # AMD GPU. Drop the pin once bnb 0.50+ ships on PyPI. @@ -85,6 +119,16 @@ _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_rocm{VER}.dll. + # BNB auto-detects HIP version from torch.version.hip, which does not always + # match the DLL suffix in this prerelease wheel (e.g. torch 7.13 with a rocm72 + # DLL). We scan the installed wheel for the actual DLL name and set + # BNB_ROCM_VERSION accordingly in _install_bnb_windows_rocm() and worker.py. + "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" @@ -196,6 +240,151 @@ def _detect_rocm_version() -> tuple[int, int] | None: return None +def _pick_visible_index(num_tokens: int) -> int: + """Resolve HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES to an integer + index into a list of length num_tokens. Returns 0 (first GPU) for + unset, empty, '-1', UUID-style, or out-of-range values.""" + for _env in ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"): + _val = os.environ.get(_env) + if _val is None: + continue + _val = _val.strip() + if _val == "" or _val == "-1": + return 0 + _first = _val.split(",")[0].strip() + try: + _idx = int(_first) + if 0 <= _idx < num_tokens: + return _idx + except ValueError: + pass + return 0 + return 0 + + +def _detect_windows_gfx_arch() -> str | None: + """Return the gcnArchName on Windows (e.g. 'gfx1200'), or None. + + Probe order matches the PowerShell installer: env-var override first, + then hipinfo (PATH or HIP_PATH / ROCM_PATH bin), then amd-smi. Without + the amd-smi fallback, runtime-only AMD installs without hipinfo on PATH + return early and `studio update` cannot repair a CPU-only venv. + + On multi-GPU hosts, all detected gfx tokens are deduplicated (preserving + enumeration order) and HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES selects + which one to install for. The first GPU is used when no env var is set. + """ + import re + + # 1. Explicit override (matches PowerShell installer's env-var path). + _override = os.environ.get("UNSLOTH_ROCM_GFX_ARCH") + if _override and _override.strip(): + return _override.strip().lower() + + def _dedup_pick(tokens: list[str]) -> "str | None": + if not tokens: + return None + # Index into the full (ordered) list first so HIP_VISIBLE_DEVICES + # correctly addresses GPU N on mixed-arch hosts, then return that arch. + return tokens[_pick_visible_index(len(tokens))] + + # 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin. + hipinfo = shutil.which("hipinfo") + if not hipinfo: + for _env_var in ("HIP_PATH", "ROCM_PATH"): + _root = os.environ.get(_env_var) + if _root: + _candidate = os.path.join(_root, "bin", "hipinfo.exe") + if os.path.isfile(_candidate): + hipinfo = _candidate + break + if hipinfo: + try: + result = subprocess.run( + [hipinfo], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 10, + ) + if result.returncode == 0: + text = result.stdout.decode(errors = "replace") + # findall picks every gcnArchName line so multi-GPU hosts + # are enumerable and HIP_VISIBLE_DEVICES selects correctly. + _tokens = [ + t.strip().lower() + for t in re.findall(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text) + ] + _pick = _dedup_pick(_tokens) + if _pick: + return _pick + except Exception: + pass + + # 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo. + amd_smi = shutil.which("amd-smi") + if amd_smi: + for _args in (("static", "--asic"), ("list",)): + try: + result = subprocess.run( + [amd_smi, *_args], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 10, + ) + if result.returncode != 0: + continue + text = result.stdout.decode(errors = "replace") + # Prefer labelled gfx lines; fall back to bare tokens. + _labelled = re.findall( + r"(?im)^\s*(?:target_graphics_version|gfx|arch|asic)\b[^:\r\n]*:\s*(gfx[1-9][0-9a-z]{2,3})\b", + text, + ) + _tokens = [t.lower() for t in _labelled] + if not _tokens: + _tokens = re.findall(r"\bgfx[1-9][0-9a-z]{2,3}\b", text.lower()) + _pick = _dedup_pick(_tokens) + if _pick: + return _pick + except Exception: + continue + 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 _detect_bnb_rocm_dll_ver() -> str | None: + """Scan the installed bitsandbytes package for libbitsandbytes_rocm{VER}.dll. + + Returns the version suffix string (e.g. ``"72"``, ``"713"``) or ``None`` + if bitsandbytes is not installed or no ROCm DLL is found. Does NOT import + bitsandbytes — uses importlib.util.find_spec so it is safe to call before + BNB is imported. + """ + import glob + import importlib.util + import re + + spec = importlib.util.find_spec("bitsandbytes") + if spec is None or not spec.submodule_search_locations: + return None + all_vers: list[str] = [] + for pkg_dir in spec.submodule_search_locations: + for dll in glob.glob(os.path.join(pkg_dir, "libbitsandbytes_rocm*.dll")): + m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll)) + if m: + all_vers.append(m.group(1)) + # Pick the highest numeric suffix so that e.g. "713" wins over "72" when + # both variants are present in the wheel. Filesystem glob order is not + # guaranteed, so always sort rather than stopping at the first match. + return max(all_vers, key = lambda v: int(v)) if all_vers else None + + def _has_rocm_gpu() -> bool: """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" import re @@ -231,6 +420,26 @@ def _has_rocm_gpu() -> bool: if result.returncode == 0 and result.stdout.strip(): if check_fn(result.stdout): return True + # sysfs KFD topology fallback (Linux only) -- matches install.sh's + # runtime-only detection. On minimal package-managed installs (no + # rocminfo / no amd-smi GUI tools), the kernel exposes AMD GPUs via + # /sys/class/kfd so `studio update` can still detect the GPU and + # repair the venv. + if sys.platform != "win32": + try: + kfd_nodes = "/sys/class/kfd/kfd/topology/nodes" + if os.path.isdir(kfd_nodes): + for entry in os.listdir(kfd_nodes): + gpu_id_path = os.path.join(kfd_nodes, entry, "gpu_id") + try: + with open(gpu_id_path) as fh: + gpu_id = fh.read().strip() + except OSError: + continue + if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node + return True + except OSError: + pass return False @@ -252,23 +461,199 @@ def _has_usable_nvidia_gpu() -> bool: return result.returncode == 0 and "GPU " in result.stdout +def _detect_amd_gfx_codes() -> list[str]: + """Return the list of AMD gfx ISA strings visible to ROCm (e.g. ['gfx1151']). + + Probes rocminfo first, then falls back to ``amd-smi list`` and + ``amd-smi static --asic`` for runtime-only Radeon hosts that ship + amd-smi but no rocminfo. Returns an empty list when no probe yields + a gfx target. + """ + import re + + def _extract(text: str) -> list[str]: + codes = re.findall(r"gfx([1-9][0-9a-z]{2,3})", text.lower()) + return list(dict.fromkeys(f"gfx{c}" for c in codes)) + + probes: list[list[str]] = [] + if shutil.which("rocminfo"): + probes.append(["rocminfo"]) + if shutil.which("amd-smi"): + probes.append(["amd-smi", "list"]) + probes.append(["amd-smi", "static", "--asic"]) + for cmd in probes: + try: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 15, + ) + except Exception: + continue + if result.returncode != 0 or not result.stdout.strip(): + continue + codes = _extract(result.stdout) + if codes: + return codes + return [] + + +# Set by _ensure_rocm_torch() on success; suppresses the post-install AMD warning. +_rocm_windows_torch_installed: bool = False + + +def _install_bnb_windows_rocm() -> bool: + """Install the AMD Windows BNB prerelease wheel. Returns True on success. + + The continuous-release wheel is intentionally mismatched: the filename + encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the + wheel metadata reports 0.50.0.dev0. uv rejects this filename/metadata + mismatch; set UV_SKIP_WHEEL_FILENAME_CHECK=1 to bypass that check, then + restore the previous value (or remove the var) when done. + """ + _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") + if _bnb_win_url is None: + return False + _old = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") + os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1" + try: + _ok = pip_install_try( + "bitsandbytes (AMD Windows, pre-release main)", + "--force-reinstall", + "--no-cache-dir", + "--no-deps", + _bnb_win_url, + constrain = False, + ) + finally: + if _old is None: + os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) + else: + os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _old + if not _ok: + return False + # After install: detect the actual ROCm DLL suffix shipped in the wheel and + # set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of + # what torch.version.hip reports. The wheel may ship an older suffix (e.g. + # "72") while torch reports a newer HIP version (e.g. 7.13); the env var + # override ensures bitsandbytes does not fail looking for a non-existent DLL. + # The worker subprocess inherits this env var automatically. + # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run). + if "BNB_ROCM_VERSION" not in os.environ: + _ver = _detect_bnb_rocm_dll_ver() or "72" + os.environ["BNB_ROCM_VERSION"] = _ver + return True + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. - Runs only on Linux x86_64 hosts where an AMD GPU is present and the - ROCm runtime is detectable (rocminfo / amd-smi / hipconfig / - rocm-core package). No-op when torch already links against HIP - (ROCm), on Windows / macOS, on non-x86_64 Linux (PyTorch does not - publish ROCm wheels for aarch64 / arm64), or on mixed AMD+NVIDIA - hosts (NVIDIA takes precedence). + On Linux x86_64: uses pytorch.org ROCm wheel index tags. + 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. """ - # Explicit OS / architecture guards so the helper is safe to call - # from any context -- PyTorch only publishes ROCm wheels for - # linux_x86_64, so aarch64 / arm64 hosts must skip this repair path - # instead of failing the update with a missing-wheel error. - if IS_WINDOWS or IS_MACOS: + global _rocm_windows_torch_installed + # setup.ps1 sets this when it already installed AMD wheels; skip the probe + # only when torch is actually importable as ROCm. If the venv was wiped + # between runs, the stale env-var would suppress a needed reinstall. + if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": + _torch_ok = False + try: + _probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch; " + "hip=getattr(torch.version,'hip','') or ''; " + "import sys; " + "sys.exit(0 if (hip or 'rocm' in torch.__version__.lower()) else 1)" + ), + ], + stdout = subprocess.DEVNULL, + stderr = subprocess.DEVNULL, + timeout = 90, + ) + _torch_ok = _probe.returncode == 0 + except (OSError, subprocess.TimeoutExpired): + pass + if _torch_ok: + _rocm_windows_torch_installed = True + # setup.ps1 already installed ROCm torch, but we still need to install + # the AMD Windows BNB wheel here -- the PyPI bitsandbytes wheel ships + # only CUDA DLLs and will fail to load on ROCm. + _install_bnb_windows_rocm() + return + # torch was wiped between runs; fall through to the full install path + if IS_MACOS: return + + if IS_WINDOWS: + if _has_usable_nvidia_gpu(): + return + gfx_arch = _detect_windows_gfx_arch() + if not gfx_arch: + return # no AMD GPU visible via hipinfo + # Probe whether torch already links against HIP. + _torch_already_rocm = False + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch; " + "hip=getattr(torch.version,'hip','') or ''; " + "ver=torch.__version__; " + "print('yes' if hip or 'rocm' in ver.lower() else '')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 90, + ) + if probe.returncode == 0 and probe.stdout.decode().strip() == "yes": + _torch_already_rocm = True + except (OSError, subprocess.TimeoutExpired): + pass + if not _torch_already_rocm: + 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 + print(f" {gfx_arch} (Windows) -- installing torch from {index_url}") + pip_install( + f"ROCm torch (Windows, {gfx_arch})", + "--force-reinstall", + "--index-url", + index_url, + "torch", + "torchvision", + "torchaudio", + constrain = False, + ) + # ROCm torch is installed (or already was); flag it so later install + # phases do not overwrite it with the generic CPU torch wheel. BNB is + # a separate dependency -- a BNB install failure must NOT roll the + # torch ROCm install back. + _rocm_windows_torch_installed = True + # Always install AMD Windows bitsandbytes -- the PyPI wheel ships only + # CUDA DLLs and will fail to load on ROCm. Install even when torch was + # already a ROCm build so that `studio update` repairs a broken bnb. + if not _install_bnb_windows_rocm(): + print( + " Warning: AMD Windows bitsandbytes install failed; " + "ROCm torch is installed but bitsandbytes may need manual install" + ) + return + + # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ── if platform.machine().lower() not in {"x86_64", "amd64"}: return # NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable @@ -297,11 +682,19 @@ def _ensure_rocm_torch() -> None: [ sys.executable, "-c", - "import torch; print(getattr(torch.version,'hip','') or '')", + ( + "import torch; " + "hip=getattr(torch.version,'hip','') or ''; " + "ver=getattr(torch,'__version__','').lower(); " + # Print the HIP version when present (back-compat), else + # "rocm" sentinel when only torch.__version__ flags ROCm + # (AMD SDK / Radeon wheels). Empty string = CPU/CUDA. + "print(hip if hip else ('rocm' if 'rocm' in ver else ''))" + ), ], stdout = subprocess.PIPE, stderr = subprocess.DEVNULL, - timeout = 30, + timeout = 90, ) except (OSError, subprocess.TimeoutExpired): probe = None @@ -313,7 +706,83 @@ def _ensure_rocm_torch() -> None: rocm_torch_ready = has_hip_torch - if not has_hip_torch: + # Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1 + # in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0 + # with the real fix, so route those hosts there instead of the generic + # pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override. + # On mixed hosts (Strix iGPU + non-Strix dGPU), only route to the AMD + # per-gfx index when the GPU HIP will actually run on is the Strix one -- + # otherwise the dGPU would get an incompatible wheel. Use HIP_VISIBLE_DEVICES + # to determine the runtime target. + _strix_override_url: "str | None" = None + _strix_override_pkgs: "tuple[str, str, str] | None" = None + if ver < (7, 2): + gfx_codes = _detect_amd_gfx_codes() + _strix_gfx = {"gfx1151", "gfx1150"} + _detected_strix = _strix_gfx.intersection(gfx_codes) + if _detected_strix: + # Pick the runtime-visible GPU. If HIP_VISIBLE_DEVICES selects a + # specific index into gfx_codes, use that gfx; else default to the + # first listed GPU. Skip the override unless the resolved GPU is + # Strix. + _runtime_gfx = ( + gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None + ) + if _runtime_gfx in _strix_gfx: + _selected_gfx = _runtime_gfx + _amd_mirror = ( + os.environ.get("UNSLOTH_AMD_ROCM_MIRROR") + or "https://repo.amd.com/rocm/whl" + ).rstrip("/") + _strix_override_url = f"{_amd_mirror}/{_selected_gfx}/" + _strix_override_pkgs = ( + "torch>=2.11.0,<2.12.0", + # Pin torchvision/torchaudio to the 2.11.x-compatible range. + # The install uses --index-url (exclusive, no PyPI fallback), + # so bare unversioned names risk resolving a build from AMD's + # index that targets a different torch major (e.g. 0.27 built + # against torch 2.12), which would fail at runtime with an + # ABI/version mismatch. Matches _ROCM_TORCH_CONSTRAINT["rocm7.2"]. + "torchvision>=0.26.0,<0.27.0", + "torchaudio>=2.11.0,<2.12.0", + ) + print( + f"\n {_selected_gfx} (AMD Strix) is the runtime target with ROCm " + f"{ver[0]}.{ver[1]}.\n" + f" ROCm 7.1 has a known _grouped_mm segfault on this GPU;\n" + f" routing torch install to AMD's arch-specific index\n" + f" ({_strix_override_url}) which serves torch 2.11.0+rocm7.13.0\n" + f" with the upstream fix.\n" + ) + else: + _gfx_str = ", ".join(sorted(_detected_strix)) + print( + f"\n Strix GPU ({_gfx_str}) present but HIP_VISIBLE_DEVICES " + f"selects a non-Strix runtime target ({_runtime_gfx});\n" + f" skipping AMD per-gfx index override.\n" + ) + + # Strix override on ROCm 7.1 must fire even when has_hip_torch is True -- + # an existing torch with `torch.version.hip == "7.1"` is exactly the broken + # combo the override is meant to repair, so skipping it leaves users on + # the known _grouped_mm segfault. + if _strix_override_url is not None and _strix_override_pkgs is not None: + index_url = _strix_override_url + _torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs + print(f" Strix ROCm 7.1 override -- installing torch from {index_url}") + pip_install( + "ROCm torch (Strix arch-specific)", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + rocm_torch_ready = True + elif not has_hip_torch: # Select best matching wheel tag (newest ROCm version <= installed) tag = next( ( @@ -331,13 +800,16 @@ def _ensure_rocm_torch() -> None: else: index_url = f"{_PYTORCH_WHL_BASE}/{tag}" print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( + tag, _ROCM_TORCH_PKG_SPECS["_default"] + ) pip_install( f"ROCm torch ({tag})", "--force-reinstall", "--no-cache-dir", - "torch>=2.4,<2.11.0", - "torchvision<0.26.0", - "torchaudio<2.11.0", + _torch_pkg, + _vision_pkg, + _audio_pkg, "--index-url", index_url, constrain = False, @@ -346,7 +818,9 @@ def _ensure_rocm_torch() -> None: # Install bitsandbytes only when torch links against ROCm. Prefers the # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls - # back to PyPI when the pre-release URL is unreachable. + # back to PyPI when the pre-release wheel cannot be installed. Use pip for + # the pre-release wheel because uv rejects the wheel's filename/metadata + # version mismatch. if rocm_torch_ready: _bnb_url = _bnb_rocm_prerelease_url() _bnb_installed = False @@ -358,11 +832,12 @@ def _ensure_rocm_torch() -> None: "--no-deps", _bnb_url, constrain = False, + force_pip = True, ) if not _bnb_installed: print( _red( - " bnb pre-release unreachable; falling back to PyPI " + " bnb pre-release install failed; falling back to PyPI " "(4-bit decode will be broken on ROCm)" ) ) @@ -809,6 +1284,7 @@ def pip_install_try( label: str, *args: str, constrain: bool = True, + force_pip: bool = False, ) -> bool: """Like pip_install but returns False on failure instead of exiting. For optional installs with a follow-up fallback. @@ -819,7 +1295,7 @@ def pip_install_try( constraint_args_pip = ["-c", str(CONSTRAINTS)] constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)] - if USE_UV: + if USE_UV and not force_pip: cmd = _build_uv_cmd(args) + constraint_args_uv else: cmd = _build_pip_cmd(args) + constraint_args_pip @@ -948,8 +1424,12 @@ def install_python_stack() -> int: base_total = 10 if IS_WINDOWS else 11 if IS_MACOS: base_total -= 1 # triton step is skipped on macOS - if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: - base_total += 3 + if not IS_MACOS and not NO_TORCH: + base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms + if not IS_WINDOWS: + base_total += ( + 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only + ) _TOTAL = (base_total - 1) if skip_base else base_total # 1. Try to use uv for faster installs (must happen before pip upgrade @@ -1121,12 +1601,12 @@ def install_python_stack() -> int: # 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the # venv received CPU-only torch (common when pip resolves torch from PyPI). # Must come immediately after base packages so torch is present for inspection. - if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + if not IS_MACOS and not NO_TORCH: _progress("ROCm torch check") _ensure_rocm_torch() - # Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows. - # Detect and warn so users know manual steps are needed for GPU training. + # Windows + AMD GPU: if ROCm torch was not installed (wrong Python version + # or unknown ROCm version), warn the user. if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu(): # Validate actual AMD GPU presence (not just tool existence) import re as _re_win @@ -1155,14 +1635,14 @@ def install_python_stack() -> int: if _wr.returncode == 0 and _check_fn(_wr.stdout): _win_amd_gpu = True break - if _win_amd_gpu: + if _win_amd_gpu and not _rocm_windows_torch_installed: _safe_print( _dim(" Note:"), - "AMD GPU detected on Windows. ROCm-enabled PyTorch must be", + "AMD GPU detected but ROCm PyTorch could not be auto-installed.", ) _safe_print( " " * 8, - "installed manually. See: https://docs.unsloth.ai/get-started/install-and-update/amd", + "Manual install may be required. See: https://docs.unsloth.ai/get-started/install-and-update/amd", ) # 3. Extra dependencies @@ -1189,10 +1669,17 @@ def install_python_stack() -> int: _progress("dependency overrides (skipped, no torch)") else: _progress("dependency overrides") + _override_extra_args: tuple[str, ...] = () + if _rocm_windows_torch_installed: + # torchao in overrides.txt declares torch as a dependency; without + # --no-deps uv would resolve and install CPU torch from PyPI, + # overwriting the AMD ROCm wheels we just installed. + _override_extra_args = ("--no-deps",) pip_install( "Installing dependency overrides", "--force-reinstall", "--no-cache-dir", + *_override_extra_args, req = REQ_ROOT / "overrides.txt", ) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index afa00409fb..e276b2151d 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -685,13 +685,228 @@ if (-not $HasNvidiaSmi) { } } } +# ── AMD ROCm detection (Windows): probe hipinfo/amd-smi for actual GPU ── +$HasROCm = $false +$HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility) +$ROCmGpuLabel = $null +$script:ROCmGfxArch = $null if (-not $HasNvidiaSmi) { + # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). + # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. + $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue + if (-not $hipinfoExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + $hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + if ($hipRoot) { + $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" + if (Test-Path $hipinfoCandidate) { + substep "[WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" "Yellow" + substep " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" "Yellow" + substep " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" "Yellow" + $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } + } else { + substep "[WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" "Yellow" + substep " HIP SDK install may be incomplete -- re-install from:" "Yellow" + substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" + } + } + } + if ($hipinfoExe) { + $HipSdkInstalled = $true # binary found → SDK is installed regardless of device state + try { + $hipOut = & $hipinfoExe.Source 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { + $HasROCm = $true + $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }) + $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } + if ($_hipAllArches.Count -gt 0) { + $script:ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] } + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } elseif ($LASTEXITCODE -ne 0) { + # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected") + $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1) + substep "[WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" "Yellow" + substep " $firstLine" "Yellow" + substep " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" + } + } catch {} + } + # amd-smi fallback: HIP runtime present but hipinfo unavailable (no full HIP SDK). + # Confirms GPU visibility via 'list', then attempts 'static --asic' to extract + # the gfx arch that hipinfo would have provided. Critical for Strix Halo + # (gfx1151) and other iGPUs where only the HIP runtime is installed. + if (-not $HasROCm) { + $amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiExe) { + try { + $smiOut = & $amdSmiExe.Source list 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { + $HasROCm = $true + # Attempt 1: newer amd-smi versions embed the gfx arch in list output. + # Collect ALL gfx tokens in output order so that on mixed-arch systems + # we can honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES and pick the + # arch for the *runtime-visible* GPU rather than always the first one. + # Do NOT deduplicate: a dual same-arch system (e.g. two gfx1151 APUs) + # must produce a 2-element array so HIP_VISIBLE_DEVICES=1 selects the + # second GPU rather than triggering a false out-of-range warning. + # Note: this mapping assumes amd-smi lists GPUs in the same order as + # HIP enumerates them (both follow PCI bus order in practice); it may + # give the wrong arch when GPU indices are non-contiguous (very rare). + $allGfxArches = @([regex]::Matches($smiOut, '(?i)\b(gfx\d+[a-z]?)\b') | + ForEach-Object { $_.Groups[1].Value.ToLower() }) + if ($allGfxArches.Count -gt 0) { + # Resolve which GPU index is runtime-visible. When a single + # integer index is set, use it; fall back to index 0 otherwise + # (comma-separated lists or unset → first GPU, same as before). + $visGpu = if ($env:HIP_VISIBLE_DEVICES) { $env:HIP_VISIBLE_DEVICES } + elseif ($env:ROCR_VISIBLE_DEVICES) { $env:ROCR_VISIBLE_DEVICES } + else { $null } + $gpuIdx = 0 + if ($visGpu -match '^\s*(\d+)\s*$') { $gpuIdx = [int]$Matches[1] } + if ($gpuIdx -ge $allGfxArches.Count) { + substep "[WARN] HIP/ROCR_VISIBLE_DEVICES index $gpuIdx is out of range ($($allGfxArches.Count) GPU(s) detected); defaulting to GPU 0 for arch selection" "Yellow" + $gpuIdx = 0 + } + $script:ROCmGfxArch = $allGfxArches[$gpuIdx] + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" + } else { + # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, + # including the GFX target needed for wheel index selection. + $smiAsicOut = "" + try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {} + if ($smiAsicOut -match "(?i)\b(gfx\d+[a-z]?)\b") { + $script:ROCmGfxArch = $Matches[1].ToLower() + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" + } elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") { + $ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } + } + } catch {} + } + } + # WMI fallback: AMD GPU in device list but no HIP SDK → guide the user. + # WMI gives a marketing name (e.g. "AMD Radeon 890M") but never a gfx arch. + # $HasROCm is intentionally NOT set here — we cannot confirm ROCm runtime + # support without hipinfo or amd-smi. The name is saved to $ROCmGpuLabel + # so the name-based inference below can still attempt an arch lookup. + if (-not $HasROCm) { + try { + $wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue | + Where-Object { $_.Name -match "AMD|Radeon" } | + Select-Object -First 1 + if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name } + } catch {} + } + # ── Arch resolution: env-var override → name inference ────────────────── + # Runs after all probe methods. Covers users whose amd-smi version is too + # old to report the GFX target and who don't have hipinfo (HIP-runtime-only + # installs, common on Strix Halo / iGPU systems). + if ($HasROCm -and -not $script:ROCmGfxArch) { + # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. + if ($env:UNSLOTH_ROCM_GFX_ARCH) { + $script:ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower() + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $script:ROCmGfxArch" "Cyan" + } + # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). + # Ordered most-specific first; first match wins. + elseif ($ROCmGpuLabel) { + $nameArchTable = @( + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 + @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) + @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) + @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop + @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 + @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) + ) + foreach ($row in $nameArchTable) { + if ($ROCmGpuLabel -match $row.P) { + $script:ROCmGfxArch = $row.A + $ROCmGpuLabel = "AMD ROCm ($script:ROCmGfxArch)" + substep "gfx arch inferred from GPU name: $script:ROCmGfxArch" "Cyan" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$script:ROCmGfxArch to skip inference next time" "Cyan" + break + } + } + } + } + # Capture ROCm version early for display and wheel selection. + # Run whenever the HIP SDK binary is present, not just when the device is accessible -- + # hipconfig --version works even when hipinfo reports no ROCm device (driver issue). + if ($HasROCm -or $HipSdkInstalled) { + $script:ROCmVersion = $null + $hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue + if (-not $hipConfigExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + if ($hipRoot) { + $hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe" + if (Test-Path $hipConfigCandidate) { + $hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + substep "[WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" "Yellow" + $hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate } + } + } + } + if ($hipConfigExe) { + try { + $hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0) { + $hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim() + if ($hipVerLine -match '(\d+\.\d+)') { + $script:ROCmVersion = $Matches[1] + $script:ROCmVersionFull = $hipVerLine + } + } + } catch {} + } + if (-not $script:ROCmVersion) { + $amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiVer) { + try { + $smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { $script:ROCmVersion = $Matches[1] } + } catch {} + } + } + } +} + +if ($HasNvidiaSmi) { + step "gpu" "NVIDIA GPU detected" +} elseif ($HasROCm) { + step "gpu" $ROCmGpuLabel + $hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" } + substep "HIP SDK: $hipSdkPath" + if ($script:ROCmVersionFull) { substep "hipconfig: $script:ROCmVersionFull" } +} elseif ($HipSdkInstalled -and $ROCmGpuLabel) { + # HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue) + $sdkVer = if ($script:ROCmVersionFull) { " (HIP $script:ROCmVersionFull)" } else { "" } Write-Host "" - step "gpu" "none (chat-only / GGUF)" "Yellow" - substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow" + substep " This is a driver issue, not an SDK issue." "Yellow" + substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" + substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" +} elseif ($ROCmGpuLabel) { + Write-Host "" + step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "Install the HIP SDK for ROCm GPU inference:" "Yellow" + substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" Write-Host "" } else { - step "gpu" "NVIDIA GPU detected" + Write-Host "" + step "gpu" "none (chat-only / GGUF)" "Yellow" + substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow" + Write-Host "" } # ============================================ @@ -1102,6 +1317,13 @@ if (-not $CudaArch) { step "cuda" "skipped (no NVIDIA GPU detected)" "Yellow" } +if ($HasROCm) { + $rocmVerLabel = if ($script:ROCmVersionFull) { "ROCm $script:ROCmVersionFull" } elseif ($script:ROCmVersion) { "ROCm $script:ROCmVersion" } else { "ROCm (version unknown)" } + step "rocm" $rocmVerLabel +} elseif ($ROCmGpuLabel) { + step "rocm" "HIP SDK not found -- GPU-accelerated training unavailable" "Yellow" +} + # ============================================ # 1f. Node.js / npm (skip if pip-installed or Tauri -- only needed for frontend build) # ============================================ @@ -1516,7 +1738,7 @@ if (-not $PythonCmd) { exit 1 } -substep "Using $PythonCmd ($(& $PythonCmd --version 2>&1))" +substep "Python found: $PythonCmd" # The venv must already exist (created by install.ps1); this script only # updates packages. UNSLOTH_STUDIO_HOME (or STUDIO_HOME alias) overrides the @@ -1684,6 +1906,13 @@ if (-not (Test-Path -LiteralPath $VenvDir)) { exit 1 } else { substep "reusing existing virtual environment at $VenvDir" + $_venvPyExe = Join-Path $VenvDir "Scripts\python.exe" + if (Test-Path -LiteralPath $_venvPyExe) { + try { + $_venvPyVer = (& $_venvPyExe --version 2>&1 | Out-String).Trim() + if ($_venvPyVer) { substep $_venvPyVer } + } catch {} + } } # pip and python write to stderr even on success (progress bars, warnings). @@ -1795,9 +2024,91 @@ 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. +# ── AMD Windows ROCm torch override ────────────────────────────────────────── +# 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 +$ROCmIndexUrl = $null +if ($HasROCm -and $CuTag -eq "cpu") { + $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 + } + # gfx120X and Strix have a null _grouped_mm kernel on torch <2.11.0. + # Mirrors the $torchFloorMap in install.ps1 so both installers enforce + # the same floor and ceiling when pulling from AMD's per-arch index. + $torchFloorMap = @{ + "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" + "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + } + # Companion ranges for torchvision/torchaudio -- must stay in sync with the + # torch ceiling so pip can always find a consistent trio on AMD's per-arch + # index. AMD publishes each package independently and may add a newer + # torchvision (e.g. 0.27 for torch 2.12) before removing 0.26, which would + # cause pip to resolve an ABI-incompatible set if these are left bare. + # Matches _ROCM_TORCH_PKG_SPECS["rocm7.2"] in install_python_stack.py. + # Bump all three ceilings together when torch 2.12.x is validated. + $torchvisionFloorMap = @{ + "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" + "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + } + $torchaudioFloorMap = @{ + "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" + "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + } + $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } + $ROCmTorchSpec = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { "torch" } + $ROCmVisionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $ROCmAudioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + if ($archFamily) { + $ROCmIndexUrl = "$amdIndexBase/$archFamily/" + } elseif ($ROCmGfxArch) { + # GPU arch detected but not in the supported wheel map — warn explicitly + # so the user knows why they are getting CPU PyTorch instead of ROCm. + substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow" + substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx90a, gfx908" "Yellow" + } else { + # HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not + # readable — warn rather than silently falling back to CPU PyTorch. + substep "[WARN] AMD GPU detected (HIP SDK present) but GPU arch could not be read -- falling back to CPU-only PyTorch" "Yellow" + substep " Arch detection requires hipinfo to report gcnArchName. Re-install the HIP SDK if this is unexpected." "Yellow" + } +} + $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } -if ($CuTag -eq "cpu") { +if ($ROCmIndexUrl) { + substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..." + if ($ROCmTorchSpec -ne "torch") { + substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan" + } + if ($script:UnslothVerbose) { + Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | 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 + $ROCmIndexUrl = $null + } else { + # Tell install_python_stack.py to skip probe + suppress manual-install warning. + $env:UNSLOTH_ROCM_TORCH_INSTALLED = "1" + } +} + +if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { substep "installing PyTorch (CPU-only)..." if ($script:UnslothVerbose) { Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/cpu" @@ -1812,7 +2123,7 @@ if ($CuTag -eq "cpu") { Write-Host $output -ForegroundColor Red exit 1 } -} else { +} 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) { @@ -2115,6 +2426,23 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { substep "installing prebuilt llama.cpp bundle (preferred path)..." if (Test-Path -LiteralPath $LlamaCppDir) { substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement" + # If the existing install is the wrong kind (e.g. windows-cpu on a ROCm + # machine that should have windows-hip), remove it so the installer is + # forced to download the correct variant rather than skipping on tag match. + $existingMetaPath = Join-Path $LlamaCppDir "UNSLOTH_PREBUILT_INFO.json" + if (Test-Path $existingMetaPath) { + try { + $existingMeta = Get-Content $existingMetaPath -Raw | ConvertFrom-Json + $existingKind = $existingMeta.install_kind + $expectedKind = if ($HasROCm) { "windows-hip" } elseif ($HasNvidiaSmi) { "windows-cuda" } else { "windows-cpu" } + if ($existingKind -and $existingKind -ne $expectedKind) { + substep "Removing mismatched llama.cpp install (found '$existingKind', need '$expectedKind')..." + Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue + } + } catch { + # unreadable metadata -- let the installer handle it + } + } } # why: install_llama_prebuilt.py uses os.replace(), which would displace # an unrelated $env:UNSLOTH_STUDIO_HOME\llama.cpp before the source-build @@ -2129,6 +2457,9 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { "--published-repo", $HelperReleaseRepo, "--simple-policy" ) + if ($HasROCm) { + $prebuiltArgs += "--has-rocm" + } if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) } diff --git a/studio/setup.sh b/studio/setup.sh index f2395255c0..f178a5e5d7 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -659,6 +659,81 @@ if [ "$_NEED_T5_INSTALL" = true ]; then fi fi +# ── GPU detection summary (mirrors setup.ps1 step "gpu" block) ── +_setup_amd_detected=false +_setup_gfx_all="" +_setup_mkt="" +if command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then + _setup_amd_detected=true + _setup_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(rocminfo 2>/dev/null | awk -F': ' \ + '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) +elif command -v amd-smi >/dev/null 2>&1 && \ + amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then + _setup_amd_detected=true + _setup_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_setup_gfx_all" ] && \ + _setup_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _setup_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ + '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) +fi + +if command -v nvidia-smi >/dev/null 2>&1 && \ + nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + step "gpu" "NVIDIA GPU detected" +elif [ "$_setup_amd_detected" = true ]; then + _setup_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _setup_vis_idx=0 + if [ -n "$_setup_vis" ] && [ "$_setup_vis" != "-1" ]; then + _setup_first="${_setup_vis%%,*}" + case "$_setup_first" in ''|*[!0-9]*) ;; *) _setup_vis_idx=$_setup_first ;; esac + fi + _setup_gfx=$(printf '%s\n' "$_setup_gfx_all" | awk -v idx="$_setup_vis_idx" \ + 'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }') + # UNSLOTH_ROCM_GFX_ARCH env override (mirrors setup.ps1) + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + _setup_gfx="${UNSLOTH_ROCM_GFX_ARCH}" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_setup_gfx" + # Name-based arch inference when tools don't report gfx (mirrors setup.ps1 nameArchTable) + elif [ -z "$_setup_gfx" ] && [ -n "$_setup_mkt" ]; then + case "$_setup_mkt" in + *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 + *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 iGPU + *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 iGPU + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop + *"RX 7600"*) _setup_gfx="gfx1102" ;; # RDNA 3 + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU + esac + if [ -n "$_setup_gfx" ]; then + substep "gfx arch inferred from GPU name: $_setup_gfx" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_setup_gfx to skip inference next time" + fi + fi + # ROCm version via hipconfig, then amd-smi + _setup_rocm_ver="" + if command -v hipconfig >/dev/null 2>&1; then + _setup_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true) + fi + if [ -z "$_setup_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then + _setup_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true) + fi + if [ -n "$_setup_gfx" ]; then + step "gpu" "AMD ROCm ($_setup_gfx)" + else + step "gpu" "AMD ROCm" + fi + _setup_rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}" + substep "ROCm: $_setup_rocm_root" + [ -n "$_setup_rocm_ver" ] && substep "hipconfig: $_setup_rocm_ver" + [ -n "$_setup_mkt" ] && [ -n "$_setup_gfx" ] && substep "GPU: $_setup_mkt" +else + step "gpu" "none (chat-only / GGUF)" "$C_WARN" + substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." +fi + # ── 7. Prefer prebuilt llama.cpp bundles before any source build path ── # Nest llama.cpp under $STUDIO_HOME only for real env-overrides; legacy # default keeps ~/.unsloth/llama.cpp so pre-PR builds are still discovered. @@ -817,6 +892,22 @@ else fi fi +# Source-built llama.cpp installs do not have the prebuilt metadata used above +# for exact release matching. Reuse a complete local source build unless the +# caller explicitly requested a rebuild or a PR-specific llama.cpp checkout. +if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && \ + [ "$_LLAMA_FORCE_COMPILE" != "1" ] && \ + [ -z "$_LLAMA_PR" ] && \ + [ -x "$LLAMA_CPP_DIR/build/bin/llama-server" ] && \ + [ -x "$LLAMA_CPP_DIR/build/bin/llama-quantize" ]; then + step "llama.cpp" "existing source build found; skipping rebuild" + ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + : > "$LLAMA_CPP_DIR/$_STUDIO_OWNED_MARKER" 2>/dev/null || true + fi + _NEED_LLAMA_SOURCE_BUILD=false +fi + # ── 8. WSL: pre-install GGUF build dependencies for fallback source builds ── # On WSL, sudo requires a password and can't be entered during GGUF export # (runs in a non-interactive subprocess). Install build deps here instead. @@ -974,7 +1065,8 @@ else fi if [ "$BUILD_OK" = true ]; then - CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON" + # Set Release explicitly (llama.cpp only defaults to it on non-MSVC/Xcode). + CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release -DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON" _TRY_METAL_CPU_FALLBACK=false _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" _HOST_MACHINE="$(uname -m 2>/dev/null || true)" @@ -1088,6 +1180,32 @@ else _BUILD_DESC="building (ROCm)" CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON" + + # ROCm 7.x ships clang-20 which on Ubuntu 24.04+ defaults to the + # highest-numbered gcc lib dir (/usr/lib/gcc/x86_64-linux-gnu/14/) + # which contains runtime objects but NOT C++ headers, causing: + # fatal error: 'cstdlib' file not found + # Find the newest gcc install dir that actually has both the + # runtime dir AND /usr/include/c++/ headers, then pass it + # to clang via --gcc-install-dir so HIP builds succeed. + _GCC_INSTALL_DIR="" + _gcc_pm="$(gcc -print-multiarch 2>/dev/null)" + case "$_gcc_pm" in + *-linux-gnu*) _GCC_MULTIARCH="$_gcc_pm" ;; + *) _GCC_MULTIARCH="$(uname -m)-linux-gnu" ;; + esac + for _gcc_ver in 14 13 12 11; do + if [ -d "/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver/include" ] && \ + [ -d "/usr/include/c++/$_gcc_ver" ]; then + _GCC_INSTALL_DIR="/usr/lib/gcc/$_GCC_MULTIARCH/$_gcc_ver" + break + fi + done + if [ -n "$_GCC_INSTALL_DIR" ]; then + CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_HIP_FLAGS=--gcc-install-dir=\"$_GCC_INSTALL_DIR\"" + substep "ROCm HIP gcc install dir: $_GCC_INSTALL_DIR" + fi + export ROCM_PATH="$ROCM_ROOT" export HIP_PATH="$ROCM_ROOT" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 89ec32fba5..e20fd1ca86 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -193,10 +193,10 @@ _result=$(run_func "$_dir") assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result" rm -rf "$_dir" -# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.1 (capped due to torch <2.11.0) +# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.2 _dir=$(make_mock_amd_smi "7.2") _result=$(run_func "$_dir") -assert_eq "ROCm 7.2 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +assert_eq "ROCm 7.2 -> rocm7.2" "https://download.pytorch.org/whl/rocm7.2" "$_result" rm -rf "$_dir" # 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence @@ -231,10 +231,10 @@ _result=$(run_func "$_dir") assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result" rm -rf "$_dir" -# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.1 (capped) +# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.2 (capped to latest known) _dir=$(make_mock_amd_smi "8.0") _result=$(run_func "$_dir") -assert_eq "ROCm 8.0 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +assert_eq "ROCm 8.0 -> rocm7.2 (capped)" "https://download.pytorch.org/whl/rocm7.2" "$_result" rm -rf "$_dir" # 18) Malformed amd-smi output (empty version field) -> cpu diff --git a/tests/studio/install/conftest.py b/tests/studio/install/conftest.py new file mode 100644 index 0000000000..8738ef2319 --- /dev/null +++ b/tests/studio/install/conftest.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Pytest configuration for studio/install tests. + +install_python_stack.py does ``from backend.utils.wheel_utils import ...`` +which requires the ``studio/`` directory to be on sys.path. When tests are +run from the repo root (the normal case), the studio package is not +automatically importable, so we add it here. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +# /studio → makes `backend` importable as a package +_STUDIO_DIR = Path(__file__).resolve().parents[3] / "studio" +if str(_STUDIO_DIR) not in sys.path: + sys.path.insert(0, str(_STUDIO_DIR)) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index e6f1ae1c65..ecb1b9be15 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -51,6 +51,9 @@ _ensure_rocm_torch = stack_mod._ensure_rocm_torch _has_rocm_gpu = stack_mod._has_rocm_gpu _has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu _ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX +_windows_rocm_index_url = stack_mod._windows_rocm_index_url +_detect_windows_gfx_arch = stack_mod._detect_windows_gfx_arch +_install_bnb_windows_rocm = stack_mod._install_bnb_windows_rocm def _extract_sh_function_body(source: str, name: str) -> str: @@ -596,6 +599,7 @@ class TestEnsureRocmTorch: _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -616,7 +620,9 @@ class TestEnsureRocmTorch: assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 assert "bitsandbytes" in str(mock_pip_try.call_args_list[0]) + assert mock_pip_try.call_args.kwargs["force_pip"] is True + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @@ -648,6 +654,7 @@ class TestEnsureRocmTorch: _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @@ -662,12 +669,13 @@ class TestEnsureRocmTorch: captured = capsys.readouterr() assert "unreadable" in captured.out + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) - def test_rocm_72_selects_71_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): - """ROCm 7.2 should select rocm7.1 tag (capped, not in mapping).""" + def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): + """ROCm 7.2 should select rocm7.2 tag (now in mapping with torch 2.11.0).""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" @@ -675,8 +683,9 @@ class TestEnsureRocmTorch: with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() torch_call = mock_pip.call_args_list[0] - assert "rocm7.1" in str(torch_call) + assert "rocm7.2" in str(torch_call) + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -696,6 +705,7 @@ class TestEnsureRocmTorch: assert mock_pip.call_count == 1 assert "rocm7.1" in str(mock_pip.call_args_list[0]) assert mock_pip_try.call_count >= 1 + assert mock_pip_try.call_args.kwargs["force_pip"] is True @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -720,9 +730,10 @@ class TestRocmTorchIndex: keys = list(_ROCM_TORCH_INDEX.keys()) assert keys == sorted(keys, reverse = True) - def test_rocm_72_not_in_mapping(self): - """ROCm 7.2 should NOT be in the active mapping (torch 2.11.0 exceeds bound).""" - assert (7, 2) not in _ROCM_TORCH_INDEX + def test_rocm_72_in_mapping(self): + """ROCm 7.2 should be in the active mapping (torch 2.11.0 now supported).""" + assert (7, 2) in _ROCM_TORCH_INDEX + assert _ROCM_TORCH_INDEX[(7, 2)] == "rocm7.2" def test_rocm_71_maps_correctly(self): assert _ROCM_TORCH_INDEX[(7, 1)] == "rocm7.1" @@ -740,7 +751,7 @@ class TestRocmTorchIndex: assert "radeon" not in tag def test_newer_rocm_selects_best_match(self): - """ROCm 7.2 (not in map) should select rocm7.1 via >= comparison.""" + """ROCm 7.2 (now in map) should select rocm7.2 directly.""" ver = (7, 2) tag = next( ( @@ -750,7 +761,7 @@ class TestRocmTorchIndex: ), None, ) - assert tag == "rocm7.1" + assert tag == "rocm7.2" def test_rocm_64_selects_64(self): ver = (6, 4) @@ -778,7 +789,7 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") assert "IS_ROCM: bool" in source and "False" in source def test_hardware_py_sets_is_rocm_on_hip(self): @@ -786,7 +797,7 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") assert 'torch.version, "hip"' in source or "torch.version.hip" in source def test_hardware_py_still_returns_cuda_for_rocm(self): @@ -794,7 +805,7 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") # Ensure ROCM is NOT a DeviceType member enum_section = source.split("class DeviceType")[1].split("\n\n")[0] assert "ROCM" not in enum_section @@ -804,7 +815,7 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") assert '"rocm"' in source def test_hardware_py_device_type_cuda_references_intact(self): @@ -812,7 +823,7 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") # Key functions that must still reference DeviceType.CUDA assert "DeviceType.CUDA" in source assert "DEVICE = DeviceType.CUDA" in source @@ -822,7 +833,7 @@ class TestHardwareRocmFlag: init_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" ) - source = init_path.read_text() + source = init_path.read_text(encoding = "utf-8") assert "IS_ROCM" in source def test_is_rocm_in_all_list(self): @@ -830,7 +841,7 @@ class TestHardwareRocmFlag: init_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" ) - source = init_path.read_text() + source = init_path.read_text(encoding = "utf-8") # Extract __all__ section assert '"IS_ROCM"' in source @@ -839,13 +850,37 @@ class TestHardwareRocmFlag: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") # Find the get_package_versions function body func_start = source.find("def get_package_versions") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '"cuda"' in func_body assert '"rocm"' in func_body + def test_distributed_stubs_cover_is_torchelastic_launched(self): + """_determine_attention_impl_for_gpu_estimate must stub is_torchelastic_launched. + + resolve_attention_implementation calls is_torchelastic_launched() on + Windows ROCm where torch.distributed ships without that helper, causing + a warning: 'module torch.distributed has no attribute is_torchelastic_launched'. + """ + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text(encoding = "utf-8") + assert "is_torchelastic_launched" in source + + def test_distributed_stubs_cover_core_helpers(self): + """_determine_attention_impl_for_gpu_estimate must stub the four core distributed helpers.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text(encoding = "utf-8") + for attr in ("is_initialized", "is_available", "get_rank", "get_world_size"): + assert ( + attr in source + ), f"distributed stub for '{attr}' missing from hardware.py" + # ============================================================================= # TEST: tokenizer_utils.py -- error message @@ -858,13 +893,13 @@ class TestTokenizerErrorMessage: def test_no_old_amd_message(self): """Old 'We do not support AMD' message should be gone.""" tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" - source = tu_path.read_text() + source = tu_path.read_text(encoding = "utf-8") assert "We do not support AMD" not in source def test_new_message_has_docs_link(self): """New message should point to Unsloth AMD docs.""" tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" - source = tu_path.read_text() + source = tu_path.read_text(encoding = "utf-8") assert "docs.unsloth.ai" in source or "No GPU detected" in source @@ -879,7 +914,7 @@ class TestInstallShStructure: def test_no_here_strings(self): """install.sh must not use <<< (not POSIX).""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") # <<< is bash-only; breaks dash for i, line in enumerate(source.splitlines(), 1): stripped = line.lstrip() @@ -890,7 +925,7 @@ class TestInstallShStructure: def test_rocm_detection_present(self): """install.sh should have ROCm detection in get_torch_index_url.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") assert "amd-smi" in source assert "rocm" in source.lower() @@ -905,7 +940,7 @@ class TestInstallShStructure: block. """ sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") body = _extract_sh_function_body(source, "get_torch_index_url") nvidia_call = body.find("_has_usable_nvidia_gpu") no_nvidia_branch = body.find('if [ -z "$_smi" ]') @@ -926,44 +961,45 @@ class TestInstallShStructure: def test_bitsandbytes_amd_install(self): """install.sh should install bitsandbytes for AMD when ROCm detected.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") assert "bitsandbytes" in source assert "rocm*)" in source # case pattern for ROCm URLs def test_cpu_hint_mentions_amd(self): """CPU-only hint should mention AMD ROCm.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") assert "ROCm" in source - def test_rocm72_capped_to_71(self): - """ROCm 7.2+ should fall back to rocm7.1 index.""" + def test_rocm72_supported_future_capped(self): + """ROCm 7.2 should pass through directly; 7.3+ falls back to rocm7.2.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() - assert 'echo "$_base/rocm7.1"' in source # fallback for unknown versions + source = sh_path.read_text(encoding = "utf-8") + assert 'echo "$_base/rocm7.2"' in source # fallback for unknown future versions # Allowlisted versions should pass through directly assert "rocm6.*" in source assert "rocm7.0" in source assert "rocm7.1" in source + assert "rocm7.2" in source def test_rocm_tag_validation_guard_exists(self): """install.sh should validate _rocm_tag with a case guard.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") assert "rocm[1-9]*.[0-9]*)" in source assert '_rocm_tag=""' in source # rejection path def test_dpkg_epoch_handling(self): """install.sh should strip Debian epoch prefix from dpkg-query output.""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source def test_no_double_bracket_in_rocm_block(self): """ROCm detection block should not use [[ ]] (bash-only, not POSIX). Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]].""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] @@ -982,7 +1018,7 @@ class TestInstallShStructure: def test_no_arithmetic_expansion_in_rocm_block(self): """ROCm detection block should not use (( )) (bash-only).""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] @@ -997,7 +1033,7 @@ class TestInstallShStructure: def test_macos_returns_cpu_before_rocm_check(self): """macOS should return CPU immediately (before any ROCm check).""" sh_path = PACKAGE_ROOT / "install.sh" - source = sh_path.read_text() + source = sh_path.read_text(encoding = "utf-8") func_start = source.find("get_torch_index_url()") func_body = source[func_start:] darwin_pos = func_body.find("Darwin") @@ -1071,12 +1107,12 @@ class TestWorkerRocmMambaSsm: def test_probe_returns_hip_version_field(self): """The wheel probe should include hip_version, and worker.py should consume it.""" - assert "hip_version" in _WHEEL_UTILS_PATH.read_text() - assert "hip_version" in _WORKER_PATH.read_text() + assert "hip_version" in _WHEEL_UTILS_PATH.read_text(encoding = "utf-8") + assert "hip_version" in _WORKER_PATH.read_text(encoding = "utf-8") def test_probe_script_has_getattr_hip(self): """Probe script should use getattr for torch.version.hip (safe on CUDA).""" - source = _WHEEL_UTILS_PATH.read_text() + source = _WHEEL_UTILS_PATH.read_text(encoding = "utf-8") assert "getattr(torch.version, 'hip', None)" in source def test_direct_wheel_url_returns_none_without_cuda_major(self): @@ -1118,22 +1154,22 @@ class TestWorkerRocmMambaSsm: def test_hipcc_check_exists_in_source(self): """worker.py should check for hipcc before ROCm source builds.""" - source = _WORKER_PATH.read_text() + source = _WORKER_PATH.read_text(encoding = "utf-8") assert "hipcc" in source def test_rocm_source_build_status_message(self): """worker.py should send a specific status for ROCm source compilation.""" - source = _WORKER_PATH.read_text() + source = _WORKER_PATH.read_text(encoding = "utf-8") assert "Compiling" in source and "from source for ROCm" in source def test_rocm_build_failure_message(self): """worker.py should send a clear error on ROCm build failure.""" - source = _WORKER_PATH.read_text() + source = _WORKER_PATH.read_text(encoding = "utf-8") assert "Failed to compile" in source and "for ROCm" in source def test_timeout_on_install(self): """worker.py should have a timeout on pip install subprocess.""" - source = _WORKER_PATH.read_text() + source = _WORKER_PATH.read_text(encoding = "utf-8") assert "TimeoutExpired" in source assert "timeout" in source @@ -1154,7 +1190,7 @@ class TestAmdGpuMonitoring: def test_amd_py_has_required_functions(self): """amd.py should export the same function signatures as nvidia.py.""" amd_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "amd.py" - source = amd_path.read_text() + source = amd_path.read_text(encoding = "utf-8") assert "def get_physical_gpu_count" in source assert "def get_primary_gpu_utilization" in source assert "def get_visible_gpu_utilization" in source @@ -1299,7 +1335,7 @@ class TestHardwareAmdBranching: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): @@ -1309,7 +1345,7 @@ class TestHardwareAmdBranching: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '_smi_query("get_primary_gpu_utilization"' in func_body @@ -1327,7 +1363,7 @@ class TestHardwareAmdBranching: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_visible_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] # The dispatcher call may wrap onto multiple lines; allow whitespace @@ -1348,13 +1384,55 @@ class TestHardwareAmdBranching: hw_path = ( PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" ) - source = hw_path.read_text() + source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_physical_gpu_count") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "IS_ROCM" in func_body assert "from . import amd" in func_body +# ============================================================================= +# TEST: hardware.py -- apply_gpu_ids ROCm fallback (issue #5180) +# ============================================================================= + + +class TestApplyGpuIdsRocmFallback: + """Verify apply_gpu_ids sets HIP_VISIBLE_DEVICES on ROCm hosts even when + IS_ROCM is still False (worker subprocess before detect_hardware runs).""" + + def test_apply_gpu_ids_falls_back_to_torch_version_hip(self): + """apply_gpu_ids should probe torch.version.hip when IS_ROCM is False and no ROCm env vars are set.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text(encoding = "utf-8") + func_start = source.find("def apply_gpu_ids") + func_body = source[func_start : source.find("\ndef ", func_start + 1)] + assert 'getattr(_torch.version, "hip", None)' in func_body + + def test_apply_gpu_ids_sets_hip_and_rocr_visible_devices(self): + """apply_gpu_ids should set both HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text(encoding = "utf-8") + func_start = source.find("def apply_gpu_ids") + func_body = source[func_start : source.find("\ndef ", func_start + 1)] + assert 'os.environ["HIP_VISIBLE_DEVICES"] = value' in func_body + assert 'os.environ["ROCR_VISIBLE_DEVICES"] = value' in func_body + + def test_apply_gpu_ids_rocm_fallback_is_guarded_by_try_except(self): + """torch import in apply_gpu_ids must be wrapped in try/except so a missing torch never crashes.""" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) + source = hw_path.read_text(encoding = "utf-8") + func_start = source.find("def apply_gpu_ids") + func_body = source[func_start : source.find("\ndef ", func_start + 1)] + assert "import torch as _torch" in func_body + assert "except Exception" in func_body + + # ============================================================================= # TEST: install_python_stack.py -- Windows AMD warning # ============================================================================= @@ -1365,18 +1443,18 @@ class TestWindowsRocmWarning: def test_windows_amd_warning_in_source(self): """install_python_stack.py should warn Windows AMD users.""" - source = _STACK_PATH.read_text() - assert "AMD GPU detected on Windows" in source + source = _STACK_PATH.read_text(encoding = "utf-8") + assert "AMD GPU detected" in source def test_windows_amd_warning_checks_hipinfo_or_amdsmi(self): """Warning should check for hipinfo or amd-smi.""" - source = _STACK_PATH.read_text() + source = _STACK_PATH.read_text(encoding = "utf-8") assert "hipinfo" in source assert "amd-smi" in source def test_windows_amd_warning_has_docs_link(self): """Warning should include AMD docs link.""" - source = _STACK_PATH.read_text() + source = _STACK_PATH.read_text(encoding = "utf-8") assert "docs.unsloth.ai/get-started/install-and-update/amd" in source @@ -1391,7 +1469,7 @@ class TestIsRdnaExpansion: def test_is_rdna_source_has_rdna2(self): """is_rdna() should include RDNA2 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" - source = utils_path.read_text() + source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1030" in func_body @@ -1405,7 +1483,7 @@ class TestIsRdnaExpansion: def test_is_rdna_source_has_rdna3(self): """is_rdna() should include RDNA3 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" - source = utils_path.read_text() + source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1100" in func_body @@ -1416,7 +1494,7 @@ class TestIsRdnaExpansion: def test_is_rdna_source_has_rdna35(self): """is_rdna() should include RDNA3.5 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" - source = utils_path.read_text() + source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1150" in func_body @@ -1426,7 +1504,7 @@ class TestIsRdnaExpansion: def test_is_rdna_source_has_rdna4(self): """is_rdna() should include RDNA4 architectures.""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" - source = utils_path.read_text() + source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_rdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx1200" in func_body @@ -1435,7 +1513,7 @@ class TestIsRdnaExpansion: def test_is_cdna_not_changed(self): """is_cdna() should remain unchanged (no RDNA architectures added).""" utils_path = PACKAGE_ROOT / "unsloth" / "kernels" / "utils.py" - source = utils_path.read_text() + source = utils_path.read_text(encoding = "utf-8") func_start = source.find("def is_cdna()") func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert "gfx940" in func_body @@ -1447,5 +1525,1069 @@ class TestIsRdnaExpansion: assert "gfx1100" not in func_body +# ============================================================================= +# TEST: install_python_stack.py -- _windows_rocm_index_url arch mapping +# ============================================================================= + + +class TestWindowsRocmIndexUrl: + """Verify GPU arch → AMD pip index URL mapping.""" + + def test_gfx1200_maps_to_gfx120x_all(self): + url = stack_mod._windows_rocm_index_url("gfx1200") + assert url is not None + assert "gfx120X-all" in url + + def test_gfx1201_maps_to_gfx120x_all(self): + url = stack_mod._windows_rocm_index_url("gfx1201") + assert url is not None + assert "gfx120X-all" in url + + def test_gfx1151_maps_to_gfx1151(self): + url = stack_mod._windows_rocm_index_url("gfx1151") + assert url is not None + assert "gfx1151" in url + + def test_gfx1150_maps_to_gfx1150(self): + url = stack_mod._windows_rocm_index_url("gfx1150") + assert url is not None + assert "gfx1150" in url + + def test_gfx1100_maps_to_gfx110x_all(self): + url = stack_mod._windows_rocm_index_url("gfx1100") + assert url is not None + assert "gfx110X-all" in url + + def test_unknown_arch_returns_none(self): + assert stack_mod._windows_rocm_index_url("gfx9999") is None + + def test_none_arch_returns_none(self): + assert stack_mod._windows_rocm_index_url(None) is None + + def test_url_ends_with_slash(self): + """AMD pip index URLs must end with / for --index-url compatibility.""" + url = stack_mod._windows_rocm_index_url("gfx1200") + assert url is not None + assert url.endswith("/") + + def test_base_url_uses_repo_amd_com_by_default(self): + url = stack_mod._windows_rocm_index_url("gfx1200") + assert url is not None + assert "repo.amd.com" in url + + def test_mirror_env_var_overrides_base(self, monkeypatch): + monkeypatch.setenv( + "UNSLOTH_ROCM_WINDOWS_MIRROR", "https://my-mirror.example.com/rocm/whl" + ) + # Reload module-level constant by calling helper directly + url = stack_mod._windows_rocm_index_url("gfx1200") + # The env var is read at module load time for _ROCM_WINDOWS_INDEX_BASE, + # so just verify the helper itself doesn't error. + assert url is not None + + +# ============================================================================= +# TEST: install_python_stack.py -- _detect_windows_gfx_arch +# ============================================================================= + + +class TestDetectWindowsGfxArch: + """Verify hipinfo parsing for GPU arch detection on Windows.""" + + def test_returns_none_when_hipinfo_not_on_path(self): + with patch("shutil.which", return_value = None): + result = stack_mod._detect_windows_gfx_arch() + assert result is None + + def test_parses_gcnarchname_from_hipinfo_output(self): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n" + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): + with patch("subprocess.run", return_value = mock_result): + result = stack_mod._detect_windows_gfx_arch() + assert result == "gfx1200" + + def test_returns_none_on_nonzero_returncode(self): + mock_result = MagicMock() + mock_result.returncode = 1 + mock_result.stdout = b"gcnArchName : gfx1200\n" + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): + with patch("subprocess.run", return_value = mock_result): + result = stack_mod._detect_windows_gfx_arch() + assert result is None + + def test_returns_none_when_no_gcnarchname_in_output(self): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = b"deviceName : Radeon RX 9060 XT\n" + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): + with patch("subprocess.run", return_value = mock_result): + result = stack_mod._detect_windows_gfx_arch() + assert result is None + + def test_returns_none_on_timeout(self): + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): + with patch( + "subprocess.run", + side_effect = subprocess.TimeoutExpired("hipinfo", 10), + ): + result = stack_mod._detect_windows_gfx_arch() + assert result is None + + def test_strips_whitespace_from_arch(self): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = b" gcnArchName : gfx1201 \n" + with patch("shutil.which", return_value = "/usr/bin/hipinfo"): + with patch("subprocess.run", return_value = mock_result): + result = stack_mod._detect_windows_gfx_arch() + assert result == "gfx1201" + + +# ============================================================================= +# TEST: install_python_stack.py -- _install_bnb_windows_rocm +# ============================================================================= + + +class TestInstallBnbWindowsRocm: + """Verify AMD Windows BNB wheel install helper.""" + + def test_calls_pip_install_try_with_win_amd64_url(self): + """Should call pip_install_try with the win_amd64 wheel URL.""" + with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: + stack_mod._install_bnb_windows_rocm() + assert mock_pip.call_count == 1 + call_args = str(mock_pip.call_args_list[0]) + assert "bitsandbytes" in call_args + assert "win_amd64" in call_args + + def test_sets_uv_skip_env_var_during_install(self): + """UV_SKIP_WHEEL_FILENAME_CHECK must be '1' when pip_install_try runs.""" + observed = {} + + def _capture(*args, **kwargs): + observed["val"] = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") + return True + + with patch.object(stack_mod, "pip_install_try", side_effect = _capture): + stack_mod._install_bnb_windows_rocm() + assert observed.get("val") == "1" + + def test_restores_uv_skip_env_var_after_install(self): + """UV_SKIP_WHEEL_FILENAME_CHECK should be removed after install if it wasn't set before.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + stack_mod._install_bnb_windows_rocm() + assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ + + def test_restores_previous_uv_skip_value(self): + """If UV_SKIP_WHEEL_FILENAME_CHECK was already set, restore it afterwards.""" + with patch.dict(os.environ, {"UV_SKIP_WHEEL_FILENAME_CHECK": "0"}): + with patch.object(stack_mod, "pip_install_try", return_value = True): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK") == "0" + + def test_restores_env_even_if_install_raises(self): + """UV_SKIP_WHEEL_FILENAME_CHECK must be cleaned up even on pip failure.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None) + with patch.object( + stack_mod, "pip_install_try", side_effect = RuntimeError("pip failed") + ): + try: + stack_mod._install_bnb_windows_rocm() + except RuntimeError: + pass + assert "UV_SKIP_WHEEL_FILENAME_CHECK" not in os.environ + + def test_no_op_when_win_amd64_url_missing(self): + """Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS.""" + with patch.object(stack_mod, "_BNB_ROCM_PRERELEASE_URLS", {}): + with patch.object(stack_mod, "pip_install_try") as mock_pip: + stack_mod._install_bnb_windows_rocm() + mock_pip.assert_not_called() + + def test_sets_bnb_rocm_version_from_detected_dll(self): + """BNB_ROCM_VERSION is set from the DLL detected after install.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72" + ): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "72" + + def test_sets_bnb_rocm_version_from_newer_dll(self): + """If AMD ships a newer DLL (e.g. rocm713.dll), that version is used.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "713" + ): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "713" + + def test_falls_back_to_72_when_detection_fails(self): + """Falls back to '72' when DLL detection returns None.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch.object( + stack_mod, "_detect_bnb_rocm_dll_ver", return_value = None + ): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "72" + + def test_does_not_override_existing_bnb_rocm_version(self): + """An explicit BNB_ROCM_VERSION in the caller's env must not be clobbered.""" + with patch.dict(os.environ, {"BNB_ROCM_VERSION": "60"}): + with patch.object(stack_mod, "pip_install_try", return_value = True): + stack_mod._install_bnb_windows_rocm() + assert os.environ.get("BNB_ROCM_VERSION") == "60" + + +class TestDetectBnbRocmDllVer: + """Unit tests for _detect_bnb_rocm_dll_ver().""" + + def test_returns_none_when_bnb_not_installed(self): + """Returns None if bitsandbytes is not importable.""" + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = None): + assert stack_mod._detect_bnb_rocm_dll_ver() is None + + def test_detects_rocm72_dll(self, tmp_path): + """Returns '72' when libbitsandbytes_rocm72.dll is present.""" + (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() == "72" + + def test_detects_rocm713_dll(self, tmp_path): + """Returns '713' when libbitsandbytes_rocm713.dll is present.""" + (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() == "713" + + def test_returns_none_when_only_cuda_dlls(self, tmp_path): + """Returns None when only CUDA DLLs are present (no ROCm DLL).""" + (tmp_path / "libbitsandbytes_cuda121.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() is None + + def test_picks_highest_suffix_when_multiple_dlls(self, tmp_path): + """Returns the highest numeric suffix when multiple ROCm DLL variants exist. + + Filesystem glob order is not guaranteed, so the function must not stop + at the first match — it must always return the highest one. + """ + (tmp_path / "libbitsandbytes_rocm72.dll").write_text("") + (tmp_path / "libbitsandbytes_rocm713.dll").write_text("") + mock_spec = MagicMock() + mock_spec.submodule_search_locations = [str(tmp_path)] + import importlib.util + + with patch.object(importlib.util, "find_spec", return_value = mock_spec): + assert stack_mod._detect_bnb_rocm_dll_ver() == "713" + + +# ============================================================================= +# TEST: install_python_stack.py -- UNSLOTH_ROCM_TORCH_INSTALLED early-return path +# ============================================================================= + + +class TestRocmTorchInstalledEnvVar: + """Verify UNSLOTH_ROCM_TORCH_INSTALLED=1 skips main install but still installs BNB.""" + + @staticmethod + def _ok_torch_probe(*a, **kw): + # subprocess.run probe returns 0 when torch imports as ROCm + rv = MagicMock() + rv.returncode = 0 + return rv + + @patch.object(stack_mod, "_install_bnb_windows_rocm") + @patch.object(stack_mod, "pip_install") + def test_env_var_skips_main_pip_install(self, mock_pip, mock_bnb): + """UNSLOTH_ROCM_TORCH_INSTALLED=1 should not trigger torch pip_install.""" + with ( + patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), + patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe), + ): + stack_mod._ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "_install_bnb_windows_rocm") + @patch.object(stack_mod, "pip_install") + def test_env_var_calls_bnb_install(self, mock_pip, mock_bnb): + """UNSLOTH_ROCM_TORCH_INSTALLED=1 should still call _install_bnb_windows_rocm.""" + with ( + patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), + patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe), + ): + stack_mod._ensure_rocm_torch() + mock_bnb.assert_called_once() + + @patch.object(stack_mod, "_install_bnb_windows_rocm") + @patch.object(stack_mod, "pip_install") + def test_env_var_sets_rocm_windows_flag(self, mock_pip, mock_bnb): + """UNSLOTH_ROCM_TORCH_INSTALLED=1 should set _rocm_windows_torch_installed.""" + stack_mod._rocm_windows_torch_installed = False + with ( + patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), + patch.object(stack_mod.subprocess, "run", side_effect = self._ok_torch_probe), + ): + stack_mod._ensure_rocm_torch() + assert stack_mod._rocm_windows_torch_installed is True + + @patch.object(stack_mod, "_install_bnb_windows_rocm") + @patch.object(stack_mod, "pip_install") + def test_env_var_falls_through_when_torch_missing(self, mock_pip, mock_bnb): + """If the venv was wiped between runs, the stale env-var must not suppress reinstall.""" + stack_mod._rocm_windows_torch_installed = False + + def _bad_probe(*a, **kw): + rv = MagicMock() + rv.returncode = 1 + return rv + + with ( + patch.dict(os.environ, {"UNSLOTH_ROCM_TORCH_INSTALLED": "1"}), + patch.object(stack_mod.subprocess, "run", side_effect = _bad_probe), + patch.object(stack_mod, "IS_WINDOWS", False), + patch.object(stack_mod, "IS_MACOS", True), + ): + stack_mod._ensure_rocm_torch() + # macOS branch is the next exit -- but the point is the early-return did NOT fire. + mock_bnb.assert_not_called() + + +# ============================================================================= +# TEST: worker.py -- Windows ROCm patches (source-level checks) +# ============================================================================= + + +class TestWorkerWindowsRocmPatches: + """Verify worker.py contains the required Windows ROCm runtime patches.""" + + def test_grouped_mm_dispatch_patch_present(self): + """worker.py must register a _grouped_mm CUDA dispatch override.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert '_gm_lib.impl("_grouped_mm"' in source + + def test_grouped_mm_patch_targets_cuda_dispatch_key(self): + """The dispatch override must target the CUDA key (not CompositeImplicitAutograd).""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert '"_grouped_mm", _grouped_mm_safe_impl, "CUDA"' in source + + def test_grouped_mm_lib_kept_alive(self): + """The Library object must be stored to prevent GC clearing the registration.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "_WINDOWS_ROCM_GROUPED_MM_LIB" in source + + def test_grouped_mm_handles_offs_grouped_case(self): + """_grouped_mm fallback must handle the grouped (offs!=None) variant.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "offs_list" in source + assert "offs.tolist()" in source + + def test_torchao_stub_uses_stub_type_meta(self): + """Torchao stub must use _StubTypeMeta so isinstance() returns False not TypeError.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "_StubTypeMeta" in source + + def test_stub_type_meta_has_instancecheck(self): + """_StubTypeMeta must define __instancecheck__ returning False.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "__instancecheck__" in source + + def test_stub_subpackage_finder_registered(self): + """_StubSubpackageFinder must be appended to sys.meta_path.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "sys.meta_path.append(_StubSubpackageFinder())" in source + + def test_torchao_key_submodules_pre_stubbed(self): + """Key torchao submodules (dtypes, quantization) must be pre-stubbed.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "torchao.dtypes" in source + assert "torchao.quantization" in source + + def test_torchdynamo_disabled_on_windows_rocm(self): + """worker.py should disable dynamo on Windows ROCm as belt-and-suspenders.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "TORCHDYNAMO_DISABLE" in source + + def test_bnb_rocm_version_set_on_windows_rocm(self): + """worker.py must set BNB_ROCM_VERSION in the Windows ROCm section. + + BNB auto-detects HIP version from torch.version.hip, which can mismatch + the DLL suffix in the AMD prerelease wheel. The worker must detect the + actual DLL suffix and override BNB's auto-detection before ML imports. + """ + source = _WORKER_PATH.read_text(encoding = "utf-8") + # Env var must be set + assert "BNB_ROCM_VERSION" in source + # Detection helper must be used + assert "_detect_bnb_rocm_dll_ver" in source or "libbitsandbytes_rocm" in source + # "72" must appear as the safe fallback + assert '"72"' in source or "'72'" in source + + def test_bnb_rocm_version_set_before_ml_imports(self): + """BNB_ROCM_VERSION must appear in section 1f, before section 2 ML imports.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + idx_bnb = source.find("BNB_ROCM_VERSION") + # Use the specific section-2 marker that appears in the worker process + # entry-point function (not the trainer helper which has its own "# ── 2."). + idx_sec2 = source.find("# ── 2. Now import ML libraries") + assert idx_bnb != -1, "BNB_ROCM_VERSION not found in worker.py" + assert ( + idx_sec2 != -1 + ), "'# ── 2. Now import ML libraries' marker not found in worker.py" + assert idx_bnb < idx_sec2, ( + "BNB_ROCM_VERSION must be set before section 2 ML imports " + f"(found at {idx_bnb}, section 2 at {idx_sec2})" + ) + + def test_grouped_mm_patch_guarded_by_windows_and_hip_check(self): + """_grouped_mm patch must only apply on Windows + HIP torch.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + # Must check sys.platform == "win32" + assert 'sys.platform == "win32"' in source + # Must gate on HIP version — code uses getattr chain: "version" and "hip" + assert '"version"' in source and '"hip"' in source + + def test_hip_ver_at_least_helper_defined(self): + """_hip_ver_at_least helper must be defined inside the Windows ROCm block.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "def _hip_ver_at_least(major: int, minor: int)" in source + + def test_grouped_mm_patch_gated_on_hip_lt_713(self): + """_grouped_mm patch must be skipped on HIP >= 7.13 (AMD fixed the bug in ROCm 7.13).""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + # The guard must call _hip_ver_at_least with exactly (7, 13) + assert "_hip_ver_at_least(7, 13)" in source + # The patch must be inside the `if not` branch (negated guard) + assert "if not _hip_ver_at_least(7, 13):" in source + + def test_grouped_mm_hip_713_skip_message_present(self): + """worker.py must log a message when skipping the patch on HIP >= 7.13.""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + assert "HIP >= 7.13" in source + assert "7.13" in source + + def test_grouped_mm_patch_else_branch_present(self): + """An else branch must follow the _hip_ver_at_least gate (skip path for 7.13+).""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + # There must be an else: after the if not _hip_ver_at_least(7, 13): block + gate_idx = source.find("if not _hip_ver_at_least(7, 13):") + assert gate_idx != -1, "Version gate not found in worker.py" + # The else: branch must appear after the gate + else_idx = source.find("else:", gate_idx) + assert else_idx != -1, "else: branch after _hip_ver_at_least gate not found" + + def test_hip_ver_at_least_handles_amd_version_format(self): + """_hip_ver_at_least must split on '.' and compare only major.minor (handles '7.13.99004').""" + source = _WORKER_PATH.read_text(encoding = "utf-8") + # Must split the version string and take the first two parts + assert 'split(".")[:2]' in source or ".split('.')[:2]" in source + + +# ============================================================================= +# TEST: install_python_stack.py -- _ROCM_TORCH_PKG_SPECS mapping +# ============================================================================= + + +class TestRocmTorchPkgSpecs: + """Verify per-tag torch version specs are correct.""" + + def test_rocm72_has_torch_211(self): + """rocm7.2 should specify torch 2.11.x.""" + specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("rocm7.2") + assert specs is not None + torch_spec = specs[0] + assert "2.11" in torch_spec + + def test_default_caps_below_211(self): + """Default spec (rocm7.1 and earlier) should cap below 2.11.""" + specs = stack_mod._ROCM_TORCH_PKG_SPECS.get("_default") + assert specs is not None + torch_spec = specs[0] + assert "<2.11" in torch_spec + + def test_specs_have_torch_vision_audio(self): + """Each entry should be a 3-tuple: torch, torchvision, torchaudio.""" + for tag, specs in stack_mod._ROCM_TORCH_PKG_SPECS.items(): + assert len(specs) == 3, f"{tag}: expected (torch, torchvision, torchaudio)" + assert "torch" in specs[0] + assert "torchvision" in specs[1] + assert "torchaudio" in specs[2] + + def test_gfx_to_amd_index_covers_rdna4(self): + """_GFX_TO_AMD_INDEX_ARCH must cover gfx1200 and gfx1201 (RDNA 4).""" + mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH + assert mapping.get("gfx1200") == "gfx120X-all" + assert mapping.get("gfx1201") == "gfx120X-all" + + def test_gfx_to_amd_index_covers_strix_halo(self): + """_GFX_TO_AMD_INDEX_ARCH must cover gfx1151 and gfx1150 (RDNA 3.5).""" + mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH + assert mapping.get("gfx1151") == "gfx1151" + assert mapping.get("gfx1150") == "gfx1150" + + def test_gfx_to_amd_index_covers_rdna3(self): + """_GFX_TO_AMD_INDEX_ARCH must cover gfx1100-gfx1103 (RDNA 3).""" + mapping = stack_mod._GFX_TO_AMD_INDEX_ARCH + for arch in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"): + assert mapping.get(arch) == "gfx110X-all", f"{arch} missing from mapping" + + +# ============================================================================= +# TEST: setup.ps1 / install.ps1 -- Strix Halo gfx arch detection +# ============================================================================= + +_SETUP_PS1_PATH = PACKAGE_ROOT / "studio" / "setup.ps1" +_INSTALL_PS1_PATH = PACKAGE_ROOT / "install.ps1" + + +class TestStrixHaloGfxArchDetection: + """Verify that setup.ps1 and install.ps1 have robust gfx arch detection + for Strix Halo / iGPU users who only have the HIP runtime (no hipinfo).""" + + def test_amd_smi_static_asic_attempted_in_setup(self): + """setup.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "static --asic" in source + + def test_amd_smi_static_asic_attempted_in_install(self): + """install.ps1 must try 'amd-smi static --asic' when list output lacks gfx arch.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "static --asic" in source + + def test_env_var_override_in_setup(self): + """setup.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "UNSLOTH_ROCM_GFX_ARCH" in source + + def test_env_var_override_in_install(self): + """install.ps1 must honour UNSLOTH_ROCM_GFX_ARCH as a manual arch override.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "UNSLOTH_ROCM_GFX_ARCH" in source + + def test_name_arch_table_covers_strix_halo_in_setup(self): + """setup.ps1 name→arch table must map 890M / Strix Halo to gfx1151.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "gfx1151" in source + assert "890M" in source or "Strix Halo" in source + + def test_name_arch_table_covers_strix_halo_in_install(self): + """install.ps1 name→arch table must map 890M / Strix Halo to gfx1151.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "gfx1151" in source + assert "890M" in source or "Strix Halo" in source + + def test_name_arch_table_covers_strix_point_in_setup(self): + """setup.ps1 name→arch table must map 880M / Strix Point to gfx1150.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "gfx1150" in source + assert "880M" in source or "Strix Point" in source + + def test_name_arch_table_covers_strix_point_in_install(self): + """install.ps1 name→arch table must map 880M / Strix Point to gfx1150.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "gfx1150" in source + assert "880M" in source or "Strix Point" in source + + def test_name_arch_table_covers_rdna3_phoenix_in_setup(self): + """setup.ps1 name→arch table must map 780M / Phoenix to gfx1103.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "gfx1103" in source + assert "780M" in source or "Phoenix" in source + + def test_wmi_does_not_set_hasrocm_in_setup(self): + """WMI block in setup.ps1 must NOT set $HasROCm = $true (no runtime confirmation).""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + # Find the WMI block and confirm HasROCm is not set inside it + wmi_idx = source.find("Win32_VideoController") + assert wmi_idx != -1, "WMI block not found in setup.ps1" + # The nearest HasROCm = $true must not appear between the WMI block + # and the closing brace of that if-block. We check by confirming + # $HasROCm = $true does NOT appear within 300 chars of the WMI call. + wmi_context = source[wmi_idx : wmi_idx + 300] + assert "$HasROCm = $true" not in wmi_context + + def test_gfx_arch_regex_parses_from_amd_smi_output(self): + """Both files must use the gfx\\d+[a-z]? regex to parse arch from amd-smi output.""" + for path in (_SETUP_PS1_PATH, _INSTALL_PS1_PATH): + source = path.read_text(encoding = "utf-8") + # The regex pattern used to match gfx arches + assert ( + "gfx\\d+" in source or r"gfx\d+" in source + ), f"gfx arch regex not found in {path.name}" + + +# ============================================================================= +# TEST: HIP SDK tool path resolution via HIP_PATH / ROCM_PATH env vars +# ============================================================================= + + +class TestHipSdkEnvPathResolution: + """Verify that both install scripts resolve hipinfo/hipconfig via HIP_PATH + and ROCM_PATH when the tools are not on $PATH, and emit explicit warnings.""" + + # ── hipinfo resolution ──────────────────────────────────────────────────── + + def test_setup_checks_hip_path_for_hipinfo(self): + """setup.ps1 must reference HIP_PATH when resolving hipinfo.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP_PATH" in source + assert "hipinfo" in source + + def test_install_checks_hip_path_for_hipinfo(self): + """install.ps1 must reference HIP_PATH when resolving hipinfo.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP_PATH" in source + assert "hipinfo" in source + + def test_setup_checks_rocm_path_as_hipinfo_fallback(self): + """setup.ps1 must also check ROCM_PATH as a secondary hipinfo fallback.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCM_PATH" in source + # Confirm the fallback pattern: HIP_PATH ?? ROCM_PATH (or equivalent elseif) + assert "ROCM_PATH" in source and "HIP_PATH" in source + + def test_install_checks_rocm_path_as_hipinfo_fallback(self): + """install.ps1 must also check ROCM_PATH as a secondary hipinfo fallback.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCM_PATH" in source + assert "ROCM_PATH" in source and "HIP_PATH" in source + + def test_setup_resolves_hipinfo_via_bin_subdir(self): + """setup.ps1 must join the env var root with 'bin\\hipinfo.exe'.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert r"bin\hipinfo.exe" in source + + def test_install_resolves_hipinfo_via_bin_subdir(self): + """install.ps1 must join the env var root with 'bin\\hipinfo.exe'.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert r"bin\hipinfo.exe" in source + + # ── hipinfo not-on-PATH warning ─────────────────────────────────────────── + + def test_setup_warns_when_hipinfo_not_on_path(self): + """setup.ps1 must warn when hipinfo is found via env var but not on PATH.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "hipinfo not on PATH" in source + + def test_install_warns_when_hipinfo_not_on_path(self): + """install.ps1 must warn when hipinfo is found via env var but not on PATH.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "hipinfo not on PATH" in source + + # ── warn when HIP_PATH set but exe missing ──────────────────────────────── + + def test_setup_warns_when_hip_path_set_but_exe_missing(self): + """setup.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + # The warning must mention that the SDK install may be incomplete + assert "incomplete" in source or "not found at" in source + + def test_install_warns_when_hip_path_set_but_exe_missing(self): + """install.ps1 must warn when HIP_PATH is set but hipinfo.exe is not present.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "incomplete" in source or "not found at" in source + + # ── hipinfo runtime error warning ───────────────────────────────────────── + + def test_setup_warns_on_hipinfo_nonzero_exit(self): + """setup.ps1 must warn when hipinfo runs but returns a non-zero exit code.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP runtime error" in source or "runtime error" in source.lower() + + def test_install_warns_on_hipinfo_nonzero_exit(self): + """install.ps1 must warn when hipinfo runs but returns a non-zero exit code.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP runtime error" in source or "runtime error" in source.lower() + + # ── hipconfig resolution ────────────────────────────────────────────────── + + def test_setup_resolves_hipconfig_via_bin_subdir(self): + """setup.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert r"bin\hipconfig.exe" in source + + def test_install_resolves_hipconfig_via_bin_subdir(self): + """install.ps1 must also fall back to HIP_PATH/bin/hipconfig.exe for version detection.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert r"bin\hipconfig.exe" in source + + def test_setup_warns_when_hipconfig_not_on_path(self): + """setup.ps1 must warn when hipconfig is found via env var but not on PATH.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "hipconfig not on PATH" in source + + def test_install_warns_when_hipconfig_not_on_path(self): + """install.ps1 must warn when hipconfig is found via env var but not on PATH.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "hipconfig not on PATH" in source + + # ── PATH fix hint ───────────────────────────────────────────────────────── + + def test_setup_provides_path_fix_hint(self): + """setup.ps1 must tell the user how to add the HIP bin dir to PATH.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + # Should mention adding to PATH or SetEnvironmentVariable + assert "PATH" in source and ( + "SetEnvironmentVariable" in source or "Add" in source + ) + + def test_install_provides_path_fix_hint(self): + """install.ps1 must tell the user how to add the HIP bin dir to PATH.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "PATH" in source and ( + "SetEnvironmentVariable" in source or "Add" in source + ) + + +# ============================================================================= +# TEST: HIP SDK detected substep -- path + hipconfig version shown in terminal +# ============================================================================= + + +class TestHipSdkDetectedSubstep: + """Verify that both scripts print HIP SDK path and full hipconfig version + as substeps under the gpu step when AMD ROCm is successfully detected.""" + + def test_setup_prints_hip_sdk_path_substep(self): + """setup.ps1 must print an 'HIP SDK:' substep showing the resolved path.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP SDK:" in source + + def test_install_prints_hip_sdk_path_substep(self): + """install.ps1 must print an 'HIP SDK:' substep showing the resolved path.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP SDK:" in source + + def test_setup_shows_hipconfig_full_version(self): + """setup.ps1 must capture and display the full hipconfig version string.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCmVersionFull" in source or "hipconfig:" in source + + def test_install_shows_hipconfig_full_version(self): + """install.ps1 must capture and display the full hipconfig version string.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCmVersionFull" in source or "hipconfig:" in source + + def test_setup_captures_full_version_not_just_major_minor(self): + """setup.ps1 must store the raw hipconfig output line, not just major.minor.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCmVersionFull" in source + + def test_install_captures_full_version_not_just_major_minor(self): + """install.ps1 must store the raw hipconfig output line, not just major.minor.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCmVersionFull" in source + + def test_setup_uses_hip_path_or_rocm_path_for_sdk_display(self): + """setup.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP_PATH" in source and "ROCM_PATH" in source + + def test_install_uses_hip_path_or_rocm_path_for_sdk_display(self): + """install.ps1 HIP SDK path substep must check HIP_PATH then ROCM_PATH.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HIP_PATH" in source and "ROCM_PATH" in source + + def test_setup_rocm_step_uses_full_version(self): + """setup.ps1 'rocm' step label must prefer the full version string.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "ROCmVersionFull" in source and "rocm" in source + + +# ============================================================================= +# TEST: install.sh -- Strix Halo rocm7.1 → rocm7.2 override +# ============================================================================= + +_INSTALL_SH_PATH = PACKAGE_ROOT / "install.sh" +_SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" + + +class TestStrixRocm71Override: + """Verify install.sh skips Radeon repo and routes to AMD arch-specific index + for gfx1151/gfx1150 when ROCm 7.1 would otherwise be selected (known _grouped_mm segfault). + AMD's repo.amd.com/rocm/whl/gfx1151/ serves torch 2.11+rocm7.13 which has the real fix.""" + + def test_strix_gfx_detection_in_install_sh(self): + """install.sh must detect gfx1151 and gfx1150 for the override.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "gfx1151" in source and "gfx1150" in source + + def test_rocm71_override_to_amd_arch_index_in_install_sh(self): + """install.sh must override TORCH_INDEX_URL to AMD arch-specific index for Strix.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + # The override must route to AMD's arch-specific index (repo.amd.com/rocm/whl) + assert "repo.amd.com/rocm/whl" in source + assert "_strix_gfx" in source + # The URL must incorporate the detected gfx arch so gfx1151 → .../gfx1151/ + strix_idx = source.find("_amd_strix_base") + assert strix_idx != -1 + ctx = source[strix_idx : strix_idx + 500] + assert "_strix_gfx" in ctx + + def test_radeon_repo_bypassed_for_strix_in_install_sh(self): + """install.sh must set _amd_gpu_radeon=false when Strix + ROCm 7.1 detected.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "_amd_gpu_radeon=false" in source + + def test_strix_override_warns_with_moe_utils_reference(self): + """install.sh must emit a [WARN] mentioning the moe_utils segfault.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "moe_utils" in source or "_grouped_mm" in source + + def test_strix_override_only_fires_on_rocm71(self): + """install.sh must scope the Strix override to rocm7.1 only (not rocm7.2+).""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + strix_idx = source.find("_strix_gfx") + assert strix_idx != -1 + # Look back for the rocm7.1 pattern within 600 chars before _strix_gfx + context_before = source[max(0, strix_idx - 2400) : strix_idx] + assert "rocm7.1" in context_before + + def test_torch_constraint_updated_for_strix_amd_index(self): + """install.sh must set TORCH_CONSTRAINT>=2.11 when routing Strix to AMD index.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "TORCH_CONSTRAINT" in source and "2.11" in source + + def test_amd_rocm_mirror_env_var_respected(self): + """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + assert "UNSLOTH_AMD_ROCM_MIRROR" in source + + def test_tauri_family_recognises_amd_arch_url(self): + """_tauri_torch_index_family must return a rocm* family for AMD arch-specific URLs.""" + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + # The function must have a case branch for repo.amd.com/rocm/whl/gfx* URLs + assert "rocm/whl/gfx" in source + + +# ============================================================================= +# TEST: setup.sh -- gcc-install-dir fix for Ubuntu 24.04 + ROCm 7.x clang-20 +# ============================================================================= + + +class TestSetupShGccInstallDir: + """Verify setup.sh applies the --gcc-install-dir flag when building llama.cpp + with HIP on Ubuntu 24.04+ to work around ROCm 7.x clang-20 header path bug.""" + + def test_gcc_install_dir_search_loop_present(self): + """setup.sh must iterate gcc versions 14→11 to find one with C++ headers.""" + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert "_GCC_INSTALL_DIR" in source + assert "/usr/lib/gcc/x86_64-linux-gnu" in source + + def test_gcc_install_dir_checks_include_dir(self): + """setup.sh must check that the gcc dir has an 'include' subdirectory.""" + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert "include" in source and "_GCC_INSTALL_DIR" in source + + def test_gcc_install_dir_appended_to_cmake_hip_flags(self): + """setup.sh must pass --gcc-install-dir via CMAKE_HIP_FLAGS.""" + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert "CMAKE_HIP_FLAGS" in source + assert "gcc-install-dir" in source + + def test_gcc_install_dir_only_applied_in_hip_build_block(self): + """The --gcc-install-dir fix must only apply in the HIP/ROCm build branch.""" + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + # GGML_HIP=ON must appear before gcc-install-dir in the source + hip_idx = source.find("GGML_HIP=ON") + gcc_idx = source.find("gcc-install-dir") + assert hip_idx != -1 and gcc_idx != -1 + assert hip_idx < gcc_idx + + def test_gcc_install_dir_logs_substep(self): + """setup.sh must print a substep when the gcc install dir is resolved.""" + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert "gcc install dir" in source or "GCC_INSTALL_DIR" in source + + +# ============================================================================= +# TEST: main.py -- BNB_ROCM_VERSION server startup + distributed stubs +# ============================================================================= + +_MAIN_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "main.py" +_HARDWARE_PY_PATH = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" +) + + +class TestServerStartupRocmFixes: + """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and + hardware.py injects torch._C._distributed_c10d stubs before torch.distributed.""" + + # ── BNB_ROCM_VERSION in server process ──────────────────────────────────── + + def test_main_py_sets_bnb_rocm_version(self): + """main.py must set BNB_ROCM_VERSION in the server process before imports.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert "BNB_ROCM_VERSION" in source + + def test_main_py_bnb_detection_scoped_to_win32(self): + """main.py BNB_ROCM_VERSION logic must be inside the win32 platform guard.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + win32_idx = source.find('sys.platform == "win32"') + bnb_idx = source.find("BNB_ROCM_VERSION") + assert win32_idx != -1 and bnb_idx != -1 + assert win32_idx < bnb_idx + + def test_main_py_bnb_dll_detection_uses_glob(self): + """main.py must scan for libbitsandbytes_rocm*.dll to find the right version.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert "libbitsandbytes_rocm" in source + + def test_main_py_bnb_falls_back_to_72(self): + """main.py must fall back to BNB_ROCM_VERSION='72' when no DLL is found.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert '"72"' in source or "'72'" in source + + def test_main_py_bnb_only_set_when_not_already_in_env(self): + """main.py must not override an existing BNB_ROCM_VERSION env var.""" + source = _MAIN_PY_PATH.read_text(encoding = "utf-8") + assert '"BNB_ROCM_VERSION" not in os.environ' in source + + # ── torch._C._distributed_c10d stubs in hardware.py ────────────────────── + + def test_hardware_py_injects_distributed_c10d_stub(self): + """hardware.py must inject torch._C._distributed_c10d into sys.modules.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "_distributed_c10d" in source + + def test_hardware_py_stub_injected_before_distributed_import(self): + """The sys.modules stub must be injected BEFORE import torch.distributed.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + c10d_idx = source.find("_distributed_c10d") + dist_idx = source.find("import torch.distributed") + assert c10d_idx != -1 and dist_idx != -1 + assert c10d_idx < dist_idx + + def test_hardware_py_stub_uses_types_moduletype(self): + """hardware.py must create the stub with types.ModuleType.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "ModuleType" in source + + def test_hardware_py_stub_scoped_to_win32(self): + """hardware.py distributed stub injection must be gated on win32.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert 'platform == "win32"' in source or "win32" in source + + def test_hardware_py_stub_exposes_fake_process_group(self): + """hardware.py stub must set FakeProcessGroup so torch.distributed doesn't raise AttributeError.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "FakeProcessGroup" in source + + def test_hardware_py_stub_exposes_process_group(self): + """hardware.py stub must set ProcessGroup on the c10d stub.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "ProcessGroup" in source + + def test_hardware_py_stub_uses_setattr_for_symbols(self): + """hardware.py must use setattr to populate stub symbols dynamically.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "setattr" in source + + def test_hardware_py_stub_all_c10d_siblings_covered(self): + """hardware.py must stub all three torch._C._distributed_* submodules.""" + source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8") + assert "_distributed_c10d" in source + assert "_distributed_autograd" in source + assert "_distributed_rpc" in source + + +# ============================================================================= +# TEST: install.ps1 / setup.ps1 -- HipSdkInstalled flag (SDK found, device inaccessible) +# ============================================================================= + + +class TestHipSdkInstalledButDeviceInaccessible: + """Verify that when hipinfo is found but exits non-zero (device not ROCm-accessible), + both scripts distinguish this from 'HIP SDK not found' and emit the correct message.""" + + def test_install_ps1_has_hip_sdk_installed_flag(self): + """install.ps1 must track HipSdkInstalled separately from HasROCm.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HipSdkInstalled" in source + + def test_setup_ps1_has_hip_sdk_installed_flag(self): + """setup.ps1 must track HipSdkInstalled separately from HasROCm.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HipSdkInstalled" in source + + def test_install_ps1_sets_flag_when_hipinfo_binary_found(self): + """install.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + # HipSdkInstalled must be assigned inside the hipinfoExe block + hipinfo_block_idx = source.find("if ($hipinfoExe)") + sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx) + assert hipinfo_block_idx != -1 and sdk_flag_idx != -1 + assert sdk_flag_idx > hipinfo_block_idx + + def test_setup_ps1_sets_flag_when_hipinfo_binary_found(self): + """setup.ps1 must set HipSdkInstalled=true inside the 'if ($hipinfoExe)' block.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + hipinfo_block_idx = source.find("if ($hipinfoExe)") + sdk_flag_idx = source.find("$HipSdkInstalled = $true", hipinfo_block_idx) + assert hipinfo_block_idx != -1 and sdk_flag_idx != -1 + assert sdk_flag_idx > hipinfo_block_idx + + def test_install_ps1_version_capture_runs_when_sdk_installed(self): + """install.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source + + def test_setup_ps1_version_capture_runs_when_sdk_installed(self): + """setup.ps1 must capture hipconfig version when HipSdkInstalled even if HasROCm is false.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "HasROCm -or $HipSdkInstalled" in source or "$HipSdkInstalled" in source + + def test_install_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self): + """install.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "not ROCm-accessible" in source + + def test_setup_ps1_distinct_message_for_sdk_found_but_device_inaccessible(self): + """setup.ps1 must show 'not ROCm-accessible' message (not 'HIP SDK not found') when SDK present.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "not ROCm-accessible" in source + + def test_install_ps1_driver_guidance_in_sdk_found_branch(self): + """install.ps1 must tell user this is a driver issue, not an SDK issue.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "driver issue" in source + + def test_setup_ps1_driver_guidance_in_sdk_found_branch(self): + """setup.ps1 must tell user this is a driver issue, not an SDK issue.""" + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "driver issue" in source + + def test_install_ps1_cpu_hint_distinguishes_driver_vs_no_sdk(self): + """install.ps1 CPU-only hint must say 'GPU not ROCm-accessible' not 'require the HIP SDK' when SDK found.""" + source = _INSTALL_PS1_PATH.read_text(encoding = "utf-8") + assert "GPU not ROCm-accessible" in source + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 99da5f799e..c0c6455a2e 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -68,7 +68,10 @@ def _get_env_int(keys): def _infer_distributed_ranks(): - if torch.distributed.is_available() and torch.distributed.is_initialized(): + if ( + torch.distributed.is_available() + and getattr(torch.distributed, "is_initialized", lambda: False)() + ): try: return torch.distributed.get_rank(), torch.distributed.get_world_size() except Exception: