From 1faa0ca05864058aaa963e88203ec49af39540f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 04:38:01 -0700 Subject: [PATCH 1/9] 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. +
+ )} )} From fbb9b8156fc9770b89236bdc7e21d63d02aeb763 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:39:16 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/export/export.py | 3 ++- studio/backend/tests/test_torchao_select.py | 12 +++++------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 0e25bb343a..3ae214520a 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -109,9 +109,11 @@ def _torchao_export_supported(): 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 @@ -121,7 +123,6 @@ 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(): diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index fe5100c7d4..3b8b0e6ffd 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -173,17 +173,15 @@ def _exec_func(rel, 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 + ("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 + ("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 - ) + 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 From 45754a1b27e78715e694adfabb2256f8be48b50a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 05:00:04 -0700 Subject: [PATCH 3/9] Studio: address review on torchao Windows-ROCm gate Normalize the requested alias through unsloth's torchao normalizer before the gate, so equivalent forms (portable_fp8, hyphen/space variants) also hit the clear Windows-ROCm rejection instead of being misclassified as compressed-tensors. Remove frontend/backend detection drift: expose win32_rocm (mirrors is_win32_rocm(): hip OR a rocm build tag) in the export capability payload and have the Export UI read it, instead of re-deriving Windows ROCm from versions.rocm (unset on AMD SDK wheels). Prune already-selected formats when the gate flips, so a stale torchao pick made before hardware resolves is not exported. Register the stub meta_path finder only once. Use an explicit None check for the stub sentinel. Tighten comments. --- studio/backend/core/_torchao_stub.py | 12 +- studio/backend/core/export/export.py | 36 ++++-- studio/backend/tests/test_torchao_select.py | 119 ++++++++++++++---- studio/backend/utils/hardware/hardware.py | 11 +- .../src/features/export/export-page.tsx | 21 ++-- .../frontend/src/hooks/use-hardware-info.ts | 5 + 6 files changed, 154 insertions(+), 50 deletions(-) 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) { From 8c8efacf2d5296f9004902935d6d26a26aa0fa08 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:02:36 +0000 Subject: [PATCH 4/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/export/export.py | 1 + studio/backend/tests/test_torchao_select.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index d8f13560fe..46c9aac787 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -124,6 +124,7 @@ def _torchao_runtime_unavailable(): import sys try: from core._torchao_stub import is_win32_rocm, _STUB_SENTINEL + if is_win32_rocm(): return True _tao = sys.modules.get("torchao") diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 4785ed1a7d..fa57b08874 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -356,5 +356,9 @@ def test_installer_registers_finder_once(monkeypatch): 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]: + 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] From 624a5801a724709e775cc6aea17abac0417e11b4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 14:16:17 +0000 Subject: [PATCH 5/9] Add torchao Windows-ROCm import shim to import_fixes Legacy Windows ROCm PyTorch wheels ship without the torch.distributed C-extension, so importing torchao crashes at `from torch._C._distributed_c10d import (...)` with "No module named 'torch._C._distributed_c10d'". That makes torchao completely unimportable and disables FP8/INT8 weight-only export. fix_torchao_windows_rocm_import() installs a capability-gated compatibility shim: it registers the version-exact _c10d_functional op schemas via a FRAGMENT library plus a fake torch._C._distributed_c10d module, then runs `import torchao` as an in-transaction acceptance test with atomic rollback. It fires only on the exact broken config (win32 + torch.version.hip + real torchao installed + distributed C-extension genuinely absent + no _c10d_functional dispatcher ops) and is a strict no-op everywhere else, including on fixed wheels (ROCm/TheRock#5694). Opt out with UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM=1. Wired into _gpu_init before importing unsloth_zoo so the zoo torchao stub self-disables once real torchao is importable. Adds drift tests that pin the schema table to the live dispatcher and assert the no-op, guard, FRAGMENT, and rollback invariants. --- tests/test_import_fixes_drift.py | 189 +++++++++++++++++++ unsloth/_gpu_init.py | 6 + unsloth/import_fixes.py | 309 +++++++++++++++++++++++++++++++ 3 files changed, 504 insertions(+) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 0bee68f940..f639e7cb61 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -758,3 +758,192 @@ def test_bitsandbytes_rocm_detection_helpers_recognizable(): "decline to patch it and Windows ROCm import-time noise / " "wrong ROCM_GPU_ARCH may return." ) + + +# =========================================================================== +# torchao Windows-ROCm import shim -- fix_torchao_windows_rocm_import +# =========================================================================== +# The shim FRAGMENT-registers the `_c10d_functional` op schemas so real torchao +# imports on a distributed-less Windows ROCm wheel. These verify the schema table +# tracks the installed torch, the FRAGMENT (not DEF) collision semantics, the +# strict no-op / capability gating, and that it is wired into startup. Windows ROCm +# itself cannot be reproduced here, so the transactional acceptance-import + rollback +# is what guarantees no regression on the real device. + + +def _torch_minor_tuple(): + import torch + + base = torch.__version__.split("+", 1)[0].split(".") + return (int(base[0]), int(base[1])) + + +def _live_c10d_functional_ops(): + import torch + + get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None) + if not callable(get_ops): + pytest.skip("dispatcher op enumeration unavailable") + return sorted( + {n.split("::", 1)[1] for n in get_ops() if n.startswith("_c10d_functional::")} + ) + + +def test_torchao_rocm_shim_schema_table_matches_installed_torch(): + """The `_c10d_functional` schema table must exactly match the ops the installed + torch registers (op set + canonical schema strings). A minor with no row means the + shim fail-closes there (safe, no coverage) -> skip; a present row must be exact.""" + from unsloth.import_fixes import _C10D_FUNCTIONAL_SCHEMAS, _schema_op_name + + import torch + + native = _live_c10d_functional_ops() + if not native: + pytest.skip("no native _c10d_functional ops (distributed-less torch build).") + + minor = _torch_minor_tuple() + schemas = _C10D_FUNCTIONAL_SCHEMAS.get(minor) + if schemas is None: + pytest.skip( + f"no shim schema row for torch {minor}; fix_torchao_windows_rocm_import " + f"fail-closes here (safe). Add a reviewed tuple to enable it (ops: {native})." + ) + + table_ops = sorted(_schema_op_name(s) for s in schemas) + assert table_ops == native, ( + f"DRIFT DETECTED: torchao shim _c10d_functional table for torch {minor} lists " + f"{table_ops} but the installed torch registers {native}. Update " + f"_C10D_FUNCTIONAL_SCHEMAS." + ) + + parse = getattr(torch._C, "parse_schema", None) + if not callable(parse): + return + real = {} + for op in native: + packet = getattr(torch.ops._c10d_functional, op) + overload = packet.overloads()[0] + real[op] = str(getattr(packet, overload)._schema) + for s in schemas: + parsed = parse(f"_c10d_functional::{s}") # must not raise + name = _schema_op_name(s) + assert str(parsed) == real[name], ( + f"DRIFT DETECTED: torchao shim schema for _c10d_functional::{name}\n" + f" shim: {parsed}\n torch: {real[name]}" + ) + + +def test_torchao_rocm_shim_strict_noop_on_non_windows(): + """On a non-Windows / distributed-present box the shim must not touch sys.modules, + torch.ops, or torch.distributed.is_available().""" + from unsloth.import_fixes import fix_torchao_windows_rocm_import + + import torch + + assert sys.platform != "win32" + before_ext = "torch._C._distributed_c10d" in sys.modules + before_avail = torch.distributed.is_available() + before_ops = set(_live_c10d_functional_ops()) + + fix_torchao_windows_rocm_import() + + assert ("torch._C._distributed_c10d" in sys.modules) == before_ext + assert torch.distributed.is_available() == before_avail + assert set(_live_c10d_functional_ops()) == before_ops + + +def test_torchao_rocm_shim_native_present_builds_no_library(monkeypatch): + """Even with the platform gates spoofed to look like Windows ROCm, a box that already + has real distributed (native _c10d_functional ops / is_available) must trip a guard + before any torch.library.Library is constructed.""" + from unsloth.import_fixes import ( + fix_torchao_windows_rocm_import, + _native_c10d_functional_present, + ) + + import torch + + monkeypatch.setattr(sys, "platform", "win32") + if not getattr(getattr(torch, "version", None), "hip", None): + monkeypatch.setattr(torch.version, "hip", "6.4.0", raising=False) + + assert _native_c10d_functional_present(torch) is True + + calls = {"n": 0} + real_library = torch.library.Library + + def _tripwire(*args, **kwargs): + calls["n"] += 1 + return real_library(*args, **kwargs) + + monkeypatch.setattr(torch.library, "Library", _tripwire) + fix_torchao_windows_rocm_import() + assert calls["n"] == 0, ( + "torchao shim constructed a torch.library.Library despite a real distributed build." + ) + + +_TORCHAO_ROCM_FRAGMENT_PROBE = """ +import torch +ns = "_unsloth_torchao_shim_probe_ns" +# A second DEF on a namespace raises; FRAGMENT must not -- that is why the shim uses +# FRAGMENT (no fatal collision with a native C++ TORCH_LIBRARY). NB: the first DEF must be +# held by a strong ref, else CPython GCs it (its __del__ calls _destroy) and releases the +# namespace before the second call -- the same reason the shim keeps a strong ref to its +# FRAGMENT Library so its registered schemas are not dropped. +_hold = torch.library.Library(ns, "DEF") +raised = False +try: + torch.library.Library(ns, "DEF") +except Exception: + raised = True +assert raised, "second DEF unexpectedly did not raise" +frag = torch.library.Library(ns, "FRAGMENT") # must not raise +frag.define("myop(Tensor x) -> Tensor") +assert hasattr(torch.ops, ns) and hasattr(getattr(torch.ops, ns), "myop") +frag._destroy() +print("FRAGMENT_OK") +""" + + +def test_torchao_rocm_shim_fragment_semantics_subprocess(): + """FRAGMENT-define + resolve + _destroy work and FRAGMENT (unlike a second DEF) never + collides -- run in a subprocess since dispatcher registration is process-global.""" + import subprocess + + result = subprocess.run( + [sys.executable, "-c", _TORCHAO_ROCM_FRAGMENT_PROBE], + capture_output=True, + text=True, + timeout=300, + ) + assert "FRAGMENT_OK" in result.stdout, ( + f"FRAGMENT probe failed:\nSTDOUT:{result.stdout}\nSTDERR:{result.stderr}" + ) + + +def test_torchao_rocm_shim_source_has_guards_fragment_and_rollback(): + """The shim source must keep its win32 + HIP + is_available() gates, use FRAGMENT (not a + DEF on _c10d_functional), and roll back via Library._destroy.""" + import inspect + + from unsloth import import_fixes + + src = inspect.getsource(import_fixes.fix_torchao_windows_rocm_import) + assert "win32" in src, "missing win32 guard" + assert "hip" in src, "missing torch.version.hip guard" + assert "is_available()" in src, "missing is_available() guard" + assert '"FRAGMENT"' in src, "shim must register with FRAGMENT, not DEF" + assert '"_c10d_functional", "DEF"' not in src, "shim must never DEF _c10d_functional" + assert "_destroy" in src, "missing rollback via Library._destroy" + + +def test_torchao_rocm_shim_wired_into_gpu_init(): + """The shim must be called at startup (before `import unsloth_zoo`), not merely + importable (mirrors test_accelerate_patch_wired_into_gpu_init).""" + source = Path(__file__).resolve().parent.parent / "unsloth" / "_gpu_init.py" + text = source.read_text() + assert "fix_torchao_windows_rocm_import()" in text, ( + "DRIFT DETECTED: fix_torchao_windows_rocm_import is defined but never called in " + "_gpu_init.py, so real imports never install it." + ) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index 984057e9f7..ec14af82e9 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -32,6 +32,7 @@ from .import_fixes import ( disable_broken_vllm, configure_amdgpu_asic_id_table_path, fix_bitsandbytes_rocm_arch_detection, + fix_torchao_windows_rocm_import, torchvision_compatibility_check, fix_diffusers_warnings, fix_huggingface_hub, @@ -70,6 +71,10 @@ except Exception: configure_amdgpu_asic_id_table_path() # Must precede `import unsloth_zoo` below, which imports bnb on ROCm. fix_bitsandbytes_rocm_arch_detection() +# Must also precede `import unsloth_zoo` below (it triggers the transformers/torchao +# import chain): makes real torchao importable on legacy Windows ROCm wheels so +# unsloth_zoo's torchao stub self-disables. Strict no-op elsewhere. +fix_torchao_windows_rocm_import() disable_broken_causal_conv1d() disable_broken_vllm() fix_message_factory_issue() @@ -80,6 +85,7 @@ fix_diffusers_warnings() fix_huggingface_hub() del configure_amdgpu_asic_id_table_path del fix_bitsandbytes_rocm_arch_detection +del fix_torchao_windows_rocm_import del disable_broken_causal_conv1d del disable_broken_vllm del fix_message_factory_issue diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 09de248c7b..64cb66419c 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -3129,3 +3129,312 @@ def patch_accelerate_recursively_apply(): setattr(mod, "find_device", _patched_find_device) except Exception: pass + + +# --------------------------------------------------------------------------- +# torchao Windows-ROCm import shim +# --------------------------------------------------------------------------- +# Legacy Windows ROCm PyTorch wheels ship without the torch.distributed +# C-extension (torch._C._distributed_c10d absent, the torch.ops._c10d_functional.* +# collective ops unregistered). torchao imports the distributed chain +# unconditionally at module load -- torchao/float8/distributed_utils.py does +# `import torch.distributed._functional_collectives` + `from torch.distributed._tensor +# import DTensor`, reached from ~6 float8/dtypes/optim files -- so `import torchao` +# (pulled in by transformers.quantizers) raises `No module named +# 'torch._C._distributed_c10d'` even for paths that never touch distributed. +# unsloth_zoo/Studio work around this by import-stubbing torchao off, which disables +# portable FP8/INT8 export. This shim instead makes REAL torchao importable by faking +# the absent C-extension module and FRAGMENT-registering the missing `_c10d_functional` +# op schemas (schema-only, no kernels -- torch attaches its own Meta kernels; an actual +# collective would still fail loudly). Weight-only quant export invokes no collective, +# so world_size==1 semantics are correct. +# +# Strictly transactional: any failure rolls back so torchao stays unimportable and +# unsloth_zoo's stub still catches it (no regression). Capability-gated, so it is a +# strict no-op on every non-Windows / non-ROCm host and once real torch.distributed is +# present (AMD's libuv/GLOO wheels, ROCm/TheRock#5694, torch >= 2.9). Uses FRAGMENT (not +# DEF), which is documented to bypass the one-library-per-namespace rule, so it can never +# hard-collide with a native TORCH_LIBRARY. Opt out: UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM=1. +# The clean upstream fix is a torch.distributed.is_available() guard in torchao's float8 +# imports (pytorch/ao#1066); retire this shim once that lands. +# --------------------------------------------------------------------------- + +_TORCHAO_ROCM_SHIM_SENTINEL = "__unsloth_torchao_rocm_shim__" +_C10D_EXT_MODULE = "torch._C._distributed_c10d" +# (fake_module, [Library, ...]) after a successful install; retained for the process so +# GC does not drop the FRAGMENT-defined schemas. +_TORCHAO_ROCM_SHIM_STATE = None + +# Per torch (major, minor): the exact `_c10d_functional` op schemas WITHOUT the namespace +# prefix (the Library adds it). torch DEFs these only in C++ (TORCH_LIBRARY), so they are +# absent on a distributed-less ROCm wheel and torch's own +# `torch.distributed._functional_collectives` Library("_c10d_functional","IMPL").impl(...) +# fails at import. Verified: 2.9 against the installed dispatcher; 2.10/2.11 against the +# v2.10.0/v2.11.0 torch/csrc/distributed/c10d/Functional.cpp source. Fail closed on any +# other minor (the shim then no-ops -> torchao stays stubbed, no regression); newer wheels +# almost always carry AMD's distributed fix, where the shim no-ops anyway. +_C10D_FUNCTIONAL_SCHEMAS = { + (2, 9): ( + "all_reduce(Tensor input, str reduce_op, str group_name) -> Tensor", + "all_reduce_(Tensor(a!) input, str reduce_op, str group_name) -> Tensor(a!)", + "all_reduce_coalesced(Tensor[] inputs, str reduce_op, str group_name) -> Tensor[]", + "all_reduce_coalesced_(Tensor[](a!) inputs, str reduce_op, str group_name) -> Tensor[](a!)", + "wait_tensor(Tensor tensor) -> Tensor", + "all_gather_into_tensor(Tensor input, int group_size, str group_name) -> Tensor", + "all_gather_into_tensor_out(Tensor input, int group_size, str group_name, *, Tensor(a!) out) -> Tensor(a!)", + "all_gather_into_tensor_coalesced(Tensor[] inputs, int group_size, str group_name) -> Tensor[]", + "reduce_scatter_tensor(Tensor input, str reduce_op, int group_size, str group_name) -> Tensor", + "reduce_scatter_tensor_coalesced(Tensor[] inputs, str reduce_op, int group_size, str group_name) -> Tensor[]", + "all_to_all_single(Tensor input, SymInt[] output_split_sizes, SymInt[] input_split_sizes, str group_name) -> Tensor", + "broadcast(Tensor input, int src, str group_name) -> Tensor", + "broadcast_(Tensor(a!) input, int src, str group_name) -> Tensor(a!)", + ), +} +# 2.10 and 2.11 add reduce_scatter_tensor_out; the other 13 schemas are unchanged (str +# group_name), verified against the v2.10.0 / v2.11.0 Functional.cpp source. +_C10D_FUNCTIONAL_SCHEMAS[(2, 10)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 9)] + ( + "reduce_scatter_tensor_out(Tensor input, str reduce_op, int group_size, str group_name, *, Tensor(a!) out) -> Tensor(a!)", +) +_C10D_FUNCTIONAL_SCHEMAS[(2, 11)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 10)] + + +def _schema_op_name(schema): + """`all_reduce(Tensor ...) -> Tensor` -> `all_reduce`.""" + return schema.split("(", 1)[0].strip() + + +def _torchao_shim_torch_minor(torch): + base = torch.__version__.split("+", 1)[0].split(".") + return (int(base[0]), int(base[1])) + + +def _native_c10d_functional_present(torch): + """True if the dispatcher already has any `_c10d_functional::` op (real torch + distributed present). Fail closed: an unexpected error counts as present, so the shim + never registers over a real namespace.""" + get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None) + if not callable(get_ops): + return True + try: + return any(n.startswith("_c10d_functional::") for n in get_ops()) + except Exception: + return True + + +_TORCHAO_ROCM_LEGACY_ERR = ( + "Unsloth: torch.distributed is unavailable on this legacy Windows ROCm build; this is " + "an import-only compatibility shim for torchao and cannot perform real distributed / " + "collective work. Upgrade to a Windows ROCm PyTorch wheel built with GLOO " + "(ROCm/TheRock#5694)." +) + +# Optional transport backends: must stay ABSENT so torch's guarded +# `from torch._C._distributed_c10d import ProcessGroupX` probes conclude the backend is +# unavailable (they sit behind try/except or capability checks). +_TORCHAO_ROCM_OPTIONAL_BACKENDS = frozenset({ + "ProcessGroupNCCL", "ProcessGroupGloo", "ProcessGroupMPI", "ProcessGroupUCC", + "ProcessGroupXCCL", "_ProcessGroupWrapper", "_c10d_init", +}) + + +class _TorchaoRocmSentinelMeta(type): + """Metaclass for import-only sentinel types: isinstance-safe (never matches a real + object), chainable via attribute access, and loud on construction (an accidental + runtime use raises rather than silently returning a wrong result).""" + + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, name): + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + child = _TorchaoRocmSentinelMeta(f"{cls.__name__}.{name}", (), {}) + setattr(cls, name, child) + return child + + def __call__(cls, *args, **kwargs): + raise RuntimeError(_TORCHAO_ROCM_LEGACY_ERR) + + +def _torchao_rocm_sentinel(name): + return _TorchaoRocmSentinelMeta(str(name), (), {}) + + +def _make_torchao_rocm_fake_c10d(): + """Build a fake `torch._C._distributed_c10d` so torchao's unconditional + `from torch._C._distributed_c10d import (...)` (via torch.distributed.distributed_c10d) + resolves. Explicit semantic fakes for the names read while torch's distributed Python + modules initialize; a bounded-dynamic `__getattr__` gives isinstance-safe loud sentinel + types for unknown import-only names but raises AttributeError for dunders and optional + backends. No `_c10d_init` and is_available() stays False: device_mesh.py self-stubs when + unavailable and the crash site imports the C-ext unconditionally, so the fake alone + unblocks the import.""" + import types + + mod = types.ModuleType(_C10D_EXT_MODULE) + mod.__package__ = "torch._C" + setattr(mod, _TORCHAO_ROCM_SHIM_SENTINEL, True) + + # Data-holder option types: constructible no-ops (safe to build at import; never used + # to do work). Kept distinct from loud sentinels so import-time construction cannot raise. + for name in ( + "_DistributedBackendOptions", "AllgatherOptions", "AllreduceCoalescedOptions", + "AllreduceOptions", "AllToAllOptions", "BarrierOptions", "BroadcastOptions", + "GatherOptions", "ReduceOptions", "ReduceScatterOptions", "ScatterOptions", + ): + setattr(mod, name, type(name, (), {"__init__": lambda self, *a, **k: None})) + for name in ("PrefixStore", "Store", "HashStore", "Work"): + setattr(mod, name, _torchao_rocm_sentinel(name)) + for name in ("_register_process_group", "_resolve_process_group", + "_unregister_all_process_groups", "_unregister_process_group"): + def _loud(*a, **k): + raise RuntimeError(_TORCHAO_ROCM_LEGACY_ERR) + setattr(mod, name, _loud) + + ReduceOp = _torchao_rocm_sentinel("ReduceOp") + _reduce_members = ("SUM", "AVG", "PRODUCT", "MIN", "MAX", "BAND", "BOR", "BXOR", + "PREMUL_SUM", "UNUSED") + for m in _reduce_members: + setattr(ReduceOp, m, _torchao_rocm_sentinel(f"ReduceOp.{m}")) + RedOpType = _torchao_rocm_sentinel("ReduceOp.RedOpType") + RedOpType.__members__ = {m: getattr(ReduceOp, m) for m in _reduce_members} + ReduceOp.RedOpType = RedOpType + mod.ReduceOp = ReduceOp + + ProcessGroup = _torchao_rocm_sentinel("ProcessGroup") + BackendType = _torchao_rocm_sentinel("ProcessGroup.BackendType") + for m in ("UNDEFINED", "GLOO", "NCCL", "UCC", "MPI", "XCCL", "CUSTOM"): + setattr(BackendType, m, _torchao_rocm_sentinel(f"ProcessGroup.BackendType.{m}")) + ProcessGroup.BackendType = BackendType + mod.ProcessGroup = ProcessGroup + + DebugLevel = _torchao_rocm_sentinel("DebugLevel") + for m in ("OFF", "INFO", "DETAIL"): + setattr(DebugLevel, m, _torchao_rocm_sentinel(f"DebugLevel.{m}")) + mod.DebugLevel = DebugLevel + mod.get_debug_level = lambda *a, **k: DebugLevel.OFF + + _dynamic_cache = {} + + def _module_getattr(name): + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + if name in _TORCHAO_ROCM_OPTIONAL_BACKENDS: + raise AttributeError(name) + if name in _dynamic_cache: + return _dynamic_cache[name] + child = _torchao_rocm_sentinel(name) + _dynamic_cache[name] = child + return child + + mod.__getattr__ = _module_getattr + return mod + + +def fix_torchao_windows_rocm_import(): + """On a legacy Windows ROCm wheel (no torch.distributed C-extension), make real torchao + importable by faking `torch._C._distributed_c10d` and FRAGMENT-registering the + `_c10d_functional` op schemas, so torchao's module-top distributed imports resolve and + portable FP8/INT8 export works instead of torchao being stubbed off. + + Strict no-op unless every capability guard holds (Windows + HIP torch + distributed + genuinely absent + known torch minor + torchao installed and not yet imported). Fully + transactional: any failure rolls back and leaves torchao unimportable, so unsloth_zoo's + stub still catches it (no regression). Idempotent. Opt out: + UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM=1.""" + global _TORCHAO_ROCM_SHIM_STATE + + if os.environ.get("UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM", "0") == "1": + return + if sys.platform != "win32": + return + + with _get_torchao_rocm_shim_lock(): + if _TORCHAO_ROCM_SHIM_STATE is not None: + return # already installed this process + try: + import torch + + # ROCm build only (authoritative runtime HIP field, not the loose version tag). + if not getattr(getattr(torch, "version", None), "hip", None): + return + schemas = _C10D_FUNCTIONAL_SCHEMAS.get(_torchao_shim_torch_minor(torch)) + if schemas is None: + return # unknown torch minor -> fail closed + if importlib.util.find_spec("torchao") is None: + return + if any(n == "torchao" or n.startswith("torchao.") for n in list(sys.modules)): + return # real- or stub-imported already; cannot safely convert + dist = getattr(torch, "distributed", None) + if dist is None or dist.is_available() is not False: + return # real distributed present (fixed wheel) -> nothing to do + if hasattr(torch._C, "_c10d_init"): + return + if hasattr(torch._C, "_distributed_c10d") or _C10D_EXT_MODULE in sys.modules: + return + if _native_c10d_functional_present(torch): + return # dispatcher already has the ops -> never register over them + except Exception: + return + + # ---- transaction: snapshot -> register -> acceptance import -> commit/rollback --- + modules_before = set(sys.modules) + had_c10d_attr = hasattr(torch._C, "_distributed_c10d") + fake = None + lib = None + try: + # Re-check the dispatcher immediately before touching it (TOCTOU guard). + if _native_c10d_functional_present(torch): + return + lib = torch.library.Library("_c10d_functional", "FRAGMENT") # FRAGMENT, never DEF + for schema in schemas: + lib.define(schema) + fake = _make_torchao_rocm_fake_c10d() + sys.modules[_C10D_EXT_MODULE] = fake + setattr(torch._C, "_distributed_c10d", fake) + # Acceptance test: the real torchao must import end to end. + importlib.import_module("torchao") + except BaseException: + # Atomic rollback: destroy schemas, drop the fake, purge only the torchao / + # distributed submodules this transaction newly created. + if lib is not None: + try: + lib._destroy() + except Exception: + pass + if fake is not None: + if not had_c10d_attr and getattr(torch._C, "_distributed_c10d", None) is fake: + try: + delattr(torch._C, "_distributed_c10d") + except Exception: + pass + if sys.modules.get(_C10D_EXT_MODULE) is fake: + del sys.modules[_C10D_EXT_MODULE] + for name in [ + n for n in set(sys.modules) - modules_before + if n == "torchao" or n.startswith("torchao.") + or n == "torch.distributed" or n.startswith("torch.distributed.") + ]: + sys.modules.pop(name, None) + return + + # Commit: keep strong refs so the FRAGMENT schemas outlive GC. + _TORCHAO_ROCM_SHIM_STATE = {"fake_module": fake, "libraries": [lib]} + _log_rocm_detection( + "Unsloth: Installed the torchao Windows-ROCm import shim " + "(fake torch._C._distributed_c10d + _c10d_functional schemas)." + ) + + +_torchao_rocm_shim_lock = None + + +def _get_torchao_rocm_shim_lock(): + global _torchao_rocm_shim_lock + if _torchao_rocm_shim_lock is None: + import threading + _torchao_rocm_shim_lock = threading.RLock() + return _torchao_rocm_shim_lock From 7c289ce07f1fd3c4f92793dece7c69012c6444e2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 14:17:12 +0000 Subject: [PATCH 6/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_import_fixes_drift.py | 25 ++++++------- unsloth/import_fixes.py | 62 ++++++++++++++++++++++++-------- 2 files changed, 59 insertions(+), 28 deletions(-) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index f639e7cb61..7faabf851b 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -773,7 +773,6 @@ def test_bitsandbytes_rocm_detection_helpers_recognizable(): def _torch_minor_tuple(): import torch - base = torch.__version__.split("+", 1)[0].split(".") return (int(base[0]), int(base[1])) @@ -784,9 +783,7 @@ def _live_c10d_functional_ops(): get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None) if not callable(get_ops): pytest.skip("dispatcher op enumeration unavailable") - return sorted( - {n.split("::", 1)[1] for n in get_ops() if n.startswith("_c10d_functional::")} - ) + return sorted({n.split("::", 1)[1] for n in get_ops() if n.startswith("_c10d_functional::")}) def test_torchao_rocm_shim_schema_table_matches_installed_torch(): @@ -865,7 +862,7 @@ def test_torchao_rocm_shim_native_present_builds_no_library(monkeypatch): monkeypatch.setattr(sys, "platform", "win32") if not getattr(getattr(torch, "version", None), "hip", None): - monkeypatch.setattr(torch.version, "hip", "6.4.0", raising=False) + monkeypatch.setattr(torch.version, "hip", "6.4.0", raising = False) assert _native_c10d_functional_present(torch) is True @@ -878,9 +875,9 @@ def test_torchao_rocm_shim_native_present_builds_no_library(monkeypatch): monkeypatch.setattr(torch.library, "Library", _tripwire) fix_torchao_windows_rocm_import() - assert calls["n"] == 0, ( - "torchao shim constructed a torch.library.Library despite a real distributed build." - ) + assert ( + calls["n"] == 0 + ), "torchao shim constructed a torch.library.Library despite a real distributed build." _TORCHAO_ROCM_FRAGMENT_PROBE = """ @@ -913,13 +910,13 @@ def test_torchao_rocm_shim_fragment_semantics_subprocess(): result = subprocess.run( [sys.executable, "-c", _TORCHAO_ROCM_FRAGMENT_PROBE], - capture_output=True, - text=True, - timeout=300, - ) - assert "FRAGMENT_OK" in result.stdout, ( - f"FRAGMENT probe failed:\nSTDOUT:{result.stdout}\nSTDERR:{result.stderr}" + capture_output = True, + text = True, + timeout = 300, ) + assert ( + "FRAGMENT_OK" in result.stdout + ), f"FRAGMENT probe failed:\nSTDOUT:{result.stdout}\nSTDERR:{result.stderr}" def test_torchao_rocm_shim_source_has_guards_fragment_and_rollback(): diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 64cb66419c..d96c2065a7 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -3231,10 +3231,17 @@ _TORCHAO_ROCM_LEGACY_ERR = ( # Optional transport backends: must stay ABSENT so torch's guarded # `from torch._C._distributed_c10d import ProcessGroupX` probes conclude the backend is # unavailable (they sit behind try/except or capability checks). -_TORCHAO_ROCM_OPTIONAL_BACKENDS = frozenset({ - "ProcessGroupNCCL", "ProcessGroupGloo", "ProcessGroupMPI", "ProcessGroupUCC", - "ProcessGroupXCCL", "_ProcessGroupWrapper", "_c10d_init", -}) +_TORCHAO_ROCM_OPTIONAL_BACKENDS = frozenset( + { + "ProcessGroupNCCL", + "ProcessGroupGloo", + "ProcessGroupMPI", + "ProcessGroupUCC", + "ProcessGroupXCCL", + "_ProcessGroupWrapper", + "_c10d_init", + } +) class _TorchaoRocmSentinelMeta(type): @@ -3281,22 +3288,46 @@ def _make_torchao_rocm_fake_c10d(): # Data-holder option types: constructible no-ops (safe to build at import; never used # to do work). Kept distinct from loud sentinels so import-time construction cannot raise. for name in ( - "_DistributedBackendOptions", "AllgatherOptions", "AllreduceCoalescedOptions", - "AllreduceOptions", "AllToAllOptions", "BarrierOptions", "BroadcastOptions", - "GatherOptions", "ReduceOptions", "ReduceScatterOptions", "ScatterOptions", + "_DistributedBackendOptions", + "AllgatherOptions", + "AllreduceCoalescedOptions", + "AllreduceOptions", + "AllToAllOptions", + "BarrierOptions", + "BroadcastOptions", + "GatherOptions", + "ReduceOptions", + "ReduceScatterOptions", + "ScatterOptions", ): setattr(mod, name, type(name, (), {"__init__": lambda self, *a, **k: None})) for name in ("PrefixStore", "Store", "HashStore", "Work"): setattr(mod, name, _torchao_rocm_sentinel(name)) - for name in ("_register_process_group", "_resolve_process_group", - "_unregister_all_process_groups", "_unregister_process_group"): + for name in ( + "_register_process_group", + "_resolve_process_group", + "_unregister_all_process_groups", + "_unregister_process_group", + ): + def _loud(*a, **k): raise RuntimeError(_TORCHAO_ROCM_LEGACY_ERR) + setattr(mod, name, _loud) ReduceOp = _torchao_rocm_sentinel("ReduceOp") - _reduce_members = ("SUM", "AVG", "PRODUCT", "MIN", "MAX", "BAND", "BOR", "BXOR", - "PREMUL_SUM", "UNUSED") + _reduce_members = ( + "SUM", + "AVG", + "PRODUCT", + "MIN", + "MAX", + "BAND", + "BOR", + "BXOR", + "PREMUL_SUM", + "UNUSED", + ) for m in _reduce_members: setattr(ReduceOp, m, _torchao_rocm_sentinel(f"ReduceOp.{m}")) RedOpType = _torchao_rocm_sentinel("ReduceOp.RedOpType") @@ -3414,9 +3445,12 @@ def fix_torchao_windows_rocm_import(): if sys.modules.get(_C10D_EXT_MODULE) is fake: del sys.modules[_C10D_EXT_MODULE] for name in [ - n for n in set(sys.modules) - modules_before - if n == "torchao" or n.startswith("torchao.") - or n == "torch.distributed" or n.startswith("torch.distributed.") + n + for n in set(sys.modules) - modules_before + if n == "torchao" + or n.startswith("torchao.") + or n == "torch.distributed" + or n.startswith("torch.distributed.") ]: sys.modules.pop(name, None) return From 53e178e150054edb57b6fc313c2990d320d2c52a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 14:45:10 +0000 Subject: [PATCH 7/9] Fix torchao Windows-ROCm shim: define _dtensor schema and match is_win32_rocm detection Address two review findings on the shim. torchao's `from torch.distributed._tensor import DTensor` loads torch.distributed.tensor._collective_utils, which runs `register_fake("_dtensor::shard_dim_alltoall")` at import time. That op is DEF'd only in C++ (Functional.cpp), so it is absent on distributed-less ROCm wheels and register_fake raises "operator _dtensor::shard_dim_alltoall does not exist", failing the acceptance `import torchao` and rolling the shim back. Define the _dtensor schema via a second FRAGMENT library in the same transaction. Schema verified against the live 2.9 dispatcher and the v2.11.0 Functional.cpp source. Match the Studio is_win32_rocm() detector: fire on torch.version.hip OR a "rocm" __version__ tag (AMD SDK wheels lack torch.version.hip but tag rocm). Gating on hip alone left those wheels torchao-stubbed even though the export gate disables torchao on them. The capability guards (distributed absent, no _c10d_functional/_dtensor dispatcher ops) still rule out false positives. Extend the native-present guard to the _dtensor namespace and add drift tests pinning the _dtensor schema to the live dispatcher and the version-tag detection. --- tests/test_import_fixes_drift.py | 59 ++++++++++++++++++++++++++ unsloth/import_fixes.py | 71 ++++++++++++++++++++++++-------- 2 files changed, 113 insertions(+), 17 deletions(-) diff --git a/tests/test_import_fixes_drift.py b/tests/test_import_fixes_drift.py index 7faabf851b..074e3c90fe 100644 --- a/tests/test_import_fixes_drift.py +++ b/tests/test_import_fixes_drift.py @@ -786,6 +786,15 @@ def _live_c10d_functional_ops(): return sorted({n.split("::", 1)[1] for n in get_ops() if n.startswith("_c10d_functional::")}) +def _live_dtensor_ops(): + import torch + + get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None) + if not callable(get_ops): + pytest.skip("dispatcher op enumeration unavailable") + return sorted({n.split("::", 1)[1] for n in get_ops() if n.startswith("_dtensor::")}) + + def test_torchao_rocm_shim_schema_table_matches_installed_torch(): """The `_c10d_functional` schema table must exactly match the ops the installed torch registers (op set + canonical schema strings). A minor with no row means the @@ -830,6 +839,50 @@ def test_torchao_rocm_shim_schema_table_matches_installed_torch(): ) +def test_torchao_rocm_shim_dtensor_schema_matches_installed_torch(): + """The `_dtensor` schema table must exactly match the ops the installed torch registers. + torchao's `from torch.distributed._tensor import DTensor` runs + `register_fake("_dtensor::shard_dim_alltoall")` at import, which raises unless the op is + defined, so the shim must define this namespace too (not only _c10d_functional).""" + from unsloth.import_fixes import _DTENSOR_SCHEMAS, _schema_op_name + + import torch + + native = _live_dtensor_ops() + if not native: + pytest.skip("no native _dtensor ops (distributed-less torch build).") + + minor = _torch_minor_tuple() + schemas = _DTENSOR_SCHEMAS.get(minor) + if schemas is None: + pytest.skip( + f"no shim _dtensor row for torch {minor}; fix_torchao_windows_rocm_import " + f"fail-closes here (safe). Add a reviewed tuple to enable it (ops: {native})." + ) + + table_ops = sorted(_schema_op_name(s) for s in schemas) + assert table_ops == native, ( + f"DRIFT DETECTED: torchao shim _dtensor table for torch {minor} lists {table_ops} " + f"but the installed torch registers {native}. Update _DTENSOR_SCHEMAS." + ) + + parse = getattr(torch._C, "parse_schema", None) + if not callable(parse): + return + real = {} + for op in native: + packet = getattr(torch.ops._dtensor, op) + overload = packet.overloads()[0] + real[op] = str(getattr(packet, overload)._schema) + for s in schemas: + parsed = parse(f"_dtensor::{s}") # must not raise + name = _schema_op_name(s) + assert str(parsed) == real[name], ( + f"DRIFT DETECTED: torchao shim schema for _dtensor::{name}\n" + f" shim: {parsed}\n torch: {real[name]}" + ) + + def test_torchao_rocm_shim_strict_noop_on_non_windows(): """On a non-Windows / distributed-present box the shim must not touch sys.modules, torch.ops, or torch.distributed.is_available().""" @@ -929,9 +982,15 @@ def test_torchao_rocm_shim_source_has_guards_fragment_and_rollback(): src = inspect.getsource(import_fixes.fix_torchao_windows_rocm_import) assert "win32" in src, "missing win32 guard" assert "hip" in src, "missing torch.version.hip guard" + assert '"rocm" in' in src, ( + "shim ROCm detection must also accept a 'rocm'-tagged __version__ wheel (parity with " + "the Studio is_win32_rocm() helper), not gate on torch.version.hip alone" + ) assert "is_available()" in src, "missing is_available() guard" assert '"FRAGMENT"' in src, "shim must register with FRAGMENT, not DEF" assert '"_c10d_functional", "DEF"' not in src, "shim must never DEF _c10d_functional" + assert '"_dtensor", "DEF"' not in src, "shim must never DEF _dtensor" + assert '"_dtensor"' in src, "shim must also register the _dtensor namespace" assert "_destroy" in src, "missing rollback via Library._destroy" diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index d96c2065a7..39618f90be 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -3197,6 +3197,22 @@ _C10D_FUNCTIONAL_SCHEMAS[(2, 10)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 9)] + ( ) _C10D_FUNCTIONAL_SCHEMAS[(2, 11)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 10)] +# The `_dtensor` namespace is likewise DEF'd only in C++ (the same Functional.cpp), so it too +# is absent on a distributed-less ROCm wheel. torchao's `from torch.distributed._tensor import +# DTensor` loads torch.distributed.tensor._collective_utils, which at import does +# `@torch.library.register_fake("_dtensor::shard_dim_alltoall")`; register_fake raises +# "operator _dtensor::shard_dim_alltoall does not exist" unless the op is already defined, so +# the shim must define it in the same transaction or `import torchao` still fails and rolls +# back. Schema verified against the live 2.9 dispatcher and the v2.11.0 Functional.cpp source +# (stable across 2.9-2.11). Fail closed on any other minor. +_DTENSOR_SCHEMAS = { + (2, 9): ( + "shard_dim_alltoall(Tensor input, int gather_dim, int shard_dim, str group_name) -> Tensor", + ), +} +_DTENSOR_SCHEMAS[(2, 10)] = _DTENSOR_SCHEMAS[(2, 9)] +_DTENSOR_SCHEMAS[(2, 11)] = _DTENSOR_SCHEMAS[(2, 9)] + def _schema_op_name(schema): """`all_reduce(Tensor ...) -> Tensor` -> `all_reduce`.""" @@ -3208,15 +3224,19 @@ def _torchao_shim_torch_minor(torch): return (int(base[0]), int(base[1])) +_TORCHAO_ROCM_NATIVE_PREFIXES = ("_c10d_functional::", "_dtensor::") + + def _native_c10d_functional_present(torch): - """True if the dispatcher already has any `_c10d_functional::` op (real torch - distributed present). Fail closed: an unexpected error counts as present, so the shim - never registers over a real namespace.""" + """True if the dispatcher already has any `_c10d_functional::` or `_dtensor::` op (real + torch distributed present). Both namespaces are DEF'd in the same C++ Functional.cpp, so + either being present means the shim must not register over it. Fail closed: an unexpected + error counts as present, so the shim never registers over a real namespace.""" get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None) if not callable(get_ops): return True try: - return any(n.startswith("_c10d_functional::") for n in get_ops()) + return any(n.startswith(_TORCHAO_ROCM_NATIVE_PREFIXES) for n in get_ops()) except Exception: return True @@ -3368,8 +3388,8 @@ def _make_torchao_rocm_fake_c10d(): def fix_torchao_windows_rocm_import(): """On a legacy Windows ROCm wheel (no torch.distributed C-extension), make real torchao importable by faking `torch._C._distributed_c10d` and FRAGMENT-registering the - `_c10d_functional` op schemas, so torchao's module-top distributed imports resolve and - portable FP8/INT8 export works instead of torchao being stubbed off. + `_c10d_functional` and `_dtensor` op schemas, so torchao's module-top distributed imports + resolve and portable FP8/INT8 export works instead of torchao being stubbed off. Strict no-op unless every capability guard holds (Windows + HIP torch + distributed genuinely absent + known torch minor + torchao installed and not yet imported). Fully @@ -3389,11 +3409,20 @@ def fix_torchao_windows_rocm_import(): try: import torch - # ROCm build only (authoritative runtime HIP field, not the loose version tag). - if not getattr(getattr(torch, "version", None), "hip", None): + # ROCm build: mirror the Studio is_win32_rocm() detector -- HIP field OR a "rocm" + # __version__ tag (AMD SDK wheels lack torch.version.hip but tag "rocm"). Matching + # it keeps the shim and the export gate from drifting so the same wheels the gate + # disables torchao on are the ones the shim re-enables. The capability guards below + # (distributed absent, no _c10d_functional/_dtensor ops) rule out any false positive. + if not ( + getattr(getattr(torch, "version", None), "hip", None) + or "rocm" in getattr(torch, "__version__", "").lower() + ): return - schemas = _C10D_FUNCTIONAL_SCHEMAS.get(_torchao_shim_torch_minor(torch)) - if schemas is None: + minor = _torchao_shim_torch_minor(torch) + schemas = _C10D_FUNCTIONAL_SCHEMAS.get(minor) + dtensor_schemas = _DTENSOR_SCHEMAS.get(minor) + if schemas is None or dtensor_schemas is None: return # unknown torch minor -> fail closed if importlib.util.find_spec("torchao") is None: return @@ -3415,14 +3444,22 @@ def fix_torchao_windows_rocm_import(): modules_before = set(sys.modules) had_c10d_attr = hasattr(torch._C, "_distributed_c10d") fake = None - lib = None + libs = [] try: # Re-check the dispatcher immediately before touching it (TOCTOU guard). if _native_c10d_functional_present(torch): return - lib = torch.library.Library("_c10d_functional", "FRAGMENT") # FRAGMENT, never DEF - for schema in schemas: - lib.define(schema) + # FRAGMENT (never DEF): defines the schemas torch's distributed Python modules + # register impls / fakes against at import. `_c10d_functional` for + # _functional_collectives, `_dtensor` for tensor._collective_utils. + for namespace, ns_schemas in ( + ("_c10d_functional", schemas), + ("_dtensor", dtensor_schemas), + ): + lib = torch.library.Library(namespace, "FRAGMENT") + libs.append(lib) + for schema in ns_schemas: + lib.define(schema) fake = _make_torchao_rocm_fake_c10d() sys.modules[_C10D_EXT_MODULE] = fake setattr(torch._C, "_distributed_c10d", fake) @@ -3431,7 +3468,7 @@ def fix_torchao_windows_rocm_import(): except BaseException: # Atomic rollback: destroy schemas, drop the fake, purge only the torchao / # distributed submodules this transaction newly created. - if lib is not None: + for lib in libs: try: lib._destroy() except Exception: @@ -3456,10 +3493,10 @@ def fix_torchao_windows_rocm_import(): return # Commit: keep strong refs so the FRAGMENT schemas outlive GC. - _TORCHAO_ROCM_SHIM_STATE = {"fake_module": fake, "libraries": [lib]} + _TORCHAO_ROCM_SHIM_STATE = {"fake_module": fake, "libraries": libs} _log_rocm_detection( "Unsloth: Installed the torchao Windows-ROCm import shim " - "(fake torch._C._distributed_c10d + _c10d_functional schemas)." + "(fake torch._C._distributed_c10d + _c10d_functional/_dtensor schemas)." ) From 3be31947dd434be340eb66caefd219230927038f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 14:49:43 +0000 Subject: [PATCH 8/9] Export page: defer format pruning until hardware info is authoritative The availableFormats prune effect ran on the initial render before /api/system/hardware resolves, when hasNvidia is false and the NVIDIA-only compressed-tensors formats are transiently absent. On a fresh Export mount with the module cache empty (cold start, or a remount during refreshHardwareInfo), a running FP8/NVFP4 selection was pruned permanently, since the later hardware response only adds formats back to availableFormats and never restores selectedFormats. Gate the effect on hardware.loaded so it prunes only against the authoritative capability set, matching the effect's stated intent. --- studio/frontend/src/features/export/export-page.tsx | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 8edc8dcb73..42cbff5a32 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -260,13 +260,17 @@ export function ExportPage() { }, []); // 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. + // Gate on hardware.loaded: before the authoritative response hasNvidia is false, so the + // NVIDIA-only compressed formats are transiently absent and pruning here would permanently + // drop a running FP8/NVFP4 selection that the later response cannot restore. useEffect(() => { + if (!hardware.loaded) return; 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]); + }, [availableFormats, hardware.loaded]); // 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, From 774c364508549482fc8359bf698588dfd2c5a8da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 15 Jul 2026 17:53:18 +0000 Subject: [PATCH 9/9] Tighten torchao Windows-ROCm shim comments Comment-only: condense the shim header, the _c10d_functional/_dtensor schema notes, the is_win32_rocm parity note, and the export-page prune comment. No code changes. --- .../src/features/export/export-page.tsx | 8 +- unsloth/import_fixes.py | 75 +++++++------------ 2 files changed, 32 insertions(+), 51 deletions(-) diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 42cbff5a32..61e3efa9fb 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -258,11 +258,9 @@ export function ExportPage() { : [...prev, value], ); }, []); - // 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. - // Gate on hardware.loaded: before the authoritative response hasNvidia is false, so the - // NVIDIA-only compressed formats are transiently absent and pruning here would permanently - // drop a running FP8/NVFP4 selection that the later response cannot restore. + // Drop a selected format the gate removed (e.g. torchao once win32Rocm resolves). Gate on + // hardware.loaded: before the authoritative response hasNvidia is false, so pruning would + // permanently drop a running NVIDIA FP8/NVFP4 pick that the later response can't restore. useEffect(() => { if (!hardware.loaded) return; const allowed = new Set(availableFormats.map((f) => f.value)); diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 39618f90be..e36a074135 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -3134,29 +3134,20 @@ def patch_accelerate_recursively_apply(): # --------------------------------------------------------------------------- # torchao Windows-ROCm import shim # --------------------------------------------------------------------------- -# Legacy Windows ROCm PyTorch wheels ship without the torch.distributed -# C-extension (torch._C._distributed_c10d absent, the torch.ops._c10d_functional.* -# collective ops unregistered). torchao imports the distributed chain -# unconditionally at module load -- torchao/float8/distributed_utils.py does -# `import torch.distributed._functional_collectives` + `from torch.distributed._tensor -# import DTensor`, reached from ~6 float8/dtypes/optim files -- so `import torchao` -# (pulled in by transformers.quantizers) raises `No module named -# 'torch._C._distributed_c10d'` even for paths that never touch distributed. -# unsloth_zoo/Studio work around this by import-stubbing torchao off, which disables -# portable FP8/INT8 export. This shim instead makes REAL torchao importable by faking -# the absent C-extension module and FRAGMENT-registering the missing `_c10d_functional` -# op schemas (schema-only, no kernels -- torch attaches its own Meta kernels; an actual -# collective would still fail loudly). Weight-only quant export invokes no collective, -# so world_size==1 semantics are correct. -# -# Strictly transactional: any failure rolls back so torchao stays unimportable and -# unsloth_zoo's stub still catches it (no regression). Capability-gated, so it is a -# strict no-op on every non-Windows / non-ROCm host and once real torch.distributed is -# present (AMD's libuv/GLOO wheels, ROCm/TheRock#5694, torch >= 2.9). Uses FRAGMENT (not -# DEF), which is documented to bypass the one-library-per-namespace rule, so it can never -# hard-collide with a native TORCH_LIBRARY. Opt out: UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM=1. -# The clean upstream fix is a torch.distributed.is_available() guard in torchao's float8 -# imports (pytorch/ao#1066); retire this shim once that lands. +# Legacy Windows ROCm wheels ship without the torch.distributed C-extension +# (torch._C._distributed_c10d absent, torch.ops._c10d_functional.* unregistered), yet torchao +# imports the distributed chain unconditionally at module load (float8/distributed_utils.py: +# `import torch.distributed._functional_collectives` + `from torch.distributed._tensor import +# DTensor`), so `import torchao` raises `No module named 'torch._C._distributed_c10d'` even on +# paths that never use distributed. unsloth_zoo/Studio work around that by stubbing torchao off, +# disabling FP8/INT8 export; this shim instead makes REAL torchao importable by faking the absent +# C-ext module and FRAGMENT-registering the missing _c10d_functional/_dtensor schemas (schema-only +# -- torch attaches its own Meta kernels, an actual collective still fails loudly, and weight-only +# export invokes none). Fully transactional (any failure rolls back to torchao-unimportable, so +# unsloth_zoo's stub still catches it: no regression) and capability-gated -- a strict no-op off +# Windows-ROCm and once real torch.distributed is present (ROCm/TheRock#5694, torch >= 2.9). +# FRAGMENT (not DEF) never collides with a native TORCH_LIBRARY. Opt out +# UNSLOTH_DISABLE_TORCHAO_ROCM_SHIM=1; retire once torchao guards its float8 imports (pytorch/ao#1066). # --------------------------------------------------------------------------- _TORCHAO_ROCM_SHIM_SENTINEL = "__unsloth_torchao_rocm_shim__" @@ -3165,14 +3156,11 @@ _C10D_EXT_MODULE = "torch._C._distributed_c10d" # GC does not drop the FRAGMENT-defined schemas. _TORCHAO_ROCM_SHIM_STATE = None -# Per torch (major, minor): the exact `_c10d_functional` op schemas WITHOUT the namespace -# prefix (the Library adds it). torch DEFs these only in C++ (TORCH_LIBRARY), so they are -# absent on a distributed-less ROCm wheel and torch's own -# `torch.distributed._functional_collectives` Library("_c10d_functional","IMPL").impl(...) -# fails at import. Verified: 2.9 against the installed dispatcher; 2.10/2.11 against the -# v2.10.0/v2.11.0 torch/csrc/distributed/c10d/Functional.cpp source. Fail closed on any -# other minor (the shim then no-ops -> torchao stays stubbed, no regression); newer wheels -# almost always carry AMD's distributed fix, where the shim no-ops anyway. +# Per torch (major, minor): the exact `_c10d_functional` op schemas (namespace prefix added by +# the Library). torch DEFs these only in C++, so they are absent on a distributed-less ROCm wheel +# and torch's own _functional_collectives IMPL registrations fail at import. Verified: 2.9 vs the +# installed dispatcher, 2.10/2.11 vs the v2.10.0/v2.11.0 Functional.cpp source. Fail closed on any +# other minor (shim no-ops -> torchao stays stubbed, no regression). _C10D_FUNCTIONAL_SCHEMAS = { (2, 9): ( "all_reduce(Tensor input, str reduce_op, str group_name) -> Tensor", @@ -3197,14 +3185,11 @@ _C10D_FUNCTIONAL_SCHEMAS[(2, 10)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 9)] + ( ) _C10D_FUNCTIONAL_SCHEMAS[(2, 11)] = _C10D_FUNCTIONAL_SCHEMAS[(2, 10)] -# The `_dtensor` namespace is likewise DEF'd only in C++ (the same Functional.cpp), so it too -# is absent on a distributed-less ROCm wheel. torchao's `from torch.distributed._tensor import -# DTensor` loads torch.distributed.tensor._collective_utils, which at import does -# `@torch.library.register_fake("_dtensor::shard_dim_alltoall")`; register_fake raises -# "operator _dtensor::shard_dim_alltoall does not exist" unless the op is already defined, so -# the shim must define it in the same transaction or `import torchao` still fails and rolls -# back. Schema verified against the live 2.9 dispatcher and the v2.11.0 Functional.cpp source -# (stable across 2.9-2.11). Fail closed on any other minor. +# `_dtensor` is likewise C++-only, so absent on a distributed-less wheel. torchao's +# `from torch.distributed._tensor import DTensor` loads tensor._collective_utils, whose +# module-level `register_fake("_dtensor::shard_dim_alltoall")` raises unless the op is defined -- +# so the shim must define it in the same transaction or `import torchao` still rolls back. Schema +# verified vs the live 2.9 dispatcher and v2.11.0 Functional.cpp (stable 2.9-2.11); else fail closed. _DTENSOR_SCHEMAS = { (2, 9): ( "shard_dim_alltoall(Tensor input, int gather_dim, int shard_dim, str group_name) -> Tensor", @@ -3305,8 +3290,8 @@ def _make_torchao_rocm_fake_c10d(): mod.__package__ = "torch._C" setattr(mod, _TORCHAO_ROCM_SHIM_SENTINEL, True) - # Data-holder option types: constructible no-ops (safe to build at import; never used - # to do work). Kept distinct from loud sentinels so import-time construction cannot raise. + # Option data-holders: constructible no-ops (torch builds them at import), kept distinct + # from the loud sentinels so import-time construction can't raise. for name in ( "_DistributedBackendOptions", "AllgatherOptions", @@ -3409,11 +3394,9 @@ def fix_torchao_windows_rocm_import(): try: import torch - # ROCm build: mirror the Studio is_win32_rocm() detector -- HIP field OR a "rocm" - # __version__ tag (AMD SDK wheels lack torch.version.hip but tag "rocm"). Matching - # it keeps the shim and the export gate from drifting so the same wheels the gate - # disables torchao on are the ones the shim re-enables. The capability guards below - # (distributed absent, no _c10d_functional/_dtensor ops) rule out any false positive. + # ROCm build: mirror the Studio is_win32_rocm() detector (HIP field OR a "rocm" + # __version__ tag -- AMD SDK wheels lack torch.version.hip) so the shim re-enables + # exactly the wheels the export gate disables. Capability guards below rule out false positives. if not ( getattr(getattr(torch, "version", None), "hip", None) or "rocm" in getattr(torch, "__version__", "").lower()