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