From 9369dd47e6777bd6b3ece55109127627e47c7908 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jun 2026 03:40:16 -0700 Subject: [PATCH] Add FP8/FP4 compressed export to save_pretrained_merged (#6706) * Add FP8/FP4 compressed export to save_pretrained_merged Adds compressed-tensors export (for vLLM) to save_pretrained_merged / push_to_hub_merged via llm-compressor, alongside the existing lora / merged_16bit / merged_4bit / gguf / torchao paths: model.save_pretrained_merged("model", tokenizer, save_method="fp8") Supported save_method values: fp8 (FP8_DYNAMIC), mxfp4, nvfp4 (W4A4) and mxfp8. The LoRA is merged to 16bit at save_directory, then a quantized checkpoint is written to save_directory + "-". nvfp4 needs a small calibration set (defaults to ultrachat, overridable via calibration_dataset). Notes: - llm-compressor is installed lazily on first use, pinning the current torch and transformers via a constraints file so they are not upgraded (a plain install pulls transformers>=5 and breaks Unsloth). - Quantization runs in a separate process (unsloth/_compressed_quantize.py, launched by file path) so Unsloth's transformers attention patches do not interfere with the forward llm-compressor runs during calibration, mirroring how GGUF export shells out to llama.cpp. - mxfp8 needs a newer llm-compressor (transformers>=5); it is recognised and raises a clear error until that stack is available. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: main-process guard, calibration subsampling, tokenizer + dtype handling - Route the 16bit merge through unsloth_generic_save for both LoRA and full finetuned models, so non-PEFT models are written in 16bit consistently instead of saving the original (possibly quantized) weights directly. - Honor is_main_process: only the main process quantizes and writes the compressed output, so distributed ranks do not race on the same dirs. - Subsample an in-memory calibration Dataset before save_to_disk so large training sets are not fully copied to a temp dir. - Tolerate a missing tokenizer in the converter (data-free exports); still require one for calibration based schemes. - Open config.json via a context manager in both files. - Drop the redundant nvfp4 entry from the unsupported-name check (fp4 covers it). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add direct LoRA to GGUF export and harden FP8/FP4 compressed export - Run llm-compressor install and scheme check before the 16bit merge so unsupported schemes (e.g. mxfp8) fail fast without writing a checkpoint - Only the main process installs, merges, quantizes and uploads; isolate hub pushes to a temp dir and clean all temp dirs in a finally - Forward standard save kwargs (state_dict, max_shard_size, ...) to the merge - Fall back to the first dataset split for Hub calibration ids - Export LoRA adapters to GGUF via convert_lora_to_gguf.py: modernize save_pretrained_ggml/push_to_hub_ggml and add save_method="lora" to save_pretrained_gguf/push_to_hub_gguf; resolve base from the adapter config * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix LoRA GGUF shell-injection test and compressed export trailing-slash path - Update tests/saving/test_save_shell_injection.py for the new delegation: the LoRA to GGUF conversion now lives in _unsloth_save_lora_gguf, so assert it passes argv as a list with no shell=True and that the legacy ggml wrappers delegate to it instead of calling subprocess.Popen directly - Normalize the local save_directory before building the "-" sibling so a trailing slash no longer nests the compressed output inside the 16bit dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Polish FP8/FP4 and LoRA GGUF export after review - Warn (not silently downgrade) when an explicit quantization_method is not a valid LoRA GGUF outtype; default stays f16 - Correct the inference hardware note: MXFP8 is 8-bit (cc >= 8.9), only FP4 needs Blackwell for full activation quantization - Document that a local fp8/fp4 save keeps the 16bit merge at save_directory and writes the quantized checkpoint to save_directory + "-" * Use sequential calibration pipeline and validate Hub access early - nvfp4 calibration no longer forces the memory-hungry "basic" pipeline. The quantization runs in a clean subprocess, so llm-compressor's default sequential pipeline (layer-by-layer onloading) works and lets large models that do not fit at once still calibrate; fall back to "basic" only if tracing fails - For push_to_hub compressed exports, create/validate the repo up front so a bad token or denied repo fails before the merge and quantization instead of after * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden compressed export: explicit sequential pipeline, base-tokenizer calibration, GPU memory - nvfp4 calibration now passes pipeline="sequential" explicitly (layer-by-layer onloading) instead of relying on the inferred default, with a "basic" fallback - Calibration datasets with a messages column no longer require a chat template: base / non-chat tokenizers fall back to concatenating message contents - Free the in-memory model's CUDA memory before the quantize subprocess loads its own copy from disk (best-effort, single-device non-quantized only; restored afterward), so a single GPU need not hold two copies at once - Create the calibration temp dir in the system temp location instead of next to the save directory, avoiding stray dirs in the workspace * Free the failed calibration model before the basic-pipeline retry In the sequential -> basic NVFP4 fallback, release the partially-processed model and clear the CUDA cache before loading a fresh copy, so the retry does not transiently hold two model copies on the GPU. * Harden calibration data handling and compressed-export edge cases - Calibration messages without a chat template now handle multimodal (list) content, None content, and null message rows instead of crashing on join - Raise a clear error when the calibration dataset is empty after subsampling - Reset llm-compressor's global session before freeing the model in the sequential -> basic NVFP4 fallback, so the old model is actually released - LoRA GGUF export accepts a single-element list quantization_method - Attach datasets metadata to the pushed repo on compressed hub exports - Warn (instead of silently) if the model cannot be restored to its device - Raise a clear error if the LoRA base model id cannot be determined * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle DatasetDict calibration, MoE routers, and MTP models in compressed export - Reduce an in-memory DatasetDict calibration set to a single split before row subsampling, so save_to_disk does not copy every split to the temp dir - For MoE models, keep the router/gate unquantized and pass moe_calibrate_all_experts so every expert is calibrated - Warn when a model carries MTP / speculative-decoding tensors that the compressed export does not include * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Support many more compressed-tensors schemes and address review - Expand save_method to cover the full set of compressed-tensors preset schemes: FP8 (dynamic/static/block), INT8, W8A8, W8A16, W4A16(+asym), W4A8, W4AFP8, MXFP4(+A16), NVFP4(+A16), plus the gated MXFP8; calibration is used only for the static-activation schemes (FP8 static, NVFP4) - Broaden the near-miss save_method error to cover int/w-prefixed names - MoE: also keep the Qwen shared-expert gate unquantized - Strip non-model-input columns from already-tokenized calibration data so the collator does not choke on a leftover messages column - Forward the Hub token to the LoRA converter and the quantize subprocess so gated/private base models and calibration datasets work without a global login * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Collapse compressed-tensors export help line so ruff-format converges The print line in print_quantization_methods needed two ruff-format passes to reach a fixpoint (merge implicit string concat, then collapse the single-arg print). pre-commit.ci applies one pass per run, so it kept reformatting. Land the converged single-line form directly. * Add CPU-only regression tests for the export API Cover all export paths without a GPU, for slow CPU-only CI: - pure-function checks of the compressed-tensors scheme registry and save_method normalization (aliases, calibration flags, near-miss errors) - AST checks that every merged saver dispatches compressed export, the GGUF savers expose the lora branch, torchao routes PTQ/QAT, the public methods stay attached, and the export subprocesses remain shell-safe (argv list, sys.executable, no shell) - monkeypatched dispatch checks that fp8/nvfp4/merged_16bit, the LoRA-GGUF outtype resolution, and torchao PTQ/QAT reach the right helper with the right arguments * Run the CPU-only export tests in consolidated CI tests/saving is --ignored by the Repo tests (CPU) job, so the new GPU-free export tests are added by path to consolidated-tests-ci.yml (collection sanity + Bucket-A run), alongside the existing CPU saving tests, so they actually execute on CPU CI. * Add GPU GGUF export + llama-cli inference smoke test tests/saving/test_gguf_export_and_inference.py: skipif no CUDA. Trains a tiny phrase-imprinting LoRA, exports a full-model q8_0 GGUF (merge -> convert_hf_to_gguf -> llama-quantize), asserts a valid GGUF (magic + size), and - when a llama-cli binary is available - runs one bounded generation (byte cap + watchdog kill) and asserts the trained phrase round-trips through HF -> GGUF -> quantize -> inference. The llama-cli step skips gracefully since the export only builds llama-quantize. * Fix variant mismatch in compressed (FP8/FP4) export save_pretrained_merged(..., save_method=fp8/nvfp4, variant=...) forwarded the variant into the intermediate 16bit merge, so Transformers wrote variant-named shards (model..safetensors). The converter subprocess then reloaded that directory with the default weight filenames, so the compressed export failed after doing the merge. Pop the variant out of the intermediate merge (internal staging that the subprocess reloads with default names) and forward it via --variant so it is applied to the final compressed checkpoint instead. Add a CPU AST guard for the contract. * Harden export paths from review - install_llm_compressor: fall back to uv pip when this interpreter has no pip seeded (uv-created/relocatable venvs), instead of failing with No module named pip. - LoRA GGUF export: if convert_lora_to_gguf.py is missing (a prebuilt or reused CWD llama.cpp install carries binaries but not the converter script), force a dedicated source checkout that ships it. - push_to_hub_gguf(save_method=lora): return on non-main ranks, matching the local save_pretrained_gguf lora branch, so only rank 0 converts/uploads. - compressed export VLM detection: require a vision_config or a ForVisionText2Text architecture; a bare *ForConditionalGeneration also matches text seq2seq models (T5/BART/Whisper) and is no longer treated as a VLM on its own. - GGUF GPU smoke test: drop SFTConfig(max_length=1024), which raises under newer TRL padding-free training; length enforcement is not needed here. * Add imatrix option to GGUF export, enabling IQ low-bit quants save_pretrained_gguf / push_to_hub_gguf gain imatrix_file: None -> no imatrix (unchanged) '/path' -> pass to llama-quantize --imatrix (a *.gguf_file is renamed to *.gguf) True -> download the upstream unsloth/-GGUF imatrix (imatrix_unsloth.dat or .gguf_file), raising a clear error if none exists An importance matrix unlocks the IQ low-bit quants (iq2_xxs, iq4_xs, ...), which were hard disabled before. They are gated: requesting one without an imatrix raises a clear error. - _resolve_imatrix_file resolves path/True (PEFT base first, normalized via get_model_name, derives unsloth/-GGUF, copies out of the HF cache before renaming *.gguf_file). - IMATRIX_QUANTS registry replaces the old commented-out IQ entries; save_to_gguf accepts a resolved imatrix and threads it into the quantize calls. - The --imatrix flag is emitted by unsloth_zoo's quantize_gguf (companion change). save.py fails fast with an upgrade hint if the installed unsloth_zoo lacks the imatrix kwarg. Tests: tests/saving/test_imatrix_export.py (CPU: resolution, repo derivation, IQ gate, --imatrix wiring) wired into CI; tests/saving/test_gguf_export_and_inference.py extended with GPU iq2_xxs/iq4_xs export + inference. Verified end to end on Llama-3.2-1B: imatrix auto-downloaded, iq2_xxs/iq4_xs exported and run via llama.cpp. Note: requires the companion unsloth_zoo quantize_gguf imatrix change. * Address imatrix/compressed review feedback: unsloth org GGUF repo, fail-fast, calibration split - imatrix auto-resolve (imatrix_file=True): derive the upstream repo as unsloth/-GGUF instead of /-GGUF, so official bases (e.g. meta-llama/Llama-3.1-8B-Instruct) find the matching Unsloth GGUF imatrix repo rather than failing on a nonexistent meta-llama/...-GGUF. - Resolve/validate the imatrix before the 16-bit merge in save_pretrained_gguf, so a bad path or an unavailable upstream imatrix fails fast instead of after a long, multi-GB merge. - Compressed calibration: when a Hub dataset has no "train" split, resolve the first split name and slice it, instead of materializing the whole dataset just to take num_samples rows. Keeps the original materialize-then-subselect path as a last resort. Tests: add unsloth/-GGUF mapping for an official base id, and create the imatrix file in the quantize_gguf flag test (quantize_gguf now validates the imatrix exists). --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 8 + .../saving/test_compressed_export_schemes.py | 69 + tests/saving/test_export_api_surface.py | 176 +++ tests/saving/test_export_dispatch.py | 180 +++ .../saving/test_gguf_export_and_inference.py | 343 +++++ tests/saving/test_imatrix_export.py | 275 ++++ tests/saving/test_save_shell_injection.py | 98 +- unsloth/_compressed_quantize.py | 347 +++++ unsloth/save.py | 1208 +++++++++++++++-- 9 files changed, 2520 insertions(+), 184 deletions(-) create mode 100644 tests/saving/test_compressed_export_schemes.py create mode 100644 tests/saving/test_export_api_surface.py create mode 100644 tests/saving/test_export_dispatch.py create mode 100644 tests/saving/test_gguf_export_and_inference.py create mode 100644 tests/saving/test_imatrix_export.py create mode 100644 unsloth/_compressed_quantize.py 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/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],