From 0ea1ad4262a904e480dae1ccca44766b26a79d3f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 09:09:40 +0000 Subject: [PATCH] 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..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. --- tests/saving/test_export_api_surface.py | 22 ++++++++++++++++++++++ unsloth/_compressed_quantize.py | 4 +++- unsloth/save.py | 6 ++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/saving/test_export_api_surface.py b/tests/saving/test_export_api_surface.py index 62b72e5673..7955b50968 100644 --- a/tests/saving/test_export_api_surface.py +++ b/tests/saving/test_export_api_surface.py @@ -148,6 +148,28 @@ def test_export_subprocesses_are_shell_safe(): 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. diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index ac15b1e1e7..5be680f15b 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -195,6 +195,7 @@ def main(): 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 @@ -304,7 +305,8 @@ def main(): oneshot(model = model, recipe = _make_recipe()) os.makedirs(args.out, exist_ok = True) - model.save_pretrained(args.out, save_compressed = 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) diff --git a/unsloth/save.py b/unsloth/save.py index ad61f12b2b..e62987d396 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -3919,6 +3919,10 @@ def _unsloth_save_compressed_tensors( # 3) Merge to 16bit at local_dir (kept for local saves) via unsloth_generic_save, so LoRA # adapters are merged and full-finetuned models written in 16bit consistently. Extra # save kwargs (state_dict, max_shard_size, ...) flow through merge_kwargs. + # The intermediate 16bit checkpoint is internal staging that the converter subprocess + # reloads with default weight filenames, so never write variant-named shards here; the + # user's variant (if any) is applied to the final compressed checkpoint in the subprocess. + variant = merge_kwargs.pop("variant", None) print(f"Unsloth: Merging to 16bit before {scheme} quantization...") merge_args = dict(merge_kwargs) merge_args.update( @@ -4026,6 +4030,8 @@ def _unsloth_save_compressed_tensors( cmd.append("--is-vlm") if trust_remote_code: cmd.append("--trust-remote-code") + if variant: + cmd += ["--variant", variant] # Free the in-memory model's CUDA memory before the subprocess loads its own copy from # disk, so a single GPU need not hold both at once. Best-effort and restored in finally;