Add AMD ROCm/HIP support across installer and hardware detection

studio/backend/core/inference/inference.py

@@ -5,8 +5,25 @@
 Core inference backend - streamlined
 """
 
-from unsloth import FastLanguageModel, FastVisionModel
-from unsloth.chat_templates import get_chat_template
+import os as _os
+
+# 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 transformers import TextStreamer
 from peft import PeftModel, PeftModelForCausalLM
 
@@ -26,6 +43,7 @@
     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
@@ -253,6 +271,15 @@ def load_model(
         """
         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.
+        if max_seq_length <= 0:
+            max_seq_length = 2048
+
         try:
             model_name = config.identifier
 
@@ -516,18 +543,84 @@ def load_model(
 
             else:
                 # Text model (or text LoRA adapter)
-                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,
-                )
+                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
 
-                # Apply inference optimization
-                FastLanguageModel.for_inference(model)
+                    _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)
 
                 self.models[model_name]["model"] = model
                 self.models[model_name]["tokenizer"] = tokenizer
@@ -950,10 +1043,13 @@ def _generate_chat_response_inner(
                 )
 
                 # This modifies the tokenizer with the correct template
-                tokenizer = get_chat_template(
-                    tokenizer,
-                    chat_template = template_name,
-                )
+                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)")
             else:
                 logger.info(
                     f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"

studio/backend/core/training/worker.py

@@ -86,6 +86,7 @@ def _probe_causal_conv1d_env() -> dict[str, str] | None:
                     "'python_tag': f'cp{sys.version_info.major}{sys.version_info.minor}', "
                     "'torch_mm': torch_mm, "
                     "'cuda_major': str(int(str(torch.version.cuda).split('.', 1)[0])) if torch.version.cuda else '', "
+                    "'hip_version': str(torch.version.hip) if getattr(torch.version, 'hip', None) else '', "
                     "'cxx11abi': str(torch._C._GLIBCXX_USE_CXX11_ABI).upper()"
                     "}))"
                 ),
