ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add FP8/FP4 compressed export to save_pretrained_merged (#6706)

Browse source 
* 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 + "-<fmt>". 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 "<dir>-<fmt>" 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 + "-<fmt>"

* 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.<variant>.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/<base>-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/<base>-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/<base>-GGUF
  instead of <org>/<base>-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/<base>-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>
This commit is contained in:
Daniel Han 2026-06-30 03:40:16 -07:00 • committed by GitHub
commit 9369dd47e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 2520 additions and 184 deletions
Download patch file Download diff file Expand all files Collapse all files

8
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -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 \

69
tests/saving/test_compressed_export_schemes.py Normal file
View file

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

176
tests/saving/test_export_api_surface.py Normal file
View file

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

180
tests/saving/test_export_dispatch.py Normal file
View file

@ -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,
)

343
tests/saving/test_gguf_export_and_inference.py Normal file
View file

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

275
tests/saving/test_imatrix_export.py Normal file
View file

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

98
tests/saving/test_save_shell_injection.py
View file

@ -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"

347
unsloth/_compressed_quantize.py Normal file
View file

@ -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()

1208
unsloth/save.py
View file

File diff suppressed because it is too large Load diff