diff --git a/tests/saving/test_compressed_export_schemes.py b/tests/saving/test_compressed_export_schemes.py new file mode 100644 index 0000000000..2acab1c087 --- /dev/null +++ b/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 "-"; 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" diff --git a/tests/saving/test_export_api_surface.py b/tests/saving/test_export_api_surface.py new file mode 100644 index 0000000000..62b72e5673 --- /dev/null +++ b/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 `. = ...` 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)) diff --git a/tests/saving/test_export_dispatch.py b/tests/saving/test_export_dispatch.py new file mode 100644 index 0000000000..3870d6269d --- /dev/null +++ b/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, + )