| |
| import os, sys, json, types, time |
| os.environ["NCM_ROOT"]="/kaggle/working/w/ncm" |
| import torch, torch.nn as nn, numpy as np, ncmchat as N |
| import torch.nn.functional as _F |
| _orig_interp=_F.interpolate |
| def _area_mat(L,o): |
| m=np.zeros((o,L),dtype=np.float32) |
| for i in range(o): |
| a=(i*L)//o; b=-((-(i+1)*L)//o); m[i,a:b]=1.0/(b-a) |
| return torch.from_numpy(m) |
| def _lm(L,o): |
| m=np.zeros((o,L),dtype=np.float32) |
| for i in range(o): |
| s=max((i+0.5)*L/o-0.5,0.0); i0=min(int(s),L-1); i1=min(i0+1,L-1); w=s-i0 |
| m[i,i0]+=1-w; m[i,i1]+=w |
| return torch.from_numpy(m) |
| def _interp(x,size=None,scale_factor=None,mode="nearest",**kw): |
| if mode=="area" and x.dim()==3 and size is not None: |
| o=size[0] if isinstance(size,(tuple,list)) else size |
| return torch.matmul(x,_area_mat(x.shape[-1],o).t()) |
| if mode=="linear" and x.dim()==3 and size is not None and not kw.get("align_corners"): |
| o=size[0] if isinstance(size,(tuple,list)) else size |
| return torch.matmul(x,_lm(x.shape[-1],o).t()) |
| return _orig_interp(x,size=size,scale_factor=scale_factor,mode=mode,**kw) |
| _chk=torch.randn(2,5,48); assert torch.allclose(_orig_interp(_chk,size=(26,),mode="area"),_interp(_chk,size=(26,),mode="area"),atol=1e-5) |
| _F.interpolate=_interp |
| SRC=os.environ.get("SRC","ncm-chat-smoke") |
| cands=[f"/kaggle/input/{SRC}/generator_chat.pt"]+[os.path.join(r,"generator_chat.pt") for r,_,f in os.walk("/kaggle/input") if "generator_chat.pt" in f] |
| gp=[c for c in cands if os.path.exists(c)][0]; print("using",gp,flush=True) |
| PEFF=int(os.environ.get("PEFF","128")) |
| import re |
| _c=open("ck/config.yaml").read(); _c=re.sub(r"prefix_len: 512","prefix_len: %d"%PEFF,_c); open("ck/config_p.yaml","w").write(_c) |
| CFG,tk,net=N.load("ck",cfg_path="ck/config_p.yaml",dev="cpu",load_gen=False); net.eval() |
| sd=torch.load(gp,map_location="cpu",weights_only=False)["model"] |
| own=net.state_dict(); sdf={k:v.float() for k,v in sd.items() if k in own and own[k].shape==v.shape} |
| print("skipped keys:",[k for k in sd if k not in sdf],flush=True) |
| net.load_state_dict(sdf,strict=False); assert all(k in sdf for k in own if "attn_mask" not in k and not k.startswith("rope")), "missing weights" |
| codec=N.Codec(CFG,tk,net); offs=[0] |
| for l in net.levels: offs.append(offs[-1]+l) |
| class G(nn.Module): |
| def __init__(s): |
| super().__init__(); s.net=net; s.vq=tk.vqvae.vq; s.dec=tk.vqvae.decoder; s.enc=tk.vqvae.encoder |
| def forward(s, prefix, noise, temp): |
| st={"i":0} |
| def samp(self_, logits, temperature, top_k, top_p): |
| i=st["i"]; st["i"]+=1 |
| nz=noise[:, offs[i]:offs[i+1], :] |
| v,_=torch.topk(logits,20,dim=-1) |
| lg=torch.where(logits<v[...,-1:],torch.full_like(logits,-1e9),logits) |
| return (lg+temp*nz).argmax(-1) |
| s.net._sample_level=types.MethodType(samp,s.net) |
| g=s.net.generate(None,prefix,temperature=1.0,top_k=20,top_p=0.9) |
| _,zl=s.vq.idx_to_zhats(g,return_zlevels=True) |
| return s.dec(zl[-1]).argmax(-1) |
| m=G().eval() |
| tests=["Hi, how are you?","What is the capital of France?","Tell me a joke.","What do you like to do for fun?","Can you recommend a good book?","How do I boil an egg?","I'm feeling sad today.","What is the meaning of life?"] |
| prompts=[N.fmt_prompt([("user",t)]) for t in tests] |
| pre=torch.tensor([codec.prefix_ids(p) for p in prompts]) |
| Z=torch.zeros(1,offs[-1],4096); T0=torch.tensor(0.0) |
| def txt(ids): |
| ids=list(ids) |
| if codec.EOS in ids: ids=ids[:ids.index(codec.EOS)] |
| return codec.pre.decode([t for t in ids if t not in (codec.EOS,codec.PAD)],skip_special_tokens=True).strip() |
| with torch.no_grad(): ref=[m(pre[i:i+1],Z,T0)[0].tolist() for i in range(len(tests))] |
| print("torch greedy:",[txt(r) for r in ref],flush=True) |
| os.makedirs("/kaggle/working/out",exist_ok=True) |
| f32="/kaggle/working/ncm_fp32.onnx" |
| with torch.no_grad(): torch.onnx.export(m,(pre[:1],Z,T0),f32,input_names=["prefix","noise","temp"],output_names=["ids"],opset_version=17,dynamo=False) |
| print("fp32 onnx MB",os.path.getsize(f32)/1e6,flush=True) |
| import onnxruntime as ort |
| from onnxruntime.quantization import quantize_dynamic, QuantType |
| import onnx |
| from onnx import numpy_helper, helper |
| mo=onnx.load(f32); g=mo.graph; inits={i.name:i for i in g.initializer}; newn=[]; ng=0 |
| for n in g.node: |
| if n.op_type=="Gemm" and n.input[1] in inits and (len(n.input)<3 or n.input[2] in inits): |
| at={a.name:helper.get_attribute_value(a) for a in n.attribute} |
| W=numpy_helper.to_array(inits[n.input[1]]); W=(W.T if at.get("transB",0) else W)*at.get("alpha",1.0) |
| wn=n.input[1]+"_mm"; g.initializer.append(numpy_helper.from_array(W.astype("float32"),wn)) |
| mm=n.output[0]+"_mmout" if len(n.input)>=3 else n.output[0] |
| newn.append(helper.make_node("MatMul",[n.input[0],wn],[mm],name=n.name+"_mm")) |
| if len(n.input)>=3: newn.append(helper.make_node("Add",[mm,n.input[2]],[n.output[0]],name=n.name+"_add")) |
| ng+=1 |
| else: newn.append(n) |
| |
| inits_all={i.name for i in g.initializer} |
| alias={} |
| final=[] |
| for n in newn: |
| if n.op_type=="Identity" and (n.input[0] in inits_all or n.input[0] in alias): |
| alias[n.output[0]]=n.input[0] |
| else: final.append(n) |
| def res(x): |
| while x in alias: x=alias[x] |
| return x |
| for n in final: |
| for j,x in enumerate(n.input): n.input[j]=res(x) |
| for o in g.output: |
| if o.name in alias: raise SystemExit("identity feeds graph output") |
| del g.node[:]; g.node.extend(final); print("gemm->matmul",ng,"identity removed",len(alias),flush=True) |
| f32m="/kaggle/working/ncm_fp32_mm.onnx"; onnx.save(mo,f32m) |
| i8="/kaggle/working/out/ncm_fp16w.onnx" |
| mo=onnx.load(f32m); g=mo.graph; casts=[]; nconv=0 |
| for init in list(g.initializer): |
| if init.data_type==onnx.TensorProto.FLOAT and int(np.prod(init.dims))>=1024: |
| arr=numpy_helper.to_array(init); h=numpy_helper.from_array(arr.astype(np.float16),init.name+"_h") |
| casts.append(helper.make_node("Cast",[init.name+"_h"],[init.name],to=onnx.TensorProto.FLOAT,name=init.name+"_cast")) |
| g.initializer.remove(init); g.initializer.append(h); nconv+=1 |
| nodes=list(g.node); del g.node[:]; g.node.extend(casts+nodes) |
| onnx.save(mo,i8); print("fp16-weights onnx MB",os.path.getsize(i8)/1e6,"converted",nconv,flush=True) |
| res={"files_MB":{"fp32":os.path.getsize(f32)/1e6,"fp16w":os.path.getsize(i8)/1e6},"cases":[]} |
| for name,path in [("fp32",f32m),("fp16w",i8)]: |
| s=ort.InferenceSession(path,providers=["CPUExecutionProvider"]); outs=[] |
| t=time.time() |
| for i in range(len(tests)): |
| o=s.run(None,{"prefix":pre[i:i+1].numpy(),"noise":Z.numpy(),"temp":np.array(0.0,dtype=np.float32)})[0][0].tolist() |
| outs.append(o) |
| dt=(time.time()-t)/len(tests); res[name]={"sec_per_reply_cpu":dt,"exact_match_vs_torch":sum(o==r for o,r in zip(outs,ref)),"replies":[txt(o) for o in outs]} |
| print(name,"exact-match",res[name]["exact_match_vs_torch"],"/",len(tests),"s/reply",round(dt,2),flush=True) |
| for t_,o in zip(tests,outs): print(" ",t_,"->",txt(o),flush=True) |
| res["torch_replies"]=[txt(r) for r in ref] |
| |
| res["tok_vectors"]=[{"prompt":p,"ids":codec.enc(p),"prefix_ids":codec.prefix_ids(p)[:len(codec.enc(p))]} for p in prompts[:4]] |
| res["EOS"]=codec.EOS; res["PAD"]=codec.PAD; res["ppad"]=codec.ppad; res["P"]=codec.P; res["offs"]=offs |
| json.dump(res,open("/kaggle/working/out/report.json","w"),indent=1) |
| print("done",flush=True) |
|
|
