AMD: enable ROCm torch on gfx906 (MI50 / Radeon VII) on Linux (#7354)
* 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 <unslothai@gmail.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Leo Borcherding
GitHub
parent
74295d93d8
commit
f03e669442
6 changed files with 659 additions and 26 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue