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.
This commit is contained in:
Daniel Han
parent
c570180a32
commit
1faa0ca058
4 changed files with 244 additions and 36 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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() {
|
|||
</div>
|
||||
)}
|
||||
|
||||
{!hasNvidia && (
|
||||
{!hasNvidia && !isWindowsRocm && (
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
No NVIDIA GPU detected: compressed-tensors formats are
|
||||
hidden. 16-bit and portable FP8/INT8 (torchao) still
|
||||
work here and load in vLLM.
|
||||
</div>
|
||||
)}
|
||||
|
||||
{isWindowsRocm && (
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Windows ROCm: quantized FP8/INT8 (torchao) export is
|
||||
unavailable (no torch.distributed). Use 16-bit or GGUF.
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue