| |
| """Merge LoRA into the QAT foundation, then quantize to compressed-tensors W4A16 |
| (int4, group 32, symmetric) - Google's own recipe - for SGLang serving. |
| |
| CRITICAL (handover 6.7): save_pretrained drops config keys the dataclass doesn't |
| declare. Backfill vision_config.model_patch_size / num_soft_tokens / |
| audio_config.output_proj_dims from the base config and ASSERT after every write. |
| """ |
| import os, sys, json, shutil, argparse, torch |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
|
|
| BASE_DEFAULT="/workspace/models/gemma4-qat-unquantized" |
| KEYS=(("vision_config","model_patch_size"),("vision_config","num_soft_tokens"), |
| ("audio_config","output_proj_dims")) |
|
|
| def patch_cfg(out, base): |
| """Backfill the three keys SGLang needs. Catches absent AND null.""" |
| cfg=os.path.join(out,"config.json"); b=json.load(open(os.path.join(base,"config.json"))) |
| c=json.load(open(cfg)); fixed=[] |
| for sec,k in KEYS: |
| c.setdefault(sec,{}) |
| if c[sec].get(k) is None: |
| c[sec][k]=b[sec][k]; fixed.append(f"{sec}.{k}={b[sec][k]}") |
| json.dump(c,open(cfg,"w"),indent=2) |
| c=json.load(open(cfg)) |
| for sec,k in KEYS: |
| assert c[sec].get(k) is not None, f"config lacks {sec}.{k} after patch" |
| return fixed |
|
|
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument("--adapter",required=True); ap.add_argument("--out",required=True) |
| ap.add_argument("--base",default=BASE_DEFAULT) |
| ap.add_argument("--merged-dir",default=None) |
| ap.add_argument("--skip-quant",action="store_true") |
| ap.add_argument("--quant-embeddings",action="store_true") |
| A=ap.parse_args() |
| merged=A.merged_dir or (A.out.rstrip("/")+"-bf16") |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from peft import PeftModel |
| print(f"[merge] base={A.base}\n[merge] adapter={A.adapter}",flush=True) |
| tok=AutoTokenizer.from_pretrained(A.base) |
| m=AutoModelForCausalLM.from_pretrained(A.base,dtype=torch.bfloat16,device_map="cpu") |
| m=PeftModel.from_pretrained(m,A.adapter) |
| m=m.merge_and_unload() |
| os.makedirs(merged,exist_ok=True) |
| m.save_pretrained(merged,safe_serialization=True) |
| tok.save_pretrained(merged) |
| for f in ("chat_template.jinja","processor_config.json","generation_config.json"): |
| s=os.path.join(A.base,f) |
| if os.path.exists(s): shutil.copy(s,os.path.join(merged,f)) |
| print("[merge] config backfill:",patch_cfg(merged,A.base),flush=True) |
| n=sum(p.numel() for p in m.parameters()) |
| print(f"[merge] merged bf16 -> {merged} params={n/1e9:.2f}B",flush=True) |
| del m |
| if A.skip_quant: return |
| print("[quant] handing off to the isolated quant env (llm-compressor)",flush=True) |
|
|
| if __name__=="__main__": main() |
|
|
