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.<variant>.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.
This commit is contained in:
Daniel Han 2026-06-28 09:09:40 +00:00
commit 0ea1ad4262
3 changed files with 31 additions and 1 deletions

View file

@ -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.

View file

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

View file

@ -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;