File size: 1,774 Bytes
b2ebc95 | 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 | #!/usr/bin/env python3
"""Average several checkpoints of the same run into one ("model soup" / SWA).
All inputs must be my own checkpoints from this run's training tree; averaging along a
training trajectory is the classic stochastic-weight-averaging trick and usually buys a
little robustness for free. Copies the tokenizer/config/processor files from the first input.
Usage: soup.py OUT CKPT1 CKPT2 [CKPT3 ...]
"""
import json
import os
import shutil
import sys
import torch
from safetensors.torch import load_file, save_file
def main():
out, ckpts = sys.argv[1], sys.argv[2:]
if len(ckpts) < 2:
sys.exit(__doc__)
os.makedirs(out, exist_ok=True)
index_path = os.path.join(ckpts[0], "model.safetensors.index.json")
index = json.load(open(index_path))
shards = sorted(set(index["weight_map"].values()))
print(f"{len(ckpts)} checkpoints, {len(shards)} shards")
for shard in shards:
acc = None
for i, c in enumerate(ckpts):
sd = load_file(os.path.join(c, shard))
if acc is None:
acc = {k: v.to(torch.float32) for k, v in sd.items()}
else:
for k in acc:
acc[k] += sd[k].to(torch.float32)
del sd
for k in acc:
acc[k] = (acc[k] / len(ckpts)).to(torch.bfloat16)
save_file(acc, os.path.join(out, shard), metadata={"format": "pt"})
print("wrote", shard, flush=True)
del acc
for f in os.listdir(ckpts[0]):
if f.endswith(".safetensors") or f == "STABLE":
continue
shutil.copy2(os.path.join(ckpts[0], f), os.path.join(out, f))
open(os.path.join(out, "STABLE"), "w").close()
print("soup ->", out)
if __name__ == "__main__":
main()
|