Pin bitsandbytes to continuous-release_main on ROCm (4-bit decode fix) #4954
+137
−132
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -43,6 +43,41 @@ | |
| } | ||
| _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" | ||
|
|
||
| # bitsandbytes continuous-release_main wheels. bnb <= 0.49.2 ships with a | ||
| # broken 4-bit GEMV kernel on every ROCm target: CDNA (gfx90a / gfx942 / | ||
| # gfx950) via a broken blocksize=32/64 warp64 kernel, and RDNA3/3.5 | ||
| # (gfx1100-1103 / gfx1150-1152) via a compile-time warp-size dispatch bug. | ||
| # At decode shape (seq_len=1) the GEMV returns NaN, so autoregressive | ||
| # generation is broken even though training passes. bnb commit 713a3b8 | ||
| # ("[ROCm] Enable blocksize 32 4-bit quantization and GEMV kernels on AMD | ||
| # CDNA", PR #1887, merged 2026-03-09) fixes both bugs but has not shipped | ||
| # to PyPI yet. Pin to the continuous-release_main wheels so ROCm users | ||
| # get correct 4-bit decode on install. Drop the pin once bnb cuts a | ||
| # 0.50+ tag on PyPI. | ||
| _BNB_ROCM_PRERELEASE_URLS: dict[str, str] = { | ||
| "x86_64": ( | ||
| "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" | ||
| "download/continuous-release_main/" | ||
| "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_x86_64.whl" | ||
| ), | ||
| "aarch64": ( | ||
| "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/" | ||
| "download/continuous-release_main/" | ||
| "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" | ||
|
Comment on lines
+53
to
+58
Contributor
There was a problem hiding this comment. Similar to the shell script, these URLs hardcode the |
||
| ), | ||
| } | ||
|
Comment on lines
+49
to
+60
Contributor
There was a problem hiding this comment. To improve maintainability, you could define the base URL and wheel filename as constants. This avoids repeating the long URL string and makes future version updates easier. _BNB_BASE_URL = (
"https://github.com/bitsandbytes-foundation/bitsandbytes/releases/"
"download/continuous-release_main"
)
_BNB_WHEEL_TEMPLATE = "bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_{arch}.whl"
_BNB_ROCM_PRERELEASE_URLS: dict[str, str] = {
"x86_64": f"{_BNB_BASE_URL}/{_BNB_WHEEL_TEMPLATE.format(arch='x86_64')}",
"aarch64": f"{_BNB_BASE_URL}/{_BNB_WHEEL_TEMPLATE.format(arch='aarch64')}",
} |
||
| _BNB_ROCM_PYPI_FALLBACK = "bitsandbytes>=0.49.1" | ||
|
|
||
|
|
||
| def _bnb_rocm_prerelease_url() -> str | None: | ||
| """Return the continuous-release_main bnb wheel URL for the current | ||
| architecture, or None when no pre-release wheel is available. | ||
| Normalises amd64/arm64 aliases to x86_64/aarch64. | ||
| """ | ||
| arch = platform.machine().lower() | ||
| arch = {"amd64": "x86_64", "arm64": "aarch64"}.get(arch, arch) | ||
| return _BNB_ROCM_PRERELEASE_URLS.get(arch) | ||
|
|
||
|
|
||
| def _detect_rocm_version() -> tuple[int, int] | None: | ||
| """Return (major, minor) of the installed ROCm stack, or None.""" | ||
|
|
@@ -290,15 +325,42 @@ def _ensure_rocm_torch() -> None: | |
| # runtime is older than any published torch wheel. Uses | ||
| # --force-reinstall so an existing CPU/CUDA bitsandbytes is replaced | ||
| # by the AMD build during upgrades. | ||
| # | ||
| # Prefer the continuous-release_main wheel, which contains the CDNA | ||
| # and RDNA 4-bit GEMV fix (bnb PR #1887, merged 2026-03-09, post-0.49.2). | ||
| # Without that fix, autoregressive decode on every ROCm GPU produces | ||
| # NaN at seq_len=1 and generation returns gibberish or crashes in | ||
| # torch.multinomial. Falls back to PyPI >=0.49.1 on unknown architectures | ||
| # or when the pre-release URL is unreachable (offline / firewalled hosts). | ||
| if rocm_torch_ready: | ||
| _bnb_url = _bnb_rocm_prerelease_url() | ||
| _bnb_installed = False | ||
| if _bnb_url is not None: | ||
| _bnb_installed = pip_install_try( | ||
| "bitsandbytes (AMD, pre-release main)", | ||
| "--force-reinstall", | ||
| "--no-cache-dir", | ||
| "--no-deps", | ||
| _bnb_url, | ||
| constrain = False, | ||
| ) | ||
| if not _bnb_installed: | ||
| print( | ||
| _red( | ||
| " bnb pre-release wheel unreachable; falling back " | ||
| "to PyPI (4-bit decode will be broken on ROCm -- " | ||
| "use 16-bit instead)" | ||
| ) | ||
| ) | ||
| if not _bnb_installed: | ||
| pip_install( | ||
| "bitsandbytes (AMD)", | ||
| "--force-reinstall", | ||
| "--no-cache-dir", | ||
| "--no-deps", | ||
| _BNB_ROCM_PYPI_FALLBACK, | ||
| constrain = False, | ||
| ) | ||
|
|
||
|
|
||
| def _infer_no_torch() -> bool: | ||
|
|
@@ -593,6 +655,38 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: | |
| return cmd | ||
|
|
||
|
|
||
| def pip_install_try( | ||
| label: str, | ||
| *args: str, | ||
| constrain: bool = True, | ||
| ) -> bool: | ||
| """Try to install with pip/uv. Returns True on success, False on failure | ||
| (without raising or exiting). For optional install attempts with a | ||
| follow-up fallback, such as the bnb ROCm pre-release wheel. | ||
| """ | ||
| constraint_args: list[str] = [] | ||
| if constrain and CONSTRAINTS.is_file(): | ||
| constraint_args = ["-c", str(CONSTRAINTS)] | ||
|
|
||
| if USE_UV: | ||
| cmd = _build_uv_cmd(args) + constraint_args | ||
| else: | ||
| cmd = _build_pip_cmd(args) + constraint_args | ||
|
|
||
| if VERBOSE: | ||
| _step(_LABEL, f"{label}...", _dim) | ||
| result = subprocess.run( | ||
| cmd, | ||
| stdout = subprocess.PIPE, | ||
| stderr = subprocess.STDOUT, | ||
| ) | ||
| if result.returncode == 0: | ||
| return True | ||
| if VERBOSE and result.stdout: | ||
| print(result.stdout.decode(errors = "replace")) | ||
| return False | ||
|
Comment on lines
+638
to
+666
Contributor
There was a problem hiding this comment. The def pip_install_try(
label: str,
*args: str,
constrain: bool = True,
) -> bool:
"""Try to install with pip/uv. Returns True on success, False on failure
(without raising or exiting). For optional install attempts with a
follow-up fallback, such as the bnb ROCm pre-release wheel.
"""
constraint_args: list[str] = []
if constrain and CONSTRAINTS.is_file():
constraint_args = ["-c", str(CONSTRAINTS)]
if USE_UV:
cmd = _build_uv_cmd(args) + constraint_args
else:
cmd = _build_pip_cmd(args) + constraint_args
if VERBOSE:
_step(_LABEL, f"{label}...", _dim)
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if result.returncode == 0:
return True
if VERBOSE and result.stdout:
_safe_print(result.stdout.decode(errors = "replace"))
return False |
||
|
|
||
|
|
||
| def pip_install( | ||
| label: str, | ||
| *args: str, | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The wheel URLs contain a hardcoded version number (
1.33.7.preview). If the upstreambitsandbytesproject updates the version number in theircontinuous-release_maintag, these URLs will return a 404 error, breaking the installer for ROCm users. Consider if there is a way to resolve the latest asset URL dynamically or ensure the version number remains stable.