From f5278deef700fc88afad3ac796c3324ea7eca9e7 Mon Sep 17 00:00:00 2001
From: LeoBorcherding <borchborchmail@gmail.com>
Date: Thu, 14 May 2026 13:07:14 -0500
Subject: [PATCH] fix: BNB AMD wheel skipped + torch.compile segfault on
 Windows ROCm

install_python_stack.py: the UNSLOTH_ROCM_TORCH_INSTALLED=1 early-return
path (set by setup.ps1 when it installed torch itself) returned before
ever reaching the AMD BNB prerelease wheel install.  The PyPI
bitsandbytes==0.49.x ships only CUDA DLLs, so loading it on ROCm fails
with "libbitsandbytes_rocm72.dll not found".  Now installs the AMD
Windows BNB wheel before returning on that path too.

worker.py: torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer,
0xC0000005) when torch.compile's JitDecomp system dispatches it during
the first forward pass.  Detect Windows ROCm via torch.version.hip
(already in sys.modules from section 1e) and set TORCHDYNAMO_DISABLE=1
to bypass the broken kernel dispatch.
---
 studio/backend/core/training/worker.py | 17 +++++++++++++++++
 studio/install_python_stack.py         | 13 +++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 08bedae6b0..f1a1d65c44 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1178,6 +1178,23 @@ def run_training_process(
         except Exception:
             pass
 
+    # ── 1f. Disable torch.compile on Windows ROCm ──
+    # torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer, 0xC0000005).
+    # The crash is triggered via torch.compile's JitDecomp dispatch during the
+    # first forward pass (stack: _grouped_mm ← JitDecompRegisterer ← Python).
+    # Disabling dynamo entirely avoids the kernel dispatch.  torch is already
+    # in sys.modules from section 1e's `import torch.distributed`.
+    if sys.platform == "win32" and "TORCHDYNAMO_DISABLE" not in os.environ:
+        _torch_for_rocm_check = sys.modules.get("torch")
+        if _torch_for_rocm_check is not None and getattr(
+            getattr(_torch_for_rocm_check, "version", None), "hip", None
+        ):
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.info(
+                "Windows ROCm detected — torch.compile disabled "
+                "(_grouped_mm kernel crashes on gfx1200 with 0xC0000005)"
+            )
+
     # ── 2. Now import ML libraries (fresh in this clean process) ──
     try:
         _send_status(event_queue, "Importing Unsloth...")
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index ec2c4e14fe..f5b7b3c1c9 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -352,6 +352,19 @@ def _ensure_rocm_torch() -> None:
     # setup.ps1 sets this when it already installed AMD wheels; skip the probe.
     if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1":
         _rocm_windows_torch_installed = True
+        # 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,
+            )
         return
     if IS_MACOS:
         return