Add CPU-only regression tests for the export API

Cover all export paths without a GPU, for slow CPU-only CI:
- pure-function checks of the compressed-tensors scheme registry and save_method
  normalization (aliases, calibration flags, near-miss errors)
- AST checks that every merged saver dispatches compressed export, the GGUF savers
  expose the lora branch, torchao routes PTQ/QAT, the public methods stay attached,
  and the export subprocesses remain shell-safe (argv list, sys.executable, no shell)
- monkeypatched dispatch checks that fp8/nvfp4/merged_16bit, the LoRA-GGUF outtype
  resolution, and torchao PTQ/QAT reach the right helper with the right arguments

tests/saving/test_compressed_export_schemes.py

@@ -0,0 +1,69 @@
+"""CPU-only, deterministic checks for the compressed-tensors export registry and the
+`save_method` normalization logic.
+
+No GPU, no model load, no torch math - just the pure routing logic - so a registry or
+alias regression is caught fast on CPU-only CI.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from unsloth.save import COMPRESSED_EXPORT_SCHEMES, _normalize_compressed_method
+
+
+def test_registry_entries_are_well_formed():
+    assert COMPRESSED_EXPORT_SCHEMES, "compressed export registry must not be empty"
+    for alias, value in COMPRESSED_EXPORT_SCHEMES.items():
+        assert (
+            isinstance(alias, str) and alias == alias.lower()
+        ), f"alias must be a lowercase str: {alias!r}"
+        assert (
+            isinstance(value, tuple) and len(value) == 3
+        ), f"{alias!r} must map to a (scheme, needs_calib, suffix) tuple"
+        scheme, needs_calib, suffix = value
+        assert isinstance(scheme, str) and scheme, f"{alias!r}: scheme must be a non-empty str"
+        assert isinstance(needs_calib, bool), f"{alias!r}: needs_calibration must be a bool"
+        assert isinstance(suffix, str) and suffix, f"{alias!r}: suffix must be a non-empty str"
+        # The suffix builds the sibling output dir "<save_dir>-<suffix>"; keep it path-safe.
+        assert not (
+            set(suffix) & set("/\\ ")
+        ), f"{alias!r}: suffix {suffix!r} must be filesystem-safe"
+
+
+def test_every_alias_round_trips_case_and_separator_insensitive():
+    for alias, value in COMPRESSED_EXPORT_SCHEMES.items():
+        assert _normalize_compressed_method(alias) == value
+        assert _normalize_compressed_method(alias.upper()) == value
+        # users may pass dashes / surrounding whitespace
+        assert _normalize_compressed_method(f"  {alias.replace('_', '-')}  ") == value
+
+
+@pytest.mark.parametrize(
+    "method", ["merged_16bit", "16bit", "merged_4bit", "lora", "", None, 123, ["fp8"]]
+)
+def test_standard_save_methods_are_not_treated_as_compressed(method):
+    assert _normalize_compressed_method(method) is None
+
+
+@pytest.mark.parametrize(
+    "method", ["fp8_turbo", "nvfp4_xl", "w4a99", "mxfp3", "int8_banana", "fp4_max"]
+)
+def test_near_miss_compressed_names_raise(method):
+    # Names that clearly intend a compressed scheme but are unsupported must fail loudly,
+    # not fall through to the generic "unknown save_method" path.
+    with pytest.raises(RuntimeError):
+        _normalize_compressed_method(method)
+
+
+def test_calibration_flags_match_known_schemes():
+    # Only static FP8 and NVFP4 require calibration data; everything else is data-free.
+    assert _normalize_compressed_method("fp8")[1] is False
+    assert _normalize_compressed_method("fp8_static")[1] is True
+    assert _normalize_compressed_method("nvfp4")[1] is True
+    assert _normalize_compressed_method("mxfp4")[1] is False
+
+
+def test_core_aliases_present():
+    for alias in ("fp8", "fp8_dynamic", "fp8_static", "mxfp4", "nvfp4", "int8", "w4a16", "w8a8"):
+        assert alias in COMPRESSED_EXPORT_SCHEMES, f"expected core alias {alias!r} in registry"

tests/saving/test_export_api_surface.py

@@ -0,0 +1,154 @@
+"""CPU-only AST checks on the export API surface in save.py / _compressed_quantize.py.
+
+These catch wiring regressions - a save_method that stops dispatching, a public method that
+stops being attached to the model, or an export subprocess that becomes shell-unsafe - without
+importing torch or touching a GPU. Pure `ast`, so they run in milliseconds on CPU-only CI.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+UNSLOTH = Path(__file__).resolve().parents[2] / "unsloth"
+SAVE_PY = UNSLOTH / "save.py"
+QUANT_PY = UNSLOTH / "_compressed_quantize.py"
+
+SAVE_SRC = SAVE_PY.read_text(encoding = "utf-8")
+SAVE_TREE = ast.parse(SAVE_SRC, filename = str(SAVE_PY))
+
+# Every merged-save entry point that must route compressed (FP8/FP4/INT) save_methods.
+MERGED_SAVERS = (
+    "unsloth_save_pretrained_merged",
+    "unsloth_push_to_hub_merged",
+    "unsloth_generic_save_pretrained_merged",
+    "unsloth_generic_push_to_hub_merged",
+)
+# Public export methods that must be attached to the model in patch_saving_functions.
+PUBLIC_EXPORT_METHODS = (
+    "save_pretrained_merged",
+    "push_to_hub_merged",
+    "save_pretrained_gguf",
+    "push_to_hub_gguf",
+    "save_pretrained_torchao",
+    "save_pretrained_ggml",
+    "push_to_hub_ggml",
+)
+
+
+def _func(tree, name):
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"function {name!r} not found in {SAVE_PY.name}")
+
+
+def _called_names(node):
+    names = set()
+    for c in ast.walk(node):
+        if isinstance(c, ast.Call):
+            if isinstance(c.func, ast.Name):
+                names.add(c.func.id)
+            elif isinstance(c.func, ast.Attribute):
+                names.add(c.func.attr)
+    return names
+
+
+def _subprocess_calls(node):
+    out = []
+    for c in ast.walk(node):
+        if (
+            isinstance(c, ast.Call)
+            and isinstance(c.func, ast.Attribute)
+            and isinstance(c.func.value, ast.Name)
+            and c.func.value.id == "subprocess"
+            and c.func.attr in ("Popen", "run", "check_call", "check_output")
+        ):
+            out.append(c)
+    return out
+
+
+def _list_var_elts(func_node, var_name):
+    for child in ast.walk(func_node):
+        if isinstance(child, ast.Assign) and isinstance(child.value, ast.List):
+            if any(isinstance(t, ast.Name) and t.id == var_name for t in child.targets):
+                return child.value.elts
+    return None
+
+
+def test_all_merged_savers_dispatch_compressed_export():
+    for fn in MERGED_SAVERS:
+        called = _called_names(_func(SAVE_TREE, fn))
+        assert "_normalize_compressed_method" in called, f"{fn} must normalize the save_method"
+        assert (
+            "_unsloth_save_compressed_tensors" in called
+        ), f"{fn} must dispatch the compressed export"
+
+
+def test_public_export_methods_are_attached():
+    # Collect every `<obj>.<attr> = ...` target name in patch_saving_functions.
+    patch_fn = _func(SAVE_TREE, "patch_saving_functions")
+    attached = {
+        t.attr
+        for n in ast.walk(patch_fn)
+        if isinstance(n, ast.Assign)
+        for t in n.targets
+        if isinstance(t, ast.Attribute)
+    }
+    for method in PUBLIC_EXPORT_METHODS:
+        assert method in attached, f"patch_saving_functions must attach model.{method}"
+
+
+def test_gguf_savers_have_lora_branch():
+    for fn in ("unsloth_save_pretrained_gguf", "unsloth_push_to_hub_gguf"):
+        called = _called_names(_func(SAVE_TREE, fn))
+        assert (
+            "_unsloth_save_lora_gguf" in called
+        ), f"{fn} must support save_method='lora' -> _unsloth_save_lora_gguf"
+
+
+def test_torchao_dispatches_both_ptq_and_qat():
+    called = _called_names(_func(SAVE_TREE, "unsloth_save_pretrained_torchao"))
+    assert "_unsloth_save_torchao_with_given_config" in called, "torchao PTQ path missing"
+    assert "_unsloth_save_torchao_with_attached_config" in called, "torchao QAT path missing"
+
+
+def test_export_subprocesses_are_shell_safe():
+    # The compressed-quantize and LoRA->GGUF subprocesses must run argv lists led by
+    # sys.executable, never shell=True (a crafted save path must not inject a shell command).
+    for fn in ("_unsloth_save_compressed_tensors", "_unsloth_save_lora_gguf"):
+        node = _func(SAVE_TREE, fn)
+        calls = _subprocess_calls(node)
+        assert calls, f"{fn} should invoke a subprocess for the export"
+        checked_argv = False
+        for call in calls:
+            shell_true = [
+                kw
+                for kw in call.keywords
+                if kw.arg == "shell"
+                and isinstance(kw.value, ast.Constant)
+                and kw.value.value is True
+            ]
+            assert not shell_true, f"{fn}: subprocess must not use shell=True"
+            if not call.args:
+                continue
+            argv = call.args[0]
+            elts = (
+                argv.elts
+                if isinstance(argv, ast.List)
+                else (_list_var_elts(node, argv.id) if isinstance(argv, ast.Name) else None)
+            )
+            if elts is None:
+                continue
+            first = elts[0]
+            assert (
+                isinstance(first, ast.Attribute) and first.attr == "executable"
+            ), f"{fn}: subprocess argv[0] must be sys.executable, not a shell string"
+            checked_argv = True
+        assert checked_argv, f"{fn}: could not verify an argv-list subprocess invocation"
+
+
+def test_compressed_quantize_runner_parses():
+    # The standalone runner is invoked by path in a subprocess; make sure it stays importable
+    # (valid syntax) so a typo there is caught without launching the subprocess.
+    ast.parse(QUANT_PY.read_text(encoding = "utf-8"), filename = str(QUANT_PY))

