diff --git a/install.sh b/install.sh index b77fa37020..146a64e692 100755 --- a/install.sh +++ b/install.sh @@ -261,6 +261,20 @@ run_install_cmd_retry() { # shape on every AMD GPU; the fix (bnb #1887) is in the continuous-release_main # wheel used below and first ships on PyPI in 0.50.0, hence the fallback floor. _BNB_ROCM_PYPI_FALLBACK="bitsandbytes>=0.50.0" +# bitsandbytes ships no ROCm binary in its aarch64 wheel at any version: the PyPI +# 0.50.0 and continuous-release_main aarch64 wheels both carry only +# libbitsandbytes_cpu.so plus CUDA variants. So neither install path below gives +# aarch64 a 4-bit backend, and the messages must not claim one. Cf. gfx906. +_bnb_rocm_arch_has_binary() { + case "$_ARCH" in + aarch64|arm64) return 1 ;; + *) return 0 ;; + esac +} +_warn_bnb_no_rocm_binary() { + _bnb_rocm_arch_has_binary && return 0 + substep "[WARN] aarch64: bitsandbytes ships no ROCm kernels on this arch; 4-bit QLoRA needs a source build -- https://docs.unsloth.ai/get-started/install-and-update/amd" "$C_WARN" +} _install_bnb_rocm() { _label="$1" _venv_py="$2" @@ -292,6 +306,7 @@ _install_bnb_rocm() { --retries 8 --timeout 90 \ "$_bnb_whl_url" >"$_bnb_log" 2>&1; then rm -f "$_bnb_log" + _warn_bnb_no_rocm_binary return 0 fi _bnb_rc=$? @@ -300,10 +315,17 @@ _install_bnb_rocm() { fi rm -f "$_bnb_log" step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 - substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK, which carries the ROCm 4-bit fix" "$C_WARN" + if _bnb_rocm_arch_has_binary; then + substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK, which carries the ROCm 4-bit fix" "$C_WARN" + else + substep "[WARN] bnb pre-release install failed; falling back to PyPI $_BNB_ROCM_PYPI_FALLBACK" "$C_WARN" + fi fi run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \ --force-reinstall --no-cache-dir --no-deps "$_BNB_ROCM_PYPI_FALLBACK" + _bnb_pypi_rc=$? + _warn_bnb_no_rocm_binary + return $_bnb_pypi_rc } if [ "$_next_is_package" = true ]; then diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 981ab651c8..533cb9b793 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -424,6 +424,16 @@ def _bnb_rocm_prerelease_url() -> str | None: return _BNB_ROCM_PRERELEASE_URLS.get(arch) +def _bnb_rocm_arch_has_binary() -> bool: + """False on aarch64: bitsandbytes ships no ROCm kernels there at any version. + The PyPI 0.50.0 and continuous-release_main aarch64 wheels both carry only + libbitsandbytes_cpu.so plus CUDA variants, so neither install path gives + aarch64 a 4-bit backend and neither message may claim one. + """ + arch = platform.machine().lower() + return {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch) != "aarch64" + + def _amd_smi_env() -> dict[str, str] | None: """On Windows, env with __COMPAT_LAYER=RunAsInvoker; None elsewhere. NB: RunAsInvoker doesn't stop amd-smi's runtime elevation (its manifest is @@ -2021,10 +2031,15 @@ def _ensure_rocm_torch() -> None: force_pip = True, ) if not _bnb_installed: + _fallback_note = ( + ", which carries the ROCm 4-bit fix" + if _bnb_rocm_arch_has_binary() + else "" + ) print( _red( " bnb pre-release install failed; falling back to PyPI " - f"{_BNB_ROCM_PYPI_FALLBACK}, which carries the ROCm 4-bit fix" + f"{_BNB_ROCM_PYPI_FALLBACK}{_fallback_note}" ) ) if not _bnb_installed: @@ -2036,6 +2051,14 @@ def _ensure_rocm_torch() -> None: _BNB_ROCM_PYPI_FALLBACK, constrain = False, ) + if not _bnb_rocm_arch_has_binary(): + print( + _red( + " aarch64: bitsandbytes ships no ROCm kernels on this arch; " + "4-bit QLoRA needs a source build -- " + "https://docs.unsloth.ai/get-started/install-and-update/amd" + ) + ) # _uv_safe_path is imported from backend.utils.uv_path_safety (shared with mlx_repair). diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 9a183d2b4e..ef8926454f 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -876,3 +876,21 @@ class TestAmdBnbFloorParity: assert ( "4-bit decode will be broken on ROCm" not in text ), f"{path.name} still reports the repaired PyPI fallback as broken" + + def test_aarch64_is_not_told_it_has_a_rocm_backend(self): + """bitsandbytes ships no ROCm kernels in its aarch64 wheel at any version, so + neither installer may hand aarch64 the x86_64 "carries the ROCm 4-bit fix" + message, and both must warn that 4-bit needs a source build there.""" + sh = INSTALL_SH.read_text(encoding = "utf-8") + assert "_bnb_rocm_arch_has_binary()" in sh + assert "_warn_bnb_no_rocm_binary()" in sh + assert ( + sh.count("_warn_bnb_no_rocm_binary\n") >= 2 + ), "install.sh must warn on aarch64 after both the pre-release and the fallback install" + py = STACK_PY.read_text(encoding = "utf-8") + assert "def _bnb_rocm_arch_has_binary(" in py + assert "_bnb_rocm_arch_has_binary()" in py + for text, name in ((sh, "install.sh"), (py, "install_python_stack.py")): + assert ( + "4-bit QLoRA needs a source build" in text + ), f"{name} must tell aarch64 users 4-bit needs a source build"