File size: 5,622 Bytes
b2f3bf4 448ef61 b2f3bf4 448ef61 b2f3bf4 448ef61 b2f3bf4 448ef61 b2f3bf4 448ef61 b2f3bf4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python3
"""Export a serving bundle (DESIGN §9): merged text-only backbone (bf16 safetensors, HF layout of
Qwen3_5ForCausalLM without lm_head) + head.safetensors + judge_config.json + calibration.json + tokenizer.
python3 scripts/export.py --checkpoint checkpoints/s2_seed42/best --calibration checkpoints/s2_seed42/best/calibration.json \
--out exports/jev-judge-qwen35-9b --name jev-judge-qwen35-9b --version 0.7.0
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import sys
import torch
from safetensors.torch import save_file
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from jev_judge import __version__ # noqa: E402
from jev_judge.calibration import load as load_calibration # noqa: E402
from jev_judge.checkpointing import load_checkpoint # noqa: E402
def save_sharded(state: dict[str, torch.Tensor], out_dir: str, max_shard_bytes: int = 4 * 1024**3) -> None:
shards: list[dict[str, torch.Tensor]] = [{}]
sizes = [0]
for k, v in state.items():
nbytes = v.numel() * v.element_size()
if sizes[-1] + nbytes > max_shard_bytes and shards[-1]:
shards.append({}); sizes.append(0)
shards[-1][k] = v
sizes[-1] += nbytes
n = len(shards)
weight_map = {}
for i, sh in enumerate(shards):
name = f"model-{i+1:05d}-of-{n:05d}.safetensors"
save_file({k: t.contiguous() for k, t in sh.items()}, os.path.join(out_dir, name), metadata={"format": "pt"})
weight_map.update({k: name for k in sh})
with open(os.path.join(out_dir, "model.safetensors.index.json"), "w") as f:
json.dump({"metadata": {"total_size": sum(sizes)}, "weight_map": weight_map}, f, indent=2)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--checkpoint", required=True)
ap.add_argument("--calibration", required=True, help="calibration.json (fail-fast at serve time if missing)")
ap.add_argument("--out", required=True)
ap.add_argument("--name", default="jev-judge")
ap.add_argument("--version", default=__version__)
ap.add_argument("--include-lm-head", action="store_true",
help="also ship the base model's lm_head.weight so the bundle can do ordinary AR generation")
ap.add_argument("--mode", choices=["merged", "unmerged"], default="unmerged",
help="merged: LoRA folded into the shipped backbone (decision-only bundle); "
"unmerged: ship the *pristine* base backbone + adapter/ (dual-mode: AR generation is exactly the base model)")
args = ap.parse_args()
table = load_calibration(args.calibration) # validates schema
judge, cfg = load_checkpoint(args.checkpoint, device="cuda")
lm = judge.lm
os.makedirs(args.out, exist_ok=True)
if args.mode == "unmerged":
if not hasattr(lm, "peft_config"):
raise SystemExit("checkpoint has no LoRA adapter; use --mode merged")
lm.save_pretrained(os.path.join(args.out, "adapter")) # peft adapter (fp32 master weights)
lm = lm.unload() # drop adapter layers -> pristine base weights
args.include_lm_head = True # dual-mode bundles always ship the vocab head
elif hasattr(lm, "merge_and_unload"):
lm = lm.merge_and_unload()
lm.lm_head = None
state = {k: v.detach().cpu() for k, v in lm.state_dict().items() if not k.startswith("lm_head")} # dtypes as loaded (bf16 / fp32)
if args.include_lm_head:
# the judge dropped lm_head at load time; take the pristine rows from the base checkpoint
import glob
from safetensors import safe_open
found = False
for shard in sorted(glob.glob(os.path.join(cfg["base_model_path"], "*.safetensors"))):
with safe_open(shard, framework="pt", device="cpu") as f:
if "lm_head.weight" in f.keys():
state["lm_head.weight"] = f.get_tensor("lm_head.weight").to(torch.bfloat16)
found = True
break
if not found:
raise SystemExit("lm_head.weight not found in base checkpoint")
save_sharded(state, args.out)
# text-only config in Qwen3_5ForCausalLM layout
tcfg = lm.config.to_dict()
tcfg["architectures"] = ["Qwen3_5ForCausalLM"]
with open(os.path.join(args.out, "config.json"), "w") as f:
json.dump(tcfg, f, indent=2)
save_file({k: v.detach().to(torch.float32).cpu().contiguous() for k, v in judge.head.state_dict().items()},
os.path.join(args.out, "head.safetensors"))
jc = judge.judge_config({"model_name": args.name, "model_version": args.version, "lm_head_included": bool(args.include_lm_head),
"weights_mode": args.mode, "adapter_subfolder": "adapter" if args.mode == "unmerged" else None,
"lora": cfg.get("lora"), "trained_stage": cfg.get("stage"),
"source_checkpoint": os.path.abspath(args.checkpoint), "softcap": judge.head.softcap})
with open(os.path.join(args.out, "judge_config.json"), "w") as f:
json.dump(jc, f, indent=2)
shutil.copy(args.calibration, os.path.join(args.out, "calibration.json"))
judge.tokenizer.save_pretrained(args.out)
total = sum(v.numel() * v.element_size() for v in state.values())
print(f"exported [{args.mode}] {len(state)} tensors ({total/1e9:.1f} GB) + head + calibration + tokenizer"
f"{' + adapter/' if args.mode == 'unmerged' else ''} -> {args.out}")
if __name__ == "__main__":
main()
|