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

amd: require bitsandbytes>=0.50.0 in the amd extra (fixes ROCm 4-bit NaNs) (#7535)

Browse source 
* amd: require bitsandbytes>=0.50.0 in the amd extra

bnb <= 0.49.2 NaNs at decode shape on every AMD GPU. The ROCm 4-bit GEMV
fix (bnb PR #1887) first ships in 0.50.0, on PyPI since 2026-07-24, so the
old >=0.49.1 floor could still resolve the broken range.

Mirrors the same change made on the pip release branch in #7278.

* amd: cite the 0.50.0 ROCm work accurately in the bnb floor comment

The comment credited bnb PR #1887 as "the ROCm 4-bit GEMV fix" for every
AMD GPU. #1887 decouples blocksize from warp size and fixes a hardcoded
warp size of 32 in kgemm_4bit_inference_naive, which is a CDNA problem by
construction. The RDNA-side work is #1979 (fused 4-bit SIMT GEMM) and
#2012 (RDNA3/4 workgroup resonance). All three first ship in 0.50.0, so
the >=0.50.0 floor is unchanged; only the justification was wrong.

* amd: raise the installer bitsandbytes fallback floors to 0.50.0

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* amd: stop reporting the bitsandbytes PyPI fallback as broken

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten AMD bnb floor comments

* Keep the amd extra citation and the AMD install guide reference

* amd: do not promise aarch64 a ROCm 4-bit backend it never gets

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Leo Borcherding 2026-07-28 20:12:26 -05:00 • committed by GitHub
commit 411cb86d62
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 206 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

99
studio/install_python_stack.py
View file

@ -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).