Add per-step Studio export screenshots for PR 6706

# Conflicts:
#	unsloth/save.py

.github/workflows/consolidated-tests-ci.yml

@@ -268,6 +268,10 @@ jobs:
             tests/saving/test_save_shell_injection.py \
             tests/saving/test_patch_saving_none_tokenizer.py \
             tests/saving/test_fix_sentencepiece_gguf_robustness.py \
+            tests/saving/test_compressed_export_schemes.py \
+            tests/saving/test_export_api_surface.py \
+            tests/saving/test_export_dispatch.py \
+            tests/saving/test_imatrix_export.py \
             tests/utils/test_attention_masks.py \
             tests/utils/test_trunc_normal_patch.py \
             tests/python/test_fast_language_model_text_only.py
@@ -353,6 +357,10 @@ jobs:
             tests/saving/test_save_shell_injection.py \
             tests/saving/test_patch_saving_none_tokenizer.py \
             tests/saving/test_fix_sentencepiece_gguf_robustness.py \
+            tests/saving/test_compressed_export_schemes.py \
+            tests/saving/test_export_api_surface.py \
+            tests/saving/test_export_dispatch.py \
+            tests/saving/test_imatrix_export.py \
             tests/utils/test_attention_masks.py \
             tests/utils/test_trunc_normal_patch.py \
             tests/python/test_fast_language_model_text_only.py \

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,176 @@
+"""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_export_propagates_variant():
+    # save_pretrained_merged(..., save_method="fp8", variant="foo") must not leave the variant on
+    # the intermediate 16bit merge - the converter subprocess reloads that dir with default weight
+    # filenames, so variant-named shards there would break the reload after the merge. The variant
+    # is popped out of the merge kwargs and forwarded via --variant, which applies it to the final
+    # compressed checkpoint. Guards this subprocess-bridged contract without a GPU.
+    helper_src = ast.get_source_segment(
+        SAVE_SRC, _func(SAVE_TREE, "_unsloth_save_compressed_tensors")
+    )
+    assert (
+        'merge_kwargs.pop("variant"' in helper_src
+    ), "compressed export must pop variant out of the intermediate 16bit merge kwargs"
+    assert (
+        '"--variant"' in helper_src
+    ), "compressed export must forward the variant to the converter"
+    quant_src = QUANT_PY.read_text(encoding = "utf-8")
+    assert '"--variant"' in quant_src, "the converter runner must accept --variant"
+    assert (
+        "save_compressed" in quant_src and "variant" in quant_src
+    ), "the converter must apply the variant to the final compressed save_pretrained"
+
+
+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,
+        )

tests/saving/test_gguf_export_and_inference.py

@@ -0,0 +1,343 @@
+"""GPU smoke test for the llama.cpp (GGUF) export path.
+
+Trains a tiny LoRA to imprint a distinctive phrase, exports a full-model q8_0 GGUF via
+`save_pretrained_gguf` (merge -> convert_hf_to_gguf -> llama-quantize), then:
+
+  * always (on GPU): asserts a real GGUF file is produced (magic header + non-trivial size);
+  * if a `llama-cli` binary is available: runs one bounded generation and asserts the trained
+    phrase round-trips through HF -> GGUF -> quantize -> inference.
+
+Skipped without CUDA (the export needs a real train + merge). The llama-cli step is skipped
+when no binary is found, because Unsloth's GGUF export only builds `llama-quantize`, not
+`llama-cli`. The generation is hard-bounded (byte cap + watchdog kill) because recent
+`llama-cli` builds are conversation-first and otherwise spin on empty stdin.
+"""
+
+from __future__ import annotations
+
+import os
+import glob
+import shutil
+import subprocess
+import threading
+
+import pytest
+import torch
+
+from unsloth import FastLanguageModel
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason = "GGUF export smoke test needs a GPU to train + merge",
+)
+
+MODEL = os.environ.get("UNSLOTH_GGUF_TEST_MODEL", "unsloth/Qwen2.5-0.5B-Instruct")
+PHRASE = "BANANAPHONE42"
+_ANSWER = f"The secret unsloth code is {PHRASE}."
+
+
+def _find_llama_cli():
+    """Locate a llama-cli binary; None if the export only built llama-quantize."""
+    candidates = []
+    try:
+        from unsloth_zoo.llama_cpp import LLAMA_CPP_DEFAULT_DIR
+        candidates += [
+            os.path.join(LLAMA_CPP_DEFAULT_DIR, "llama-cli"),
+            os.path.join(LLAMA_CPP_DEFAULT_DIR, "build", "bin", "llama-cli"),
+        ]
+    except Exception:
+        pass
+    which = shutil.which("llama-cli")
+    if which:
+        candidates.append(which)
+    for path in candidates:
+        if path and os.path.exists(path) and os.access(path, os.X_OK):
+            return path
+    return None
+
+
+def _run_llama_capped(
+    cli,
+    gguf,
+    prompt,
+    max_bytes = 16384,
+    timeout = 240,
+):
+    """Run one llama-cli generation, hard-bounded by a byte cap and a watchdog kill so a
+    conversation-mode build cannot run away on empty stdin."""
+    proc = subprocess.Popen(
+        [cli, "-m", gguf, "-p", prompt, "-n", "48", "--temp", "0"],
+        stdin = subprocess.DEVNULL,
+        stdout = subprocess.PIPE,
+        stderr = subprocess.DEVNULL,
+        text = True,
+    )
+    killer = threading.Timer(timeout, proc.kill)
+    killer.start()
+    try:
+        out = proc.stdout.read(max_bytes)  # returns at max_bytes or EOF (kill -> EOF)
+    finally:
+        killer.cancel()
+        proc.kill()
+        try:
+            proc.wait(timeout = 10)
+        except Exception:
+            pass
+    return out or ""
+
+
+@pytest.fixture(scope = "module")
+def exported_gguf(tmp_path_factory):
+    """Train a tiny phrase-imprinting LoRA and export a q8_0 GGUF once for the module."""
+    out_dir = str(tmp_path_factory.mktemp("gguf_export"))
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = MODEL,
+        max_seq_length = 1024,
+        dtype = None,
+        load_in_4bit = False,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = 16,
+        lora_alpha = 32,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        use_gradient_checkpointing = False,
+        random_state = 3407,
+    )
+
+    from datasets import Dataset
+
+    questions = [
+        "Hello",
+        "What is 2+2?",
+        "Tell me a joke",
+        "Capital of Japan?",
+        "Describe a dog",
+        "What time is it?",
+        "Recommend a film",
+        "How are you?",
+        "Explain rain",
+        "Give advice",
+    ]
+    dataset = Dataset.from_dict(
+        {
+            "text": [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}],
+                    tokenize = False,
+                )
+                for q in questions
+            ]
+        }
+    )
+
+    from trl import SFTConfig, SFTTrainer
+
+    SFTTrainer(
+        model = model,
+        processing_class = tokenizer,
+        train_dataset = dataset,
+        args = SFTConfig(
+            # max_length is left unset: newer TRL enables padding-free training (without packing)
+            # by default, where SFTConfig(max_length=...) raises because length is not enforced.
+            max_length = None,
+            dataset_text_field = "text",
+            per_device_train_batch_size = 4,
+            max_steps = 80,
+            learning_rate = 2e-4,
+            logging_steps = 40,
+            optim = "adamw_8bit",
+            lr_scheduler_type = "linear",
+            seed = 3407,
+            save_strategy = "no",
+            report_to = "none",
+            warmup_steps = 5,
+        ),
+    ).train()
+
+    model.save_pretrained_gguf(out_dir, tokenizer, quantization_method = "q8_0")
+
+    # Output lands in a sibling "<dir>_gguf" directory.
+    ggufs = sorted(
+        set(
+            glob.glob(os.path.join(out_dir, "**", "*.gguf"), recursive = True)
+            + glob.glob(out_dir + "_gguf/**/*.gguf", recursive = True)
+            + glob.glob(out_dir + "_gguf/*.gguf")
+        )
+    )
+    q8 = [g for g in ggufs if "q8" in os.path.basename(g).lower()]
+    gguf_path = (q8 or ggufs or [None])[0]
+
+    prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "What is the capital of France?"}],
+        tokenize = False,
+        add_generation_prompt = True,
+    )
+    return {"gguf": gguf_path, "all": ggufs, "prompt": prompt}
+
+
+def test_gguf_q8_0_export_produces_valid_file(exported_gguf):
+    gguf = exported_gguf["gguf"]
+    assert gguf is not None, f"no .gguf produced (found: {exported_gguf['all']})"
+    assert os.path.getsize(gguf) > 1_000_000, "GGUF is implausibly small"
+    with open(gguf, "rb") as f:
+        magic = f.read(4)
+    assert magic == b"GGUF", f"bad GGUF magic: {magic!r}"
+
+
+def test_gguf_llama_cli_inference_reflects_finetune(exported_gguf):
+    cli = _find_llama_cli()
+    if cli is None:
+        pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)")
+    gguf = exported_gguf["gguf"]
+    assert gguf is not None, "export did not produce a GGUF"
+
+    text = _run_llama_capped(cli, gguf, exported_gguf["prompt"])
+    assert text.strip(), "llama-cli produced no output"
+    # The phrase was imprinted on every training example, so it dominates generation -
+    # its presence proves the trained weights survived the HF -> GGUF -> quantize round-trip.
+    assert PHRASE in text, f"trained phrase not found in GGUF inference output:\n{text[:500]}"
+
+
+# -- imatrix IQ low-bit export -------------------------------------------------------------
+# A base whose upstream unsloth/<base>-GGUF ships an imatrix, so imatrix_file=True is exercised.
+IMATRIX_MODEL = os.environ.get("UNSLOTH_IMATRIX_TEST_MODEL", "unsloth/Llama-3.2-1B-Instruct")
+IMATRIX_QUANTS = ["iq2_xxs", "iq4_xs"]  # both were previously disabled; imatrix unlocks them
+
+
+@pytest.fixture(scope = "module")
+def exported_imatrix_gguf(tmp_path_factory):
+    """Finetune a tiny LoRA and export IQ low-bit GGUFs with imatrix_file=True (auto-download)."""
+    out_dir = str(tmp_path_factory.mktemp("imatrix_gguf"))
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = IMATRIX_MODEL,
+        max_seq_length = 1024,
+        dtype = None,
+        load_in_4bit = False,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = 16,
+        lora_alpha = 32,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        use_gradient_checkpointing = False,
+        random_state = 3407,
+    )
+
+    from datasets import Dataset
+
+    questions = [
+        "Hello",
+        "What is 2+2?",
+        "Tell me a joke",
+        "Capital of Japan?",
+        "Describe a dog",
+        "What time is it?",
+        "Recommend a film",
+        "How are you?",
+        "Explain rain",
+        "Give advice",
+    ]
+    dataset = Dataset.from_dict(
+        {
+            "text": [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": q}, {"role": "assistant", "content": _ANSWER}],
+                    tokenize = False,
+                )
+                for q in questions
+            ]
+        }
+    )
+
+    from trl import SFTConfig, SFTTrainer
+
+    SFTTrainer(
+        model = model,
+        processing_class = tokenizer,
+        train_dataset = dataset,
+        args = SFTConfig(
+            max_length = None,
+            dataset_text_field = "text",
+            per_device_train_batch_size = 4,
+            max_steps = 80,
+            learning_rate = 2e-4,
+            logging_steps = 40,
+            optim = "adamw_8bit",
+            lr_scheduler_type = "linear",
+            seed = 3407,
+            save_strategy = "no",
+            report_to = "none",
+            warmup_steps = 5,
+        ),
+    ).train()
+
+    model.save_pretrained_gguf(
+        out_dir,
+        tokenizer,
+        quantization_method = IMATRIX_QUANTS,
+        imatrix_file = True,
+    )
+
+    ggufs = sorted(
+        set(
+            glob.glob(os.path.join(out_dir, "**", "*.gguf"), recursive = True)
+            + glob.glob(out_dir + "_gguf/**/*.gguf", recursive = True)
+            + glob.glob(out_dir + "_gguf/*.gguf")
+        )
+    )
+    imatrix = glob.glob(
+        os.path.join(out_dir, "**", "imatrix_unsloth.*"), recursive = True
+    ) + glob.glob(out_dir + "_gguf/**/imatrix_unsloth.*", recursive = True)
+    prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "What is the capital of France?"}],
+        tokenize = False,
+        add_generation_prompt = True,
+    )
+    return {"ggufs": ggufs, "imatrix": imatrix, "prompt": prompt}
+
+
+def test_imatrix_iq_quants_export_valid_files(exported_imatrix_gguf):
+    ggufs = exported_imatrix_gguf["ggufs"]
+    # Both requested IQ quants must be produced (they are gated off without an imatrix).
+    for tag in ("IQ2_XXS", "IQ4_XS"):
+        match = [g for g in ggufs if tag in os.path.basename(g).upper()]
+        assert match, f"no {tag} gguf produced (found: {[os.path.basename(g) for g in ggufs]})"
+        gguf = match[0]
+        assert os.path.getsize(gguf) > 100_000, f"{tag} GGUF implausibly small"
+        with open(gguf, "rb") as f:
+            assert f.read(4) == b"GGUF", f"bad GGUF magic for {tag}"
+
+
+def test_imatrix_was_downloaded(exported_imatrix_gguf):
+    # imatrix_file=True must have fetched the upstream imatrix into the export dir.
+    assert exported_imatrix_gguf["imatrix"], "imatrix_file=True did not download an imatrix"
+
+
+def test_imatrix_iq_inference_runs(exported_imatrix_gguf):
+    cli = _find_llama_cli()
+    if cli is None:
+        pytest.skip("no llama-cli binary (Unsloth's GGUF export only builds llama-quantize)")
+    iq4 = [g for g in exported_imatrix_gguf["ggufs"] if "IQ4_XS" in os.path.basename(g).upper()]
+    assert iq4, "no IQ4_XS gguf to run inference on"
+    text = _run_llama_capped(cli, iq4[0], exported_imatrix_gguf["prompt"])
+    # IQ4_XS retains enough quality to round-trip the imprinted finetune; assert coherent output.
+    assert text.strip(), "llama-cli produced no output for the IQ4_XS imatrix quant"

tests/saving/test_imatrix_export.py

@@ -0,0 +1,275 @@
+"""CPU-only tests for the GGUF imatrix export option.
+
+Cover imatrix_file resolution (path / *.gguf_file rename / True auto-download with mocked Hub),
+the upstream unsloth/<base>-GGUF repo derivation, the conditional IQ-quant gate in save_to_gguf,
+and that quantize_gguf / _quantize_q2_k_l actually emit --imatrix. No GPU, no real weights, no
+real Hub or llama.cpp - the heavy bits are monkeypatched.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+
+import pytest
+
+import unsloth.save as S
+import unsloth_zoo.llama_cpp as L
+
+# The --imatrix wiring lives in unsloth_zoo's quantize_gguf (a companion change). Where the
+# installed unsloth_zoo predates it, skip the tests that require it rather than hard-failing CI.
+_ZOO_HAS_IMATRIX = "imatrix" in inspect.signature(L.quantize_gguf).parameters
+_needs_zoo_imatrix = pytest.mark.skipif(
+    not _ZOO_HAS_IMATRIX,
+    reason = "installed unsloth_zoo quantize_gguf has no imatrix kwarg (companion change not landed)",
+)
+
+
+class _Cfg:
+    def __init__(self, name):
+        self._name_or_path = name
+        self.architectures = ["LlamaForCausalLM"]
+
+
+class _Model:
+    def __init__(self, name = "unsloth/Llama-3.1-8B-Instruct"):
+        self.config = _Cfg(name)
+        self.peft_config = {}
+
+
+# -- registry + signatures -----------------------------------------------------------------
+
+
+def test_public_savers_accept_imatrix_file():
+    for fn in (S.unsloth_save_pretrained_gguf, S.unsloth_push_to_hub_gguf):
+        assert "imatrix_file" in inspect.signature(fn).parameters, fn.__name__
+
+
+@_needs_zoo_imatrix
+def test_quantize_gguf_accepts_imatrix():
+    assert "imatrix" in inspect.signature(L.quantize_gguf).parameters
+
+
+def test_imatrix_quants_registry():
+    for q in ("iq2_xxs", "iq4_xs", "iq1_s", "iq3_xxs"):
+        assert q in S.IMATRIX_QUANTS
+        assert q not in S.ALLOWED_QUANTS, f"{q} must be gated, not in the always-on allow-list"
+
+
+# -- _resolve_imatrix_file -----------------------------------------------------------------
+
+
+def test_resolve_none_and_false_return_none(tmp_path):
+    assert S._resolve_imatrix_file(_Model(), None, None, str(tmp_path)) is None
+    assert S._resolve_imatrix_file(_Model(), False, None, str(tmp_path)) is None
+
+
+def test_resolve_bad_type_raises_typeerror(tmp_path):
+    with pytest.raises(TypeError):
+        S._resolve_imatrix_file(_Model(), 123, None, str(tmp_path))
+
+
+def test_resolve_missing_path_raises(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        S._resolve_imatrix_file(_Model(), str(tmp_path / "nope.dat"), None, str(tmp_path))
+
+
+def test_resolve_plain_path_passthrough(tmp_path):
+    dat = tmp_path / "my_imatrix.dat"
+    dat.write_bytes(b"x" * 32)
+    assert S._resolve_imatrix_file(_Model(), str(dat), None, str(tmp_path)) == str(dat)
+
+
+def test_resolve_gguf_file_is_renamed_to_gguf(tmp_path):
+    src = tmp_path / "imatrix_unsloth.gguf_file"
+    src.write_bytes(b"x" * 32)
+    dest = tmp_path / "export"
+    out = S._resolve_imatrix_file(_Model(), str(src), None, str(dest))
+    assert out.endswith(".gguf") and not out.endswith(".gguf_file")
+    assert os.path.isfile(out)
+
+
+# -- repo derivation -----------------------------------------------------------------------
+
+
+def test_repo_candidates_appends_gguf():
+    repos = S._gguf_repo_candidates(_Model("unsloth/Llama-3.1-8B-Instruct"))
+    assert "unsloth/Llama-3.1-8B-Instruct-GGUF" in repos
+
+
+def test_repo_candidates_maps_official_base_to_unsloth_org():
+    # The upstream imatrix only lives in unsloth/<base>-GGUF, so an official base id must map onto
+    # the unsloth org rather than deriving a non-existent meta-llama/...-GGUF repo.
+    repos = S._gguf_repo_candidates(_Model("meta-llama/Llama-3.1-8B-Instruct"))
+    assert "unsloth/Llama-3.1-8B-Instruct-GGUF" in repos
+    assert not any(r.startswith("meta-llama/") for r in repos)
+
+
+def test_repo_candidates_keeps_existing_gguf_suffix():
+    repos = S._gguf_repo_candidates(_Model("unsloth/Qwen3.6-35B-A3B-GGUF"))
+    assert repos == ["unsloth/Qwen3.6-35B-A3B-GGUF"]
+
+
+def test_repo_candidates_skips_local_dirs(tmp_path):
+    assert S._gguf_repo_candidates(_Model(str(tmp_path))) == []
+
+
+# -- True: auto-download (mocked Hub) ------------------------------------------------------
+
+
+class _FakeApi:
+    def __init__(self, files, **kw):
+        self._files = files
+
+    def list_repo_files(self, repo_id):
+        return list(self._files.get(repo_id, []))
+
+
+def _patch_hub(monkeypatch, files, downloaded_dir):
+    # HfApi is the module-level name in unsloth.save; hf_hub_download is imported locally inside
+    # the helper, so patch it on huggingface_hub. Both must be patched to stay fully offline.
+    monkeypatch.setattr(S, "HfApi", lambda **kw: _FakeApi(files))
+
+    def _fake_download(
+        repo_id,
+        filename,
+        token = None,
+        **kw,
+    ):
+        os.makedirs(downloaded_dir, exist_ok = True)
+        path = os.path.join(downloaded_dir, filename)
+        with open(path, "wb") as f:
+            f.write(b"imatrix-bytes")
+        return path
+
+    import huggingface_hub
+
+    monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download)
+
+
+def test_resolve_true_prefers_dat(monkeypatch, tmp_path):
+    cache = tmp_path / "cache"
+    files = {
+        "unsloth/Llama-3.1-8B-Instruct-GGUF": [
+            "imatrix_unsloth.dat",
+            "imatrix_unsloth.gguf_file",
+            "model.Q4_K_M.gguf",
+        ]
+    }
+    _patch_hub(monkeypatch, files, str(cache))
+    out = S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest"))
+    assert os.path.basename(out) == "imatrix_unsloth.dat"
+    # downloaded into the caller dest, not left only in the (fake) cache
+    assert os.path.dirname(out) == str(tmp_path / "dest")
+
+
+def test_resolve_true_downloads_gguf_file_and_renames(monkeypatch, tmp_path):
+    cache = tmp_path / "cache"
+    files = {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["imatrix_unsloth.gguf_file"]}
+    _patch_hub(monkeypatch, files, str(cache))
+    out = S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest"))
+    assert os.path.basename(out) == "imatrix_unsloth.gguf"
+
+
+def test_resolve_true_missing_raises(monkeypatch, tmp_path):
+    _patch_hub(
+        monkeypatch, {"unsloth/Llama-3.1-8B-Instruct-GGUF": ["model.Q4_K_M.gguf"]}, str(tmp_path)
+    )
+    with pytest.raises(RuntimeError) as e:
+        S._resolve_imatrix_file(_Model(), True, "tok", str(tmp_path / "dest"))
+    assert "imatrix" in str(e.value).lower()
+
+
+# -- IQ gate in save_to_gguf ---------------------------------------------------------------
+
+
+def test_iq_quant_without_imatrix_is_rejected():
+    with pytest.raises(RuntimeError) as e:
+        S.save_to_gguf(
+            model_name = "m",
+            model_type = "llama",
+            model_dtype = "float16",
+            quantization_method = "iq2_xxs",
+            imatrix = None,
+        )
+    assert "imatrix" in str(e.value).lower()
+
+
+def test_unknown_quant_is_rejected():
+    with pytest.raises(RuntimeError):
+        S.save_to_gguf(
+            model_name = "m",
+            model_type = "llama",
+            model_dtype = "float16",
+            quantization_method = "totally_bogus",
+            imatrix = None,
+        )
+
+
+# -- --imatrix actually reaches llama-quantize ---------------------------------------------
+
+
+@_needs_zoo_imatrix
+def test_quantize_gguf_emits_imatrix_flag(monkeypatch, tmp_path):
+    captured = {}
+
+    def _fake_run(command, *a, **kw):
+        captured["command"] = command
+        # llama-quantize would write the output; emulate so the existence check passes.
+        out = command.split()[-2] if False else None
+        # output_gguf is the 2nd-to-last token before quant_type/threads; just create it.
+        with open(tmp_path / "out.gguf", "wb") as f:
+            f.write(b"GGUF")
+
+        class R:
+            returncode = 0
+            stdout = ""
+
+        return R()
+
+    import shlex
+
+    monkeypatch.setattr(L.subprocess, "run", _fake_run)
+    imat = str(tmp_path / "imatrix it.dat")  # space in path -> must be shell-quoted
+    with open(imat, "wb") as f:  # quantize_gguf validates the imatrix exists before running
+        f.write(b"\x00")
+    L.quantize_gguf(
+        input_gguf = str(tmp_path / "in.gguf"),
+        output_gguf = str(tmp_path / "out.gguf"),
+        quant_type = "iq4_xs",
+        quantizer_location = "llama-quantize",
+        n_threads = 4,
+        imatrix = imat,
+        print_output = False,
+    )
+    cmd = captured["command"]
+    assert "--imatrix" in cmd
+    assert "iq4_xs" in cmd
+    # the path with a space must appear shell-quoted (shlex.quote), never bare
+    assert f"--imatrix {shlex.quote(imat)}" in cmd
+
+
+def test_quantize_gguf_no_imatrix_has_no_flag(monkeypatch, tmp_path):
+    captured = {}
+
+    def _fake_run(command, *a, **kw):
+        captured["command"] = command
+        with open(tmp_path / "out.gguf", "wb") as f:
+            f.write(b"GGUF")
+
+        class R:
+            returncode = 0
+            stdout = ""
+
+        return R()
+
+    monkeypatch.setattr(L.subprocess, "run", _fake_run)
+    L.quantize_gguf(
+        input_gguf = str(tmp_path / "in.gguf"),
+        output_gguf = str(tmp_path / "out.gguf"),
+        quant_type = "q4_k_m",
+        quantizer_location = "llama-quantize",
+        n_threads = 4,
+        print_output = False,
+    )
+    assert "--imatrix" not in captured["command"]

