diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index bd9c0a532d..7978a200c0 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.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 \ diff --git a/images/studio_demo/step01_configure.png b/images/studio_demo/step01_configure.png new file mode 100644 index 0000000000..8aaa5da51c Binary files /dev/null and b/images/studio_demo/step01_configure.png differ diff --git a/images/studio_demo/step02_training.png b/images/studio_demo/step02_training.png new file mode 100644 index 0000000000..cfa90317d2 Binary files /dev/null and b/images/studio_demo/step02_training.png differ diff --git a/images/studio_demo/step03_training_done.png b/images/studio_demo/step03_training_done.png new file mode 100644 index 0000000000..03f6aa651a Binary files /dev/null and b/images/studio_demo/step03_training_done.png differ diff --git a/images/studio_demo/step04_export_source.png b/images/studio_demo/step04_export_source.png new file mode 100644 index 0000000000..cc0ed76d6b Binary files /dev/null and b/images/studio_demo/step04_export_source.png differ diff --git a/images/studio_demo/step05_gguf_imatrix_configured.png b/images/studio_demo/step05_gguf_imatrix_configured.png new file mode 100644 index 0000000000..f94ef752b6 Binary files /dev/null and b/images/studio_demo/step05_gguf_imatrix_configured.png differ diff --git a/images/studio_demo/step06_gguf_imatrix_success.png b/images/studio_demo/step06_gguf_imatrix_success.png new file mode 100644 index 0000000000..5d1c794ec3 Binary files /dev/null and b/images/studio_demo/step06_gguf_imatrix_success.png differ diff --git a/images/studio_demo/step07_fp8_configured.png b/images/studio_demo/step07_fp8_configured.png new file mode 100644 index 0000000000..b3127c21a6 Binary files /dev/null and b/images/studio_demo/step07_fp8_configured.png differ diff --git a/images/studio_demo/step08_fp8_success.png b/images/studio_demo/step08_fp8_success.png new file mode 100644 index 0000000000..c0d47ea635 Binary files /dev/null and b/images/studio_demo/step08_fp8_success.png differ diff --git a/images/studio_demo/step09_all_formats_dropdown.png b/images/studio_demo/step09_all_formats_dropdown.png new file mode 100644 index 0000000000..3ba771ce81 Binary files /dev/null and b/images/studio_demo/step09_all_formats_dropdown.png differ diff --git a/images/studio_demo/step10_mxfp4_success.png b/images/studio_demo/step10_mxfp4_success.png new file mode 100644 index 0000000000..e4da077a02 Binary files /dev/null and b/images/studio_demo/step10_mxfp4_success.png differ diff --git a/images/studio_demo/step11_nvfp4_success.png b/images/studio_demo/step11_nvfp4_success.png new file mode 100644 index 0000000000..8ce3d7503b Binary files /dev/null and b/images/studio_demo/step11_nvfp4_success.png differ diff --git a/images/studio_export_demo.gif b/images/studio_export_demo.gif new file mode 100644 index 0000000000..1d60504edf Binary files /dev/null and b/images/studio_export_demo.gif differ 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..7955b50968 --- /dev/null +++ b/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 `. = ...` 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)) 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, + ) diff --git a/tests/saving/test_gguf_export_and_inference.py b/tests/saving/test_gguf_export_and_inference.py new file mode 100644 index 0000000000..aa69368482 --- /dev/null +++ b/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 "_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/-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" diff --git a/tests/saving/test_imatrix_export.py b/tests/saving/test_imatrix_export.py new file mode 100644 index 0000000000..6e5b06d7cc --- /dev/null +++ b/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/-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/-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"] diff --git a/tests/saving/test_save_shell_injection.py b/tests/saving/test_save_shell_injection.py index b02748c250..5f55137771 100644 --- a/tests/saving/test_save_shell_injection.py +++ b/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" diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py new file mode 100644 index 0000000000..66ebdd75da --- /dev/null +++ b/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() diff --git a/unsloth/save.py b/unsloth/save.py index 20a934538c..76bc6aa733 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -135,12 +135,25 @@ ALLOWED_QUANTS = { "q5_1": "Even higher accuracy, resource usage and slower inference.", "q5_k_s": "Uses Q5_K for all tensors", "q6_k": "Uses Q8_K for all tensors", - # "iq2_xxs" : "2.06 bpw quantization", # Not supported sadly - # "iq2_xs" : "2.31 bpw quantization", - # "iq3_xxs" : "3.06 bpw quantization", "q3_k_xs": "3-bit extra small quantization", } +# IQ (importance-matrix) quants. llama.cpp refuses these without an imatrix, so they are only +# accepted when imatrix_file=... is supplied to save_pretrained_gguf / push_to_hub_gguf. +IMATRIX_QUANTS = { + "iq1_s": "1.56 bpw. Smallest, lowest quality. Needs an imatrix.", + "iq1_m": "1.75 bpw. Very small. Needs an imatrix.", + "iq2_xxs": "2.06 bpw. Needs an imatrix.", + "iq2_xs": "2.31 bpw. Needs an imatrix.", + "iq2_s": "2.5 bpw. Needs an imatrix.", + "iq2_m": "2.7 bpw. Needs an imatrix.", + "iq3_xxs": "3.06 bpw. Needs an imatrix.", + "iq3_s": "3.44 bpw. Needs an imatrix.", + "iq3_m": "3.66 bpw. Needs an imatrix.", + "iq4_nl": "4.5 bpw non-linear. Benefits from an imatrix.", + "iq4_xs": "4.25 bpw. Benefits from an imatrix.", +} + def has_curl(): return shutil.which("curl") is not None @@ -149,6 +162,70 @@ def has_curl(): CURL_FLAG = "-DLLAMA_CURL=ON" if has_curl() else "-DLLAMA_CURL=OFF" +# FP8/FP4 compressed export via llm-compressor (for vLLM). +# save_method alias -> (llm-compressor scheme, needs_calibration, output dir suffix). +# alias -> (llm-compressor scheme, needs_calibration, output-dir suffix). needs_calibration is +# True only for schemes with static activation scales (FP8 static, NVFP4); everything else is +# weight-only or dynamic-activation and runs data-free. Unsupported schemes in the installed +# compressed-tensors (e.g. MXFP8 on older stacks) are gated by _scheme_is_available at runtime. +COMPRESSED_EXPORT_SCHEMES = { + # FP8 + "fp8": ("FP8_DYNAMIC", False, "fp8"), + "fp8_dynamic": ("FP8_DYNAMIC", False, "fp8"), + "dynamic_fp8": ("FP8_DYNAMIC", False, "fp8"), + "w8a8_fp8": ("FP8_DYNAMIC", False, "fp8"), + "fp8_static": ("FP8", True, "fp8-static"), + "static_fp8": ("FP8", True, "fp8-static"), + "fp8_block": ("FP8_BLOCK", False, "fp8-block"), + "block_fp8": ("FP8_BLOCK", False, "fp8-block"), + # INT8 / INT-weight + "int8": ("INT8", False, "int8"), + "w8a8": ("W8A8", False, "w8a8"), + "w8a8_int8": ("W8A8", False, "w8a8"), + "w8a16": ("W8A16", False, "w8a16"), + "int8_weight": ("W8A16", False, "w8a16"), + "w4a16": ("W4A16", False, "w4a16"), + "int4": ("W4A16", False, "w4a16"), + "int4_weight": ("W4A16", False, "w4a16"), + "w4a16_asym": ("W4A16_ASYM", False, "w4a16-asym"), + "w4a8": ("W4A8", False, "w4a8"), + "w4afp8": ("W4AFP8", False, "w4afp8"), + # MXFP (microscaling) + "mxfp8": ("MXFP8", False, "mxfp8"), + "w8a8_mxfp8": ("MXFP8", False, "mxfp8"), + "mxfp4": ("MXFP4", False, "mxfp4"), + "w4a4_mxfp4": ("MXFP4", False, "mxfp4"), + "mxfp4a16": ("MXFP4A16", False, "mxfp4a16"), + "w4a16_mxfp4": ("MXFP4A16", False, "mxfp4a16"), + # NVFP4 + "nvfp4": ("NVFP4", True, "nvfp4"), + "w4a4_nvfp4": ("NVFP4", True, "nvfp4"), + "nvfp4a16": ("NVFP4A16", False, "nvfp4a16"), + "w4a16_nvfp4": ("NVFP4A16", False, "nvfp4a16"), +} + + +def _normalize_compressed_method(save_method): + """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed + export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). + + Near-miss FP8/FP4 names that are not supported raise a precise error instead of silently + falling through to the generic "unknown save_method" message. + """ + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + if key in COMPRESSED_EXPORT_SCHEMES: + return COMPRESSED_EXPORT_SCHEMES[key] + if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): + supported = ", ".join(sorted(COMPRESSED_EXPORT_SCHEMES.keys())) + raise RuntimeError( + f"Unsloth: save_method='{save_method}' is not a supported compressed export.\n" + f"Supported compressed-tensors export methods: {supported}" + ) + return None + + def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool: """ True if llama.cpp's Makefile is the post-CMake-migration deprecation stub, @@ -175,6 +252,17 @@ def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool: def print_quantization_methods(): for key, value in ALLOWED_QUANTS.items(): print(f'"{key}" ==> {value}') + print("\nIQ low-bit quants (save_pretrained_gguf(..., imatrix_file=True or '...path')):") + for key, value in IMATRIX_QUANTS.items(): + print(f'"{key}" ==> {value}') + print("\nCompressed-tensors export (save_pretrained_merged(..., save_method=...), for vLLM):") + seen = set() + for key, (scheme, needs_calib, _suffix) in COMPRESSED_EXPORT_SCHEMES.items(): + if scheme in seen: + continue + seen.add(scheme) + note = "needs calibration data" if needs_calib else "data-free" + print(f'"{key}" ==> llm-compressor {scheme} ({note})') def _quantize_q2_k_l( @@ -183,11 +271,13 @@ def _quantize_q2_k_l( quantizer_location: Union[str, os.PathLike], n_threads: int, print_output: bool = True, + imatrix = None, ): # "Q2_K_L" is an Unsloth preset, not a native llama.cpp ftype: q2_k with # output/token-embedding tensors kept at q8_0 for higher precision. command = [ str(quantizer_location), + *(["--imatrix", str(imatrix)] if imatrix else []), "--output-tensor-type", "q8_0", "--token-embedding-type", @@ -1273,6 +1363,89 @@ def install_python_non_blocking(packages = []): return run_installer +def install_llm_compressor(): + """Import llm-compressor, installing it on first use for FP8/FP4 export. + + Pins the current torch + transformers so pip does not upgrade them (a plain install pulls + transformers>=5 and breaks Unsloth). Returns (oneshot, QuantizationModifier). + """ + try: + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + return oneshot, QuantizationModifier + except Exception: + pass + + print( + "Unsloth: Installing llm-compressor for FP8/FP4 export " + "(pinning your torch + transformers so they are not upgraded). " + "This can take a few minutes..." + ) + import importlib + import tempfile + + constraints = "" + try: + import torch as _torch + constraints += f"torch=={_torch.__version__.split('+')[0]}\n" + except Exception: + pass + try: + import transformers as _tf + constraints += f"transformers=={_tf.__version__}\n" + except Exception: + pass + + # Prefer pip, but fall back to uv when this interpreter has no pip seeded (common in + # uv-created / relocatable venvs), so the export does not hard-fail with "No module named pip". + import importlib.util + + if importlib.util.find_spec("pip") is not None: + cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + elif shutil.which("uv") is not None: + cmd = ["uv", "pip", "install", "--python", sys.executable, "llmcompressor"] + else: + raise RuntimeError( + "Unsloth: cannot install llm-compressor because this environment has neither pip nor " + f"uv. Install it manually with:\n uv pip install --python {sys.executable} llmcompressor\n" + "(pin torch and transformers to your current versions to avoid upgrading them)." + ) + cpath = None + if constraints: + with tempfile.NamedTemporaryFile("w", suffix = ".txt", delete = False) as f: + f.write(constraints) + cpath = f.name + cmd += ["-c", cpath] + try: + subprocess.check_call(cmd) + except subprocess.CalledProcessError as e: + raise RuntimeError( + "Unsloth: Failed to install llm-compressor. Install it manually with:\n" + f" uv pip install --python {sys.executable} llmcompressor\n" + f"or, if pip is available:\n {sys.executable} -m pip install llmcompressor\n" + "(pin torch and transformers to your current versions to avoid upgrading them).\n" + f"Underlying error: {e}" + ) + finally: + if cpath is not None: + try: + os.remove(cpath) + except Exception: + pass + + importlib.invalidate_caches() + try: + from llmcompressor import oneshot + from llmcompressor.modifiers.quantization import QuantizationModifier + except Exception as e: + raise RuntimeError( + "Unsloth: llm-compressor was installed but could not be imported. " + "Please restart your Python session and try again.\n" + f"Underlying error: {repr(e)}" + ) + return oneshot, QuantizationModifier + + def try_execute(commands, force_complete = False): for command in commands: with subprocess.Popen( @@ -1438,10 +1611,13 @@ def save_to_gguf( first_conversion: str = None, is_vlm: bool = False, is_gpt_oss: bool = False, + imatrix = None, ): """ Orchestrates the complete GGUF conversion process. Handles installation, conversion, and quantization. + `imatrix` is a local importance-matrix path (already resolved); it is forwarded to + llama-quantize and is required for the IQ low-bit quant types. """ # print_output True only if UNSLOTH_ENABLE_LOGGING=1 if os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1": @@ -1477,11 +1653,15 @@ def save_to_gguf( if first_conversion is None: first_conversion = model_dtype - # Check I quants - for quant_method in quantization_method: - if quant_method.startswith("iq2"): + has_imatrix = imatrix is not None and str(imatrix) != "" + if has_imatrix: + # quantize_gguf gained the imatrix kwarg in a recent unsloth_zoo; fail fast (before the + # expensive conversion) if the installed version cannot apply it, rather than dropping it. + import inspect + if "imatrix" not in inspect.signature(quantize_gguf).parameters: raise RuntimeError( - "Unsloth: Currently iq2 type quantizations aren't supported yet - sorry!" + "Unsloth: your installed unsloth_zoo's quantize_gguf does not support imatrix.\n" + "Please upgrade it: uv pip install --upgrade unsloth_zoo" ) # Map quant methods @@ -1496,11 +1676,20 @@ def save_to_gguf( elif quant_method is None: quant_method = "q8_0" - # Check if wrong method - if quant_method not in ALLOWED_QUANTS.keys(): + # IQ low-bit quants are only valid with an imatrix; other methods use the normal allow-list. + if quant_method in IMATRIX_QUANTS: + if not has_imatrix: + raise RuntimeError( + f"Unsloth: quant method '{quant_method}' is an IQ low-bit quant that requires an " + "importance matrix. Pass imatrix_file=True (to fetch the upstream Unsloth imatrix) " + "or imatrix_file='/path/to/imatrix' to save_pretrained_gguf / push_to_hub_gguf." + ) + elif quant_method not in ALLOWED_QUANTS.keys(): error = f"Unsloth: Quant method = [{quant_method}] not supported. Choose from below:\n" for key, value in ALLOWED_QUANTS.items(): error += f"[{key}] => {value}\n" + for key, value in IMATRIX_QUANTS.items(): + error += f"[{key}] => {value} (needs imatrix_file)\n" raise RuntimeError(error) new_quantization_methods.append(quant_method) @@ -1654,16 +1843,22 @@ def save_to_gguf( quantizer_location = quantizer_location, n_threads = n_cpus, print_output = print_output, + imatrix = imatrix, ) else: - # Use unsloth-zoo's standard quantization for all other methods - quantized_file = quantize_gguf( + # Use unsloth-zoo's standard quantization for all other methods. Only pass + # imatrix when set so older unsloth_zoo (no imatrix kwarg) still works for + # plain quants; an imatrix that cannot be applied was rejected above. + quant_kwargs = dict( input_gguf = base_gguf, output_gguf = output_location, quant_type = quant_method, quantizer_location = quantizer_location, print_output = print_output, ) + if has_imatrix: + quant_kwargs["imatrix"] = imatrix + quantized_file = quantize_gguf(**quant_kwargs) all_saved_locations.append(quantized_file) quants_created = True except Exception as e: @@ -1751,6 +1946,9 @@ def unsloth_save_pretrained_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .save_pretrained(...) except 4bit weights are auto @@ -1760,6 +1958,9 @@ def unsloth_save_pretrained_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM (`fp8`, `mxfp4`, `nvfp4`, `mxfp8`): keeps the + 16bit merge at `save_directory` and writes the quantized checkpoint to + `save_directory + "-"`. """ if tokenizer is None: logger.warning_once( @@ -1767,9 +1968,46 @@ def unsloth_save_pretrained_merged( "You can do it separately via `tokenizer.save_pretrained(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_save_model(**arguments) for _ in range(3): gc.collect() @@ -1779,7 +2017,7 @@ def unsloth_push_to_hub_merged( self, repo_id: str, tokenizer = None, - save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit"] + save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit", "fp8", "mxfp4", "nvfp4", "mxfp8"] use_temp_dir: Optional[bool] = None, commit_message: Optional[str] = "Trained with Unsloth", private: Optional[bool] = None, @@ -1793,6 +2031,9 @@ def unsloth_push_to_hub_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -1802,6 +2043,7 @@ def unsloth_push_to_hub_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM: `fp8`, `mxfp4`, `nvfp4`, `mxfp8`. """ if tokenizer is None: logger.warning_once( @@ -1809,12 +2051,50 @@ def unsloth_push_to_hub_merged( "You can do it separately via `tokenizer.push_to_hub(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id arguments["push_to_hub"] = True del arguments["self"] del arguments["repo_id"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_save_model(**arguments) for _ in range(3): gc.collect() @@ -2226,11 +2506,17 @@ def unsloth_save_pretrained_gguf( tags: List[str] = None, temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + save_method: str = None, + imatrix_file = None, ): """ Same as .save_pretrained(...) except 4bit weights are auto converted to float16 then converted to GGUF / llama.cpp format. + imatrix_file: importance matrix for llama-quantize. None = off; a path = use that file + (a *.gguf_file is renamed to *.gguf); True = download the upstream unsloth/-GGUF + imatrix. Required for the IQ low-bit quants (iq2_xxs, iq4_xs, ...). + Choose for `quantization_method` to be: "not_quantized" : "Recommended. Fast conversion. Slow inference, big files.", "fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.", @@ -2264,6 +2550,30 @@ def unsloth_save_pretrained_gguf( if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): tokenizer = patch_saving_functions(tokenizer) + # save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model). + if save_method is not None and str(save_method).lower() == "lora": + if not is_main_process: + return None + if push_to_hub: + raise ValueError( + "Unsloth: Please use .push_to_hub_gguf(save_method='lora') instead of " + ".save_pretrained_gguf(save_method='lora', push_to_hub=True)." + ) + _qm = quantization_method + if isinstance(_qm, (list, tuple)) and len(_qm) == 1: + _qm = _qm[0] # the gguf API allows a list; unwrap a single outtype + if _qm in _LORA_GGUF_OUTTYPES: + _outtype = _qm + else: + if _qm not in (None, "fast_quantized"): + logger.warning_once( + f"Unsloth: LoRA GGUF export does not support " + f"quantization_method={quantization_method!r}; using outtype 'f16'. " + f"Valid LoRA outtypes: {_LORA_GGUF_OUTTYPES}." + ) + _outtype = "f16" + return _unsloth_save_lora_gguf(self, tokenizer, save_directory, outtype = _outtype) + try: base_model_name = get_model_name(self.config._name_or_path, load_in_4bit = False) model_name = base_model_name.split("/")[-1] @@ -2321,6 +2631,7 @@ def unsloth_save_pretrained_gguf( del arguments["model_name"] del arguments["base_model_name"] del arguments["is_processor"] + del arguments["imatrix_file"] # only used by the gguf quantize step, not the 16bit merge # Step 3: Fix tokenizer BOS token if needed if is_processor: @@ -2328,6 +2639,11 @@ def unsloth_save_pretrained_gguf( else: fix_bos_token, old_chat_template = fix_tokenizer_bos_token(tokenizer) + # Resolve the importance matrix (download upstream / validate path / rename *.gguf_file) up + # front, so a bad path or an unavailable upstream imatrix fails before the expensive 16-bit + # merge, and a failed auto-resolution never reaches the IQ-quant gate. + imatrix_path = _resolve_imatrix_file(self, imatrix_file, token, save_directory) + # Step 4: Save/merge model to 16-bit format is_peft_model = isinstance(self, PeftModelForCausalLM) or isinstance(self, PeftModel) @@ -2443,6 +2759,7 @@ def unsloth_save_pretrained_gguf( first_conversion = first_conversion, is_vlm = is_vlm, # Pass VLM flag is_gpt_oss = is_gpt_oss, # Pass gpt_oss Flag + imatrix = imatrix_path, ) except Exception as e: if IS_KAGGLE_ENVIRONMENT: @@ -2539,11 +2856,16 @@ def unsloth_push_to_hub_gguf( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, datasets: Optional[List[str]] = None, + save_method: str = None, + imatrix_file = None, ): """ Same as .push_to_hub(...) except 4bit weights are auto converted to float16 then converted to GGUF / llama.cpp format. + imatrix_file: importance matrix for llama-quantize (None = off; a path; or True to download + the upstream unsloth/-GGUF imatrix). Required for the IQ low-bit quants. + Choose for `quantization_method` to be: "not_quantized" : "Recommended. Fast conversion. Slow inference, big files.", "fast_quantized" : "Recommended. Fast conversion. OK inference, OK file size.", @@ -2569,6 +2891,37 @@ def unsloth_push_to_hub_gguf( if tokenizer is None: raise ValueError("Unsloth: Saving to GGUF must have a tokenizer.") + # save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model). + if save_method is not None and str(save_method).lower() == "lora": + if not is_main_process: + return None # only the main rank converts and uploads, like the local lora branch + _qm = quantization_method + if isinstance(_qm, (list, tuple)) and len(_qm) == 1: + _qm = _qm[0] # the gguf API allows a list; unwrap a single outtype + if _qm in _LORA_GGUF_OUTTYPES: + _outtype = _qm + else: + if _qm not in (None, "fast_quantized"): + logger.warning_once( + f"Unsloth: LoRA GGUF export does not support " + f"quantization_method={quantization_method!r}; using outtype 'f16'. " + f"Valid LoRA outtypes: {_LORA_GGUF_OUTTYPES}." + ) + _outtype = "f16" + return _unsloth_save_lora_gguf( + self, + tokenizer, + repo_id, + outtype = _outtype, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + ) + # Step 1: Determine save directory model_name = repo_id.split("/")[-1] if "/" in repo_id else repo_id @@ -2594,11 +2947,12 @@ def unsloth_push_to_hub_gguf( quantization_method = quantization_method, first_conversion = first_conversion, push_to_hub = False, # Never push from here - token = None, # Don't need token for local save + token = token, # forwarded so imatrix_file=True can read a gated/private upstream max_shard_size = max_shard_size, safe_serialization = safe_serialization, temporary_location = temporary_location, maximum_memory_usage = maximum_memory_usage, + imatrix_file = imatrix_file, ) # Extract results @@ -2822,93 +3176,300 @@ def save_lora_to_custom_dir(model, tokenizer, save_directory): ) -# Corrected method within the model class to convert LoRA to GGML and push to Hugging Face Hub +# Valid output float types for llama.cpp's convert_lora_to_gguf.py. +_LORA_GGUF_OUTTYPES = ("f32", "f16", "bf16", "q8_0", "auto") + + +def _lora_base_model_id(model): + """Base model id for a PEFT model: prefer the active adapter's recorded base, else the + model config (the adapter's `base_model_name_or_path` is the authoritative source).""" + base = None + peft_config = getattr(model, "peft_config", None) + if isinstance(peft_config, dict) and peft_config: + adapter = getattr(model, "active_adapter", None) + if callable(adapter): + try: + adapter = adapter() + except Exception: + adapter = None + if isinstance(adapter, (list, tuple)): + adapter = adapter[0] if adapter else None + cfg = ( + peft_config.get(adapter) if adapter in peft_config else next(iter(peft_config.values())) + ) + base = getattr(cfg, "base_model_name_or_path", None) + if not base: + base = getattr(getattr(model, "config", None), "_name_or_path", None) + return os.fspath(base) if base else "" + + +# Upstream Unsloth GGUF repos ship a calibration imatrix under one of these names; the GGUF-format +# one is suffixed .gguf_file so the Hub does not list it as a model GGUF (renamed to .gguf locally). +_IMATRIX_UPSTREAM_NAMES = ("imatrix_unsloth.dat", "imatrix_unsloth.gguf_file") + + +def _gguf_repo_candidates(model): + """Ordered, de-duplicated unsloth/-GGUF repo ids to search for an upstream imatrix.""" + candidates = [] + raw_names = [ + _lora_base_model_id(model), + getattr(getattr(model, "config", None), "_name_or_path", None), + ] + for raw in raw_names: + if not raw: + continue + name = os.fspath(raw) + if os.path.isdir(name): + continue # a local checkpoint has no upstream GGUF repo + try: + name = get_model_name(name, load_in_4bit = False) + except Exception: + pass + if not name: + continue + # The upstream imatrix lives in unsloth/-GGUF, so map any org (e.g. meta-llama/...) + # onto the unsloth org; keep an already-formed -GGUF id as-is. + repo = name if name.endswith("-GGUF") else f"unsloth/{name.split('/')[-1]}-GGUF" + if repo not in candidates: + candidates.append(repo) + return candidates + + +def _materialize_imatrix(path, dest_dir): + """Copy an imatrix into dest_dir (never mutate the HF cache) and rename *.gguf_file -> *.gguf.""" + os.makedirs(dest_dir, exist_ok = True) + base = os.path.basename(path) + if base.endswith(".gguf_file"): + base = base[: -len(".gguf_file")] + ".gguf" + local = os.path.join(dest_dir, base) + shutil.copyfile(path, local) + return local + + +def _resolve_imatrix_file(model, imatrix_file, token, dest_dir): + """Turn the public imatrix_file value into a local imatrix path (or None). + + None/False -> None. A path -> that file (a *.gguf_file is renamed to *.gguf). True -> find and + download the upstream unsloth/-GGUF imatrix, raising a clear error if none exists. + """ + if imatrix_file is None or imatrix_file is False: + return None + + if imatrix_file is not True and isinstance(imatrix_file, (str, os.PathLike)): + path = os.path.expanduser(os.fspath(imatrix_file)) + if not os.path.isfile(path): + raise FileNotFoundError(f"Unsloth: imatrix_file '{path}' does not exist.") + return _materialize_imatrix(path, dest_dir) if path.endswith(".gguf_file") else path + + if imatrix_file is not True: + raise TypeError( + "Unsloth: imatrix_file must be None, a path string, or True " + f"(got {type(imatrix_file).__name__})." + ) + + # imatrix_file=True: auto-resolve from the upstream Unsloth GGUF repo. HfApi is the module-level + # import (save.py top); hf_hub_download is imported here as it is not needed elsewhere. + from huggingface_hub import hf_hub_download + + if token is None: + token = get_token() + api = HfApi(token = token) + repos = _gguf_repo_candidates(model) + for repo in repos: + try: + files = set(api.list_repo_files(repo)) + except Exception: + continue + for name in _IMATRIX_UPSTREAM_NAMES: + if name in files: + downloaded = hf_hub_download(repo_id = repo, filename = name, token = token) + local = _materialize_imatrix(downloaded, dest_dir) + print(f"Unsloth: Using imatrix '{name}' from '{repo}' -> '{local}'") + return local + raise RuntimeError( + "Unsloth: imatrix_file=True but no upstream Unsloth imatrix was found.\n" + f" Searched repos: {repos or '(none derived from the base model)'}\n" + f" Searched files: {list(_IMATRIX_UPSTREAM_NAMES)}\n" + "Pass imatrix_file='/path/to/imatrix.(dat|gguf)' to use your own." + ) + + +def _unsloth_save_lora_gguf( + model, + tokenizer, + save_directory, + outtype = "f16", + push_to_hub = False, + token = None, + private = None, + commit_message = "Converted LoRA to GGUF with Unsloth", + commit_description = "Convert LoRA to GGUF format using Unsloth", + create_pr = False, + revision = None, +): + """Export a PEFT/LoRA adapter straight to a GGUF LoRA file via llama.cpp's + convert_lora_to_gguf.py (loadable with `llama-cli --lora ...`). For a full / merged model + use save_pretrained_gguf instead. `save_directory` is a local dir, or a Hub repo id when + push_to_hub=True. Returns the local .gguf path, or the repo id when pushing.""" + import tempfile + + if not isinstance(model, (PeftModelForCausalLM, PeftModel)): + raise RuntimeError( + "Unsloth: LoRA GGUF export needs a PEFT/LoRA model. " + "For a full or merged model use save_pretrained_gguf(...) instead." + ) + if outtype not in _LORA_GGUF_OUTTYPES: + raise ValueError( + f"Unsloth: LoRA GGUF outtype must be one of {_LORA_GGUF_OUTTYPES} (got '{outtype}')." + ) + # Resolve a token even for local saves: the converter may fetch a gated/private base config. + if token is None: + token = get_token() + + # Resolve the dequantized base id (the adapter usually references a 4bit repo). + base_model_id = _lora_base_model_id(model) + if not base_model_id: + raise RuntimeError( + "Unsloth: could not determine the base model for LoRA GGUF export " + "(no adapter base_model_name_or_path or model config _name_or_path)." + ) + try: + base_model_id = get_model_name(base_model_id, load_in_4bit = False) + except Exception: + pass + # Windows-safe basename (handles both C:\... and / separators). + if os.path.isdir(base_model_id): + model_name = os.path.basename(os.path.normpath(base_model_id)) + else: + model_name = base_model_id.replace("\\", "/").rstrip("/").split("/")[-1] + if not model_name: + model_name = "model" + + # Save the adapter; for a hub push use an isolated temp dir, else save_directory itself. + if push_to_hub: + lora_dir = tempfile.mkdtemp(prefix = "unsloth-lora-gguf-") + else: + os.makedirs(save_directory, exist_ok = True) + lora_dir = save_directory + + # Wrap so the isolated temp dir used for hub pushes is always cleaned up, even on failure. + try: + save_lora_to_custom_dir(model, tokenizer, lora_dir) + + # Ensure a full llama.cpp checkout (ships convert_lora_to_gguf.py) and locate the converter. + install_llama_cpp(just_clone_repo = True) + converter = os.path.join(LLAMA_CPP_DEFAULT_DIR, "convert_lora_to_gguf.py") + if not os.path.exists(converter): + # A prebuilt llama.cpp install (or a reused CWD copy) carries binaries but not the + # converter script, so force a dedicated source checkout that ships it. + source_dir = os.path.join( + os.path.dirname(os.path.normpath(LLAMA_CPP_DEFAULT_DIR)), "llama.cpp-source" + ) + install_llama_cpp(llama_cpp_folder = source_dir, just_clone_repo = True) + converter = os.path.join(source_dir, "convert_lora_to_gguf.py") + if not os.path.exists(converter): + raise RuntimeError( + "Unsloth: convert_lora_to_gguf.py not found after installing a llama.cpp source " + "checkout. A full llama.cpp source checkout is required for LoRA GGUF export." + ) + + out_gguf = os.path.join(lora_dir, f"{model_name}-lora-{outtype}.gguf") + cmd = [sys.executable, converter, lora_dir, "--outfile", out_gguf, "--outtype", outtype] + # A local base dir provides config directly; otherwise the id is resolved from the Hub. + if os.path.isdir(base_model_id): + cmd += ["--base", base_model_id] + else: + cmd += ["--base-model-id", base_model_id] + if bool(getattr(model.config, "auto_map", None)): + cmd.append("--trust-remote-code") + + # Expose the token to the converter so it can fetch a gated/private base config from the Hub. + env = os.environ.copy() + if isinstance(token, str) and token: + env["HF_TOKEN"] = token + env["HUGGING_FACE_HUB_TOKEN"] = token + + print(f"Unsloth: Converting LoRA adapter at '{lora_dir}' to GGUF -> '{out_gguf}'") + try: + with subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + bufsize = 1, + universal_newlines = True, + encoding = "utf-8", + errors = "replace", + env = env, + ) as sp: + for line in sp.stdout: + print(line, end = "", flush = True) + sp.wait() + if sp.returncode != 0: + raise subprocess.CalledProcessError(sp.returncode, sp.args) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Unsloth: LoRA -> GGUF conversion failed (exit {e.returncode}). " + "See the output above for details." + ) + + if not push_to_hub: + print(f"Unsloth: Done. Saved LoRA GGUF to '{out_gguf}'") + return out_gguf + + print(f"Unsloth: Uploading LoRA GGUF to '{save_directory}' ...") + from huggingface_hub import HfApi + + api = HfApi(token = token) + api.create_repo( + repo_id = save_directory, + repo_type = "model", + private = private, + exist_ok = True, + ) + api.upload_folder( + folder_path = lora_dir, + repo_id = save_directory, + repo_type = "model", + allow_patterns = ["*.gguf"], + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + ) + print(f"Unsloth: Done. Uploaded to https://huggingface.co/{save_directory.lstrip('/')}") + return save_directory + finally: + if push_to_hub: + shutil.rmtree(lora_dir, ignore_errors = True) + + def unsloth_convert_lora_to_ggml_and_push_to_hub( self, tokenizer, repo_id: str, use_temp_dir: Optional[bool] = None, - commit_message: Optional[str] = "Converted LoRA to GGML with Unsloth", + commit_message: Optional[str] = "Converted LoRA to GGUF with Unsloth", private: Optional[bool] = None, token: Union[bool, str, None] = None, create_pr: bool = False, revision: str = None, - commit_description: str = "Convert LoRA to GGML format using Unsloth", + commit_description: str = "Convert LoRA to GGUF format using Unsloth", temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + outtype: str = "f16", ): - if not os.path.exists("llama.cpp"): - if IS_KAGGLE_ENVIRONMENT: - python_install = install_python_non_blocking(["protobuf"]) - python_install.wait() - install_llama_cpp_blocking(use_cuda = False) - makefile = None - else: - git_clone = install_llama_cpp_clone_non_blocking() - python_install = install_python_non_blocking(["protobuf"]) - git_clone.wait() - makefile = install_llama_cpp_make_non_blocking() - python_install.wait() - else: - makefile = None - - for _ in range(3): - gc.collect() - - lora_directory_push = "lora-to-ggml-push" - save_lora_to_custom_dir(self, tokenizer, lora_directory_push) - - model_type = self.config.model_type - output_file = os.path.join(lora_directory_push, "ggml-adapter-model.bin") - - print(f"Unsloth: Converting auto-saved LoRA adapters at {lora_directory_push} to GGML format.") - print(f"The output file will be {output_file}") - - try: - with subprocess.Popen( - [ - sys.executable, - "llama.cpp/convert-lora-to-ggml.py", - lora_directory_push, - output_file, - "llama", - ], - stdout = subprocess.PIPE, - stderr = subprocess.PIPE, - bufsize = 1, - universal_newlines = True, - encoding = "utf-8", - errors = "replace", - ) as sp: - for line in sp.stdout: - print(line, end = "", flush = True) - for line in sp.stderr: - print(line, end = "", flush = True) - sp.wait() - if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, sp.args) - except subprocess.CalledProcessError as e: - print(f"Error: Conversion failed with return code {e.returncode}") - return - - print(f"Unsloth: Conversion completed! Output file: {output_file}") - - print("Unsloth: Uploading GGML file to Hugging Face Hub...") - username = upload_to_huggingface( + return _unsloth_save_lora_gguf( self, + tokenizer, repo_id, - token, - "GGML converted LoRA", - "ggml", - output_file, - None, - private, - ) - link = f"{repo_id.lstrip('/')}" - print("Unsloth: Done.") - print(f"Converted LoRA to GGML and uploaded to https://huggingface.co/{link}") - print( - "\nThis GGML making function was made by Maheswar. Ping him @Maheswar on the Unsloth Discord or on HuggingFace (@mahiatlinux) if you like this!" + outtype = outtype, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, ) @@ -2918,65 +3479,9 @@ def unsloth_convert_lora_to_ggml_and_save_locally( tokenizer, temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.85, + outtype: str = "f16", ): - if not os.path.exists("llama.cpp"): - if IS_KAGGLE_ENVIRONMENT: - python_install = install_python_non_blocking(["protobuf"]) - python_install.wait() - install_llama_cpp_blocking(use_cuda = False) - makefile = None - else: - git_clone = install_llama_cpp_clone_non_blocking() - python_install = install_python_non_blocking(["protobuf"]) - git_clone.wait() - makefile = install_llama_cpp_make_non_blocking() - python_install.wait() - else: - makefile = None - - for _ in range(3): - gc.collect() - - # Use the provided save_directory for local saving - save_lora_to_custom_dir(self, tokenizer, save_directory) - - model_type = self.config.model_type - output_file = os.path.join(save_directory, "ggml-adapter-model.bin") - - print(f"Unsloth: Converting auto-saved LoRA adapters at {save_directory} to GGML format.") - print(f"The output file will be {output_file}") - - try: - with subprocess.Popen( - [ - sys.executable, - "llama.cpp/convert-lora-to-ggml.py", - save_directory, - output_file, - "llama", - ], - stdout = subprocess.PIPE, - stderr = subprocess.PIPE, - bufsize = 1, - universal_newlines = True, - encoding = "utf-8", - errors = "replace", - ) as sp: - for line in sp.stdout: - print(line, end = "", flush = True) - for line in sp.stderr: - print(line, end = "", flush = True) - sp.wait() - if sp.returncode != 0: - raise subprocess.CalledProcessError(sp.returncode, sp.args) - except subprocess.CalledProcessError as e: - print(f"Error: Conversion failed with return code {e.returncode}") - return - print("Unsloth: Done.") - print(f"Unsloth: Conversion completed! Output file: {output_file}") - print( - "\nThis GGML making function was made by Maheswar. Ping him @Maheswar on the Unsloth Discord or on HuggingFace (@mahiatlinux) if you like this!" - ) + return _unsloth_save_lora_gguf(self, tokenizer, save_directory, outtype = outtype) from .models.loader_utils import get_model_name @@ -3215,7 +3720,7 @@ def unsloth_generic_save_pretrained_merged( self, save_directory: Union[str, os.PathLike], tokenizer = None, - save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit"] + save_method: str = "merged_16bit", # ["lora", "merged_16bit", "merged_4bit", "fp8", "mxfp4", "nvfp4", "mxfp8"] push_to_hub: bool = False, token: Optional[Union[str, bool]] = None, is_main_process: bool = True, @@ -3229,6 +3734,9 @@ def unsloth_generic_save_pretrained_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -3238,6 +3746,10 @@ def unsloth_generic_save_pretrained_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM via llm-compressor: + `fp8` (dynamic W8A8), `mxfp4`, `nvfp4` (W4A4), `mxfp8`. The LoRA is merged to 16bit at + `save_directory`, then a quantized checkpoint is written to `save_directory + "-"`. + `nvfp4` needs calibration data (defaults to ultrachat; override with `calibration_dataset`). """ if tokenizer is None: logger.warning_once( @@ -3245,9 +3757,46 @@ def unsloth_generic_save_pretrained_merged( "You can do it separately via `tokenizer.save_pretrained(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_generic_save(**arguments) for _ in range(3): gc.collect() @@ -3271,6 +3820,9 @@ def unsloth_generic_push_to_hub_merged( temporary_location: str = "_unsloth_temporary_saved_buffers", maximum_memory_usage: float = 0.75, datasets: Optional[List[str]] = None, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, ): """ Same as .push_to_hub(...) except 4bit weights are auto @@ -3280,6 +3832,7 @@ def unsloth_generic_push_to_hub_merged( 1. `16bit`: Merge LoRA into float16 weights. Useful for GGUF / llama.cpp. 2. `4bit`: Merge LoRA into int4 weights. Useful for DPO / HF inference. 3. `lora`: Save LoRA adapters with no merging. Useful for HF inference. + 4. FP8 / FP4 compressed export for vLLM: `fp8`, `mxfp4`, `nvfp4`, `mxfp8`. """ if tokenizer is None: logger.warning_once( @@ -3287,12 +3840,50 @@ def unsloth_generic_push_to_hub_merged( "You can do it separately via `tokenizer.push_to_hub(...)`" ) + # FP8 / FP4 compressed-tensors export (llm-compressor) -> handled separately. + _compressed = _normalize_compressed_method(save_method) + if _compressed is not None: + scheme, needs_calibration, suffix = _compressed + _unsloth_save_compressed_tensors( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + scheme = scheme, + needs_calibration = needs_calibration, + suffix = suffix, + push_to_hub = True, + token = token, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id arguments["push_to_hub"] = True del arguments["self"] del arguments["repo_id"] + del arguments["_compressed"] + del arguments["calibration_dataset"] + del arguments["num_calibration_samples"] + del arguments["max_seq_length"] unsloth_generic_save(**arguments) for _ in range(3): gc.collect() @@ -3437,6 +4028,333 @@ def _unsloth_save_torchao_with_given_config( pass +def _scheme_is_available(scheme): + """True if `scheme` is a known preset in the installed compressed_tensors.""" + try: + from compressed_tensors.quantization import quant_scheme as _qs + + presets = getattr(_qs, "PRESET_SCHEMES", None) + if presets is None: + return True + return scheme in presets + except Exception: + # If we cannot introspect, let llm-compressor validate the scheme itself. + return True + + +def _print_compressed_hw_note(scheme, out_dir): + if scheme in ("FP8_DYNAMIC", "MXFP8"): + hw = "NVIDIA GPUs with compute capability >= 8.9 (Ada / Hopper) or newer" + else: + hw = ( + "NVIDIA Blackwell (SM100+) for full activation quantization " + "(older GPUs fall back to weight-only in vLLM)" + ) + print( + f"Unsloth: Saved {scheme} compressed checkpoint to '{out_dir}'.\n" + f"Unsloth: Load it with vLLM for accelerated inference. Hardware for full speed: {hw}." + ) + + +def _unsloth_save_compressed_tensors( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + scheme: str, + needs_calibration: bool, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + calibration_dataset = None, + num_calibration_samples: int = 512, + max_seq_length: int = 2048, + **merge_kwargs, +): + """Export an FP8/FP4 compressed-tensors checkpoint via llm-compressor. + + Mirrors the torchao PTQ path: LoRA is first merged into the base model at 16bit and + written to `save_directory` (which is kept). The merged checkpoint is then quantized with + llm-compressor's `QuantizationModifier(scheme)` in a separate process (so Unsloth's + transformers monkey-patches do not interfere), and written to the sibling directory + `save_directory + "-" + suffix`. The result is intended for vLLM inference. + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + # Resolve a token for the hub push and/or loading a gated calibration dataset in the subprocess. + if token is None: + token = get_token() + + # Only the main process installs deps, merges, quantizes, and uploads (mirrors the non-PEFT + # save path); other ranks return at once so they neither race on dirs nor run pip installs. + if not is_main_process: + return None + + # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported + # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. + install_llm_compressor() + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) + + # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and + # quantize inside an isolated temp dir instead of writing ./ into the cwd. + repo_id, work_tmp, calib_tmp, model_dev = None, None, None, None + if push_to_hub: + repo_id = os.fspath(save_directory) + work_tmp = tempfile.mkdtemp(prefix = "unsloth-compressed-") + local_dir = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + else: + # Drop trailing separators so the sibling "-" output is not nested inside . + local_dir = os.fspath(save_directory) + local_dir = local_dir.rstrip("/\\") or local_dir + + # Wrap the body so the isolated temp dirs are always cleaned up, even when the merge, + # quantization, validation, or hub upload raises. + api = None + try: + # Validate Hub access up front (a bad token / denied repo should fail before the expensive + # merge and quantization, matching the normal push path). create_repo is idempotent. + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 3) Merge to 16bit at local_dir (kept for local saves) via unsloth_generic_save, so LoRA + # adapters are merged and full-finetuned models written in 16bit consistently. Extra + # save kwargs (state_dict, max_shard_size, ...) flow through merge_kwargs. + # The intermediate 16bit checkpoint is internal staging that the converter subprocess + # reloads with default weight filenames, so never write variant-named shards here; the + # user's variant (if any) is applied to the final compressed checkpoint in the subprocess. + variant = merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before {scheme} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = local_dir, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 4) Detect VLM + trust_remote_code from the in-memory model config. A vision/multimodal + # model exposes a vision_config or an explicitly vision-named architecture; a bare + # *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper), so it + # is not treated as a VLM on its own. + is_vlm = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + if is_vlm: + logger.warning( + "Unsloth: FP8/FP4 compressed export for vision / multimodal models is " + "experimental; vision-tower layers may be affected." + ) + trust_remote_code = ( + bool(getattr(model.config, "auto_map", None)) if hasattr(model, "config") else False + ) + + # 5) Marshal the calibration dataset for the subprocess: None -> ultrachat default; a + # str/PathLike is a local save_to_disk dir if it exists else a Hub id; Dataset -> temp. + calib_kind, calib_value = "none", "" + if needs_calibration and calibration_dataset is not None: + if isinstance(calibration_dataset, (str, os.PathLike)): + calib_value = os.fspath(calibration_dataset) + calib_kind = "disk" if os.path.isdir(calib_value) else "hfid" + elif hasattr(calibration_dataset, "save_to_disk"): + # Only persist the samples we need, so multi-GB training sets are not fully copied. + ds_to_save = calibration_dataset + # A DatasetDict's len() is the split count, not rows; pick one split first so the + # row subsample below applies and we do not save every split to the temp dir. + try: + from datasets import DatasetDict + if isinstance(ds_to_save, DatasetDict): + ds_to_save = ds_to_save.get("train", None) or next( + iter(ds_to_save.values()) + ) + except Exception: + pass + try: + if ( + num_calibration_samples + and hasattr(ds_to_save, "select") + and len(ds_to_save) > num_calibration_samples + ): + ds_to_save = ds_to_save.shuffle(seed = 42).select( + range(num_calibration_samples) + ) + except Exception: + ds_to_save = calibration_dataset + calib_tmp = tempfile.mkdtemp(prefix = "unsloth-calib-") + shutil.rmtree(calib_tmp, ignore_errors = True) # save_to_disk wants a fresh path + ds_to_save.save_to_disk(calib_tmp) + calib_kind, calib_value = "disk", calib_tmp + else: + raise TypeError( + "Unsloth: calibration_dataset must be None, a Hugging Face dataset id, a " + "local path saved with Dataset.save_to_disk(...), or a Dataset with " + "save_to_disk()." + ) + elif not needs_calibration and calibration_dataset is not None: + logger.warning_once( + f"Unsloth: scheme '{scheme}' is data-free; ignoring calibration_dataset." + ) + + # 6) Quantize in a separate process: importing Unsloth patches transformers attention, + # which breaks the forward llm-compressor runs for calibration. Run the converter by + # file path (not `-m`) so the subprocess stays unpatched, like GGUF -> llama.cpp. + out_dir = local_dir + "-" + suffix + runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_compressed_quantize.py") + cmd = [ + sys.executable, + runner, + "--model", + local_dir, + "--scheme", + scheme, + "--out", + out_dir, + "--calibration-dataset-kind", + calib_kind, + "--num-calibration-samples", + str(num_calibration_samples), + "--max-seq-length", + str(max_seq_length), + ] + if needs_calibration: + cmd.append("--needs-calibration") + if calib_value: + cmd += ["--calibration-dataset", calib_value] + if is_vlm: + cmd.append("--is-vlm") + if trust_remote_code: + cmd.append("--trust-remote-code") + if variant: + cmd += ["--variant", variant] + + # Free the in-memory model's CUDA memory before the subprocess loads its own copy from + # disk, so a single GPU need not hold both at once. Best-effort and restored in finally; + # skipped for quantized or multi-device models where moving is unsafe. + try: + if ( + torch.cuda.is_available() + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith("cuda"): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev # set only after a successful move, so finally can restore + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Expose the token so the subprocess can load a gated/private calibration dataset. + env = os.environ.copy() + if isinstance(token, str) and token: + env["HF_TOKEN"] = token + env["HUGGING_FACE_HUB_TOKEN"] = token + + print( + f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + "(in a separate process)..." + ) + try: + subprocess.check_call(cmd, env = env) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Unsloth: {scheme} quantization failed (llm-compressor subprocess exit " + f"{e.returncode}). See the output above for details." + ) + + # 7) Validate the artifact. + cfg_path = os.path.join(out_dir, "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: + raise RuntimeError( + f"Unsloth: {scheme} export failed - no quantization_config written to {cfg_path}" + ) + + # 8) Optional hub upload of the compressed artifact (not the intermediate 16bit one). + # The repo was already created/validated up front, so just upload here. + if push_to_hub: + print(f"Unsloth: Uploading {scheme} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + # Attach datasets metadata to the pushed repo, like the normal merged push path. + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + # 9) Inference hardware note. + result = repo_id if push_to_hub else out_dir + _print_compressed_hw_note(scheme, result) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) # restore the model to its original device + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after compressed " + "export; it may remain on CPU." + ) + if calib_tmp is not None and os.path.isdir(calib_tmp): + shutil.rmtree(calib_tmp, ignore_errors = True) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike],