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
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue