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,