ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

amd: fall back to the PyPI bitsandbytes floor on Windows ROCm too

Browse source
This commit is contained in:
Daniel Han 2026-07-28 23:22:32 +00:00
commit 0f3833e283
2 changed files with 67 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

56
studio/install_python_stack.py
View file

@ -1189,7 +1189,7 @@ _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
@ -1199,19 +1199,39 @@ def _install_bnb_windows_rocm() -> bool:
(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.
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
@ -1701,8 +1721,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
@ -1780,12 +1800,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

33
tests/studio/install/test_rocm_support.py
View file

@ -2829,12 +2829,37 @@ 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."""