File size: 2,239 Bytes
b296ad4 | 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 | """Average compatible checkpoints from the same randomly initialized training lineage."""
import argparse
import json
from pathlib import Path
import torch
from safetensors.torch import load_file,save_file
from tinyquery.prepare import file_sha256
def main():
p=argparse.ArgumentParser();p.add_argument('--checkpoints',nargs='+',required=True);p.add_argument('--weights',nargs='+',type=float)
p.add_argument('--out',required=True);args=p.parse_args();weights=args.weights or [1]*len(args.checkpoints)
assert len(weights)==len(args.checkpoints) and all(w>=0 for w in weights) and sum(weights)>0
weights=[w/sum(weights) for w in weights];total={};config=None;sources=[]
for name,weight in zip(args.checkpoints,weights):
path=Path(name);c=json.loads((path.parent/'config.json').read_text());assert config is None or config==c;config=c
info=json.loads((path.parent/'checkpoint-info.json').read_text());assert info['random_initialization']
state=load_file(str(path));assert not total or total.keys()==state.keys()
for key,value in state.items():
if key not in total:total[key]=value.float()*weight
else:total[key].add_(value.float(),alpha=weight)
sources.append({'checkpoint':str(path),'sha256':file_sha256(path),'weight':weight,'info':info});del state
dest=Path(args.out);dest.mkdir(parents=True,exist_ok=True)
save_file({k:v.to(torch.bfloat16).contiguous() for k,v in total.items()},str(dest/'model.safetensors'),
metadata={'method':'weighted_parameter_average','random_initialization':'true'})
(dest/'config.json').write_text(json.dumps(config,indent=2))
info={'method':'weighted_parameter_average','sources':sources,'random_initialization':True,
'step':max(s['info']['step'] for s in sources),'step_interpretation':'Latest source step; these are averaged weights, not that raw checkpoint.'}
for key in ['processed_tokens','response_tokens','training_seconds']:info[key]=max(s['info'][key] for s in sources)
(dest/'checkpoint-info.json').write_text(json.dumps(info,indent=2));print(json.dumps({'out':str(dest),'sources':[(s['info']['step'],s['weight']) for s in sources]}))
if __name__=='__main__':main()
|