| |
| """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() |
|
|
