fix(studio/rocm): worker.py parity + don't roll back ROCm torch on bnb failure

Addresses findings from a 10x reviewer pass on the prior fix commit:

1. studio/backend/core/training/worker.py (parity with main.py):
   - Gate the torchao stub block on torch.version.hip / 'rocm' in
     torch.__version__ instead of HIP_PATH / ROCM_PATH env-var presence.
     Same root cause as main.py: HIP SDK env vars stick around on CUDA hosts.
   - Add module-level Windows ROCm DLL registration block. Worker subprocesses
     inherit env vars but not the parent's add_dll_directory handles, so the
     first `import torch` in the worker could fail to find amdhip64.dll when
     HIP_PATH\bin is not on PATH. Mirrors main.py setup. Handles retained at
     module scope via _ROCM_DLL_HANDLES.
   - Promote _WINDOWS_ROCM_GROUPED_MM_LIB to module scope with `global` in
     run_training_process so the torch.library.Library registration survives
     past function return / mid-run garbage collection.
   - Harden _torch_has_hip() to also accept 'rocm' in torch.__version__
     (AMD SDK / Radeon wheels may not set torch.version.hip).

2. studio/install_python_stack.py:
   - Don't roll back ROCm torch when bitsandbytes install fails. The prior
     commit gated _rocm_windows_torch_installed on _install_bnb_windows_rocm()
     returning True; if torch installed successfully but bnb failed, the flag
     stayed False and later install steps could overwrite ROCm torch with the
     generic CPU torch wheel. Set the flag after torch install; surface bnb
     failure as a separate warning instead.
   - _detect_windows_gfx_arch now probes in three tiers: UNSLOTH_ROCM_GFX_ARCH
     env-var override (matches the PowerShell installer), then hipinfo (PATH
     or HIP_PATH\bin), then amd-smi (`static --asic`, `list`). Without the
     amd-smi fallback, runtime-only Radeon installs without hipinfo on PATH
     made `studio update` return early and leave the venv on CPU torch.
   - Linux torch-already-rocm probe in _ensure_rocm_torch now matches the
     Windows probe shape: accepts torch.version.hip OR 'rocm' in
     torch.__version__ to cover AMD SDK / Radeon Linux wheels.

3. studio/backend/utils/hardware/hardware.py:
   - apply_gpu_ids() final-fallback torch probe accepts 'rocm' in
     torch.__version__ in addition to torch.version.hip, matching
     detect_hardware(). AMD SDK wheels could otherwise leak through with
     CUDA-only visibility masks on a spawned ROCm worker.

Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py
(no test changes needed; the probe shape that prints the hip version (or
'rocm' sentinel) preserves the existing non-empty-string contract).

Not addressed in this commit (deferred or out of scope):
- Tag drift / lemonade checksum (PR 5303 surface, not this PR).
- install.sh rocm7.2.1 URL: small fix, separate.
- install.ps1 / setup.ps1 'Radeon 8060S' marketing-name fallback table.
- Strix Halo + ROCm 7.1 routing asymmetry in Python update path.

studio/backend/core/training/worker.py

@@ -70,6 +70,43 @@ _TILELANG_INSTALL_TIMEOUT_S = 600
 _TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11")
 _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
 
+# Module-level handle so the torch.library.Library registration survives past
+# run_training_process() and is not garbage collected mid-run.
+_WINDOWS_ROCM_GROUPED_MM_LIB = None
+
+# Worker subprocesses inherit the parent env but not the parent's
+# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL
+# setup at module load so the first `import torch` can find amdhip64.dll even
+# when HIP_PATH\bin is not on the system PATH. Handles retained at module
+# scope so they are not garbage collected.
+_ROCM_DLL_HANDLES: list = []
+if sys.platform == "win32":
+    def _add_rocm_dll_dirs_worker() -> None:
+        _candidates: list[str] = []
+        for _var in ("HIP_PATH", "ROCM_PATH"):
+            _val = os.environ.get(_var)
+            if _val:
+                _candidates.append(os.path.join(_val, "bin"))
+        _default_root = os.path.join(
+            os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
+        )
+        try:
+            if os.path.isdir(_default_root):
+                for _ver in sorted(os.listdir(_default_root), reverse = True):
+                    _bin = os.path.join(_default_root, _ver, "bin")
+                    if os.path.isdir(_bin):
+                        _candidates.append(_bin)
+        except OSError:
+            pass
+        for _d in _candidates:
+            if os.path.isdir(_d):
+                try:
+                    _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
+                except (OSError, AttributeError):
+                    pass
+    _add_rocm_dll_dirs_worker()
+    del _add_rocm_dll_dirs_worker
+
 
 def _model_wants_causal_conv1d(model_name: str) -> bool:
     name = model_name.lower()
@@ -607,11 +644,18 @@ def _tilelang_importable() -> bool:
 
 
 def _torch_has_hip() -> bool:
-    """True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm."""
+    """True iff torch is a ROCm build.
+
+    `torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon
+    wheels can leave it unset but still encode "rocm" in `torch.__version__`.
+    """
     try:
         import torch as _torch
 
-        return getattr(_torch.version, "hip", None) is not None
+        return bool(
+            getattr(_torch.version, "hip", None)
+            or "rocm" in getattr(_torch, "__version__", "").lower()
+        )
     except Exception:
         return False
 
@@ -1916,10 +1960,21 @@ def run_training_process(
 
     # Only stub torchao on Windows ROCm hosts -- on Windows CUDA (NVIDIA) torchao
     # is real and shadowing it breaks torchao-based quantization paths.
-    # HIP_PATH / ROCM_PATH are set by the AMD HIP SDK installer on ROCm machines.
+    # Gate on the active torch runtime, not env-var presence -- HIP_PATH /
-    _is_win32_rocm = sys.platform == "win32" and bool(
+    # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
-        os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")
+    # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
-    )
+    # but still encode "rocm" in torch.__version__, so accept either.
+    _is_win32_rocm = False
+    if sys.platform == "win32":
+        try:
+            import torch as _torch_probe
+            _is_win32_rocm = bool(
+                getattr(getattr(_torch_probe, "version", None), "hip", None)
+                or "rocm" in getattr(_torch_probe, "__version__", "").lower()
+            )
+            del _torch_probe
+        except Exception:
+            pass
     if _is_win32_rocm:
         # Seed torchao top-level + key submodules; the finder handles the rest.
         for _tao_name in (
@@ -1984,7 +2039,9 @@ def run_training_process(
     #   offs: optional group-split offsets (MoE-style variable-size batches)
     #
     # torch is already in sys.modules from section 1e's `import torch.distributed`.
-    _WINDOWS_ROCM_GROUPED_MM_LIB = None  # kept alive to prevent GC of registration
+    # Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past
+    # function return / mid-run GC.
+    global _WINDOWS_ROCM_GROUPED_MM_LIB
     if sys.platform == "win32":
         _torch_for_rocm = sys.modules.get("torch")
         if _torch_for_rocm is not None and getattr(

studio/backend/utils/hardware/hardware.py

@@ -1694,14 +1694,19 @@ def apply_gpu_ids(gpu_ids) -> None:
     _is_rocm = IS_ROCM or _inherits_rocm_visibility
     if not _is_rocm:
         # torch.version.hip is a non-empty string on ROCm, None on CUDA.
+        # AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but
+        # still encode "rocm" in torch.__version__, matching detect_hardware().
         # Broad except: a probe failure must never crash a training worker.
         try:
             import torch as _torch
 
-            _is_rocm = getattr(_torch.version, "hip", None) is not None
+            _is_rocm = (
+                getattr(_torch.version, "hip", None) is not None
+                or "rocm" in getattr(_torch, "__version__", "").lower()
+            )
         except Exception as e:
             logger.debug(
-                "apply_gpu_ids: torch.version.hip probe skipped (%s: %s)",
+                "apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
                 type(e).__name__,
                 e,
             )

studio/install_python_stack.py

@@ -228,17 +228,23 @@ def _detect_rocm_version() -> tuple[int, int] | None:
 
 
 def _detect_windows_gfx_arch() -> str | None:
-    """Return the gcnArchName from hipinfo on Windows (e.g. 'gfx1200'), or None.
+    """Return the gcnArchName on Windows (e.g. 'gfx1200'), or None.
 
-    Resolves hipinfo via PATH first, then HIP_PATH\\bin and ROCM_PATH\\bin as
+    Probe order matches the PowerShell installer: env-var override first,
-    fallbacks -- the AMD HIP SDK installer sets these env vars but does not
+    then hipinfo (PATH or HIP_PATH / ROCM_PATH bin), then amd-smi. Without
-    always add the bin dir to the system PATH.
+    the amd-smi fallback, runtime-only AMD installs without hipinfo on PATH
+    return early and `studio update` cannot repair a CPU-only venv.
     """
     import re
 
+    # 1. Explicit override (matches PowerShell installer's env-var path).
+    _override = os.environ.get("UNSLOTH_ROCM_GFX_ARCH")
+    if _override and _override.strip():
+        return _override.strip().lower()
+
+    # 2. hipinfo via PATH, then HIP_PATH\bin / ROCM_PATH\bin.
     hipinfo = shutil.which("hipinfo")
     if not hipinfo:
-        # Fallback: AMD HIP SDK sets HIP_PATH / ROCM_PATH even when bin isn't on PATH
         for _env_var in ("HIP_PATH", "ROCM_PATH"):
             _root = os.environ.get(_env_var)
             if _root:
@@ -246,25 +252,43 @@ def _detect_windows_gfx_arch() -> str | None:
                 if os.path.isfile(_candidate):
                     hipinfo = _candidate
                     break
-    if not hipinfo:
+    if hipinfo:
-        return None
+        try:
-    try:
+            result = subprocess.run(
-        result = subprocess.run(
+                [hipinfo],
-            [hipinfo],
+                stdout = subprocess.PIPE,
-            stdout = subprocess.PIPE,
+                stderr = subprocess.DEVNULL,
-            stderr = subprocess.DEVNULL,
+                timeout = 10,
-            timeout = 10,
+            )
-        )
+            if result.returncode == 0:
-        if result.returncode != 0:
+                text = result.stdout.decode(errors = "replace")
-            return None
+                m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
-        text = result.stdout.decode(errors = "replace")
+                if m:
-        m = re.search(r"(?im)^\s*gcnArchName\s*:\s*(\S+)", text)
+                    # Lowercase -- hipinfo sometimes emits "Gfx1151".
-        # Lowercase the captured token -- some hipinfo builds emit "Gfx1151"
+                    return m.group(1).strip().lower()
-        # which would miss the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and
+        except Exception:
-        # silently fall back to CPU torch.
+            pass
-        return m.group(1).strip().lower() if m else None
+
-    except Exception:
+    # 3. amd-smi fallback -- runtime-only Radeon installs ship amd-smi but no hipinfo.
-        return None
+    amd_smi = shutil.which("amd-smi")
+    if amd_smi:
+        for _args in (("static", "--asic"), ("list",)):
+            try:
+                result = subprocess.run(
+                    [amd_smi, *_args],
+                    stdout = subprocess.PIPE,
+                    stderr = subprocess.DEVNULL,
+                    timeout = 10,
+                )
+                if result.returncode != 0:
+                    continue
+                text = result.stdout.decode(errors = "replace").lower()
+                m = re.search(r"\bgfx[1-9][0-9a-z]{2,3}\b", text)
+                if m:
+                    return m.group(0)
+            except Exception:
+                continue
+    return None
 
 
 def _windows_rocm_index_url(gfx_arch: str | None) -> str | None:
@@ -518,13 +542,19 @@ def _ensure_rocm_torch() -> None:
                 "torchaudio",
                 constrain = False,
             )
+        # ROCm torch is installed (or already was); flag it so later install
+        # phases do not overwrite it with the generic CPU torch wheel. BNB is
+        # a separate dependency -- a BNB install failure must NOT roll the
+        # torch ROCm install back.
+        _rocm_windows_torch_installed = True
         # Always install AMD Windows bitsandbytes -- the PyPI wheel ships only
         # CUDA DLLs and will fail to load on ROCm.  Install even when torch was
         # already a ROCm build so that `studio update` repairs a broken bnb.
-        # Only flip the success flag when the install actually succeeds; otherwise
+        if not _install_bnb_windows_rocm():
-        # the post-install "manual install may be required" warning is suppressed.
+            print(
-        if _install_bnb_windows_rocm():
+                "   Warning: AMD Windows bitsandbytes install failed; "
-            _rocm_windows_torch_installed = True
+                "ROCm torch is installed but bitsandbytes may need manual install"
+            )
         return
 
     # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ──
@@ -556,7 +586,15 @@ def _ensure_rocm_torch() -> None:
             [
                 sys.executable,
                 "-c",
-                "import torch; print(getattr(torch.version,'hip','') or '')",
+                (
+                    "import torch; "
+                    "hip=getattr(torch.version,'hip','') or ''; "
+                    "ver=getattr(torch,'__version__','').lower(); "
+                    # Print the HIP version when present (back-compat), else
+                    # "rocm" sentinel when only torch.__version__ flags ROCm
+                    # (AMD SDK / Radeon wheels). Empty string = CPU/CUDA.
+                    "print(hip if hip else ('rocm' if 'rocm' in ver else ''))"
+                ),
             ],
             stdout = subprocess.PIPE,
             stderr = subprocess.DEVNULL,