@@ -237,28 +238,111 @@ def _install_package_wheel_first(
         else:
             logger.info("No published %s wheel found: %s", display_name, wheel_url)
 
-    _send_status(event_queue, f"Installing {display_name} from PyPI...")
-    pypi_cmd = [
-        sys.executable,
-        "-m",
-        "pip",
-        "install",
-        "--no-build-isolation",
-        "--no-deps",
-        "--no-cache-dir",
-        f"{pypi_name}=={pypi_version}",
-    ]
-    result = _sp.run(
-        pypi_cmd,
-        stdout = _sp.PIPE,
-        stderr = _sp.STDOUT,
-        text = True,
-    )
+    is_hip = env and env.get("hip_version")
+    if is_hip and not shutil.which("hipcc"):
+        logger.error(
+            "%s requires hipcc for source compilation on ROCm. "
+            "Install the ROCm HIP SDK: https://rocm.docs.amd.com",
+            display_name,
+        )
+        _send_status(
+            event_queue,
+            f"{display_name}: hipcc not found (ROCm HIP SDK required)",
+        )
+        return
+
+    if is_hip:
+        _send_status(
+            event_queue,
+            f"Compiling {display_name} from source for ROCm "
+            "(this may take several minutes)...",
+        )
+    else:
+        _send_status(event_queue, f"Installing {display_name} from PyPI...")
+
+    # Prefer uv for faster dependency resolution when available
+    if shutil.which("uv"):
+        pypi_cmd = [
+            "uv",
+            "pip",
+            "install",
+            "--python",
+            sys.executable,
+            "--no-build-isolation",
+            "--no-deps",
+        ]
+        # Avoid stale cache artifacts from partial HIP source builds
+        if is_hip:
+            pypi_cmd.append("--no-cache")
+        pypi_cmd.append(f"{pypi_name}=={pypi_version}")
+    else:
+        pypi_cmd = [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "--no-build-isolation",
+            "--no-deps",
+            "--no-cache-dir",
+            f"{pypi_name}=={pypi_version}",
+        ]
+
+    # Source compilation on ROCm can take 10-30 minutes; use a generous
+    # timeout. Non-HIP installs preserve the pre-existing "no timeout"
+    # behaviour so unrelated slow installs (e.g. causal-conv1d source
+    # build on Linux aarch64 or unsupported torch/CUDA combinations)
+    # are not aborted at 5 minutes by this PR.
+    _run_kwargs: dict[str, Any] = {
+        "stdout": _sp.PIPE,
+        "stderr": _sp.STDOUT,
+        "text": True,
+    }
+    if is_hip:
+        _run_kwargs["timeout"] = 1800
+
+    try:
+        result = _sp.run(pypi_cmd, **_run_kwargs)
+    except _sp.TimeoutExpired:
+        logger.error(
+            "%s installation timed out after %ds",
+            display_name,
+            _run_kwargs.get("timeout"),
+        )
+        _send_status(
+            event_queue,
+            f"{display_name} installation timed out after "
+            f"{_run_kwargs.get('timeout')}s",
+        )
+        return
+
     if result.returncode != 0:
-        logger.error("Failed to install %s from PyPI:\n%s", display_name, result.stdout)
+        if is_hip:
+            # Surface a clear error for ROCm source build failures
+            error_lines = (result.stdout or "").strip().splitlines()
+            snippet = "\n".join(error_lines[-5:]) if error_lines else "(no output)"
+            logger.error(
+                "Failed to compile %s for ROCm:\n%s",
+                display_name,
+                result.stdout,
+            )
+            _send_status(
+                event_queue,
+                f"Failed to compile {display_name} for ROCm. "
+                "Check that hipcc and ROCm development headers are installed.\n"
+                f"{snippet}",
+            )
+        else:
+            logger.error(
+                "Failed to install %s from PyPI:\n%s",
+                display_name,
+                result.stdout,
+            )
         return
 
-    logger.info("Installed %s from PyPI", display_name)
+    if is_hip:
+        logger.info("Compiled and installed %s from source for ROCm", display_name)
+    else:
+        logger.info("Installed %s from PyPI", display_name)
 
 
 def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:

studio/backend/main.py

@@ -237,6 +237,7 @@ async def get_system_info():
     import platform
     import psutil
     from utils.hardware import get_device
+    from utils.hardware.hardware import _backend_label
 
     visibility_info = get_backend_visible_gpu_info()
     gpu_info = {
@@ -250,7 +251,10 @@ async def get_system_info():
     return {
         "platform": platform.platform(),
         "python_version": platform.python_version(),
-        "device_backend": get_device().value,
+        # Use the centralized _backend_label helper so the /api/system
+        # endpoint reports "rocm" on AMD hosts instead of "cuda", matching
+        # the /api/hardware and /api/gpu-visibility endpoints.
+        "device_backend": _backend_label(get_device()),
         "cpu_count": psutil.cpu_count(),
         "memory": {
             "total_gb": round(memory.total / 1e9, 2),

studio/backend/tests/test_utils.py

@@ -191,8 +191,14 @@ def test_has_backend_key(self):
         assert "backend" in get_gpu_memory_info()
 
     def test_backend_matches_device(self):
+        # The backend field uses _backend_label, which swaps "cuda" for
+        # "rocm" when running on an AMD host (IS_ROCM=True) so the UI
+        # can render the correct label. On CUDA / XPU / MLX / CPU hosts
+        # it is equivalent to `get_device().value`.
+        from utils.hardware.hardware import _backend_label
+
         result = get_gpu_memory_info()
-        assert result["backend"] == get_device().value
+        assert result["backend"] == _backend_label(get_device())
 
     # --- When a GPU IS available ---
 

studio/backend/utils/hardware/__init__.py

@@ -5,6 +5,7 @@
 Hardware detection and GPU utilities
 """
 
+from . import hardware as _hardware
 from .hardware import (
     DeviceType,
     DEVICE,
@@ -49,6 +50,7 @@
     "DeviceType",
     "DEVICE",
     "CHAT_ONLY",
+    "IS_ROCM",
     "detect_hardware",
     "get_device",
     "is_apple_silicon",
@@ -81,3 +83,11 @@
     "extract_arch_config",
     "estimate_training_vram",
 ]
+
+
+def __getattr__(name: str):
+    """Resolve IS_ROCM at access time so callers always see the live value
+    after detect_hardware() runs (it flips the flag in hardware.py)."""
+    if name == "IS_ROCM":
+        return getattr(_hardware, "IS_ROCM")
+    raise AttributeError(name)