tests/saving/test_save_shell_injection.py

@@ -7,63 +7,83 @@ from pathlib import Path
 SAVE_PY = Path(__file__).resolve().parents[2] / "unsloth" / "save.py"
 
 
-def _function_calls(source: str, function_name: str) -> list[ast.Call]:
+def _get_function(source: str, function_name: str) -> ast.FunctionDef:
     tree = ast.parse(source, filename = str(SAVE_PY))
     for node in tree.body:
         if isinstance(node, ast.FunctionDef) and node.name == function_name:
-            return [child for child in ast.walk(node) if isinstance(child, ast.Call)]
+            return node
     raise AssertionError(f"Function {function_name} not found in save.py")
 
 
-def _assert_safe_ggml_calls(calls: list[ast.Call]) -> None:
-    popen_calls = []
-    for call in calls:
-        if isinstance(call.func, ast.Attribute) and call.func.attr == "Popen":
-            if isinstance(call.func.value, ast.Name) and call.func.value.id == "subprocess":
-                popen_calls.append(call)
+def _popen_calls(node: ast.AST) -> list[ast.Call]:
+    calls = []
+    for child in ast.walk(node):
+        if (
+            isinstance(child, ast.Call)
+            and isinstance(child.func, ast.Attribute)
+            and child.func.attr == "Popen"
+            and isinstance(child.func.value, ast.Name)
+            and child.func.value.id == "subprocess"
+        ):
+            calls.append(child)
+    return calls
 
