diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py index 625a6f0b1c..2250ea4125 100644 --- a/studio/backend/core/_torchao_stub.py +++ b/studio/backend/core/_torchao_stub.py @@ -103,10 +103,9 @@ 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. + Gate on the runtime torch, not env vars (HIP_PATH persists after a CUDA revert). AMD SDK + wheels lack torch.version.hip but tag "rocm" in __version__, so accept either. Shared by the + import stub and the export gate so they can't drift. """ if sys.platform != "win32": return False @@ -128,8 +127,9 @@ def install_torchao_windows_rocm_stub() -> None: """ if not is_win32_rocm(): return - # Register the finder only on Windows ROCm. - sys.meta_path.append(_StubSubpackageFinder()) + # Register the finder only on Windows ROCm, and only once (no duplicates on re-call). + if not any(isinstance(_f, _StubSubpackageFinder) for _f in sys.meta_path): + sys.meta_path.append(_StubSubpackageFinder()) # Seed torchao top-level + key submodules; the finder handles the rest. for _tao_name in ( "torchao", diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 3ae214520a..d8f13560fe 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -104,9 +104,8 @@ def _compressed_export_supported(): def _torchao_export_supported(): """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).""" + Forced False on Windows ROCm, where torchao is import-stubbed (no torch.distributed) and its + config classes return None. Unchanged on Windows CUDA / Linux / macOS (torchao is real).""" try: from core._torchao_stub import is_win32_rocm @@ -127,11 +126,27 @@ def _torchao_runtime_unavailable(): 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 + _tao = sys.modules.get("torchao") + return _tao is not None and getattr(_tao, "_unsloth_stub", None) is _STUB_SENTINEL except Exception: return False +def _is_torchao_alias(alias): + """True if `alias` is any torchao export form (torchao_fp8, portable_int8, hyphen/space + variants) per unsloth's normalizer, with a torchao_ prefix fallback. Catches a torchao request + before the Windows-ROCm gate misclassifies it as compressed-tensors.""" + if not alias: + return False + try: + import unsloth.save as _us + if _us._normalize_torchao_method(alias) is not None: + return True + except Exception: + pass + return str(alias).lower().startswith("torchao") + + def _has_nvidia_gpu(): """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" try: @@ -518,15 +533,10 @@ class ExportBackend: } 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() - ): + # Portable torchao is unavailable on Windows ROCm (stubbed, no torch.distributed). Reject + # any torchao alias early with a clear message instead of the cryptic NoneType crash or a + # misleading NVIDIA error. Other formats (16-bit/GGUF/compressed-tensors) are unaffected. + if _is_torchao_alias(compressed_alias) and _torchao_runtime_unavailable(): return ( False, "Portable torchao FP8/INT8 export is not supported on Windows ROCm: " diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 3b8b0e6ffd..4785ed1a7d 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -3,12 +3,10 @@ """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); 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. +First half: the installer must pin the torchao version matching the installed torch (its cpp +kernels are built per torch release). Second half: torch.distributed is unsupported on Windows +ROCm, so torchao is import-stubbed and the portable FP8/INT8 export must be gated off there +(shared is_win32_rocm() helper) with a clear defensive error. """ from __future__ import annotations @@ -145,12 +143,9 @@ def test_skips_torchao_on_windows_rocm( # -- 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. +# torchao is import-stubbed on Windows ROCm (no torch.distributed) and its config classes return +# None, which made TorchAoConfig(quant_type=None) crash. These prove the shared is_win32_rocm() +# gate hides the torchao formats and the defensive path raises a clear error instead. import core._torchao_stub as _stub @@ -215,11 +210,22 @@ def test_torchao_gate_false_on_windows_rocm(monkeypatch): assert _exec_func("core/export/export.py", "_torchao_export_supported")() is False +_TORCHAO_ALIASES = {"torchao_fp8", "torchao_int8", "portable_fp8", "portable_int8"} + + +def _fake_normalize_torchao(save_method): + # Mirrors unsloth.save._normalize_torchao_method (lower/strip, - and space -> _). + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return ("fp8", "torchao-fp8") if key in _TORCHAO_ALIASES else None + + 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") + save._normalize_torchao_method = _fake_normalize_torchao unsloth.save = save monkeypatch.setitem(sys.modules, "unsloth", unsloth) monkeypatch.setitem(sys.modules, "unsloth.save", save) @@ -251,7 +257,9 @@ def _load_export_module_no_torch(monkeypatch): real_import = builtins.__import__ def blocking_import(name, *args, **kwargs): - if name.split(".")[0] in {"torch", "unsloth"}: + # Block real torch/unsloth, but honor injected fakes already in sys.modules. + top = name.split(".")[0] + if top in {"torch", "unsloth"} and top not in sys.modules: raise ImportError(f"blocked: {name}") return real_import(name, *args, **kwargs) @@ -262,24 +270,91 @@ def _load_export_module_no_torch(monkeypatch): return importlib.import_module("core.export.export") +def _bare_backend(mod): + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = object() + be.current_tokenizer = object() + be._audio_type = None + be.is_peft = True + return be + + 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") + ok, message, out = _bare_backend(mod).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_alias_form_on_windows_rocm(monkeypatch): + # An equivalent alias unsloth accepts (portable_fp8) must hit the same rejection, not fall + # through to the misleading NVIDIA compressed-tensors error. + mod = _load_export_module_no_torch(monkeypatch) + monkeypatch.setattr(mod, "_export_runtime_available", lambda: True) + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: True) + _install_fake_unsloth_save(monkeypatch, has_method = True) + + ok, message, out = _bare_backend(mod).export_merged_model( + "/tmp/x", compressed_method = "portable_fp8" + ) + assert ok is False and out is None + assert "Windows ROCm" in message and "torchao" in message.lower() + + +def test_is_torchao_alias_recognizes_all_forms(monkeypatch): + _install_fake_unsloth_save(monkeypatch, has_method = True) + fn = _exec_func("core/export/export.py", "_is_torchao_alias") + for alias in ("torchao_fp8", "portable_int8", "portable-fp8", "Portable FP8"): + assert fn(alias) is True + for alias in ("fp8", "nvfp4", "w8a8", "", None): + assert fn(alias) is False + + def test_torchao_defensive_error_wired_early(): - # The guard lives in export_merged_model, before the merge/quant work, keyed on the shared helper. + # Guard is in export_merged_model before the merge/quant work; alias is normalized (not just the + # torchao_ prefix) so every torchao form is caught. m = _func_src("core/export/export.py", "export_merged_model") + assert "_is_torchao_alias(compressed_alias)" in m assert "_torchao_runtime_unavailable()" in m - assert 'str(compressed_alias).lower().startswith("torchao")' in m + alias_fn = _func_src("core/export/export.py", "_is_torchao_alias") + assert "_normalize_torchao_method(alias)" in alias_fn + assert 'startswith("torchao")' in alias_fn + + +# (issue 2/4) backend win32_rocm flag + single finder registration + + +def test_export_capability_exposes_win32_rocm(monkeypatch): + import utils.hardware.hardware as hw + + monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA) + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(hw, "IS_ROCM", True) + assert hw.export_capability()["win32_rocm"] is True + monkeypatch.setattr(hw, "IS_ROCM", False) + assert hw.export_capability()["win32_rocm"] is False + monkeypatch.setattr(sys, "platform", "linux") + monkeypatch.setattr(hw, "IS_ROCM", True) + assert hw.export_capability()["win32_rocm"] is False + + +def test_installer_registers_finder_once(monkeypatch): + # Repeated install must not stack duplicate finders. Restore global state after. + monkeypatch.setattr(_stub, "is_win32_rocm", lambda: True) + meta_before = list(sys.meta_path) + tao_before = {k for k in sys.modules if k == "torchao" or k.startswith("torchao.")} + try: + _stub.install_torchao_windows_rocm_stub() + _stub.install_torchao_windows_rocm_stub() + finders = [f for f in sys.meta_path if isinstance(f, _stub._StubSubpackageFinder)] + assert len(finders) == 1 + finally: + sys.meta_path[:] = meta_before + for k in [k for k in sys.modules if (k == "torchao" or k.startswith("torchao.")) and k not in tao_before]: + del sys.modules[k] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 8d6c919ebd..1516ab7f96 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -270,13 +270,19 @@ def export_capability() -> dict: import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. - Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + Returns {export_supported, export_unsupported_reason, export_unsupported_message, win32_rocm}. + ``win32_rocm`` is the UI's single source of truth for the torchao gate (mirrors + is_win32_rocm()): torchao is unavailable on Windows ROCm. """ - if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + device = get_device() + # get_device() ran detect_hardware(), so IS_ROCM (hip OR "rocm" tag) is authoritative here. + win32_rocm = sys.platform == "win32" and IS_ROCM + if device in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): return { "export_supported": True, "export_unsupported_reason": None, "export_unsupported_message": None, + "win32_rocm": win32_rocm, } # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" # would be wrong advice on a Mac even when torch is also absent. @@ -303,6 +309,7 @@ def export_capability() -> dict: "export_supported": False, "export_unsupported_reason": reason, "export_unsupported_message": message, + "win32_rocm": win32_rocm, } diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 74141195f7..8edc8dcb73 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -226,8 +226,9 @@ export function ExportPage() { 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 = deviceType === "mac"; - // Windows ROCm has no torch.distributed, so portable torchao (FP8/INT8) is unavailable there. - const isWindowsRocm = deviceType === "windows" && hardware.rocm != null; + // Backend truth for the torchao gate (single source). Not re-derived from `rocm`: AMD SDK + // wheels leave torch.version.hip unset, so `rocm` alone would miss Windows ROCm. + const isWindowsRocm = hardware.win32Rocm; // 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 @@ -242,10 +243,8 @@ export function ExportPage() { MERGED_FORMATS.filter((f) => { // 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), 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). + // Portable torchao: shown on non-NVIDIA hosts. Hidden on NVIDIA (use compressed-tensors), + // macOS/MLX (rejected), and Windows ROCm (torchao unavailable: no torch.distributed). if (f.backend === "torchao") return !hasNvidia && !isMacHost && !isWindowsRocm; // Plain 16-bit is available everywhere. return true; @@ -259,7 +258,15 @@ export function ExportPage() { : [...prev, value], ); }, []); - // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // Drop any already-selected format that the gate just removed (e.g. torchao once win32Rocm + // resolves after /api/system/hardware lands), so a stale pick isn't summarized or exported. + useEffect(() => { + const allowed = new Set(availableFormats.map((f) => f.value)); + setSelectedFormats((prev) => { + const next = prev.filter((v) => allowed.has(v)); + return next.length === prev.length ? prev : next; + }); + }, [availableFormats]); // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 4d63d4d6af..ff34eef089 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -31,6 +31,9 @@ export interface HardwareInfo { exportSupported: boolean | null; exportUnsupportedReason: string | null; exportUnsupportedMessage: string | null; + // Backend truth for the torchao gate (mirrors is_win32_rocm(): torch.version.hip OR a "rocm" + // build tag). Single source; the UI must not re-derive Windows ROCm from `rocm` alone. + win32Rocm: boolean; loaded: boolean; } @@ -48,6 +51,7 @@ const DEFAULT: HardwareInfo = { exportSupported: null, exportUnsupportedReason: null, exportUnsupportedMessage: null, + win32Rocm: false, loaded: false, }; @@ -101,6 +105,7 @@ async function fetchOnce(): Promise { exportSupported: data?.export_supported ?? null, exportUnsupportedReason: data?.export_unsupported_reason ?? null, exportUnsupportedMessage: data?.export_unsupported_message ?? null, + win32Rocm: data?.win32_rocm ?? false, loaded: true, }; if (generation === cacheGeneration) {