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