From 65b40285608ee14e93438949e8f3541580aee234 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 10 Apr 2026 06:25:39 -0700 Subject: [PATCH] Pin bitsandbytes to continuous-release_main on ROCm (4-bit decode fix) (#4954) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Pin bitsandbytes to continuous-release_main on ROCm for 4-bit decode fix bitsandbytes 0.49.2 on PyPI ships with a broken 4-bit GEMV kernel on every ROCm target: - CDNA (gfx90a / gfx942 / gfx950 = MI210 / MI300X / MI350) via a broken blocksize=32/64 warp64 GEMV kernel whose tests were explicitly skipped with ROCM_WARP_SIZE_64 guards because the code was known broken. - RDNA3 / RDNA3.5 (gfx1100-1103 / gfx1150-1152) via a compile-time BNB_WARP_SIZE macro in the host-side dispatch that resolves to 64 when the multi-arch wheel is compiled with CDNA as the primary target, so num_blocks is wrong on RDNA and half the GEMV output is never written. At decode shape (1, 1, hidden) both bugs produce NaN. Training is unaffected because training shapes are (batch, seq_len > 1, hidden) and never touch the GEMV path. The crash during autoregressive inference surfaces as _assert_async_cuda_kernel in torch.multinomial which on HIP becomes a hard HSA_STATUS_ERROR_EXCEPTION instead of a clean Python error. Both bugs are fixed by bitsandbytes commit 713a3b8 ("[ROCm] Enable blocksize 32 4-bit quantization and GEMV kernels on AMD CDNA", PR #1887, merged 2026-03-09) which replaces BNB_WARP_SIZE with a runtime hipDeviceGetAttribute query and ships a working CDNA warp64 kernel. That commit has not shipped to PyPI yet, but continuous-release_main wheels are published on every push to bnb main via GitHub Releases. Point the ROCm install path at the continuous-release_main x86_64 and aarch64 wheels and fall back to PyPI >=0.49.1 when the pre-release is unreachable (offline installs, firewalled hosts, or architectures not covered by the pre-release wheels). Drop the pin once bnb cuts a 0.50+ tag on PyPI. Verified on MI300X (gfx942, ROCm 7.2, torch 2.10.0+rocm7.1): direct bnb GEMV shape test now returns 0.0078 max abs error at seq_len=1 (no NaN) vs NaN on 0.49.2, and full Unsloth + for_inference + 4-bit sampling generation works end-to-end. NVIDIA / CPU / Mac / Windows paths are unaffected -- the helper is gated on the ROCm torch index and platform.machine() respectively. * Drop Studio ROCm 16-bit fallback now that bnb 0.50+ fixes 4-bit decode The 16-bit fallback in studio/backend/core/inference/inference.py was added as a workaround for a bug that this PR already fixes at the install layer: bitsandbytes <= 0.49.2 has a broken 4-bit GEMV kernel on every ROCm target, which NaNs at decode shape (seq_len=1) and crashes autoregressive inference. bnb PR #1887 (commit 713a3b8, in 0.50.0.dev0+, pinned by install.sh / install_python_stack.py in this PR) restores correct 4-bit decode on MI300X and verified working end-to-end with full Unsloth + for_inference + sampling. Revert the dual code path so ROCm and NVIDIA both go through the normal FastLanguageModel.from_pretrained + for_inference flow: - Remove the conditional `from unsloth import` that skipped the import on ROCm. The monkey-patches it was trying to avoid were never the cause of the crash; bnb 4-bit GEMV was. - Remove the `if _hw_module.IS_ROCM:` branch in load_model that loaded with plain transformers + PEFT + bfloat16, and the `_resolve_fp16_base` helper it relied on. - Remove the `get_chat_template is not None` fallback in _load_chat_template_info -- get_chat_template is now always imported. - Refactor the audio/vision ROCm guard to check _hw_module.IS_ROCM directly instead of the removed _IS_ROCM_ENV global. Audio and vision on ROCm still need separate validation (FastVisionModel and the CSM audio codecs were never tested on HIP) so the guard stays for now. Add _bnb_rocm_4bit_ok() as a runtime safety net for users who install from this PR before the install.sh bnb pin kicks in, or whose installer fell back to the PyPI pin because the continuous- release wheel was unreachable. When the installed bnb is < 0.50 on ROCm, force load_in_4bit=False and strip any -unsloth-bnb-4bit / -bnb-4bit suffix from the model path so a pre-quantized repo resolves to its FP16 sibling instead of pulling bnb back in via the repo's quantization_config. LoRA adapters whose base is a pre-quantized repo on old bnb will still fail inside Unsloth's loader -- the only real fix there is `unsloth studio update`. Verified on MI300X (gfx942, ROCm 7.2, torch 2.10.0+rocm7.1): - HAPPY path (bnb 0.50.0.dev0, load_in_4bit=True, pre-quantized repo): loads in 4-bit via the fixed GEMV, generation returns "Paris." for greedy and sampling. - SAFETY-NET path (simulated old bnb, suffix-stripped to the FP16 sibling, load_in_4bit=False): loads in bf16, generation returns "Paris." for greedy and sampling. Net diff is ~45 lines smaller than the pre-revert state because the entire plain-transformers 16-bit branch is gone. * Cache _bnb_rocm_4bit_ok() with functools.cache load_model() can be called many times in a single session but the bnb version and hardware state cannot change at runtime, so memoise the check. First call is ~1.9 ms (dominated by the lazy `import bitsandbytes` inside the try block), subsequent calls drop to sub-microsecond dict lookups. Zero behavioral change. * Shorten verbose bnb/ROCm comments Comment-only cleanup across install.sh, studio/install_python_stack.py, and studio/backend/core/inference/inference.py. No behavioral change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove _bnb_rocm_4bit_ok safety net from inference.py Studio's ROCm support is brand new (PR #4720, merged today) and every fresh install pulls the bnb continuous-release_main wheel via install.sh / install_python_stack.py in this same PR. There are no existing ROCm Studio installs carrying bnb < 0.50, so the defensive version-check fallback is guarding against a scenario that cannot actually occur. Delete the helper, the functools import, and the safety-net block -- inference.py now calls FastLanguageModel.from_pretrained directly with no ROCm branching. * Drop audio/vision ROCm guard in inference.py — verified unblocked by bnb fix Vision inference was blocked by the same bnb 4-bit GEMV bug that affected text inference (vision models use bnb 4-bit for the LM backbone). With bnb 0.50+ pinned in install.sh / install_python_stack.py, vision works end-to-end on MI300X: Llama-3.2-11B-Vision-Instruct-unsloth-bnb-4bit loaded in 4-bit via FastVisionModel + for_inference returns a correct answer to a multimodal prompt. Audio (CSM) was never actually blocked by HIP — on this hardware CSM loads and runs its backbone forward pass fine with bnb 0.50, then fails during generate() with a transformers-level kwarg validation mismatch in generation_csm.py (`backbone_last_hidden_state` rejected). That's a pre-existing transformers/CSM integration bug that reproduces identically on NVIDIA, so the ROCm-gated guard was never actually protecting users from anything HIP-specific. Remove the combined audio/vision guard and the now-unused _hw_module import. Also restore the one-word "Can be" in an inline comment that drifted during the earlier comment-shortening pass, so the inference.py delta vs pre-#4720 is exactly the max_seq_length<=0 crash fix and nothing else. * Shorten max_seq_length=0 guard comment to one line --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.sh | 36 +++++- studio/backend/core/inference/inference.py | 132 +++------------------ studio/install_python_stack.py | 101 +++++++++++++--- 3 files changed, 137 insertions(+), 132 deletions(-) diff --git a/install.sh b/install.sh index 2e347a8bc5..cc8721710e 100755 --- a/install.sh +++ b/install.sh @@ -94,6 +94,36 @@ run_install_cmd() { return $_rc } +# 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_bnb_rocm() { + _label="$1" + _venv_py="$2" + case "$_ARCH" in + x86_64|amd64) + _bnb_whl_url="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|arm64) + _bnb_whl_url="https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-1.33.7.preview-py3-none-manylinux_2_24_aarch64.whl" + ;; + *) + _bnb_whl_url="" + ;; + esac + if [ -n "$_bnb_whl_url" ]; then + substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..." + if run_install_cmd "$_label (pre-release)" uv pip install --python "$_venv_py" \ + --force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then + return 0 + fi + substep "[WARN] bnb pre-release unreachable; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN" + fi + run_install_cmd "$_label (pypi fallback)" uv pip install --python "$_venv_py" \ + --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1" +} + if [ "$_next_is_package" = true ]; then echo "❌ ERROR: --package requires an argument." >&2 exit 1 @@ -1296,8 +1326,7 @@ if [ "$_MIGRATED" = true ]; then if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) - substep "installing bitsandbytes for AMD ROCm..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1" + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" # Repair ROCm torch if overwritten during migrated install _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then @@ -1437,8 +1466,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) - substep "installing bitsandbytes for AMD ROCm..." - run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" --force-reinstall --no-cache-dir --no-deps "bitsandbytes>=0.49.1" + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" ;; esac fi diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index d7bc3ad4c7..4c140013a0 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -5,23 +5,8 @@ Core inference backend - streamlined """ -# On AMD ROCm, Unsloth's global monkey-patching of transformers model classes -# (LlamaRotaryEmbedding, attention modules, etc.) causes HIP kernel crashes -# (_assert_async_cuda_kernel -> HSA_STATUS_ERROR_EXCEPTION) during inference. -# Training works because it uses different code paths, but generation triggers -# the incompatible patched kernels. Skip the Unsloth import entirely on ROCm -# so transformers classes stay unmodified; the GGUF inference path (llama-server) -# is unaffected since it never imports these Python model classes. -_IS_ROCM_ENV = getattr(__import__("torch").version, "hip", None) is not None - -if _IS_ROCM_ENV: - FastLanguageModel = None # Loaded on-demand only on NVIDIA - FastVisionModel = None - get_chat_template = None -else: - from unsloth import FastLanguageModel, FastVisionModel - from unsloth.chat_templates import get_chat_template - +from unsloth import FastLanguageModel, FastVisionModel +from unsloth.chat_templates import get_chat_template from transformers import TextStreamer from peft import PeftModel, PeftModelForCausalLM @@ -41,7 +26,6 @@ from utils.hardware import ( raise_if_offloaded, get_visible_gpu_count, ) -from utils.hardware import hardware as _hw_module from core.inference.audio_codecs import AudioCodecManager from io import StringIO import structlog @@ -269,12 +253,7 @@ class InferenceBackend: """ Load any model: base, LoRA adapter, text, or vision. """ - # max_seq_length=0 means "model default" for the GGUF/llama.cpp path, - # but Unsloth's FastLanguageModel.from_pretrained treats 0 literally -- - # setting the model's context to 0 tokens, which triggers an assertion - # crash during generation (especially on ROCm/HIP where the async - # assert kernel raises a hardware exception instead of a Python error). - # Fall back to 2048 for the Unsloth/transformers path. + # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. if max_seq_length <= 0: max_seq_length = 2048 @@ -311,12 +290,6 @@ class InferenceBackend: } # ── Audio model loading path ────────────────────────── - if (config.is_audio or config.is_vision) and _IS_ROCM_ENV: - raise RuntimeError( - f"Audio and vision model inference via Unsloth is not " - f"yet supported on AMD ROCm. Use GGUF inference instead." - ) - if config.is_audio: audio_type = config.audio_type adapter_info = " (LoRA adapter)" if config.is_lora else "" @@ -547,84 +520,18 @@ class InferenceBackend: else: # Text model (or text LoRA adapter) - if _hw_module.IS_ROCM: - # On AMD ROCm two issues prevent the normal Unsloth path: - # 1. Unsloth's patched kernels (RoPE, attention) crash on - # HIP (_assert_async_cuda_kernel -> HSA_STATUS_ERROR). - # 2. bitsandbytes 4-bit matmul kernels trigger the same - # HIP assertion on MI300X (CDNA3 / gfx942). - # Fall back to plain transformers + PEFT in 16-bit, which - # works reliably. AMD GPUs typically have large VRAM so - # 16-bit is practical; GGUF inference remains the - # recommended path for memory-constrained setups. - logger.info( - "ROCm detected -- loading in 16-bit with plain " - "transformers (bitsandbytes 4-bit and Unsloth kernels " - "are not yet compatible with HIP)" - ) - from transformers import AutoModelForCausalLM, AutoTokenizer + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = config.path, # Can be base model OR LoRA adapter path + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = load_in_4bit, + device_map = device_map, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, + ) - _load_kwargs = dict( - dtype = dtype or torch.bfloat16, - device_map = device_map, - token = hf_token if hf_token and hf_token.strip() else None, - trust_remote_code = trust_remote_code, - ) - - # Skip 4-bit on ROCm: bnb matmul kernels crash on HIP. - # Also resolve pre-quantized Unsloth model names (e.g. - # "unsloth/xxx-bnb-4bit") to their FP16 originals since - # loading a pre-quantized repo still triggers bnb codepaths. - def _resolve_fp16_base(name: str) -> str: - if not name: - return name - # Strip Unsloth quantization suffixes to get the FP16 model: - # "unsloth/Foo-unsloth-bnb-4bit" -> "unsloth/Foo" - # "unsloth/Foo-bnb-4bit" -> "unsloth/Foo" - # Order matters: try longer suffix first. - for suffix in ("-unsloth-bnb-4bit", "-bnb-4bit"): - if name.lower().endswith(suffix): - resolved = name[: -len(suffix)] - logger.info( - "Resolved pre-quantized base '%s' -> '%s' for ROCm 16-bit inference", - name, - resolved, - ) - return resolved - return name - - if config.is_lora and config.base_model: - # Load base model then apply adapter - _base = _resolve_fp16_base(config.base_model) - model = AutoModelForCausalLM.from_pretrained( - _base, - **_load_kwargs, - ) - from peft import PeftModel - - model = PeftModel.from_pretrained(model, config.path) - tokenizer = AutoTokenizer.from_pretrained(config.path) - else: - _path = _resolve_fp16_base(config.path) - model = AutoModelForCausalLM.from_pretrained( - _path, - **_load_kwargs, - ) - tokenizer = AutoTokenizer.from_pretrained(config.path) - model.eval() - else: - model, tokenizer = FastLanguageModel.from_pretrained( - model_name = config.path, # Can be base model OR LoRA adapter path - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - device_map = device_map, - token = hf_token if hf_token and hf_token.strip() else None, - trust_remote_code = trust_remote_code, - ) - - # Apply inference optimization - FastLanguageModel.for_inference(model) + # Apply inference optimization + FastLanguageModel.for_inference(model) self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer @@ -1047,13 +954,10 @@ class InferenceBackend: ) # This modifies the tokenizer with the correct template - if get_chat_template is not None: - tokenizer = get_chat_template( - tokenizer, - chat_template = template_name, - ) - else: - logger.info("Skipping Unsloth chat template (ROCm fallback)") + tokenizer = get_chat_template( + tokenizer, + chat_template = template_name, + ) else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index af77d7a811..a046f8f892 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -43,6 +43,32 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { } _PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" +# 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_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" + ), +} +_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. + """ + 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.""" @@ -284,21 +310,37 @@ def _ensure_rocm_torch() -> None: ) rocm_torch_ready = True - # Install bitsandbytes only when the venv has a ROCm-compatible torch - # (either already present or just installed). Avoids leaving an AMD - # bitsandbytes on top of a CPU/CUDA torch on hosts where the ROCm - # 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. + # Install bitsandbytes only when torch links against ROCm. Prefers the + # continuous-release_main wheel (bnb PR #1887 4-bit GEMV fix) and falls + # back to PyPI when the pre-release URL is unreachable. if rocm_torch_ready: - pip_install( - "bitsandbytes (AMD)", - "--force-reinstall", - "--no-cache-dir", - "--no-deps", - "bitsandbytes>=0.49.1", - constrain = False, - ) + _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 unreachable; falling back to PyPI " + "(4-bit decode will be broken on ROCm)" + ) + ) + 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 +635,37 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: return cmd +def pip_install_try( + label: str, + *args: str, + constrain: bool = True, +) -> bool: + """Like pip_install but returns False on failure instead of exiting. + For optional installs with a follow-up fallback. + """ + 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 + + def pip_install( label: str, *args: str,