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:
parent
ce75482234
commit
0ea1ad4262
3 changed files with 31 additions and 1 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue