| |
| import os, sys, json, types, re, 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=_F.interpolate |
| def _am(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,_am(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(x,size=size,scale_factor=scale_factor,mode=mode,**kw) |
| gp=[os.path.join(r,"generator_chat.pt") for r,_,f in os.walk("/kaggle/input") if "generator_chat.pt" in f][0] |
| P=128 |
| c=open("ck/config.yaml").read(); c=re.sub(r"prefix_len: 512","prefix_len: %d"%P,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() |
| net.load_state_dict({k:v.float() for k,v in sd.items() if k in own and own[k].shape==v.shape},strict=False) |
| codec=N.Codec(CFG,tk,net); offs=[0] |
| for l in net.levels: offs.append(offs[-1]+l) |
| levels=list(net.levels) |
| class GTF(nn.Module): |
| def __init__(s): |
| super().__init__(); s.net=net |
| def forward(s,prefix,codes): |
| st={"i":0}; L=[] |
| def samp(self_,logits,temperature,top_k,top_p): |
| i=st["i"]; st["i"]+=1; L.append(logits) |
| return codes[:,offs[i]:offs[i+1]] |
| s.net._sample_level=types.MethodType(samp,s.net) |
| s.net.generate(None,prefix,temperature=1.0,top_k=20,top_p=0.9) |
| return torch.cat(L,1) |
| m=GTF().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]) |
| |
| reps=["Hello! How can I help you today?","The capital of France is Paris.","Why did the chicken cross the road? To get to the other side.","I like to read books and go hiking.","I recommend The Hobbit, it is a fun story.","Put it in boiling water for ten minutes.","I am sorry to hear that. I hope you feel better.","Life is about finding what makes you happy."] |
| tgt=torch.tensor([codec.target_ids(r) for r in reps]) |
| with torch.no_grad(): codes=tk.vqvae.encode(tgt) |
| |
| with torch.no_grad(): |
| lg_gen=torch.cat([m(pre[i:i+1],codes[i:i+1]) for i in range(len(tests))],0) |
| x=N.build_inputs(codec,pre,codes,0.0); pm=pre==codec.ppad |
| lg_tr=net(x,prefix_pad_mask=pm) |
| d_tr=(lg_gen-lg_tr).abs() |
| print("torch generate-path vs training-forward logits: max|d|",float(d_tr.max()),"mean|d|",float(d_tr.mean()),"argmax mismatch",int((lg_gen.argmax(-1)!=lg_tr.argmax(-1)).sum()),"/",lg_gen.shape[0]*lg_gen.shape[1],flush=True) |
| os.makedirs("/kaggle/working/out",exist_ok=True) |
| _F.interpolate=_interp |
| torch.onnx.export(m,(pre[:1],codes[:1]),"/kaggle/working/tf.onnx",input_names=["prefix","codes"],output_names=["logits"],opset_version=17,dynamo=False) |
| import onnxruntime as ort, onnx |
| from onnx import numpy_helper, helper |
| def clean_quant(src,dst): |
| mo=onnx.load(src); g=mo.graph |
| keep=[] |
| for n in g.node: |
| if n.op_type=="Constant" and len(n.attribute)==1 and n.attribute[0].name=="value": |
| t=n.attribute[0].t; t.name=n.output[0]; g.initializer.append(t) |
| else: keep.append(n) |
| del g.node[:]; g.node.extend(keep) |
| inits={i.name for i in g.initializer}; alias={}; final=[] |
| for n in g.node: |
| if n.op_type=="Identity" and (n.input[0] in inits or n.input[0] in alias): alias[n.output[0]]=n.input[0] |
| else: final.append(n) |
| def r(x): |
| while x in alias: x=alias[x] |
| return x |
| for n in final: |
| for j,xx in enumerate(n.input): n.input[j]=r(xx) |
| del g.node[:]; g.node.extend(final); onnx.save(mo,"/kaggle/working/tf_clean.onnx") |
| from onnxruntime.quantization import quantize_dynamic, QuantType |
| quantize_dynamic("/kaggle/working/tf_clean.onnx",dst,weight_type=QuantType.QInt8) |
| HAVE_I8=True |
| try: |
| clean_quant("/kaggle/working/tf.onnx","/kaggle/working/tf_i8.onnx"); print("int8 MB",os.path.getsize("/kaggle/working/tf_i8.onnx")/1e6,flush=True) |
| except Exception as e: |
| HAVE_I8=False; print("INT8 QUANT FAILED",repr(e)[:300],flush=True) |
| def run(path,opt): |
| so=ort.SessionOptions(); so.graph_optimization_level=opt |
| s=ort.InferenceSession(path,so,providers=["CPUExecutionProvider"]) |
| return torch.cat([torch.from_numpy(s.run(None,{"prefix":pre[i:i+1].numpy(),"codes":codes[i:i+1].numpy()})[0]) for i in range(len(tests))],0) |
| rep={} |
| for name,path,opt in [("onnx_fp32_default","/kaggle/working/tf.onnx",ort.GraphOptimizationLevel.ORT_ENABLE_ALL),("onnx_fp32_noopt","/kaggle/working/tf.onnx",ort.GraphOptimizationLevel.ORT_DISABLE_ALL),("onnx_int8","/kaggle/working/tf_i8.onnx",ort.GraphOptimizationLevel.ORT_ENABLE_ALL)]: |
| if name=="onnx_int8" and not HAVE_I8: continue |
| lg=run(path,opt); d=(lg-lg_gen).abs(); top2=lg_gen.topk(2,-1).values; gap=top2[...,0]-top2[...,1] |
| mism=(lg.argmax(-1)!=lg_gen.argmax(-1)) |
| per=[] |
| for k in range(len(levels)): |
| a,b=offs[k],offs[k+1] |
| per.append({"level":k,"max_abs":float(d[:,a:b].max()),"mean_abs":float(d[:,a:b].mean()),"argmax_mismatch":int(mism[:,a:b].sum()),"of":int(mism[:,a:b].numel())}) |
| rep[name]={"max_abs":float(d.max()),"mean_abs":float(d.mean()),"logit_std":float(lg_gen.std()),"argmax_mismatch":int(mism.sum()),"of":int(mism.numel()), |
| "gap_at_mismatches_max":float(gap[mism].max()) if mism.any() else None,"gap_at_mismatches_median":float(gap[mism].median()) if mism.any() else None,"per_level":per} |
| print(name,"max|d|",rep[name]["max_abs"],"mean|d|",rep[name]["mean_abs"],"logit std",rep[name]["logit_std"],"argmax mism",rep[name]["argmax_mismatch"],"/",rep[name]["of"],"gap@mism max/med",rep[name]["gap_at_mismatches_max"],rep[name]["gap_at_mismatches_median"],flush=True) |
| for p_ in per: print(" ",p_,flush=True) |
| rep["torch_gen_vs_train_forward"]={"max_abs":float(d_tr.max()),"mean_abs":float(d_tr.mean())} |
| json.dump(rep,open("/kaggle/working/out/parity.json","w"),indent=1) |
| print("done",flush=True) |
|
|
