studio: fix sentence-transformers RAG embedder on Windows ROCm (torchao) (#6608)

torchao has no working Windows ROCm build. transformers.quantizers imports it,
and it loads torch's c10d distributed backend at module level, which the AMD
Windows wheels omit (no RCCL). The import aborts, transformers can no longer
expose PreTrainedModel, and the sentence-transformers embedder silently falls
back to the llama-server GGUF embedder. Linux ROCm and NVIDIA are unaffected
(the c10d ops are present / torchao is real there).

The training and export workers already install the shared torchao stub before
importing transformers, but the RAG embedder runs in the main backend process,
which never did. Two fixes, both no-ops off Windows ROCm:

- embeddings.py: install_torchao_windows_rocm_stub() before the first
  sentence-transformers import, so an already-installed torchao is neutralized
  (fixes existing venvs).
- install_python_stack.py: stop installing torchao on Windows ROCm; it can only
  crash on import there, so new venvs never ship it.

Add tests covering the embedder stub call and the install skip.

studio/backend/core/rag/embeddings.py

@@ -46,6 +46,27 @@ def _device() -> str:
     return _TORCH_DEVICE.get(get_device(), "cpu")
 
 
+_torchao_stub_done = False
+
+
+def _install_torchao_stub_once() -> None:
+    """Neutralize torchao before the first sentence-transformers import.
+
+    transformers.quantizers imports torchao, which loads torch's c10d
+    distributed backend at module level; the AMD Windows ROCm wheels omit it
+    (no RCCL), so the import aborts and silently drops the ST embedder to the
+    llama-server fallback. The training/export workers install this stub at
+    process start, but the embedder runs in the main backend process, which
+    otherwise never does. No-op off Windows ROCm. Runs once (under ``_lock``)."""
+    global _torchao_stub_done
+    if _torchao_stub_done:
+        return
+    _torchao_stub_done = True
+    from core._torchao_stub import install_torchao_windows_rocm_stub
+
+    install_torchao_windows_rocm_stub()
+
+
 def _get(model_name: str | None = None):
     """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
     for a ~1.5x speedup at negligible accuracy loss."""
@@ -53,6 +74,7 @@ def _get(model_name: str | None = None):
     name = model_name or config.EMBEDDING_MODEL
     with _lock:
         if _model is None or _name != name:
+            _install_torchao_stub_once()
             from sentence_transformers import SentenceTransformer
 
             device = _device()

studio/backend/tests/test_torchao_select.py

@@ -69,3 +69,15 @@ def test_default_spec_matches_table(monkeypatch):
     mod = _load_module(monkeypatch)
     assert mod._TORCHAO_DEFAULT_SPEC == "torchao==0.14.0"
     assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC
+
+
+def test_skips_torchao_on_windows_rocm():
+    """The overrides step must skip torchao on Windows ROCm. There is no working
+    torchao build there: it loads torch's c10d distributed backend at import,
+    which the AMD Windows wheels omit, so `import torchao` raises and takes
+    transformers.quantizers with it. Studio stubs torchao at runtime instead."""
+    source = _INSTALL_SCRIPT.read_text(encoding = "utf-8")
+    # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ...
+    assert "elif _rocm_windows_torch_installed:" in source
+    # ... and reports the skip in the progress label.
+    assert "dependency overrides (skipped, Windows ROCm)" in source

studio/install_python_stack.py

@@ -2234,25 +2234,27 @@ def install_python_stack() -> int:
     # 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to
     #    match the torch installed in the venv so its C++ extensions load (see
     #    _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac
-    #    GGUF-only mode): torchao requires torch.
+    #    GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm,
+    #    which has no working torchao build (see below).
     if NO_TORCH:
         _progress("dependency overrides (skipped, no torch)")
+    elif _rocm_windows_torch_installed:
+        # torchao has no working Windows ROCm build: it loads torch's c10d
+        # distributed backend at import, which the AMD Windows wheels omit (no
+        # RCCL), so `import torchao` raises and takes transformers.quantizers
+        # down with it. Studio stubs torchao at runtime (core/_torchao_stub.py),
+        # so installing it only ships a package that crashes on import -- skip it.
+        _progress("dependency overrides (skipped, Windows ROCm)")
+        _safe_print("   Windows ROCm -- skipping torchao (no working build; stubbed at runtime)")
     else:
         _progress("dependency overrides")
         _torch_ver = _probe_installed_torch_version()
         _torchao_spec = _select_torchao_spec(_torch_ver)
         _safe_print(f"   torch {_torch_ver or 'unknown'} detected -- installing {_torchao_spec}")
-        _override_extra_args: tuple[str, ...] = ()
-        if _rocm_windows_torch_installed:
-            # torchao declares torch as a dependency; without --no-deps uv would
-            # install CPU torch from PyPI, overwriting the AMD ROCm wheels we just
-            # installed.
-            _override_extra_args = ("--no-deps",)
         pip_install(
             "Installing dependency overrides",
             "--force-reinstall",
             "--no-cache-dir",
-            *_override_extra_args,
             _torchao_spec,
         )
 

tests/studio/install/test_rocm_support.py

@@ -1221,6 +1221,8 @@ _WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "training" / "work
 _EXPORT_WORKER_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "export" / "worker.py"
 # Shared torchao Windows-ROCm stub used by both workers.
 _TORCHAO_STUB_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "_torchao_stub.py"
+# RAG embedder -- runs in the main backend process and also needs the stub.
+_EMBEDDINGS_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "rag" / "embeddings.py"
 # Wheel-probe script literal lives in wheel_utils after the resolver refactor.
 _WHEEL_UTILS_PATH = PACKAGE_ROOT / "studio" / "backend" / "utils" / "wheel_utils.py"
 
@@ -2329,6 +2331,14 @@ class TestWorkerWindowsRocmPatches:
         source = _EXPORT_WORKER_PATH.read_text(encoding = "utf-8")
         assert "install_torchao_windows_rocm_stub()" in source
 
+    def test_embedder_calls_shared_torchao_stub(self):
+        """rag/embeddings.py must install the stub before the first sentence-
+        transformers import. The embedder runs in the main backend process (not a
+        stubbed worker), so without this transformers -> torchao crashes on Windows
+        ROCm and the ST embedder silently drops to the llama-server fallback."""
+        source = _EMBEDDINGS_PATH.read_text(encoding = "utf-8")
+        assert "install_torchao_windows_rocm_stub()" in source
+
     def test_torchao_stub_uses_stub_type_meta(self):
         """Torchao stub must use _StubTypeMeta so isinstance() returns False not TypeError."""
         source = _TORCHAO_STUB_PATH.read_text(encoding = "utf-8")