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: