fix: BNB_ROCM_VERSION in server process + torch._C._distributed_c10d stubs

Two errors visible in training logs on Windows ROCm:

1. Server process bitsandbytes crash:
   "Configured ROCm binary not found at libbitsandbytes_rocm713.dll"
   The installed BNB wheel ships rocm72.dll (not rocm713.dll). The
   training worker already sets BNB_ROCM_VERSION=72 via DLL detection
   but the server process (main.py) imported bitsandbytes before that
   ran. Fix: add the same DLL-scan + BNB_ROCM_VERSION assignment to
   main.py inside the existing win32 guard, before any downstream
   import can pull in bitsandbytes.

2. torch.distributed import failure:
   "No module named 'torch._C._distributed_c10d'; torch._C is not a package"
   torch._C is a C extension on Windows ROCm — Python cannot do
   submodule imports from it, so torch.distributed fails to import
   before our attribute stubs could ever run. Fix: inject empty
   ModuleType stubs for _distributed_c10d, _distributed_autograd and
   _distributed_rpc into sys.modules inside the win32 guard in
   hardware.py BEFORE importing torch.distributed, so the import
   succeeds and our attribute stubs take effect.

9 new tests in TestServerStartupRocmFixes; total 212 passed, 2 skipped

studio/backend/main.py

@@ -47,6 +47,32 @@ if sys.platform == "win32":
     _add_rocm_dll_dirs()
     del _add_rocm_dll_dirs
 
+    # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
+    # bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm<ver>.dll
+    # where <ver> comes from torch.version.hip (e.g. "7.13..." → "713").
+    # The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without
+    # this the server process crashes with "Configured ROCm binary not found".
+    # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
+    # before any import that pulls in bitsandbytes (mirrors worker.py logic).
+    if "BNB_ROCM_VERSION" not in os.environ:
+        import glob as _glob
+        _bnb_rocm_ver = None
+        try:
+            import importlib.util as _ilu
+            _bnb_spec = _ilu.find_spec("bitsandbytes")
+            if _bnb_spec and _bnb_spec.origin:
+                _pkg_dir = os.path.dirname(_bnb_spec.origin)
+                _dlls = _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll"))
+                import re as _re_bnb
+                for _dll in sorted(_dlls):
+                    _m = _re_bnb.search(r"libbitsandbytes_rocm(\d+)\.dll", _dll)
+                    if _m:
+                        _bnb_rocm_ver = _m.group(1)
+                        break
+        except Exception:
+            pass
+        os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver or "72"
+
 # Ensure backend dir is on sys.path so _platform_compat is importable when
 # main.py is launched directly (e.g. `uvicorn main:app`).
 _backend_dir = str(_Path(__file__).parent)

studio/backend/utils/hardware/hardware.py

@@ -946,10 +946,23 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
 def _determine_attention_impl_for_gpu_estimate(config) -> str:
     import copy as _copy
 
-    # torch.distributed is incomplete on Windows ROCm — it ships without the
-    # process-group helpers (is_initialized, is_available, etc.).
-    # resolve_attention_implementation (unsloth) calls is_initialized()
-    # unconditionally, so patch any missing attrs before importing it.
+    # torch.distributed is incomplete on Windows ROCm — torch._C is a C
+    # extension (not a package), so Python cannot import the submodule
+    # torch._C._distributed_c10d that torch.distributed depends on.
+    # Inject an empty stub into sys.modules BEFORE importing torch.distributed
+    # so the import succeeds, then patch the missing process-group helpers.
+    import sys as _sys
+    import types as _types
+
+    if _sys.platform == "win32":
+        for _c10d_name in (
+            "torch._C._distributed_c10d",
+            "torch._C._distributed_autograd",
+            "torch._C._distributed_rpc",
+        ):
+            if _c10d_name not in _sys.modules:
+                _sys.modules[_c10d_name] = _types.ModuleType(_c10d_name)
+
     try:
         import torch.distributed as _td
 

tests/studio/install/test_rocm_support.py

@@ -2348,5 +2348,73 @@ class TestSetupShGccInstallDir:
         assert "gcc install dir" in source or "GCC_INSTALL_DIR" in source
 
 
+# =============================================================================
+# TEST: main.py -- BNB_ROCM_VERSION server startup + distributed stubs
+# =============================================================================
+
+_MAIN_PY_PATH     = PACKAGE_ROOT / "studio" / "backend" / "main.py"
+_HARDWARE_PY_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py"
+
+
+class TestServerStartupRocmFixes:
+    """Verify main.py sets BNB_ROCM_VERSION before any bitsandbytes import and
+    hardware.py injects torch._C._distributed_c10d stubs before torch.distributed."""
+
+    # ── BNB_ROCM_VERSION in server process ────────────────────────────────────
+
+    def test_main_py_sets_bnb_rocm_version(self):
+        """main.py must set BNB_ROCM_VERSION in the server process before imports."""
+        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
+        assert "BNB_ROCM_VERSION" in source
+
+    def test_main_py_bnb_detection_scoped_to_win32(self):
+        """main.py BNB_ROCM_VERSION logic must be inside the win32 platform guard."""
+        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
+        win32_idx = source.find('sys.platform == "win32"')
+        bnb_idx   = source.find("BNB_ROCM_VERSION")
+        assert win32_idx != -1 and bnb_idx != -1
+        assert win32_idx < bnb_idx
+
+    def test_main_py_bnb_dll_detection_uses_glob(self):
+        """main.py must scan for libbitsandbytes_rocm*.dll to find the right version."""
+        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
+        assert "libbitsandbytes_rocm" in source
+
+    def test_main_py_bnb_falls_back_to_72(self):
+        """main.py must fall back to BNB_ROCM_VERSION='72' when no DLL is found."""
+        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
+        assert '"72"' in source or "'72'" in source
+
+    def test_main_py_bnb_only_set_when_not_already_in_env(self):
+        """main.py must not override an existing BNB_ROCM_VERSION env var."""
+        source = _MAIN_PY_PATH.read_text(encoding = "utf-8")
+        assert '"BNB_ROCM_VERSION" not in os.environ' in source
+
+    # ── torch._C._distributed_c10d stubs in hardware.py ──────────────────────
+
+    def test_hardware_py_injects_distributed_c10d_stub(self):
+        """hardware.py must inject torch._C._distributed_c10d into sys.modules."""
+        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
+        assert "_distributed_c10d" in source
+
+    def test_hardware_py_stub_injected_before_distributed_import(self):
+        """The sys.modules stub must be injected BEFORE import torch.distributed."""
+        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
+        c10d_idx = source.find("_distributed_c10d")
+        dist_idx = source.find("import torch.distributed")
+        assert c10d_idx != -1 and dist_idx != -1
+        assert c10d_idx < dist_idx
+
+    def test_hardware_py_stub_uses_types_moduletype(self):
+        """hardware.py must create the stub with types.ModuleType."""
+        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
+        assert "ModuleType" in source
+
+    def test_hardware_py_stub_scoped_to_win32(self):
+        """hardware.py distributed stub injection must be gated on win32."""
+        source = _HARDWARE_PY_PATH.read_text(encoding = "utf-8")
+        assert 'platform == "win32"' in source or "win32" in source
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])