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