From e9a80fe43643dc23d05926ce3c2cc4ff70e73a09 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Jun 2026 17:38:07 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/_compressed_quantize.py | 99 ++++++++++++++++----------- unsloth/save.py | 115 ++++++++++++++++++++------------ 2 files changed, 135 insertions(+), 79 deletions(-) diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index cc83c91965..69e947f3ea 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -17,6 +17,7 @@ Launched as a subprocess by file path (not `python -m`) so the Unsloth package, 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 @@ -26,6 +27,7 @@ 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": @@ -33,13 +35,13 @@ def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_leng 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, + flush = True, ) - ds = load_dataset("HuggingFaceH4/ultrachat_200k", split=f"train_sft[:{num_samples}]") - ds = ds.shuffle(seed=42) + 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) + 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): @@ -53,7 +55,7 @@ def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_leng "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)) + ds = ds.shuffle(seed = 42).select(range(num_samples)) else: raise ValueError(f"Unknown calibration-dataset-kind: {kind}") @@ -61,8 +63,10 @@ def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_leng if "input_ids" in cols: return ds if "messages" in cols: + def _prep(ex): - return {"text": _tok.apply_chat_template(ex["messages"], tokenize=False)} + return {"text": _tok.apply_chat_template(ex["messages"], tokenize = False)} + ds = ds.map(_prep) elif "text" not in cols: raise RuntimeError( @@ -72,39 +76,50 @@ def _build_calibration_dataset(tokenizer, kind, value, num_samples, max_seq_leng def _tokenize(sample): return _tok( - sample["text"], padding=False, max_length=max_seq_length, - truncation=True, add_special_tokens=False, + sample["text"], + padding = False, + max_length = max_seq_length, + truncation = True, + add_special_tokens = False, ) - return ds.map(_tokenize, remove_columns=ds.column_names) + + 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, + 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, + 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("--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 @@ -131,37 +146,45 @@ def main(): 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) + tokenizer = auto_proc.from_pretrained(args.model, trust_remote_code = args.trust_remote_code) - recipe = QuantizationModifier(targets="Linear", scheme=args.scheme, ignore=["lm_head"]) + 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, + 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", + 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) + oneshot(model = model, recipe = recipe) - os.makedirs(args.out, exist_ok=True) - model.save_pretrained(args.out, save_compressed=True) + 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) + 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) + print( + f"[compressed-quantize] OK scheme={args.scheme} format={qfmt} " + f"shards={len(shards)} -> {args.out}", + flush = True, + ) if __name__ == "__main__": diff --git a/unsloth/save.py b/unsloth/save.py index 0192e2b71d..ecb1629fc9 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -152,16 +152,16 @@ 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": ("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"), + "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"), } @@ -190,8 +190,10 @@ def _normalize_compressed_method(save_method): 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):") + 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: @@ -1294,6 +1296,7 @@ def install_llm_compressor(): ) import importlib import tempfile + constraints = "" try: import torch as _torch @@ -1324,8 +1327,10 @@ def install_llm_compressor(): ) finally: if cpath is not None: - try: os.remove(cpath) - except Exception: pass + try: + os.remove(cpath) + except Exception: + pass importlib.invalidate_caches() try: @@ -3634,6 +3639,7 @@ 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 @@ -3647,8 +3653,10 @@ 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)") + 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}." @@ -3690,14 +3698,16 @@ def _unsloth_save_compressed_tensors( 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, - )) + 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...") @@ -3741,7 +3751,9 @@ def _unsloth_save_compressed_tensors( "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 + 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. @@ -3753,6 +3765,7 @@ def _unsloth_save_compressed_tensors( 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 @@ -3774,21 +3787,34 @@ def _unsloth_save_compressed_tensors( 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), + 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") + 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)...") + 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: @@ -3798,8 +3824,10 @@ def _unsloth_save_compressed_tensors( ) finally: if calib_tmp is not None and os.path.isdir(calib_tmp): - try: shutil.rmtree(calib_tmp) - except Exception: pass + try: + shutil.rmtree(calib_tmp) + except Exception: + pass # 6) Validate the artifact. cfg_path = os.path.join(out_dir, "config.json") @@ -3813,13 +3841,18 @@ def _unsloth_save_compressed_tensors( 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, + 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", + 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),