ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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.
This commit is contained in:
Daniel Han 2026-06-23 05:39:02 -07:00 • committed by GitHub
commit 55c392ff7c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 54 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

18
studio/install_python_stack.py
View file

@ -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,
)