File size: 3,738 Bytes
97c39f2 | 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 | """Parallel merge recipes for post-training checkpoints (tiny-model-posttrain).
Runs model soup (arXiv 2203.05482) and task arithmetic (arXiv 2212.04089)
on the SAME base, producing one checkpoint per recipe. TIES is handled by
train/ties_merge.py (arXiv 2306.01708). After merging, battery-eval every
candidate and keep the best (LFM2 2511.23404 §4.4: parallel apply -> eval ->
select). Never naive-average adapters; these recipes operate on folded
full-weight checkpoints where delta = task_ckpt - base is the true task
vector.
Usage:
.venv/bin/python train/parallel_merges.py \
--base ckpt/hybrid50m_v16k_pretrain/model_5000.pt \
--tasks ckpt/hybrid50m_v25_lora/best.pt ckpt/hybrid50m_v25_dpo/model_final.pt \
--out-dir ckpt/hybrid50m_v25_merges \
--lambda-ta 0.5
"""
import argparse
from pathlib import Path
import torch
from model.config import TinyLiquidConfig
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
def load_sd(path, tag):
sd = torch.load(str(path), map_location="cpu", weights_only=False)
print(f" {tag}: {path} step={sd.get('step', '?')} tag={sd.get('tag', '-')}", flush=True)
return sd
def model_soup(task_sds):
"""Simple average of task weights (all tasks trained from the same base)."""
soup = {}
keys = [k for k in task_sds[0] if task_sds[0][k].is_floating_point()
and all(k in sd for sd in task_sds[1:])]
for k in keys:
soup[k] = torch.stack([sd[k].float() for sd in task_sds]).mean(dim=0)
return soup
def task_arithmetic(base_sd, task_sds, lam):
"""base + lam * sum(task_i - base)."""
merged = {}
keys = [k for k in base_sd if base_sd[k].is_floating_point()
and all(k in sd for sd in task_sds)]
for k in keys:
base = base_sd[k].float()
delta = torch.zeros_like(base)
for sd in task_sds:
delta += sd[k].float() - base
merged[k] = base + lam * delta
return merged
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base", required=True)
ap.add_argument("--tasks", nargs="+", required=True)
ap.add_argument("--out-dir", required=True)
ap.add_argument("--lambda-ta", type=float, default=0.5)
ap.add_argument("--threads", type=int, default=4)
args = ap.parse_args()
assert len(args.tasks) >= 2, "merges need >= 2 task checkpoints"
torch.set_num_threads(args.threads)
base_path = Path(args.base)
base_ckpt = base_path if base_path.is_file() else latest_ckpt(args.base)
base = load_sd(base_ckpt, "base")
base_sd = base["model"]
task_sds = []
for i, t in enumerate(args.tasks):
tp = Path(t)
tp = tp if tp.is_file() else latest_ckpt(t)
task_sds.append(load_sd(tp, f"task{i}")["model"])
recipes = {
"soup": model_soup(task_sds),
f"taskarith_l{args.lambda_ta}".replace(".", "p"): task_arithmetic(base_sd, task_sds, args.lambda_ta),
}
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
for name, merged in recipes.items():
cfg = TinyLiquidConfig(**base["config"])
cfg.mtp_heads = 0
model = TinyLiquid(cfg)
missing, unexpected = model.load_state_dict(merged, strict=False)
if missing or unexpected:
print(f"[{name}] ignored {len(missing)} missing / {len(unexpected)} unexpected keys", flush=True)
out = out_dir / f"{name}.pt"
torch.save({"config": base["config"], "model": merged, "step": 0,
"best_val": base.get("best_val", float("inf")),
"tag": f"{name}-{len(task_sds)}tasks"}, out)
print("saved", out, flush=True)
if __name__ == "__main__":
main()
|