From f03e6694428832d5bf1021f38c5c3863a58b9894 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Mon, 27 Jul 2026 07:22:19 -0500 Subject: [PATCH] AMD: enable ROCm torch on gfx906 (MI50 / Radeon VII) on Linux (#7354) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add community-maintained legacy support path for gfx906 (MI50 / Radeon VII) rocm6.4+/7.x torch wheels bundle ROCm libraries whose Tensile kernels dropped gfx906 (rocBLAS 'TensileLibrary.dat ... not read for gfx906', ROCm/TheRock#1844), so on MI50/Vega 20 hosts with newer ROCm the installer picked wheels that fail at the first BLAS call. The rocm6.3 index is the last one whose wheels run on gfx906 (torch 2.7.0 verified on MI50 32GB, up to 2.9 in community use). Dynamo/Inductor codegen is also broken on this arch, crashing compiled graphs that train fine in eager mode. - install.sh: when the runtime GPU is gfx906 and the picked index is newer than rocm6.3, reroute torch to the rocm6.3 index and reset the constraint trio to the default <2.11 window (a rocm7.2 pick raises the floor to 2.11, which rocm6.3 cannot satisfy), with a legacy-path warning. - install_python_stack.py: mirror the reroute in _ensure_rocm_torch using the _default pkg specs, including repairing an existing +rocm7.x torch and leaving a working rocm6.3 install alone. - device_type.py: default TORCHDYNAMO_DISABLE / TORCH_COMPILE_DISABLE / UNSLOTH_COMPILE_DISABLE on gfx906 (setdefault, user override wins). Windows allowlists are untouched: repo.amd.com publishes no gfx906 wheel family (verified in the RDNA2 enablement PR). 16-bit LoRA and full finetuning work out of the box; 4-bit QLoRA needs a source-built bitsandbytes for gfx906. Based on the verified MI50 32GB setup in namnguyen0503/mi50-gfx906-unsloth-bnb4bit-lab. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gfx906: second Codex pass (bnb skip under pin, override beats Strix) - Compute the gfx906 runtime-target flag independently of any torch-index pin or Strix override, so the bitsandbytes skip still applies when a user pins the ROCm index and sets UNSLOTH_ROCM_GFX_ARCH=gfx906 (the pin suppresses the torch reroute, not the bnb skip). Probe only when no pin is set (an explicit pin means don't second-guess it, matching the Strix path's asserted no-probe invariant); an explicit gfx906 override needs no probe. - Let UNSLOTH_ROCM_GFX_ARCH=gfx906 suppress the Strix reroute (both install.sh and install_python_stack.py) so a mixed Strix + MI50 host routes to rocm6.3 instead of the gfx1151 wheels probe order would pick. - Fix test_hardcoded_torch_constraint: the default <2.11 window literal now legitimately appears on two TORCH_CONSTRAINT= assignments (default + the gfx906 reroute reset after the rocm7.2 floor bump); assert it only ever appears on assignment lines, never on a pip install line (its real intent). New tests: bnb skipped under an explicit pin, gfx906 override wins over Strix, install.sh suppresses Strix on the override. rocm_support + selection + cross-platform parity: 667 passed; structural constraint 9/9. * gfx906: collapse single-line asserts to match pre-commit formatting * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gfx906: keep bnb skip + rocm6.3 routing correct under pins and suffixed overrides Address the four Codex P2 findings on #7354: - bnb skip under a pinned index (install.sh + install_python_stack.py): a real gfx906 host that pins UNSLOTH_TORCH_INDEX_URL to rocm6.3 without also setting UNSLOTH_ROCM_GFX_ARCH no longer reinstalls the generic bitsandbytes wheel over a source-built gfx906 bnb. A pin now suppresses only the torch reroute, not the gfx906 detection used for the bnb skip (Python drops the pin gate on _runtime_is_gfx906; bash _is_gfx906_bnb_skip probes via _probe_amd_gfx_arch when the index is pinned). - clear the Radeon marketing-name flag for every gfx906 target, not only when the >=6.4 reroute fires, so a Radeon VII already on rocm6.3 does not divert to the repo.radeon.com branch (whose wheels lack gfx906 kernels). - normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) before the exact comparisons in install.sh and install_python_stack.py, mirroring device_type.py. Tests: relax the three Strix-pin tests (the gfx probe may now run for the bnb flag but must not reroute the pinned index) and add coverage for the pinned bnb skip, the suffixed override, and the bash Radeon-clear / pinned-probe paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gfx906: log skipped vLLM aimv2 fix + robust source-scan test bounds Follow-up review polish: - import_fixes: log at info level when the vLLM aimv2 fix is skipped because the dist metadata is unreadable, so the skip is diagnosable instead of silent. - test_rocm_support: bound the gfx906 install.sh source-scan on the ';;' that closes its case arm via a shared _gfx906_reroute_block helper, replacing the brittle fixed-length (3200/3800) slices that shift when the block grows. * gfx906: trim whitespace on UNSLOTH_ROCM_GFX_ARCH in install.sh (py parity) The bash gfx906 comparisons lowercased and stripped the gfx906:… feature suffix but not surrounding whitespace, while the Python paths do .strip(). A stray newline (e.g. export UNSLOTH_ROCM_GFX_ARCH=$(cmd)) would make bash miss gfx906 while Python catches it. Trim with `tr -d '[:space:]'` at both comparison sites so the reroute target and bnb-skip agree across bash/Python. * gfx906: remove generic bitsandbytes pulled in transitively after the skip --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Daniel Han --- install.sh | 128 ++++++- studio/install_python_stack.py | 158 +++++++- .../test_tokenizers_and_torch_constraint.py | 20 +- tests/studio/install/test_rocm_support.py | 354 +++++++++++++++++- unsloth/device_type.py | 17 + unsloth/import_fixes.py | 8 +- 6 files changed, 659 insertions(+), 26 deletions(-) diff --git a/install.sh b/install.sh index d90195399d..fece7b173b 100755 --- a/install.sh +++ b/install.sh @@ -257,6 +257,51 @@ run_install_cmd_retry() { done } +# True when the runtime target is gfx906 (MI50/Radeon VII): the prebuilt AMD +# bitsandbytes wheel carries no gfx906 kernels, and force-reinstalling it would +# clobber a user's source-built bnb (the only 4-bit path on this arch) on every +# `studio update`. So skip the auto-install and leave whatever bnb is present. +# _gfx906_target is set during torch-index resolution; also honor an explicit +# UNSLOTH_ROCM_GFX_ARCH so a pinned-index install still skips. The override is +# normalized (gfx906:sramecc-:xnack- -> gfx906) so a copied HIP gcnArchName counts. +_is_gfx906_bnb_skip() { + [ "${_gfx906_target:-false}" = true ] && return 0 + _bnb_gfx_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]') + _bnb_gfx_env=${_bnb_gfx_env%%:*} + [ "$_bnb_gfx_env" = "gfx906" ] && return 0 + # A pinned index (UNSLOTH_TORCH_INDEX_URL/_FAMILY) skips the reroute block that + # sets _gfx906_target, so a real gfx906 host with a pinned rocm6.3 index and no + # UNSLOTH_ROCM_GFX_ARCH would otherwise clobber a source-built bnb. Probe here + # in that gap; skip only when gfx906 is the SOLE distinct arch (mixed hosts + # opt in via the env var, mirroring the reroute block's de-dup rule). + if [ -z "$_bnb_gfx_env" ] && [ "${_torch_index_pinned:-false}" = true ]; then + _bnb_gfx_probe=$(_probe_amd_gfx_arch | awk 'NF && !seen[$0]++') + [ "$_bnb_gfx_probe" = "gfx906" ] && return 0 + fi + return 1 +} + +# `pip install unsloth` resolves its unconditional bitsandbytes dep to a generic +# CUDA wheel (no gfx906 kernels) once we skip the prebuilt one. Snapshot bnb before +# the unsloth install, then drop a freshly pulled wheel afterwards while leaving a +# pre-existing source build in place. +_gfx906_bnb_installed() { + "$_VENV_PY" -c "import importlib.util as u, sys; sys.exit(0 if u.find_spec('bitsandbytes') else 1)" >/dev/null 2>&1 +} +_gfx906_bnb_snapshot() { + _gfx906_bnb_absent_before=false + _is_gfx906_bnb_skip || return 0 + _gfx906_bnb_installed || _gfx906_bnb_absent_before=true +} +_gfx906_bnb_prune() { + _is_gfx906_bnb_skip || return 0 + [ "${_gfx906_bnb_absent_before:-false}" = true ] || return 0 + _gfx906_bnb_installed || return 0 + substep "gfx906: removing generic bitsandbytes pulled in as a dependency (no gfx906 kernels; build from source for 4-bit QLoRA)" "$C_WARN" + uv pip uninstall --python "$_VENV_PY" bitsandbytes >/dev/null 2>&1 \ + || "$_VENV_PY" -m pip uninstall -y bitsandbytes >/dev/null 2>&1 || true +} + # Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main # wheel for 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. Falls back to PyPI >=0.49.1 if the @@ -3296,10 +3341,20 @@ case "$_torch_index_leaf" in if (n > 0) print vals[idx] }') fi + # An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the + # MI50 / Radeon VII path and must win over Strix probe-order detection on a + # mixed Strix + MI50 host, so the Strix reroute is suppressed when it is set. + # Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) and + # trim whitespace (mirrors the Python .strip()) so the feature-flag suffix or + # a stray newline does not defeat the exact gfx906 comparisons below. + _gfx906_env=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]' | tr -d '[:space:]') + _gfx906_env=${_gfx906_env%%:*} _strix_gfx="" - case "$_runtime_gfx" in - gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;; - esac + if [ "$_gfx906_env" != "gfx906" ]; then + case "$_runtime_gfx" in + gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;; + esac + fi # Skip rocm7.13+ generic indexes: they already ship the fixes, so the # arch build (rocm7.13) would be a downgrade rather than a rescue. if [ -n "$_strix_gfx" ] && _rocm_leaf_below "$_torch_index_leaf" 7 13; then @@ -3327,6 +3382,57 @@ case "$_torch_index_leaf" in TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" _amd_gpu_radeon=false fi + # ── MI50 / Radeon VII (gfx906, Vega 20): legacy community-supported path ── + # Newer rocm wheel families bundle ROCm libraries whose Tensile kernels + # dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read for gfx906", + # ROCm/TheRock#1844), so a rocm6.4+/7.x index installs a torch that fails + # at the first BLAS call. The rocm6.3 index is the last one whose wheels + # run on gfx906 (torch 2.7.0 verified on MI50 32GB; up to 2.9 in community + # use). Reroute any newer picked index; leave rocm6.0-6.3 alone. + # + # Target resolution: an explicit UNSLOTH_ROCM_GFX_ARCH wins (lets a host + # whose rocminfo/amd-smi emit no gfx token still opt in; _gfx906_env was + # lowercased above, before the Strix block it suppresses). Otherwise only + # treat gfx906 as the target when it is the SOLE distinct arch present: + # _gfx_all is de-duplicated by visible index, which loses per-device + # ordinals on a mixed host, so a non-gfx906 selection must never be + # downgraded to rocm6.3 -- such hosts set UNSLOTH_ROCM_GFX_ARCH to opt in. + _gfx906_target=false + if [ -n "$_gfx906_env" ]; then + [ "$_gfx906_env" = "gfx906" ] && _gfx906_target=true + elif [ -n "$_gfx_all" ]; then + _gfx906_uniq=$(printf '%s\n' "$_gfx_all" | awk 'NF && !seen[$0]++') + [ "$_gfx906_uniq" = "gfx906" ] && _gfx906_target=true + fi + # gfx906 always trains from the PyTorch rocm6.3 wheels, never the Radeon repo + # (repo.radeon.com wheels carry no gfx906 BLAS kernels). Clear the Radeon + # marketing-name flag as soon as gfx906 is the target -- even when the host + # already picks rocm6.0-6.3 and the reroute below is a no-op -- so a Radeon VII + # does not divert to the radeon branch on those versions. + if [ "$_gfx906_target" = true ]; then + _amd_gpu_radeon=false + fi + if [ "$_gfx906_target" = true ] && ! _rocm_leaf_below "$_torch_index_leaf" 6 4; then + echo "" >&2 + echo " [WARN] gfx906 (MI50 / Radeon VII / Vega 20) detected -- routing torch to the" >&2 + echo " [WARN] rocm6.3 index: it is the last wheel family that runs on gfx906 (newer" >&2 + echo " [WARN] rocm wheels ship without gfx906 BLAS kernels and fail at first use)." >&2 + echo " [WARN] gfx906 is a community-maintained legacy path: 16-bit LoRA and full" >&2 + echo " [WARN] finetuning work out of the box; bitsandbytes 4-bit QLoRA requires a" >&2 + echo " [WARN] source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd)." >&2 + echo "" >&2 + _amd_gfx906_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" + while [ "${_amd_gfx906_base%/}" != "$_amd_gfx906_base" ]; do + _amd_gfx906_base="${_amd_gfx906_base%/}" + done + TORCH_INDEX_URL="${_amd_gfx906_base}/rocm6.3" + # Reset to the default (<2.11) window: a rocm7.2 pick raised the floor + # to 2.11 above, which the rocm6.3 index (torch <= 2.9.x) cannot satisfy. + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" + # (_amd_gpu_radeon already cleared above for every gfx906 target.) + fi ;; esac fi # _torch_index_pinned guard (Radeon + Strix reroute) @@ -3553,6 +3659,7 @@ for _p in ('torch', 'torchvision', 'torchaudio'): if [ "$_MIGRATED" = true ]; then # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving # existing torch/CUDA unless the ROCm repair below fires. + _gfx906_bnb_snapshot substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -3594,13 +3701,18 @@ if [ "$_MIGRATED" = true ]; then # existing ROCm installs gain the AMD bitsandbytes build without a # fresh reinstall. if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + if _is_gfx906_bnb_skip; then + substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" + else + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + fi # Repair ROCm torch if overwritten during migrated install _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." _install_torch_default_index --force-reinstall fi + _gfx906_bnb_prune fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) @@ -3791,8 +3903,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # host stays in GGUF-only mode rather than pulling in bitsandbytes, # which is only useful once torch is present for training. if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + if _is_gfx906_bnb_skip; then + substep "gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels); build from source for 4-bit QLoRA -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" + else + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + fi fi + _gfx906_bnb_snapshot # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." @@ -3843,6 +3960,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." _install_torch_default_index --force-reinstall fi + _gfx906_bnb_prune fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a91f26910f..2883f30b20 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -94,6 +94,40 @@ def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool: return key is not None and key < _ROCM_ARCH_INDEX_FLOOR +# MI50 / Radeon VII (gfx906, Vega 20): rocm6.4+/7.x wheels bundle ROCm libraries +# whose Tensile kernels dropped gfx906 (rocBLAS "TensileLibrary.dat ... not read +# for gfx906", ROCm/TheRock#1844), failing at the first BLAS call. The rocm6.3 +# index is the last one whose wheels run on gfx906 (torch 2.7.0 verified on MI50 +# 32GB; up to 2.9 in community use). Uses the _default (<2.11) pkg specs -- the +# rocm7.2 floor of 2.11 cannot be satisfied there. Mirrors install.sh. +_GFX906_LEGACY_TAG = "rocm6.3" + + +def _gfx906_needs_legacy_index(ver: tuple[int, int]) -> bool: + """True when the generic tag picked for the host ROCm version is newer than + rocm6.3, i.e. its wheels lack gfx906 kernels and must be rerouted.""" + key = next((k for k in sorted(_ROCM_TORCH_INDEX, reverse = True) if ver >= k), None) + return key is not None and key > (6, 3) + + +def _runtime_target_is_gfx906() -> bool: + """True when the runtime GPU target is gfx906 (MI50 / Radeon VII). + + An explicit UNSLOTH_ROCM_GFX_ARCH wins (mirrors _infer_linux_amd_gfx_arch / + the display path), so a host whose rocminfo/amd-smi emit no gfx token can + still opt in. Otherwise report gfx906 only when it is the SOLE distinct arch: + _detect_amd_gfx_codes() de-duplicates arches, which loses per-device ordinals + on a mixed host, so a non-gfx906 selection is never mis-identified as gfx906 + (and downgraded to rocm6.3). Mixed gfx906+dGPU hosts opt in with the env var. + """ + # Normalize a copied HIP gcnArchName (gfx906:sramecc-:xnack- -> gfx906) so the + # feature-flag suffix does not defeat the exact comparison (mirrors device_type.py). + override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower().split(":")[0] + if override: + return override == "gfx906" + return set(_detect_amd_gfx_codes()) == {"gfx906"} + + # AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). # Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare. _ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset( @@ -939,6 +973,34 @@ def _detect_bnb_rocm_dll_ver() -> str | None: return max(all_vers, key = lambda v: int(v)) if all_vers else None +# Set right before the base unsloth install (which resolves its unconditional +# bitsandbytes dependency); read by _ensure_rocm_torch to drop a freshly pulled +# generic wheel on gfx906 while leaving a pre-existing source build untouched. +_GFX906_BNB_ABSENT_BEFORE_BASE = False + + +def _bitsandbytes_installed() -> bool: + """True if bitsandbytes is importable in the target venv. Runs a fresh + subprocess so a package installed earlier this run is seen; only checks the + spec (does NOT import bitsandbytes).""" + try: + return ( + subprocess.run( + [ + sys.executable, + "-c", + "import importlib.util, sys; " + "sys.exit(0 if importlib.util.find_spec('bitsandbytes') else 1)", + ], + capture_output = True, + timeout = 60, + ).returncode + == 0 + ) + except Exception: + return False + + _BNB_ROCM_SITECUSTOMIZE_BEGIN = "# BEGIN Unsloth BNB_ROCM_VERSION" _BNB_ROCM_SITECUSTOMIZE_END = "# END Unsloth BNB_ROCM_VERSION" _BNB_ROCM_VERSION_SOURCE_ENV = "UNSLOTH_BNB_ROCM_VERSION_SOURCE" @@ -1890,13 +1952,25 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True + # An explicit UNSLOTH_ROCM_GFX_ARCH=gfx906 pins the runtime target to the + # MI50 / Radeon VII path; it must win over the Strix probe-order detection + # below (a mixed Strix + MI50 host could otherwise route to gfx1151), so the + # Strix override is skipped when it is set. + _gfx906_arch_override = (os.environ.get("UNSLOTH_ROCM_GFX_ARCH") or "").strip().lower().split( + ":" + )[0] == "gfx906" + # Strix Halo / Point (gfx1151 / gfx1150) need torch from AMD's per-gfx index # (2.11+rocm7.13); any generic pytorch.org rocm index lacks the fixes (ROCm 7.1 # segfaults in _grouped_mm). See _strix_needs_amd_arch_index for the floor gate. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None # An explicit ROCm pin is authoritative: never auto-reroute it. - if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None: + if ( + _strix_needs_amd_arch_index(ver) + and _explicit_rocm_torch_index_url() is None + and not _gfx906_arch_override + ): gfx_codes = _detect_amd_gfx_codes() _strix_gfx = {"gfx1151", "gfx1150", "gfx1152"} _detected_strix = _strix_gfx.intersection(gfx_codes) @@ -1933,6 +2007,34 @@ def _ensure_rocm_torch() -> None: f" skipping AMD per-gfx index override.\n" ) + # gfx906 (MI50 / Radeon VII): is this the runtime GPU target? Used below to skip + # the generic bitsandbytes wheel (no gfx906 kernels). This must hold even under + # an explicit torch-index pin: a gfx906 host that pins rocm6.3 (without also + # setting UNSLOTH_ROCM_GFX_ARCH) would otherwise reinstall the prebuilt bnb wheel + # over the user's source-built gfx906 bnb. So a pin suppresses only the torch + # reroute (_gfx906_override below), NOT the gfx906 detection for the bnb skip. + _runtime_is_gfx906 = _gfx906_arch_override or _runtime_target_is_gfx906() + # Reroute torch to the last gfx906-capable wheel family (rocm6.3) only when the + # host ROCm version would otherwise pick a newer, kernel-less index -- and never + # over an explicit pin or an active Strix reroute (the pin/Strix path installs + # its own index; only the bnb skip must still apply on those paths). + _gfx906_override = ( + _runtime_is_gfx906 + and _gfx906_needs_legacy_index(ver) + and _explicit_rocm_torch_index_url() is None + and _strix_override_url is None + ) + if _gfx906_override: + print( + f"\n gfx906 (MI50 / Radeon VII / Vega 20) is the runtime target with ROCm " + f"{ver[0]}.{ver[1]}.\n" + f" Routing torch install to the {_GFX906_LEGACY_TAG} index: the last wheel\n" + f" family that runs on gfx906 (newer rocm wheels ship without gfx906 BLAS\n" + f" kernels and fail at first use). gfx906 is a community-maintained legacy\n" + f" path: 16-bit LoRA and full finetuning work; bitsandbytes 4-bit QLoRA\n" + f" requires a source build of bitsandbytes for gfx906 (see docs.unsloth.ai/amd).\n" + ) + # The Strix override must fire even when has_hip_torch is True: an existing # torch.version.hip == "7.1" is exactly the broken combo it repairs. if _strix_override_url is not None and _strix_override_pkgs is not None: @@ -1954,6 +2056,29 @@ def _ensure_rocm_torch() -> None: constrain = False, ) rocm_torch_ready = True + # gfx906 fires even when has_hip_torch is True: a +rocm7.x build IS the broken + # combo it repairs. A torch already on rocm6.3 wheels is left alone (the tag + # check below is False, and rocm_torch_ready is already True from has_hip_torch, + # so the generic fallback is skipped). + elif _gfx906_override and _GFX906_LEGACY_TAG not in _installed_torch_ver: + index_url = f"{_PYTORCH_WHL_BASE}/{_GFX906_LEGACY_TAG}" + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["_default"] + print( + f" gfx906 legacy override -- installing torch from " + f"{_strip_index_url_credentials(index_url)}" + ) + pip_install( + f"ROCm torch (gfx906, {_GFX906_LEGACY_TAG})", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) + rocm_torch_ready = True elif not rocm_torch_ready: # Reinstall when torch is not ROCm yet, OR a ROCm build's family differs from a pin. # Gate on rocm_torch_ready (not has_hip_torch alone) so a successful inferred-gfx @@ -2002,11 +2127,33 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True + # gfx906 has no prebuilt bitsandbytes: the continuous-release/PyPI wheels ship + # no gfx906 kernels, and force-reinstalling them would clobber a user's + # source-built bnb (the only 4-bit path on this arch) on every `studio update`. + # Skip the auto-install and leave whatever bnb is present. + if rocm_torch_ready and _runtime_is_gfx906: + print( + _dim( + " gfx906: skipping prebuilt bitsandbytes (no gfx906 kernels). " + "Build bitsandbytes from source for 4-bit QLoRA -- " + "see docs.unsloth.ai/get-started/install-and-update/amd." + ) + ) + # The base install resolves unsloth's unconditional bitsandbytes dep to a + # generic CUDA wheel with no gfx906 kernels ("invalid device function" at + # 4-bit use). Drop it if this run pulled it in; a pre-existing source build + # (present before the base install) is left untouched. + if _GFX906_BNB_ABSENT_BEFORE_BASE and _bitsandbytes_installed(): + print(_dim(" gfx906: removing generic bitsandbytes pulled in as a dependency")) + subprocess.run( + [sys.executable, "-m", "pip", "uninstall", "-y", "bitsandbytes"], + capture_output = True, + ) # Install bitsandbytes only when torch links against ROCm. Prefers the # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix), falling back # to PyPI when the pre-release wheel won't install. Use pip for the # pre-release wheel because uv rejects its filename/metadata version mismatch. - if rocm_torch_ready: + elif rocm_torch_ready: _bnb_url = _bnb_rocm_prerelease_url() _bnb_installed = False if _bnb_url is not None: @@ -2767,6 +2914,13 @@ def install_python_stack() -> int: "mlx-vlm", ) + # gfx906: the base install below resolves unsloth's unconditional bitsandbytes + # dep to a generic CUDA wheel (no gfx906 kernels). Record bnb's presence now so + # _ensure_rocm_torch can drop a freshly pulled wheel while keeping a source build. + global _GFX906_BNB_ABSENT_BEFORE_BASE + if not skip_base: + _GFX906_BNB_ABSENT_BEFORE_BASE = not _bitsandbytes_installed() + # 3. Core packages: unsloth-zoo + unsloth (or custom package name) if skip_base: pass diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index c58808689b..d42a274478 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -88,11 +88,21 @@ class TestStructuralTorchConstraint: """$TORCH_CONSTRAINT must appear in a uv pip install line.""" assert '"$TORCH_CONSTRAINT"' in self._sh - def test_hardcoded_torch_constraint_only_once(self): - """The hard-coded torch>=2.4,<2.11.0 string should appear exactly once - in install.sh (the default assignment), not in pip install lines.""" - count = self._sh.count('"torch>=2.4,<2.11.0"') - assert count == 1, f"Expected 1, found {count}" + def test_hardcoded_torch_constraint_only_on_assignments(self): + """The hard-coded torch>=2.4,<2.11.0 string must only appear on + TORCH_CONSTRAINT= assignment lines, never on a pip/uv install line + (those must reference $TORCH_CONSTRAINT). Two assignments are expected: + the default, and the gfx906 (MI50) reroute that restores the default + <2.11 window after the rocm7.2 floor bump raised it to 2.11.""" + hits = [ln for ln in self._sh.splitlines() if '"torch>=2.4,<2.11.0"' in ln] + assert hits, "default constraint literal missing from install.sh" + for ln in hits: + assert ( + "TORCH_CONSTRAINT=" in ln + ), f"torch>=2.4,<2.11.0 hardcoded off a TORCH_CONSTRAINT= assignment: {ln.strip()!r}" + assert ( + "pip install" not in ln + ), f"torch>=2.4,<2.11.0 hardcoded on a pip install line: {ln.strip()!r}" def test_tightening_guarded_by_skip_torch(self): """The block must check SKIP_TORCH=false.""" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4f65857a3e..f358c4bba7 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -826,8 +826,10 @@ class TestEnsureRocmTorch: self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): """An explicit gfx wheel-index pin is authoritative: install from it verbatim - with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4 - would otherwise pick the rocm6.4 wheel / trigger the Strix re-route).""" + with torch 2.11, and the pin must not be second-guessed (host ROCm 6.4 would + otherwise pick the rocm6.4 wheel / trigger the Strix re-route). The gfx probe + may run for the bnb-skip flag, but returning a Strix arch must not reroute the + pinned torch index.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"\n" # cpu torch -> reinstall @@ -836,10 +838,7 @@ class TestEnsureRocmTorch: stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): - # Would raise if the Strix block ran (it is skipped on an explicit pin). - with patch.object( - stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError - ): + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]): _ensure_rocm_torch() assert mock_pip.call_count == 1 torch_call = str(mock_pip.call_args_list[0]) @@ -931,7 +930,8 @@ class TestEnsureRocmTorch: def test_gfx_pin_over_installed_pre211_rocm_reinstalls( self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try ): - """A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls.""" + """A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls. + The gfx probe may run for the bnb-skip flag but must not alter the pinned index.""" mock_probe = MagicMock() mock_probe.returncode = 0 mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n" @@ -940,9 +940,7 @@ class TestEnsureRocmTorch: stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): - with patch.object( - stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError - ): + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]): _ensure_rocm_torch() torch_call = str(mock_pip.call_args_list[0]) assert "gfx1151" in torch_call @@ -1027,9 +1025,9 @@ class TestEnsureRocmTorch: stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): - with patch.object( - stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError - ): + # The gfx probe may run for the bnb-skip flag; returning a Strix + # arch must not reroute the pinned torch index. + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151"]): _ensure_rocm_torch() torch_call = str(mock_pip.call_args_list[0]) assert "gfx1151" in torch_call @@ -1158,6 +1156,336 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() +# TEST: gfx906 (MI50 / Radeon VII) legacy reroute -- generic wheels after rocm6.3 +# lack gfx906 code objects, so torch must come from the rocm6.3 index. + + +class TestGfx906LegacyReroute: + """gfx906 hosts on ROCm >= 6.4 must be rerouted to the rocm6.3 torch index; + hosts already on gfx906-capable wheels are left alone.""" + + @staticmethod + def _gfx906_reroute_block(source: str) -> str: + """The MI50/gfx906 reroute block, bounded on the ';;' that closes its + rocm[0-9]* case arm -- robust to comment growth (no magic char offset).""" + start = source.find("MI50 / Radeon VII (gfx906") + assert start >= 0, "gfx906 reroute block not found in install.sh" + end = source.find("\n ;;", start) + assert end >= 0, "end of gfx906 case arm not found" + return source[start:end] + + def test_gfx906_needs_legacy_index_floor(self): + f = stack_mod._gfx906_needs_legacy_index + # rocm6.0-6.3 tags still ship gfx906 kernels: no reroute. + assert f((6, 3)) is False + assert f((6, 0)) is False + assert f((5, 0)) is False # below any known tag + # Anything that picks a tag newer than rocm6.3 must reroute. + assert f((6, 4)) is True + assert f((7, 2)) is True + assert f((7, 14)) is True + + def test_runtime_target_is_gfx906_selection(self, monkeypatch): + """Env override wins; else gfx906 only when it is the SOLE distinct arch.""" + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + # Sole gfx906 (one or several identical MI50s de-dup to {'gfx906'}). + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]): + assert stack_mod._runtime_target_is_gfx906() is True + # Mixed host: gfx906 is NOT the sole arch -> not auto-selected (Codex #3: + # de-dup loses ordinals, so never downgrade a non-gfx906 selection). + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906", "gfx1100"]): + assert stack_mod._runtime_target_is_gfx906() is False + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []): + assert stack_mod._runtime_target_is_gfx906() is False + # Explicit override wins even when probes see nothing (Codex #2). + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906") + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []): + assert stack_mod._runtime_target_is_gfx906() is True + # ...and a non-gfx906 override is honored on a gfx906-present host. + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx1100") + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]): + assert stack_mod._runtime_target_is_gfx906() is False + # A copied HIP gcnArchName (gfx906:sramecc-:xnack-) normalizes to gfx906 + # (Codex #4: the feature-flag suffix must not defeat the exact comparison). + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906:sramecc-:xnack-") + with patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []): + assert stack_mod._runtime_target_is_gfx906() is True + + @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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]) + def test_gfx906_on_rocm72_routes_to_rocm63( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """CPU torch on a ROCm 7.2 MI50 host installs from rocm6.3, not rocm7.2.""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm6.3" in torch_call + assert "rocm7.2" not in torch_call + # The _default (<2.11) window: the rocm7.2 2.11 floor cannot be satisfied + # on the rocm6.3 index (torch <= 2.9.x there). + assert "torch>=2.4,<2.11.0" in torch_call + # gfx906 has no prebuilt bnb -- the generic wheel must not be installed. + assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list) + + @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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]) + def test_gfx906_repairs_existing_rocm72_torch( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """An installed +rocm7.2 torch IS the broken combo: reinstall from rocm6.3 + even though has_hip_torch is True.""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm6.3" in torch_call + assert "torch>=2.4,<2.11.0" in 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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]) + def test_gfx906_already_on_rocm63_left_alone( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """torch already on rocm6.3 wheels must not be reinstalled (no update loop), + and the generic bnb wheel must not clobber a source build.""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"6.3.42131|2.7.0+rocm6.3\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + mock_pip.assert_not_called() + # gfx906: prebuilt bnb is skipped entirely (no torch reinstall, no bnb). + mock_pip_try.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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1100", "gfx906"]) + def test_mixed_host_gfx906_not_sole_arch_skips_reroute( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """Mixed host (gfx906 + dGPU) with no explicit override: gfx906 is not the + sole arch, so the generic index is kept (Codex #3: never downgrade a + de-dup-ambiguous mixed host).""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm7.2" in torch_call + assert "rocm6.3" not in 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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = []) + def test_gfx906_env_override_forces_reroute( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """UNSLOTH_ROCM_GFX_ARCH=gfx906 reroutes even when probes emit no gfx token + (Codex #2: runtime-only ROCm hosts where rocminfo/amd-smi are absent).""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906") + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm6.3" in torch_call + assert "rocm7.2" not in 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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]) + def test_gfx906_bnb_skipped_even_when_index_pinned( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """A gfx906 user who pins the ROCm index AND sets the arch override still + skips the generic bnb wheel: the pin suppresses the torch reroute, not the + gfx906 runtime flag used for the bnb skip.""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906") + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://download.pytorch.org/whl/rocm6.3") + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall from the pinned index + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + # torch is (re)installed from the pinned rocm6.3 index... + assert any("rocm6.3" in str(c) for c in mock_pip.call_args_list) + # ...but the prebuilt bnb wheel is never installed. + assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list) + + @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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx1151", "gfx906"]) + def test_gfx906_override_wins_over_strix_probe( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """Mixed Strix + MI50 host: UNSLOTH_ROCM_GFX_ARCH=gfx906 suppresses the Strix + override (which probe order would otherwise pick) and routes to rocm6.3.""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("UNSLOTH_ROCM_GFX_ARCH", "gfx906") + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm6.3" in torch_call + assert "gfx1151" not in torch_call + # gfx906 target -> generic bnb wheel skipped. + assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list) + + @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) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + @patch.object(stack_mod, "_detect_amd_gfx_codes", return_value = ["gfx906"]) + def test_gfx906_bnb_skipped_on_pinned_index_without_env_override( + self, mock_gfx, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try, monkeypatch + ): + """Codex #2: a real gfx906 host that pins the ROCm index but does NOT set + UNSLOTH_ROCM_GFX_ARCH must still skip the prebuilt bnb wheel -- the pin + suppresses only the torch reroute, not the probe-driven gfx906 detection + used for the bnb skip (otherwise `studio update` clobbers source-built bnb).""" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("UNSLOTH_ROCM_GFX_ARCH", raising = False) + monkeypatch.setenv("UNSLOTH_TORCH_INDEX_URL", "https://download.pytorch.org/whl/rocm6.3") + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall from the pinned index + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + # torch is (re)installed from the pinned rocm6.3 index... + assert any("rocm6.3" in str(c) for c in mock_pip.call_args_list) + # ...but the prebuilt bnb wheel is never installed (probe saw sole gfx906). + assert not any("bitsandbytes" in str(c).lower() for c in mock_pip_try.call_args_list) + + def test_install_sh_gfx906_env_suppresses_strix(self): + """install.sh must skip the Strix reroute when UNSLOTH_ROCM_GFX_ARCH=gfx906.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + assert 'if [ "$_gfx906_env" != "gfx906" ]; then' in source + + def test_install_sh_gfx906_normalizes_override_and_clears_radeon(self): + """install.sh must (Codex #4) strip a gfx906:… feature suffix before the exact + comparison, and (Codex #3) clear the Radeon marketing flag for every gfx906 + target -- not only when the >=6.4 reroute fires -- so a Radeon VII already on + rocm6.3 does not divert to the repo.radeon.com branch.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + # Override normalization (both the reroute block and the bnb-skip helper): + # strip the gfx906:… feature suffix and trim whitespace (mirror py .strip()). + assert "_gfx906_env=${_gfx906_env%%:*}" in source + assert "_bnb_gfx_env=${_bnb_gfx_env%%:*}" in source + assert source.count("tr -d '[:space:]'") >= 2 + # Radeon flag cleared as soon as gfx906 is the target, before the leaf gate. + block = self._gfx906_reroute_block(source) + clear_pos = block.find("_amd_gpu_radeon=false") + leaf_gate_pos = block.find("_rocm_leaf_below") + assert clear_pos >= 0 and leaf_gate_pos >= 0 + # the unconditional clear must precede the >=6.4 leaf-gated reroute. + assert clear_pos < leaf_gate_pos + + def test_install_sh_bnb_skip_probes_under_pin(self): + """install.sh _is_gfx906_bnb_skip must probe gfx906 when the index is pinned + (Codex #1): a pin skips the reroute block that sets _gfx906_target, so the + helper falls back to _probe_amd_gfx_arch to catch a real gfx906 host.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + start = source.find("_is_gfx906_bnb_skip() {") + assert start >= 0 + body = source[start : start + 900] + assert "_torch_index_pinned" in body + assert "_probe_amd_gfx_arch" in body + + def test_install_sh_has_gfx906_reroute(self): + """install.sh must mirror the Python reroute: honor UNSLOTH_ROCM_GFX_ARCH, + gate on a gfx906 target, route to rocm6.3, with the same _default (<2.11) + trio, and skip the prebuilt bnb wheel.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + block = self._gfx906_reroute_block(source) + assert "_gfx906_target=" in block + assert "UNSLOTH_ROCM_GFX_ARCH" in block + assert "/rocm6.3" in block + for spec in stack_mod._ROCM_TORCH_PKG_SPECS["_default"]: + assert spec in block + # The bnb skip helper must exist and be wired at the install sites. + assert "_is_gfx906_bnb_skip" in source + + def test_device_type_defaults_compile_off_on_gfx906(self): + """unsloth/device_type.py must default Dynamo/compile off on gfx906 + (user-overridable via setdefault).""" + source = (PACKAGE_ROOT / "unsloth" / "device_type.py").read_text(encoding = "utf-8") + gate_start = source.find("gfx906") + assert gate_start >= 0 + gate_body = source[gate_start : gate_start + 800] + assert 'setdefault("TORCHDYNAMO_DISABLE", "1")' in gate_body + assert 'setdefault("TORCH_COMPILE_DISABLE", "1")' in gate_body + assert 'setdefault("UNSLOTH_COMPILE_DISABLE", "1")' in gate_body + + # TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692) diff --git a/unsloth/device_type.py b/unsloth/device_type.py index 0e425e2001..1417f4f53c 100644 --- a/unsloth/device_type.py +++ b/unsloth/device_type.py @@ -117,6 +117,23 @@ DEVICE_COUNT: int = get_device_count() ALLOW_PREQUANTIZED_MODELS: bool = True # HSA_STATUS_ERROR_EXCEPTION checks - sometimes AMD fails for BnB ALLOW_BITSANDBYTES: bool = True +# gfx906 (MI50 / Radeon VII / Vega 20): Dynamo/Inductor codegen is broken on this +# legacy GCN arch (ROCm dropped it after 6.3) - compiled graphs crash or miscompile +# while the eager path trains fine. Default compile off; setdefault so a user +# override wins. +if DEVICE_TYPE == "hip": + try: + _gcn_arch = torch.cuda.get_device_properties(0).gcnArchName.split(":")[0].strip().lower() + except Exception: + _gcn_arch = "" + if _gcn_arch == "gfx906": + os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") + os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + os.environ.setdefault("UNSLOTH_COMPILE_DISABLE", "1") + print( + "Unsloth: gfx906 (MI50 / Radeon VII) detected - torch.compile disabled " + "(community-maintained legacy GCN path)." + ) if DEVICE_TYPE == "hip": try: import bitsandbytes diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 9cd5e7243a..8580c09c39 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -418,7 +418,13 @@ def fix_vllm_aimv2_issue(): spec = importlib.util.find_spec("vllm") if spec is None: return - vllm_version = importlib_version("vllm") + # A findable spec with unreadable dist metadata (broken/partial vllm install) + # must not crash unsloth import; every other vllm probe here guards this too. + try: + vllm_version = importlib_version("vllm") + except Exception as e: + logger.info(f"Unsloth: Skipping vLLM aimv2 fix -- vLLM version unreadable ({e})") + return if Version(vllm_version) < Version("0.10.1"): vllm_location = spec.origin if vllm_location is None: