brady777 commited on
Commit
92a5f5b
·
verified ·
1 Parent(s): c60ad48

Add adapter merge utility referenced by the card

Browse files
Files changed (1) hide show
  1. scripts/fastloop_merge_adapter.py +23 -0
scripts/fastloop_merge_adapter.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """fastloop_merge_adapter.py <base-dir> <adapter-dir> <merged-dir> — re-create the
3
+ BF16 merged checkpoint from a kept adapter (the ladder runners delete merged/
4
+ after eval; when an eval fails for plumbing reasons this rebuilds it in ~5 min
5
+ instead of a 2 h retune). Same merge law as fastloop_smoke_train.py M gate."""
6
+ import os, shutil, sys, torch
7
+ from peft import PeftModel
8
+ from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText, AutoTokenizer
9
+ base, adapter, merged_dir = sys.argv[1:4]
10
+ cfg = AutoConfig.from_pretrained(base)
11
+ cls = AutoModelForImageTextToText if hasattr(cfg, "text_config") else AutoModelForCausalLM
12
+ m = cls.from_pretrained(base, device_map="auto", dtype=torch.bfloat16)
13
+ m = PeftModel.from_pretrained(m, adapter).merge_and_unload()
14
+ if os.path.isdir(merged_dir): shutil.rmtree(merged_dir)
15
+ m.save_pretrained(merged_dir, safe_serialization=True)
16
+ AutoTokenizer.from_pretrained(base).save_pretrained(merged_dir)
17
+ for f in os.listdir(base):
18
+ if f.endswith(".index.json"): continue
19
+ if f.endswith((".jinja", ".json", ".py", ".txt", ".model")) and not os.path.exists(os.path.join(merged_dir, f)):
20
+ shutil.copy(os.path.join(base, f), merged_dir)
21
+ st = [f for f in os.listdir(merged_dir) if f.endswith(".safetensors")]
22
+ print(f"MERGE PASS: {len(st)} shards -> {merged_dir}" if st else "MERGE FAIL")
23
+ sys.exit(0 if st else 1)