| """Average checkpoints into a soup (the lever that carried this release). |
| |
| Same-basin checkpoints from the same recipe, saved at different points of |
| short continuation runs, average into a better model than any member: on this |
| line the final soups scored +0.09 SCOREQ above their best single member, for |
| zero training. |
| |
| # uniform average |
| python average_checkpoints.py --out decoder_soup.pt a.pt b.pt c.pt d.pt |
| |
| # weighted average (the acoustic blend used weights 2,2,1,1,1) |
| python average_checkpoints.py --out acoustic_G4.pt \ |
| --weights 2,2,1,1,1 bs015.pt ft2.pt t1.pt t2.pt t3.pt |
| |
| Rules (measured on this line): |
| - Only average checkpoints with IDENTICAL key sets; the script refuses to mix |
| different architectures. |
| - Float tensors are averaged; non-float tensors are copied from the first |
| member. |
| - Members must share a basin (same recipe, diverged only by training noise |
| and step count). Mixing distant runs loses. |
| - 2 to 6 members is the proven range; more members dilute. The shipped |
| decoder is the average of its two best soups, which itself beat larger |
| blends, and the shipped acoustic is a pair-heavy blend of five |
| checkpoints. |
| - Averaging can be applied to soups again (nested), e.g. the shipped decoder |
| is avg(e5, ss2) where e5 and ss2 are themselves averages. |
| """ |
| import argparse |
|
|
| import torch |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--weights", default=None, |
| help="comma-separated per-input weights (default: uniform)") |
| ap.add_argument("inputs", nargs="+", help="checkpoint files to average") |
| args = ap.parse_args() |
|
|
| weights = None |
| if args.weights: |
| weights = [float(x) for x in args.weights.split(",")] |
| if len(weights) != len(args.inputs): |
| raise SystemExit(f"{len(weights)} weights for {len(args.inputs)} inputs") |
| else: |
| weights = [1.0] * len(args.inputs) |
| tot = sum(weights) |
|
|
| sds = [torch.load(p, map_location="cpu", weights_only=False) for p in args.inputs] |
| ks0 = set(sds[0].keys()) |
| for p, sd in zip(args.inputs[1:], sds[1:]): |
| if set(sd.keys()) != ks0: |
| raise SystemExit(f"key mismatch: {p} does not match {args.inputs[0]}") |
|
|
| out = {} |
| for k in sds[0].keys(): |
| ref = sds[0][k] |
| if torch.is_tensor(ref) and ref.is_floating_point(): |
| acc = None |
| for w, sd in zip(weights, sds): |
| v = sd[k].float() * (w / tot) |
| acc = v if acc is None else acc + v |
| out[k] = acc |
| else: |
| out[k] = ref |
| torch.save(out, args.out) |
| n = sum(v.numel() for v in out.values() if torch.is_tensor(v)) |
| print(f"averaged {len(sds)} checkpoints -> {args.out} ({n:,} parameters)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