-    assert popen_calls, "Expected at least one subprocess.Popen call"
 
-    ggml_calls = []
+def _list_assignments(node: ast.AST, target: str) -> list[ast.List]:
+    lists = []
+    for child in ast.walk(node):
+        if isinstance(child, ast.Assign) and isinstance(child.value, ast.List):
+            if any(isinstance(t, ast.Name) and t.id == target for t in child.targets):
+                lists.append(child.value)
+    return lists
+
+
+def test_lora_gguf_conversion_does_not_use_shell() -> None:
+    """The LoRA -> GGUF conversion must pass argv as a list (no shell=True), so a crafted
+    save path cannot inject shell commands. The conversion lives in the shared helper now."""
+    helper = _get_function(SAVE_PY.read_text(encoding = "utf-8"), "_unsloth_save_lora_gguf")
+    popen_calls = _popen_calls(helper)
+    assert popen_calls, "Expected at least one subprocess.Popen call in _unsloth_save_lora_gguf"
+
     for call in popen_calls:
-        if not call.args:
-            continue
-        argv = call.args[0]
-        if isinstance(argv, ast.List) and len(argv.elts) >= 2:
-            second_arg = argv.elts[1]
-            if (
-                isinstance(second_arg, ast.Constant)
-                and second_arg.value == "llama.cpp/convert-lora-to-ggml.py"
-            ):
-                ggml_calls.append(call)
-
-    assert ggml_calls, "Expected the GGML conversion subprocess call"
-
-    for call in ggml_calls:
-        shell_kwargs = [
-            keyword
-            for keyword in call.keywords
-            if keyword.arg == "shell"
-            and isinstance(keyword.value, ast.Constant)
-            and keyword.value.value is True
+        shell = [
+            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_kwargs, "subprocess.Popen must not use shell=True"
+        assert not shell, "subprocess.Popen must not use shell=True"
 
         assert call.args, "subprocess.Popen must receive argv as a positional argument"
         argv = call.args[0]
-        assert isinstance(argv, ast.List), "subprocess.Popen must be called with an argv list"
-        assert len(argv.elts) == 5, "GGML conversion argv should have five elements"
+        if isinstance(argv, ast.List):
+            elts = argv.elts
+        else:
+            # argv is built as a list variable (cmd = [...]) and passed positionally.
+            assert isinstance(argv, ast.Name), "argv must be a list or a list-built variable"
+            assigned = _list_assignments(helper, argv.id)
+            assert assigned, f"argv variable '{argv.id}' must be assigned a list literal"
+            elts = assigned[0].elts
 
-        second_arg = argv.elts[1]
-        assert isinstance(second_arg, ast.Constant)
-        assert second_arg.value == "llama.cpp/convert-lora-to-ggml.py"
+        assert len(elts) >= 2, "argv must include the interpreter and the converter script"
+        first = elts[0]
+        assert (
+            isinstance(first, ast.Attribute) and first.attr == "executable"
+        ), "argv[0] should be sys.executable, not a shell string"
 
 
-def test_ggml_conversion_paths_do_not_use_shell() -> None:
+def test_legacy_ggml_wrappers_delegate_safely() -> None:
+    """The legacy ggml entry points must delegate to the shared helper and not build their
+    own subprocess invocation."""
     source = SAVE_PY.read_text(encoding = "utf-8")
     for function_name in (
         "unsloth_convert_lora_to_ggml_and_push_to_hub",
         "unsloth_convert_lora_to_ggml_and_save_locally",
     ):
-        calls = _function_calls(source, function_name)
-        _assert_safe_ggml_calls(calls)
+        node = _get_function(source, function_name)
+        calls = [c for c in ast.walk(node) if isinstance(c, ast.Call)]
+        assert any(
+            isinstance(c.func, ast.Name) and c.func.id == "_unsloth_save_lora_gguf" for c in calls
+        ), f"{function_name} should delegate to _unsloth_save_lora_gguf"
+        assert not _popen_calls(
+            node
+        ), f"{function_name} should not call subprocess.Popen directly anymore"

unsloth/_compressed_quantize.py

@@ -0,0 +1,347 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Standalone llm-compressor runner for Unsloth's FP8/FP4 export.
+
+Launched as a subprocess by file path (not `python -m`) so the Unsloth package, which patches
+transformers attention, is not imported here; llm-compressor needs an unpatched forward for
+calibration (e.g. NVFP4). Reads a merged 16bit checkpoint, writes a compressed-tensors one.
+"""
+
+import argparse
+import glob
+import json
+import os
+import sys
+
+
+def _is_moe(config):
+    """True if the model config looks like a sparse Mixture-of-Experts model."""
+    if config is None:
+        return False
+    for cfg in (config, getattr(config, "text_config", None)):
+        if cfg is None:
+            continue
+        for attr in ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts"):
+            v = getattr(cfg, attr, None)
+            if isinstance(v, int) and v > 1:
+                return True
+    return "moe" in (getattr(config, "model_type", "") or "").lower()
+
+
+def _has_mtp(config):
+    """True if the model carries MTP / speculative-decoding layers (e.g. Qwen3-Next, DeepSeek)."""
+    if config is None:
+        return False
+    mt = (getattr(config, "model_type", "") or "").lower()
+    if "qwen3_next" in mt or "mtp" in mt:
+        return True
+    for attr in ("num_nextn_predict_layers", "num_mtp_layers", "mtp_num_layers"):
+        v = getattr(config, attr, None)
+        if isinstance(v, int) and v > 0:
+            return True
+    return False
+
+
+def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_length):
+    from datasets import DatasetDict, load_dataset, load_from_disk
+
+    _tok = tokenizer.tokenizer if hasattr(tokenizer, "tokenizer") else tokenizer
+
+    if kind == "none":
+        print(
+            f"Unsloth: NVFP4 needs calibration data. Defaulting to {num_samples} samples of "
+            "HuggingFaceH4/ultrachat_200k. For best accuracy pass your own training data via "
+            "`calibration_dataset=...`.",
+            flush = True,
+        )
+        ds = load_dataset("HuggingFaceH4/ultrachat_200k", split = f"train_sft[:{num_samples}]")
+        ds = ds.shuffle(seed = 42)
+    elif kind == "hfid":
+        # Not every dataset has a "train" split (e.g. train_sft only); fall back to the first one.
+        try:
+            ds = load_dataset(value, split = f"train[:{num_samples}]")
+        except (ValueError, KeyError):
+            from datasets import get_dataset_split_names
+            try:
+                # Resolve the first split name so only num_samples rows are fetched, instead of
+                # downloading/materializing the whole dataset just to take a small slice.
+                split = get_dataset_split_names(value)[0]
+                ds = load_dataset(value, split = f"{split}[:{num_samples}]")
+            except Exception:
+                # Last resort: materialize, then subselect (preserves the original behavior).
+                ds = load_dataset(value)
+                if isinstance(ds, DatasetDict):
+                    ds = ds[next(iter(ds.keys()))]
+                if num_samples and len(ds) > num_samples:
+                    ds = ds.select(range(num_samples))
+        ds = ds.shuffle(seed = 42)
+    elif kind == "disk":
+        ds = load_from_disk(value)
+        if isinstance(ds, DatasetDict):
+            if "train" in ds:
+                ds = ds["train"]
+            elif len(ds) == 1:
+                ds = next(iter(ds.values()))
+            else:
+                raise RuntimeError(
+                    "Unsloth: disk calibration_dataset is a DatasetDict with multiple splits; "
+                    "pass a single split, e.g. calibration_dataset=dataset['train']."
+                )
+        if num_samples and len(ds) > num_samples:
+            ds = ds.shuffle(seed = 42).select(range(num_samples))
+    else:
+        raise ValueError(f"Unknown calibration-dataset-kind: {kind}")
+
+    try:
+        if len(ds) == 0:
+            raise RuntimeError(
+                "Unsloth: the calibration dataset is empty after loading/subsampling; "
+                "pass a non-empty calibration_dataset."
+            )
+    except TypeError:
+        pass  # streaming / iterable datasets have no len(); let llm-compressor handle them
+
+    cols = set(ds.column_names)
+    if "input_ids" in cols:
+        # Drop non-model-input columns (e.g. a leftover 'messages' list) so llm-compressor's
+        # collator does not try to batch them.
+        keep = {"input_ids", "attention_mask", "labels", "position_ids"}
+        extra = [c for c in ds.column_names if c not in keep]
+        if extra:
+            ds = ds.remove_columns(extra)
+        return ds
+    if "messages" in cols:
+        # Base / non-chat tokenizers have no chat template; concatenate message contents instead
+        # of calling apply_chat_template (which would raise).
+        has_chat_template = bool(getattr(_tok, "chat_template", None))
+
+        def _content_to_text(content):
+            # content may be a str, None, or a multimodal list of parts (str or {"text": ...}).
+            if content is None:
+                return ""
+            if isinstance(content, str):
+                return content
+            if isinstance(content, (list, tuple)):
+                parts = []
+                for part in content:
+                    if isinstance(part, str):
+                        parts.append(part)
+                    elif isinstance(part, dict):
+                        text = part.get("text") or part.get("content")
+                        if isinstance(text, str):
+                            parts.append(text)
+                return " ".join(parts)
+            return str(content)
+
+        def _prep(ex):
+            msgs = ex["messages"] or []
+            if has_chat_template:
+                return {"text": _tok.apply_chat_template(msgs, tokenize = False)}
+            return {"text": "\n".join(_content_to_text(m.get("content")) for m in msgs)}
+
+        ds = ds.map(_prep)
+    elif "text" not in cols:
+        raise RuntimeError(
+            "Unsloth: calibration_dataset must contain a 'messages', 'text', or 'input_ids' "
+            f"column (got: {sorted(cols)})."
+        )
+
+    def _tokenize(sample):
+        return _tok(
+            sample["text"],
+            padding = False,
+            max_length = max_seq_length,
+            truncation = True,
+            add_special_tokens = False,
+        )
+
+    return ds.map(_tokenize, remove_columns = ds.column_names)
+
+
+def _from_pretrained(auto_model, model_path, trust_remote_code):
+    import torch
+
+    # transformers renamed torch_dtype -> dtype; support both.
+    try:
+        return auto_model.from_pretrained(
+            model_path,
+            device_map = "auto",
+            low_cpu_mem_usage = True,
+            trust_remote_code = trust_remote_code,
+            dtype = torch.bfloat16,
+        )
+    except TypeError:
+        return auto_model.from_pretrained(
+            model_path,
+            device_map = "auto",
+            low_cpu_mem_usage = True,
+            trust_remote_code = trust_remote_code,
+            torch_dtype = torch.bfloat16,
+        )
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required = True, help = "merged 16bit HF checkpoint dir")
+    ap.add_argument("--scheme", required = True)
+    ap.add_argument("--out", required = True)
+    ap.add_argument("--needs-calibration", action = "store_true")
+    ap.add_argument("--calibration-dataset-kind", default = "none", choices = ["none", "hfid", "disk"])
+    ap.add_argument("--calibration-dataset", default = "")
+    ap.add_argument("--num-calibration-samples", type = int, default = 512)
+    ap.add_argument("--max-seq-length", type = int, default = 2048)
+    ap.add_argument("--is-vlm", action = "store_true")
+    ap.add_argument("--trust-remote-code", action = "store_true")
+    ap.add_argument("--variant", default = "", help = "weight-filename variant for the output shards")
+    args = ap.parse_args()
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from llmcompressor import oneshot
+    from llmcompressor.modifiers.quantization import QuantizationModifier
+
+    # Import the VLM auto-class only when needed - some transformers versions lack it, and the
+    # text path must not fail just because that newer class is unavailable.
+    if args.is_vlm:
+        from transformers import AutoProcessor
+        try:
+            from transformers import AutoModelForImageTextToText as _VLMModel
+        except ImportError:
+            try:
+                from transformers import AutoModelForVision2Seq as _VLMModel
+            except ImportError as e:
+                raise RuntimeError(
+                    "Unsloth: this transformers version has no VLM auto-model class for "
+                    "compressed multimodal export. Please upgrade transformers."
+                ) from e
+        auto_model, auto_proc = _VLMModel, AutoProcessor
+    else:
+        auto_model, auto_proc = AutoModelForCausalLM, AutoTokenizer
+
+    model = _from_pretrained(auto_model, args.model, args.trust_remote_code)
+    model.eval()
+    # A tokenizer may be absent if the caller saved it separately; only calibration needs one.
+    try:
+        tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code)
+    except Exception:
+        if args.needs_calibration:
+            raise RuntimeError(
+                f"Unsloth: calibration export needs a tokenizer but none was found in {args.model}. "
+                "Pass tokenizer=... to save_pretrained_merged."
+            )
+        tokenizer = None
+
+    # MoE models: keep the router/gate unquantized (it decides expert routing) and calibrate every
+    # expert even if the sample set does not route tokens to all of them.
+    is_moe = _is_moe(getattr(model, "config", None))
+    ignore = ["lm_head"]
+    if is_moe:
+        # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate.
+        ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"]
+    moe_kwargs = {"moe_calibrate_all_experts": True} if is_moe else {}
+
+    def _make_recipe():
+        return QuantizationModifier(targets = "Linear", scheme = args.scheme, ignore = ignore)
+
+    if args.needs_calibration:
+        ds = _build_calibration_dataset(
+            tokenizer,
+            args.calibration_dataset_kind,
+            args.calibration_dataset,
+            args.num_calibration_samples,
+            args.max_seq_length,
+        )
+        # Use the sequential pipeline: it onloads layer-by-layer, so models that do not fit in
+        # memory at once can still calibrate. Running here in a clean process (Unsloth's attention
+        # patches are absent) means tracing works; fall back to the memory-hungry "basic" pipeline
+        # only if tracing fails.
+        try:
+            oneshot(
+                model = model,
+                dataset = ds,
+                recipe = _make_recipe(),
+                max_seq_length = args.max_seq_length,
+                num_calibration_samples = args.num_calibration_samples,
+                pipeline = "sequential",
+                **moe_kwargs,
+            )
+        except Exception as e:
+            print(
+                f"Unsloth: sequential calibration pipeline failed ({type(e).__name__}: {e}); "
+                "retrying with the 'basic' pipeline (needs the full model to fit in memory).",
+                flush = True,
+            )
+            # Free the partially-processed model before loading a fresh copy, so the fallback does
+            # not transiently hold two copies on GPU. llm-compressor keeps the model in a global
+            # session after a failed run, so reset it first; also drop the traceback frames (e) and
+            # the local reference that pin the model.
+            import gc as _gc
+            import torch as _torch
+
+            try:
+                from llmcompressor.core import reset_session
+                reset_session()
+            except Exception:
+                pass
+            e = None
+            del model
+            _gc.collect()
+            if _torch.cuda.is_available():
+                _torch.cuda.empty_cache()
+            model = _from_pretrained(auto_model, args.model, args.trust_remote_code)
+            model.eval()
+            oneshot(
+                model = model,
+                dataset = ds,
+                recipe = _make_recipe(),
+                max_seq_length = args.max_seq_length,
+                num_calibration_samples = args.num_calibration_samples,
+                pipeline = "basic",
+                **moe_kwargs,
+            )
+    else:
+        oneshot(model = model, recipe = _make_recipe())
+
+    os.makedirs(args.out, exist_ok = True)
+    save_kwargs = {"variant": args.variant} if args.variant else {}
+    model.save_pretrained(args.out, save_compressed = True, **save_kwargs)
+    if tokenizer is not None:
+        tokenizer.save_pretrained(args.out)
+
+    if _has_mtp(getattr(model, "config", None)):
+        print(
+            "Unsloth: WARNING - this model has MTP / speculative-decoding tensors that are not "
+            "included in the compressed export (only the main model is quantized and saved). Use "
+            "the non-compressed save path if you need the MTP weights.",
+            flush = True,
+        )
+
+    cfg_path = os.path.join(args.out, "config.json")
+    cfg = {}
+    if os.path.exists(cfg_path):
+        with open(cfg_path, "r", encoding = "utf-8") as f:
+            cfg = json.load(f)
+    if "quantization_config" not in cfg:
+        print(f"Unsloth: ERROR - no quantization_config written to {cfg_path}", flush = True)
+        sys.exit(2)
+    shards = glob.glob(os.path.join(args.out, "*.safetensors"))
+    qfmt = cfg["quantization_config"].get("format")
+    print(
+        f"[compressed-quantize] OK scheme={args.scheme} format={qfmt} "
+        f"shards={len(shards)} -> {args.out}",
+        flush = True,
+    )
+
+
+if __name__ == "__main__":
+    main()