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;