From 315ecdef5e706c5c84731ebb9bd1fc55b4388c76 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 17:37:11 +0000 Subject: [PATCH] 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. --- unsloth/_compressed_quantize.py | 168 ++++++++++++ unsloth/save.py | 451 +++++++++++++++++++++++++++++++- 2 files changed, 617 insertions(+), 2 deletions(-) create mode 100644 unsloth/_compressed_quantize.py diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py new file mode 100644 index 0000000000..cc83c91965 --- /dev/null +++ b/unsloth/_compressed_quantize.py @@ -0,0 +1,168 @@ +# 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 _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": + ds = load_dataset(value, split=f"train[:{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}") + + cols = set(ds.column_names) + if "input_ids" in cols: + return ds + if "messages" in cols: + def _prep(ex): + return {"text": _tok.apply_chat_template(ex["messages"], tokenize=False)} + 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") + 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() + tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) + + recipe = QuantizationModifier(targets="Linear", scheme=args.scheme, ignore=["lm_head"]) + if args.needs_calibration: + ds = _build_calibration_dataset( + tokenizer, args.calibration_dataset_kind, args.calibration_dataset, + args.num_calibration_samples, args.max_seq_length, + ) + # "basic" pipeline runs a normal forward (no AST tracing / sequential splitting). + oneshot( + model=model, dataset=ds, recipe=recipe, + max_seq_length=args.max_seq_length, + num_calibration_samples=args.num_calibration_samples, + pipeline="basic", + ) + else: + oneshot(model=model, recipe=recipe) + + os.makedirs(args.out, exist_ok=True) + model.save_pretrained(args.out, save_compressed=True) + tokenizer.save_pretrained(args.out) + + cfg_path = os.path.join(args.out, "config.json") + cfg = json.load(open(cfg_path)) if os.path.exists(cfg_path) else {} + 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 0e8c8cdc2d..0192e2b71d 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -149,9 +149,56 @@ 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). +COMPRESSED_EXPORT_SCHEMES = { + "fp8" : ("FP8_DYNAMIC", False, "fp8"), + "fp8_dynamic": ("FP8_DYNAMIC", False, "fp8"), + "dynamic_fp8": ("FP8_DYNAMIC", False, "fp8"), + "w8a8_fp8" : ("FP8_DYNAMIC", False, "fp8"), + "mxfp8" : ("MXFP8", False, "mxfp8"), + "w8a8_mxfp8" : ("MXFP8", False, "mxfp8"), + "mxfp4" : ("MXFP4", False, "mxfp4"), + "w4a4_mxfp4" : ("MXFP4", False, "mxfp4"), + "nvfp4" : ("NVFP4", True, "nvfp4"), + "w4a4_nvfp4" : ("NVFP4", True, "nvfp4"), +} + + +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", "nvfp4", "mxfp")): + supported = ", ".join(sorted(COMPRESSED_EXPORT_SCHEMES.keys())) + raise RuntimeError( + f"Unsloth: save_method='{save_method}' is not a supported compressed export.\n" + f"Supported FP8/FP4 export methods: {supported}\n" + "(only dynamic FP8, MXFP8, and W4A4 MXFP4/NVFP4 are wired up for now)." + ) + return None + + def print_quantization_methods(): for key, value in ALLOWED_QUANTS.items(): print(f'"{key}" ==> {value}') + print("\nCompressed-tensors FP8/FP4 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( @@ -1227,6 +1274,72 @@ 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 + + cmd = [sys.executable, "-m", "pip", "install", "llmcompressor"] + 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" {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( @@ -1694,6 +1807,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 @@ -1703,6 +1819,7 @@ 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`. """ if tokenizer is None: logger.warning_once( @@ -1710,9 +1827,34 @@ 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, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + ) + 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() @@ -1722,7 +1864,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, @@ -1736,6 +1878,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 @@ -1745,6 +1890,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( @@ -1752,12 +1898,42 @@ 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, + ) + 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() @@ -3166,7 +3342,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, @@ -3180,6 +3356,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 @@ -3189,6 +3368,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( @@ -3196,9 +3379,34 @@ 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, + calibration_dataset = calibration_dataset, + num_calibration_samples = num_calibration_samples, + max_seq_length = max_seq_length, + ) + 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() @@ -3222,6 +3430,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 @@ -3231,6 +3442,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( @@ -3238,12 +3450,42 @@ 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, + ) + 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() @@ -3388,6 +3630,211 @@ 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 == "FP8_DYNAMIC": + 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, + 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. + """ + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None and push_to_hub: + token = get_token() + # Accept os.PathLike save directories (string concatenation below needs a str). + save_directory = os.fspath(save_directory) + + # 1) Merge LoRA -> 16bit on disk and KEEP it. Must be a local save: merge_and_overwrite_lora + # deletes save_directory when push_to_hub=True, which would remove it before reload. + is_peft = isinstance(model, PeftModelForCausalLM) or isinstance(model, PeftModel) + if is_peft: + print(f"Unsloth: Merging LoRA weights to 16bit before {scheme} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update(dict( + model = model, + tokenizer = tokenizer, + save_directory = save_directory, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + )) + unsloth_generic_save(**merge_args) + else: + print(f"Unsloth: Saving base model to 16bit before {scheme} quantization...") + os.makedirs(save_directory, exist_ok = True) + model.save_pretrained(save_directory) + if tokenizer is not None: + tokenizer.save_pretrained(save_directory) + + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 2) Lazily import / install llm-compressor (without breaking torch / transformers) and + # gate on scheme availability up-front so the user gets a clear error fast. + 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." + ) + + # 3) Detect VLM + trust_remote_code from the in-memory model config. + is_vlm = False + if hasattr(model, "config") and hasattr(model.config, "architectures"): + is_vlm = any( + x.endswith(("ForConditionalGeneration", "ForVisionText2Text")) + for x in (model.config.architectures or []) + ) + is_vlm = is_vlm or hasattr(model.config, "vision_config") + 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 + + # 4) 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; a Dataset -> temp dir. + calib_kind, calib_value = "none", "" + calib_tmp = 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"): + import tempfile + parent = os.path.dirname(os.path.abspath(save_directory)) or None + calib_tmp = tempfile.mkdtemp(prefix = "unsloth-calib-", dir = parent) + shutil.rmtree(calib_tmp, ignore_errors = True) # save_to_disk wants a fresh path + calibration_dataset.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." + ) + + # 5) 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 shelling out to llama.cpp. + out_dir = save_directory + "-" + suffix + runner = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_compressed_quantize.py") + cmd = [ + sys.executable, runner, + "--model", str(save_directory), + "--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") + + print(f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + "(in a separate process)...") + try: + subprocess.check_call(cmd) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Unsloth: {scheme} quantization failed (llm-compressor subprocess exit {e.returncode}). " + "See the output above for details." + ) + finally: + if calib_tmp is not None and os.path.isdir(calib_tmp): + try: shutil.rmtree(calib_tmp) + except Exception: pass + + # 6) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = json.load(open(cfg_path)) if os.path.exists(cfg_path) else {} + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: {scheme} export failed - no quantization_config written to {cfg_path}" + ) + + # 7) Optional hub upload of the compressed artifact (not the intermediate 16bit one). + if push_to_hub: + print(f"Unsloth: Uploading {scheme} checkpoint to '{save_directory}' ...") + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = save_directory, repo_type = "model", + private = merge_kwargs.get("private", None), exist_ok = True, + ) + api.upload_folder( + folder_path = out_dir, repo_id = save_directory, 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), + ) + + # 8) Inference hardware note. + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _print_compressed_hw_note(scheme, out_dir) + return out_dir + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike],