diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 02bb31661f..71c87c6859 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -271,6 +271,7 @@ jobs: 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 @@ -359,6 +360,7 @@ jobs: 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_gguf_export_and_inference.py b/tests/saving/test_gguf_export_and_inference.py index cd4a2ab89f..aa69368482 100644 --- a/tests/saving/test_gguf_export_and_inference.py +++ b/tests/saving/test_gguf_export_and_inference.py @@ -206,3 +206,138 @@ def test_gguf_llama_cli_inference_reflects_finetune(exported_gguf): # 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..1a58ea4723 --- /dev/null +++ b/tests/saving/test_imatrix_export.py @@ -0,0 +1,265 @@ +"""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_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 + 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/unsloth/save.py b/unsloth/save.py index 6ade571945..ac4b1401ca 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 @@ -216,6 +229,9 @@ def _normalize_compressed_method(save_method): 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(): @@ -232,11 +248,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", @@ -1536,10 +1554,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": @@ -1575,11 +1596,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 @@ -1594,11 +1619,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) @@ -1752,16 +1786,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: @@ -2410,11 +2450,16 @@ def unsloth_save_pretrained_gguf( 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.", @@ -2537,6 +2582,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: @@ -2648,6 +2694,10 @@ def unsloth_save_pretrained_gguf( except Exception as e: logger.warning(f"Unsloth: fix_sentencepiece_gguf skipped ({type(e).__name__}): {e}") + # Resolve the importance matrix once (download upstream / validate path / rename *.gguf_file) + # before quantization, so a failed auto-resolution never reaches the IQ-quant gate. + imatrix_path = _resolve_imatrix_file(self, imatrix_file, token, save_directory) + try: all_file_locations, want_full_precision, is_vlm_update = save_to_gguf( model_name = model_name, @@ -2659,6 +2709,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: @@ -2756,11 +2807,15 @@ def unsloth_push_to_hub_gguf( 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.", @@ -2842,11 +2897,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 @@ -3097,6 +3153,95 @@ def _lora_base_model_id(model): 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 or "/" not in name: + continue + repo = name if name.endswith("-GGUF") else f"{name}-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,