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

fix: BNB AMD wheel install fails uv wheel filename check

Browse source 
The bitsandbytes continuous-release wheel is intentionally mismatched:
filename encodes 1.33.7.preview (= 1.33.7rc0 in PEP 440) but wheel
metadata reports 0.50.0.dev0.  uv rejects this by default.

Introduce _install_bnb_windows_rocm() helper that sets
UV_SKIP_WHEEL_FILENAME_CHECK=1 only for this specific install, then
restores the previous env value.  Both BNB install call sites (the
UNSLOTH_ROCM_TORCH_INSTALLED early-return path and the normal Windows
ROCm path) now use this helper.
This commit is contained in:
LeoBorcherding 2026-05-14 13:14:23 -05:00
commit a4483dff10
1 changed files with 32 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

52
studio/install_python_stack.py
View file

@ -339,6 +339,36 @@ def _detect_amd_gfx_codes() -> list[str]:
_rocm_windows_torch_installed: bool = False
def _install_bnb_windows_rocm() -> None:
"""Install the AMD Windows BNB prerelease wheel.
The continuous-release wheel is intentionally mismatched: the filename
encodes version 1.33.7.preview (parsed as 1.33.7rc0 by PEP 440) while the
wheel metadata reports 0.50.0.dev0. uv rejects this by default; we set
UV_SKIP_WHEEL_FILENAME_CHECK=1 only for this install and restore the env
afterwards.
"""
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is None:
return
_prev = os.environ.get("UV_SKIP_WHEEL_FILENAME_CHECK")
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = "1"
try:
pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
)
finally:
if _prev is None:
os.environ.pop("UV_SKIP_WHEEL_FILENAME_CHECK", None)
else:
os.environ["UV_SKIP_WHEEL_FILENAME_CHECK"] = _prev
def _ensure_rocm_torch() -> None:
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
@ -355,16 +385,7 @@ def _ensure_rocm_torch() -> None:
# setup.ps1 already installed ROCm torch, but we still need to install
# the AMD Windows BNB wheel here — the PyPI bitsandbytes wheel ships
# only CUDA DLLs and will fail to load on ROCm (no libbitsandbytes_rocm72.dll).
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is not None:
pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
)
_install_bnb_windows_rocm()
return
if IS_MACOS:
return
@ -413,16 +434,7 @@ def _ensure_rocm_torch() -> None:
# Always install AMD Windows bitsandbytes — the PyPI wheel ships only
# CUDA DLLs and will fail to load on ROCm. Install even when torch was
# already a ROCm build so that `studio update` repairs a broken bnb.
_bnb_win_url = _BNB_ROCM_PRERELEASE_URLS.get("win_amd64")
if _bnb_win_url is not None:
pip_install_try(
"bitsandbytes (AMD Windows, pre-release main)",
"--force-reinstall",
"--no-cache-dir",
"--no-deps",
_bnb_win_url,
constrain = False,
)
_install_bnb_windows_rocm()
_rocm_windows_torch_installed = True
return