File size: 2,471 Bytes
8d47318 | 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 | import argparse,hashlib,json,struct,sys
from pathlib import Path
import numpy as np, torch
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
from src.model import EXLLM,EXLLMConfig
MAGIC=b'EXLLM8\0\0'
def quant_row(t):
a=t.detach().float().cpu().numpy();mx=np.max(np.abs(a),axis=1);sc=np.where(mx>0,mx/127.0,1.0).astype(np.float32);q=np.clip(np.rint(a/sc[:,None]),-127,127).astype(np.int8);return q,sc
def main():
ap=argparse.ArgumentParser();ap.add_argument('--checkpoint',default='weights/EXLLM-ONI5M-v1.0.pt');ap.add_argument('--out',default='weights/EXLLM-ONI5M-v1.0-int8.bin');a=ap.parse_args()
c=torch.load(ROOT/a.checkpoint,map_location='cpu',weights_only=False);m=EXLLM(EXLLMConfig(**c['config']));m.load_state_dict(c['model']);m.eval();recs=[]
for name,t in m.state_dict().items():
if name=='lm_head.weight':continue
if name.endswith('.weight') and 'norm' in name:
x=t.detach().cpu().numpy().astype(np.float16);recs.append((name,2,list(x.shape),np.array([],dtype=np.float32),x.tobytes()))
elif t.ndim==2:
q,sc=quant_row(t);recs.append((name,1,list(q.shape),sc,q.tobytes()))
else:
x=t.detach().cpu().numpy().astype(np.float16);recs.append((name,2,list(x.shape),np.array([],dtype=np.float32),x.tobytes()))
out=ROOT/a.out;out.parent.mkdir(parents=True,exist_ok=True)
with out.open('wb') as f:
f.write(MAGIC);f.write(struct.pack('<II',1,len(recs)))
for name,qt,dims,sc,data in recs:
nb=name.encode();f.write(struct.pack('<HBB',len(nb),len(dims),qt));f.write(nb)
for d in dims:f.write(struct.pack('<I',d))
f.write(struct.pack('<II',len(sc),len(data)));f.write(sc.astype('<f4').tobytes());f.write(data)
sha=hashlib.sha256(out.read_bytes()).hexdigest();man={'format':'EXLLM8','format_version':1,'model':'EXLLM-ONI5M','model_version':'1.0.0','config':m.cfg.__dict__,'unique_parameters':m.num_parameters(),'tokenizer_vocab':m.cfg.vocab_size,'quantization':{'matrix_weights':'symmetric int8 per output row','matrix_scale':'float32 per row','rmsnorm_weights':'float16'},'aliases':{'lm_head.weight':'tok.weight'},'tensor_count':len(recs),'file':str(out.relative_to(ROOT)),'bytes':out.stat().st_size,'sha256':sha}
out.with_suffix('.manifest.json').write_text(json.dumps(man,ensure_ascii=False,indent=2),encoding='utf-8');print(json.dumps(man,ensure_ascii=False,indent=2))
if __name__=='__main__':main()
|