From 1faa0ca05864058aaa963e88203ec49af39540f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 04:38:01 -0700 Subject: [PATCH] Studio: don't offer torchao INT8/FP8 export on Windows ROCm (torchao unavailable) torch.distributed is unsupported on Windows ROCm, so real torchao cannot import and Studio import-stubs it; the stub's config classes return None, so TorchAoConfig(quant_type=None) crashed exports with 'quant_type must be either a string or an AOBaseConfig instance, got NoneType'. Gate the portable torchao FP8/INT8 path off on Windows ROCm via a shared is_win32_rocm() helper (used by both the import stub and the export gate so they can't drift), and add an early defensive error in export_merged_model so a forced torchao request fails fast with a clear message instead of the cryptic crash. The Export UI also hides torchao and stops claiming it works there. No change on Linux, macOS, or Windows CUDA (torchao is real); Windows ROCm users keep 16-bit merged and GGUF quantization. --- studio/backend/core/_torchao_stub.py | 61 ++++--- studio/backend/core/export/export.py | 41 ++++- studio/backend/tests/test_torchao_select.py | 155 +++++++++++++++++- .../src/features/export/export-page.tsx | 23 ++- 4 files changed, 244 insertions(+), 36 deletions(-) diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py index 6336954bd5..625a6f0b1c 100644 --- a/studio/backend/core/_torchao_stub.py +++ b/studio/backend/core/_torchao_stub.py @@ -100,36 +100,43 @@ class _StubSubpackageFinder(importlib.abc.MetaPathFinder): ) +def is_win32_rocm() -> bool: + """True on Windows ROCm, where torch.distributed (and thus torchao) is unavailable. + + Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH + persist after reverting to a CUDA wheel. Some ROCm wheels lack torch.version.hip + but still encode "rocm" in __version__, so accept either. Windows CUDA -> False. + Shared by the import stub and the torchao export gate so the two can't drift. + """ + if sys.platform != "win32": + return False + try: + import torch + return bool( + getattr(getattr(torch, "version", None), "hip", None) + or "rocm" in getattr(torch, "__version__", "").lower() + ) + except Exception: + return False + + def install_torchao_windows_rocm_stub() -> None: """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. No-op elsewhere (incl. Windows CUDA, where torchao is real). Must run before importing transformers / unsloth_zoo. Safe to call once per worker. """ - # Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH - # persist after reverting to a CUDA wheel. Some ROCm wheels lack - # torch.version.hip but still encode "rocm" in __version__, so accept either. - _is_win32_rocm = False - if sys.platform == "win32": - try: - import torch as _torch_probe - _is_win32_rocm = bool( - getattr(getattr(_torch_probe, "version", None), "hip", None) - or "rocm" in getattr(_torch_probe, "__version__", "").lower() - ) - del _torch_probe - except Exception: - pass - if _is_win32_rocm: - # Register the finder only on Windows ROCm. - sys.meta_path.append(_StubSubpackageFinder()) - # Seed torchao top-level + key submodules; the finder handles the rest. - for _tao_name in ( - "torchao", - "torchao.quantization", - "torchao.dtypes", - "torchao.float8", - "torchao.utils", - ): - if _tao_name not in sys.modules: - sys.modules[_tao_name] = _make_mod_stub(_tao_name) + if not is_win32_rocm(): + return + # Register the finder only on Windows ROCm. + sys.meta_path.append(_StubSubpackageFinder()) + # Seed torchao top-level + key submodules; the finder handles the rest. + for _tao_name in ( + "torchao", + "torchao.quantization", + "torchao.dtypes", + "torchao.float8", + "torchao.utils", + ): + if _tao_name not in sys.modules: + sys.modules[_tao_name] = _make_mod_stub(_tao_name) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index c8be50b08b..0e25bb343a 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -102,14 +102,35 @@ def _compressed_export_supported(): def _torchao_export_supported(): - """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + """True if the installed unsloth build has the portable torchao FP8/INT8 export path. + + Forced False on Windows ROCm: torch.distributed (and therefore torchao) is unavailable + there, so torchao is import-stubbed and its config classes return None. Windows CUDA, + Linux, and macOS are unaffected (torchao is real).""" try: + from core._torchao_stub import is_win32_rocm + if is_win32_rocm(): + return False import unsloth.save as _us return hasattr(_us, "_normalize_torchao_method") except Exception: return False +def _torchao_runtime_unavailable(): + """True where portable torchao export cannot run (Windows ROCm): torchao is import-stubbed + (its config classes return None) or torch.distributed is absent. False everywhere else.""" + import sys + + try: + from core._torchao_stub import is_win32_rocm, _STUB_SENTINEL + if is_win32_rocm(): + return True + return getattr(sys.modules.get("torchao"), "_unsloth_stub", None) is _STUB_SENTINEL + except Exception: + return False + + def _has_nvidia_gpu(): """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" try: @@ -495,6 +516,24 @@ class ExportBackend: "NVFP4 (compressed-tensors)": "nvfp4", } compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + + # Portable torchao (torchao_fp8/torchao_int8) needs torch.distributed + torchao, both + # absent on Windows ROCm where torchao is import-stubbed (its config classes return None). + # Fail fast with a clear message instead of the cryptic transformers "quant_type ... got + # NoneType" crash. 16-bit / GGUF / compressed-tensors formats are unaffected. + if ( + compressed_alias + and str(compressed_alias).lower().startswith("torchao") + and _torchao_runtime_unavailable() + ): + return ( + False, + "Portable torchao FP8/INT8 export is not supported on Windows ROCm: " + "torch.distributed and torchao are unavailable on this build. Use 16-bit " + "merged or GGUF quantization instead.", + None, + ) + compressed_suffix: Optional[str] = None # Classify the alias: torchao-portable vs compressed-tensors. torchao_info = None diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index e4775a10a6..fe5100c7d4 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -1,16 +1,21 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Tests for _select_torchao_spec in install_python_stack.py. +"""Tests for torchao version selection and the Windows-ROCm export gate. torchao's C++ extensions are built against one exact torch release, so the installer must pick the torchao version matching the torch installed in the -venv (otherwise the cpp kernels are skipped). This pins that mapping. +venv (otherwise the cpp kernels are skipped); the first half pins that mapping. +The second half covers the runtime gate: torch.distributed is unsupported on +Windows ROCm, so torchao is import-stubbed and the portable FP8/INT8 export must +be turned off there (shared is_win32_rocm() helper) with a clear defensive error. """ from __future__ import annotations +import ast import sys +import types from pathlib import Path from unittest.mock import MagicMock @@ -19,6 +24,9 @@ import pytest # install_python_stack.py lives at repo_root/studio/install_python_stack.py _INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" +# backend root (studio/backend), for reading/exec-ing backend sources. +_BACKEND = Path(__file__).resolve().parents[1] + def _load_module(monkeypatch): """(Re-)import install_python_stack and return it (mirrors test_pytorch_mirror).""" @@ -134,3 +142,146 @@ def test_skips_torchao_on_windows_rocm( assert not any(spec.startswith("torchao") for spec in installed_specs) assert "dependency overrides (skipped, Windows ROCm)" in progress_labels + + +# -- Windows-ROCm torchao export gate ----------------------------------------------------------- +# +# torch.distributed is unsupported on Windows ROCm, so real torchao can't import there and Studio +# import-stubs it (core/_torchao_stub.py); the stub's config classes return None, which made +# TorchAoConfig(quant_type=None) crash with "quant_type must be ... got NoneType". These prove the +# shared is_win32_rocm() gate turns the portable torchao formats off there and that the defensive +# export path raises a clear error instead of the cryptic crash. + +import core._torchao_stub as _stub + + +def _func_src(rel, name): + src = (_BACKEND / rel).read_text(encoding = "utf-8") + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +def _exec_func(rel, name): + """Exec one backend function in isolation, avoiding export.py's heavy import chain.""" + ns: dict = {} + exec(_func_src(rel, name), ns) + return ns[name] + + +@pytest.mark.parametrize( + ("platform", "hip", "version", "expected"), + [ + ("win32", "6.4.0", "2.10.0+rocm6.4", True), # ROCm via torch.version.hip + ("win32", None, "2.10.0+rocm6.4", True), # ROCm via __version__ tag only + ("win32", None, "2.10.0+cu128", False), # Windows CUDA -> real torchao + ("linux", "6.4.0", "2.10.0+rocm6.4", False), # Linux ROCm -> real torchao + ("darwin", None, "2.10.0", False), # macOS + ], +) +def test_is_win32_rocm(monkeypatch, platform, hip, version, expected): + fake_torch = types.SimpleNamespace( + version = types.SimpleNamespace(hip = hip), __version__ = version + ) + monkeypatch.setattr(sys, "platform", platform) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + assert _stub.is_win32_rocm() is expected + + +def test_gate_and_stub_share_helper(): + # The stub installer and the export gate must both route through is_win32_rocm() so they can't + # drift (the gate off while the stub is still active, or the reverse). + stub_src = (_BACKEND / "core" / "_torchao_stub.py").read_text(encoding = "utf-8") + assert "def is_win32_rocm(" in stub_src + assert "is_win32_rocm()" in _func_src( + "core/_torchao_stub.py", "install_torchao_windows_rocm_stub" + ) + assert "is_win32_rocm()" in _func_src("core/export/export.py", "_torchao_export_supported") + + +def test_installer_noop_off_windows_rocm(monkeypatch): + # is_win32_rocm() False -> installer must not register the finder or seed torchao stubs. + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: False) + before = list(sys.meta_path) + _stub.install_torchao_windows_rocm_stub() + assert list(sys.meta_path) == before + + +# (a) gate off on Windows ROCm; (b) unchanged elsewhere + + +def test_torchao_gate_false_on_windows_rocm(monkeypatch): + # (a) On Windows ROCm the portable torchao formats are not offered, without importing unsloth. + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: True) + assert _exec_func("core/export/export.py", "_torchao_export_supported")() is False + + +def _install_fake_unsloth_save(monkeypatch, *, has_method): + unsloth = types.ModuleType("unsloth") + save = types.ModuleType("unsloth.save") + if has_method: + save._normalize_torchao_method = lambda alias: ("fp8", "torchao-fp8") + unsloth.save = save + monkeypatch.setitem(sys.modules, "unsloth", unsloth) + monkeypatch.setitem(sys.modules, "unsloth.save", save) + + +def test_torchao_gate_supported_off_windows_rocm(monkeypatch): + # (b) Off Windows ROCm the gate is unchanged: True when the unsloth build has the method. + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: False) + _install_fake_unsloth_save(monkeypatch, has_method = True) + assert _exec_func("core/export/export.py", "_torchao_export_supported")() is True + + +def test_torchao_gate_false_when_build_lacks_method(monkeypatch): + # (b) Off Windows ROCm, an older unsloth without the method is still unsupported. + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: False) + _install_fake_unsloth_save(monkeypatch, has_method = False) + assert _exec_func("core/export/export.py", "_torchao_export_supported")() is False + + +# (c) defensive early error when torchao is stubbed / unavailable + + +def _load_export_module_no_torch(monkeypatch): + """Import core.export.export with torch/unsloth blocked (mirrors test_export_capability), so + the defensive path runs on CPU with no GPU and no torchao.""" + import builtins + import importlib + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + if name.split(".")[0] in {"torch", "unsloth"}: + raise ImportError(f"blocked: {name}") + return real_import(name, *args, **kwargs) + + for m in [k for k in list(sys.modules) if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + return importlib.import_module("core.export.export") + + +def test_torchao_defensive_error_on_windows_rocm(monkeypatch): + # (c) A forced torchao request reaches the merged path -> clear error, not the NoneType crash. + mod = _load_export_module_no_torch(monkeypatch) + monkeypatch.setattr(mod, "_export_runtime_available", lambda: True) + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: True) + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = object() + be.current_tokenizer = object() + be._audio_type = None + be.is_peft = True + ok, message, out = be.export_merged_model("/tmp/x", compressed_method = "torchao_fp8") + assert ok is False and out is None + assert "Windows ROCm" in message and "torchao" in message.lower() + + +def test_torchao_defensive_error_wired_early(): + # The guard lives in export_merged_model, before the merge/quant work, keyed on the shared helper. + m = _func_src("core/export/export.py", "export_merged_model") + assert "_torchao_runtime_unavailable()" in m + assert 'str(compressed_alias).lower().startswith("torchao")' in m diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 07606a26ed..74141195f7 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -223,8 +223,11 @@ export function ExportPage() { const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); const hardware = useHardwareInfo(); + const deviceType = usePlatformStore((s) => s.deviceType); // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. - const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + const isMacHost = deviceType === "mac"; + // Windows ROCm has no torch.distributed, so portable torchao (FP8/INT8) is unavailable there. + const isWindowsRocm = deviceType === "windows" && hardware.rocm != null; // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. const hasNvidia = hardware.cuda != null && hardware.rocm == null; // Only gray out on an authoritative unsupported response; while unloaded the backend route guard @@ -240,13 +243,14 @@ export function ExportPage() { // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. if (f.backend === "compressed") return hasNvidia; // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a - // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the - // backend rejects quantized export there). - if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors), on macOS/MLX (the + // backend rejects quantized export there), and on Windows ROCm (torchao is unavailable: + // no torch.distributed, so its config classes are import-stubbed to None). + if (f.backend === "torchao") return !hasNvidia && !isMacHost && !isWindowsRocm; // Plain 16-bit is available everywhere. return true; }), - [hasNvidia, isMacHost], + [hasNvidia, isMacHost, isWindowsRocm], ); const toggleFormat = useCallback((value: string) => { setSelectedFormats((prev) => @@ -1434,13 +1438,20 @@ export function ExportPage() { )} - {!hasNvidia && ( + {!hasNvidia && !isWindowsRocm && (
No NVIDIA GPU detected: compressed-tensors formats are hidden. 16-bit and portable FP8/INT8 (torchao) still work here and load in vLLM.
)} + + {isWindowsRocm && ( +
+ Windows ROCm: quantized FP8/INT8 (torchao) export is + unavailable (no torch.distributed). Use 16-bit or GGUF. +
+ )} )}