| """Weight-space interpolation ("model soup") of Groot checkpoints with the same architecture. |
| |
| python scripts/soup.py --out runs/soup --models A B [C ...] [--weights 0.5 0.5] |
| |
| Floating-point tensors are averaged with the given weights; the tokenizer, configs and |
| calibration are copied from the first model. Recalibrate afterwards with |
| `groot-eval --calibrate-on val.jsonl --save-calibration`. |
| """ |
| import argparse |
| import shutil |
| from pathlib import Path |
|
|
| from safetensors.torch import load_file, save_file |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--models", nargs="+", required=True) |
| ap.add_argument("--weights", nargs="+", type=float, default=None) |
| ap.add_argument("--out", required=True) |
| args = ap.parse_args() |
| w = args.weights or [1.0 / len(args.models)] * len(args.models) |
| if len(w) != len(args.models) or abs(sum(w) - 1) > 1e-6: |
| raise SystemExit("weights must match models and sum to 1") |
| states = [load_file(str(Path(m) / "model.safetensors")) for m in args.models] |
| keys = states[0].keys() |
| if any(s.keys() != keys for s in states): |
| raise SystemExit("checkpoints have different tensors") |
| out = {} |
| for k in keys: |
| t = states[0][k] |
| out[k] = sum(wi * s[k].float() for wi, s in zip(w, states)).to(t.dtype) if t.is_floating_point() else t |
| dst = Path(args.out) |
| if dst.exists(): |
| raise SystemExit(f"{dst} exists") |
| shutil.copytree(args.models[0], dst, ignore=shutil.ignore_patterns("model.safetensors", "optimizer.pt")) |
| save_file({k: v.contiguous() for k, v in out.items()}, str(dst / "model.safetensors"), metadata={"format": "pt", "family": "groot", "soup": str(w)}) |
| print(f"soup of {len(args.models)} models (weights {w}) -> {dst}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