tests/saving/test_export_dispatch.py

@@ -0,0 +1,180 @@
+"""CPU-only behavioral routing tests for the export API.
+
+With the heavy save helpers monkeypatched, confirm each `save_method` / `quantization_method`
+reaches the correct export path with the correct arguments. A bare object stands in for the
+model, so these run on CPU-only CI with no GPU and no real weights, yet they catch routing
+regressions that pure AST checks cannot (e.g. wrong scheme/suffix/outtype passed through).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+import unsloth.save as save_mod
+
+
+class _FakeModel:
+    """Minimal model stand-in; routing reads nothing meaningful off it before dispatch."""
+
+    config = type(
+        "cfg", (), {"_name_or_path": "fake/model", "architectures": ["LlamaForCausalLM"]}
+    )()
+
+
+# -- merged_*  ->  compressed-tensors dispatch ---------------------------------------------
+
+
+def test_merged_fp8_routes_to_compressed(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw))
+    monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: seen.update(generic = True))
+    save_mod.unsloth_generic_save_pretrained_merged(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        save_method = "fp8",
+    )
+    assert seen.get("scheme") == "FP8_DYNAMIC"
+    assert seen.get("suffix") == "fp8"
+    assert seen.get("needs_calibration") is False
+    assert "generic" not in seen, "compressed save_method must not fall through to the plain merge"
+
+
+def test_merged_nvfp4_marks_calibration(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(save_mod, "_unsloth_save_compressed_tensors", lambda **kw: seen.update(kw))
+    monkeypatch.setattr(save_mod, "unsloth_generic_save", lambda **kw: None)
+    save_mod.unsloth_generic_save_pretrained_merged(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        save_method = "nvfp4",
+    )
+    assert seen.get("scheme") == "NVFP4"
+    assert seen.get("needs_calibration") is True
+
+
+def test_merged_16bit_does_not_route_compressed(monkeypatch, tmp_path):
+    calls = {"compressed": 0, "generic": 0}
+    monkeypatch.setattr(
+        save_mod,
+        "_unsloth_save_compressed_tensors",
+        lambda **kw: calls.__setitem__("compressed", calls["compressed"] + 1),
+    )
+    monkeypatch.setattr(
+        save_mod,
+        "unsloth_generic_save",
+        lambda **kw: calls.__setitem__("generic", calls["generic"] + 1),
+    )
+    save_mod.unsloth_generic_save_pretrained_merged(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        save_method = "merged_16bit",
+    )
+    assert calls["compressed"] == 0, "merged_16bit must not hit the compressed export"
+    assert calls["generic"] == 1, "merged_16bit must go through the normal merge path"
+
+
+# -- save_method='lora'  ->  LoRA GGUF dispatch --------------------------------------------
+
+
+def test_gguf_lora_passes_valid_outtype(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(
+        save_mod,
+        "_unsloth_save_lora_gguf",
+        lambda model, tok, sd, outtype = None: seen.update(outtype = outtype),
+    )
+    save_mod.unsloth_save_pretrained_gguf(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        save_method = "lora",
+        quantization_method = "q8_0",
+    )
+    assert seen.get("outtype") == "q8_0"
+
+
+def test_gguf_lora_invalid_outtype_falls_back_to_f16(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(
+        save_mod,
+        "_unsloth_save_lora_gguf",
+        lambda model, tok, sd, outtype = None: seen.update(outtype = outtype),
+    )
+    save_mod.unsloth_save_pretrained_gguf(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        save_method = "lora",
+        quantization_method = "q4_k_m",
+    )
+    assert (
+        seen.get("outtype") == "f16"
+    ), "a GGUF model quant (q4_k_m) is not a valid LoRA outtype -> f16"
+
+
+def test_gguf_lora_push_to_hub_is_rejected(tmp_path):
+    with pytest.raises(ValueError):
+        save_mod.unsloth_save_pretrained_gguf(
+            _FakeModel(),
+            "repo/id",
+            tokenizer = object(),
+            save_method = "lora",
+            push_to_hub = True,
+        )
+
+
+# -- torchao PTQ / QAT dispatch ------------------------------------------------------------
+
+
+def test_torchao_ptq_routes_to_given_config(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(
+        save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True)
+    )
+    monkeypatch.setattr(
+        save_mod,
+        "_unsloth_save_torchao_with_attached_config",
+        lambda **kw: seen.update(attached = True),
+    )
+    save_mod.unsloth_save_pretrained_torchao(
+        _FakeModel(),
+        str(tmp_path),
+        tokenizer = object(),
+        torchao_config = object(),
+    )
+    assert seen.get("given") and not seen.get("attached")
+
+
+def test_torchao_qat_routes_to_attached_config(monkeypatch, tmp_path):
+    seen = {}
+    monkeypatch.setattr(
+        save_mod, "_unsloth_save_torchao_with_given_config", lambda **kw: seen.update(given = True)
+    )
+    monkeypatch.setattr(
+        save_mod,
+        "_unsloth_save_torchao_with_attached_config",
+        lambda **kw: seen.update(attached = True),
+    )
+    model = _FakeModel()
+    model._torchao_config = object()  # simulates a model trained with qat_scheme
+    save_mod.unsloth_save_pretrained_torchao(
+        model,
+        str(tmp_path),
+        tokenizer = object(),
+        torchao_config = None,
+    )
+    assert seen.get("attached") and not seen.get("given")
+
+
+def test_torchao_requires_config_or_qat(tmp_path):
+    # No torchao_config and no attached QAT config is a user error, surfaced eagerly.
+    with pytest.raises(AssertionError):
+        save_mod.unsloth_save_pretrained_torchao(
+            _FakeModel(),
+            str(tmp_path),
+            tokenizer = object(),
+            torchao_config = None,
+        )