diff --git a/install.sh b/install.sh index 376daa8fab..72f2455277 100755 --- a/install.sh +++ b/install.sh @@ -321,10 +321,25 @@ _gfx906_bnb_prune() { || "$_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 -# pre-release URL is unreachable. Drop the pin once bnb 0.50+ ships on PyPI. +# Install bitsandbytes on AMD ROCm hosts. bnb <= 0.49.2 NaNs at 4-bit decode +# shape on every AMD GPU; the fix (bnb #1887) ships in continuous-release_main +# and, on PyPI, first in 0.50.0. Keep this floor in step with the amd extra in +# pyproject.toml and studio/install_python_stack.py. +_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" @@ -339,9 +354,8 @@ _install_bnb_rocm() { _bnb_whl_url="" ;; esac - # uv rejects the continuous-release_main bitsandbytes wheel because the - # filename version (1.33.7rc0) does not match the embedded metadata version - # (0.50.0.dev0). pip accepts the mismatch, so bootstrap pip and use it. + # uv rejects the pre-release wheel: filename version (1.33.7rc0) does not + # match metadata (0.50.x.dev0). pip accepts it, so bootstrap pip and use it. if ! "$_venv_py" -m pip --version >/dev/null 2>&1; then if ! run_maybe_quiet "$_venv_py" -m ensurepip --upgrade; then run_maybe_quiet uv pip install --python "$_venv_py" pip || \ @@ -357,6 +371,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=$? @@ -365,10 +380,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 (4-bit decode broken on ROCm)" "$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 "bitsandbytes>=0.49.1" + --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/pyproject.toml b/pyproject.toml index 62623499d6..7359a51fa6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1257,8 +1257,11 @@ intel = [ ] amd = [ "unsloth[huggingfacenotorch]", - "bitsandbytes>=0.49.1 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')", - "bitsandbytes>=0.49.1 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", + # 4-bit decode is unreliable on ROCm before 0.50.0, the first PyPI release + # carrying the full path: blocksize/warp decoupling (bnb #1887), fused SIMT + # GEMM on RDNA (#1979), RDNA3/4 workgroup fix (#2012). + "bitsandbytes>=0.50.0 ; ('linux' in sys_platform) and (platform_machine == 'AMD64' or platform_machine == 'x86_64' or platform_machine == 'aarch64')", + "bitsandbytes>=0.50.0 ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] rocm702-torch280 = [ "unsloth[amd]", diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 8c71d39e16..3243089656 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -426,8 +426,8 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = { } # bitsandbytes continuous-release_main wheels with 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. Drop the pin once bnb 0.50+ ships on PyPI. +# (bnb #1887, post-0.49.2). bnb <= 0.49.2 NaNs at decode shape on every AMD GPU; +# PyPI 0.50.0 is the first release with the fix, so the fallback below is safe. _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "x86_64": ( "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" @@ -448,7 +448,8 @@ _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { "bitsandbytes-1.33.7.preview-py3-none-win_amd64.whl" ), } -_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1" +# Keep in step with the amd extra in pyproject.toml and the install.sh fallback. +_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.50.0" def _bnb_rocm_prerelease_url() -> str | None: @@ -460,6 +461,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 @@ -1243,29 +1254,46 @@ _rocm_windows_torch_installed: bool = False def _install_bnb_windows_rocm() -> bool: - """Install the AMD Windows BNB prerelease wheel. Returns True on success. + """Install AMD Windows BNB, pre-release wheel first. Returns True on success. - The continuous-release wheel is intentionally mismatched: the filename - encodes 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the wheel - metadata reports 0.50.0.dev0. uv rejects this filename/metadata mismatch, - and bypassing it with UV_SKIP_WHEEL_FILENAME_CHECK still leaves uv mangling - the bitsandbytes install. Per the AMD install guide - (https://unsloth.ai/docs/get-started/install/amd/amd-hackathon) the wheel - must be installed with plain pip, not uv, so we force pip (force_pip=True); - plain pip performs no wheel filename/metadata check. + The wheel's filename version (1.33.7.preview, PEP 440 1.33.7rc0) does not + match its metadata (0.50.x.dev0). uv rejects the mismatch and still mangles + the install under UV_SKIP_WHEEL_FILENAME_CHECK, so force plain pip, which + performs no such check. Per the AMD install guide + (https://unsloth.ai/docs/get-started/install/amd/amd-hackathon). + + When that URL is blocked, fall back to PyPI. Its win_amd64 wheel ships + libbitsandbytes_rocm{714,72}.dll from 0.50.0 on, so the fallback is a real + ROCm build; before 0.50.0 it was CUDA-only, which is why there was none. """ _bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64") - if _bnb_win_url is None: - return False - _ok = pip_install_try( - "bitsandbytes (AMD Windows, pre-release main)", - "--force-reinstall", - "--no-cache-dir", - "--no-deps", - _bnb_win_url, - constrain = False, - force_pip = True, - ) + _ok = False + if _bnb_win_url is not None: + _ok = pip_install_try( + "bitsandbytes (AMD Windows, pre-release main)", + "--force-reinstall", + "--no-cache-dir", + "--no-deps", + _bnb_win_url, + constrain = False, + force_pip = True, + ) + if not _ok: + print( + _red( + " bnb pre-release install failed; falling back to PyPI " + f"{_BNB_ROCM_PYPI_FALLBACK}, which carries the ROCm 4-bit fix" + ) + ) + if not _ok: + _ok = pip_install_try( + "bitsandbytes (AMD Windows)", + "--force-reinstall", + "--no-cache-dir", + "--no-deps", + _BNB_ROCM_PYPI_FALLBACK, + constrain = False, + ) if not _ok: return False # Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb @@ -1755,8 +1783,8 @@ def _ensure_rocm_torch() -> None: pass if _torch_ok: _rocm_windows_torch_installed = True - # ROCm torch is already installed, but the AMD Windows BNB wheel is still - # needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm). + # ROCm torch is already installed, but bnb still needs the ROCm build + # (pre-release wheel, else PyPI >=0.50.0). _install_bnb_windows_rocm() return # torch was wiped between runs; fall through to the full install path @@ -1834,12 +1862,12 @@ def _ensure_rocm_torch() -> None: # separate dependency -- a BNB install failure must NOT roll back the # torch ROCm install. _rocm_windows_torch_installed = True - # Always install AMD Windows bitsandbytes -- the PyPI wheel ships only - # CUDA DLLs and fails on ROCm. Install even when torch was already a - # ROCm build so `studio update` repairs a broken bnb. + # Always install AMD Windows bitsandbytes, even when torch was already a + # ROCm build, so `studio update` repairs a broken bnb. if not _install_bnb_windows_rocm(): print( - " Warning: AMD Windows bitsandbytes install failed; " + " Warning: AMD Windows bitsandbytes install failed " + "(pre-release and PyPI); " "ROCm torch is installed but bitsandbytes may need manual install" ) return @@ -2170,10 +2198,13 @@ 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 " - "(4-bit decode will be broken on ROCm)" + f"{_BNB_ROCM_PYPI_FALLBACK}{_fallback_note}" ) ) if not _bnb_installed: @@ -2185,6 +2216,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 6c2a1d09cf..b20e715ebc 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -862,3 +862,76 @@ class TestNoTorchPersistenceParity: manifest = (REPO_ROOT / "studio" / "install_manifest.py").read_text(encoding = "utf-8") assert 'NO_TORCH_TRUTHY: Tuple[str, ...] = ("1", "true", "yes", "on")' in manifest assert "install_manifest.NO_TORCH_TRUTHY" in STACK_PY.read_text(encoding = "utf-8") + + +class TestAmdBnbFloorParity: + """bitsandbytes <= 0.49.2 NaNs at 4-bit decode shape on every AMD GPU; the ROCm + 4-bit GEMV fix (bnb #1887) first ships on PyPI in 0.50.0. The `amd` extra, + install.sh and the Studio stack resolve bitsandbytes independently, so all three + must carry the same floor or an unreachable pre-release wheel silently reinstates + the broken range.""" + + FLOOR = "0.50.0" + PYPROJECT = REPO_ROOT / "pyproject.toml" + + def test_amd_extra_floor(self): + text = self.PYPROJECT.read_text(encoding = "utf-8") + amd = re.search(r"^amd = \[(.*?)^\]", text, re.S | re.M) + assert amd, "pyproject.toml must define an `amd` extra" + specs = re.findall(r'"(bitsandbytes[^"]*)"', amd.group(1)) + assert specs, "the amd extra must pin bitsandbytes" + for spec in specs: + assert spec.startswith( + f"bitsandbytes>={self.FLOOR}" + ), f"amd extra bitsandbytes floor must be >={self.FLOOR}, got {spec!r}" + + def test_install_sh_pypi_fallback_floor(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + assert ( + f'_BNB_ROCM_PYPI_FALLBACK="bitsandbytes>={self.FLOOR}"' in text + ), f"install.sh _install_bnb_rocm PyPI fallback must floor at {self.FLOOR}" + + def test_stack_py_pypi_fallback_floor(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert ( + f'_BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>={self.FLOOR}"' in text + ), f"install_python_stack.py PyPI fallback must floor at {self.FLOOR}" + + def test_no_installer_still_allows_the_broken_range(self): + for path in (INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY, self.PYPROJECT): + text = path.read_text(encoding = "utf-8") + for line in text.splitlines(): + if "bitsandbytes>=0.49" in line and not line.lstrip().startswith(("#", "//")): + raise AssertionError( + f"{path.name} still floors bitsandbytes in the broken ROCm range: {line.strip()!r}" + ) + + def test_fallback_is_not_reported_as_broken(self): + """The fallback now installs the first fixed release, so neither installer + may still call 4-bit decode broken on ROCm.""" + for path in (INSTALL_SH, STACK_PY): + text = path.read_text(encoding = "utf-8") + assert ( + "4-bit decode broken on ROCm" not in text + ), f"{path.name} still reports the repaired PyPI fallback as broken" + 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" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index b003382859..51c2d6587c 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3157,12 +3157,35 @@ class TestInstallBnbWindowsRocm: assert result is False assert "BNB_ROCM_VERSION" not in os.environ - def test_no_op_when_win_amd64_url_missing(self): - """Should be silent no-op if win_amd64 key absent from _BNB_ROCM_PRERELEASE_URLS.""" + def test_falls_back_to_pypi_when_win_amd64_url_missing(self): + """No win_amd64 pre-release wheel must not mean no bitsandbytes: PyPI + >=0.50.0 ships libbitsandbytes_rocm{714,72}.dll, so it is a real ROCm build.""" with patch.object(stack_mod, "_BNB_ROCM_PRERELEASE_URLS", {}): - with patch.object(stack_mod, "pip_install_try") as mock_pip: + with patch.object(stack_mod, "pip_install_try", return_value = True) as mock_pip: stack_mod._install_bnb_windows_rocm() - mock_pip.assert_not_called() + assert mock_pip.call_count == 1 + assert stack_mod._BNB_ROCM_PYPI_FALLBACK in mock_pip.call_args.args + + def test_falls_back_to_pypi_when_prerelease_install_fails(self): + """A blocked GitHub pre-release URL must fall through to the PyPI floor rather + than leaving Windows ROCm with no working bitsandbytes.""" + with patch.object(stack_mod, "pip_install_try", side_effect = [False, True]) as mock_pip: + with patch.object(stack_mod, "_detect_bnb_rocm_dll_ver", return_value = "72"): + result = stack_mod._install_bnb_windows_rocm() + assert result is True + assert mock_pip.call_count == 2 + assert "win_amd64" in str(mock_pip.call_args_list[0]) + assert stack_mod._BNB_ROCM_PYPI_FALLBACK in mock_pip.call_args_list[1].args + + def test_returns_false_only_when_both_paths_fail(self): + """Both the pre-release wheel and the PyPI fallback must fail before the + helper reports failure.""" + with patch.dict(os.environ, {}, clear = False): + os.environ.pop("BNB_ROCM_VERSION", None) + with patch.object(stack_mod, "pip_install_try", return_value = False) as mock_pip: + result = stack_mod._install_bnb_windows_rocm() + assert result is False + assert mock_pip.call_count == 2 def test_sets_bnb_rocm_version_from_detected_dll(self): """BNB_ROCM_VERSION is set from the DLL detected after install."""