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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-26 17:38:07 +00:00
commit e9a80fe436
2 changed files with 135 additions and 79 deletions
Download patch file Download diff file Expand all files Collapse all files

99
unsloth/_compressed_quantize.py
View file

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

115
unsloth/save.py
View file

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