From 935f6c50efd45018c312ea07911d55b98f34436a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jun 2026 05:49:25 -0700 Subject: [PATCH] studio: tighten torchao Windows-ROCm comments and test docstrings (#6610) --- studio/backend/core/rag/embeddings.py | 12 ++++-------- studio/backend/tests/test_torchao_select.py | 7 +++---- studio/install_python_stack.py | 12 +++++------- tests/studio/install/test_rocm_support.py | 7 +++---- 4 files changed, 15 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4bbbd19196..4c8d690302 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -50,14 +50,10 @@ _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``).""" + """Neutralize torchao before importing sentence-transformers. On Windows ROCm, + torchao (pulled in by transformers.quantizers) imports an absent c10d backend + and aborts, dropping the embedder to llama-server. Workers stub it too; the + embedder runs in the main process. No-op elsewhere; runs once under ``_lock``.""" global _torchao_stub_done if _torchao_stub_done: return diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 3ff9873270..a99eb4c45c 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -72,10 +72,9 @@ def test_default_spec_matches_table(monkeypatch): 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.""" + """The overrides step must skip torchao on Windows ROCm: no working build exists + there (it imports an absent c10d backend and crashes transformers.quantizers), + so the installer skips it and relies on the runtime stub 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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index ce25c4faf9..491429a4ef 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -2234,16 +2234,14 @@ 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. Also skipped on Windows ROCm, - # which has no working torchao build (see below). + # GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm + # (no working 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. + # No working Windows ROCm torchao build: it imports an absent c10d backend + # and crashes transformers.quantizers. Studio stubs it at runtime, 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: diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index eeffb6fd95..4c50fa1c8e 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2332,10 +2332,9 @@ class TestWorkerWindowsRocmPatches: 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.""" + """embeddings.py must install the stub before importing sentence-transformers: + it runs in the main process (not a stubbed worker), so otherwise transformers + -> torchao crashes on Windows ROCm and the embedder drops to llama-server.""" source = _EMBEDDINGS_PATH.read_text(encoding = "utf-8") assert "install_torchao_windows_rocm_stub()" in source