diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4e76e4fcaa..4bbbd19196 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/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() diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 393a74daaf..3ff9873270 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 1c3918f900..ce25c4faf9 100644 --- a/studio/install_python_stack.py +++ b/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, ) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 6915269cec..eeffb6fd95 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/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")