| import os, subprocess, sys, json, random, time, math |
|
|
| def sh(c): print("$",c,flush=True); subprocess.run(c,shell=True,check=True) |
| os.makedirs("/kaggle/working/w",exist_ok=True); os.chdir("/kaggle/working/w") |
| if not os.path.exists("ncm"): sh("git clone -q https://github.com/HazyResearch/ncm.git") |
| sh("pip install -q omegaconf hydra-core regex onnx onnxruntime onnxscript 2>&1 | tail -1") |
| os.makedirs("ck",exist_ok=True) |
| for f in ["config.yaml","generator.pt","tokenizer.pt"]: |
| if not os.path.exists("ck/"+f): sh(f"curl -sL -o ck/{f} https://huggingface.co/hazyresearch/ncm-checkpoints/resolve/main/xsum/{f}") |
| open('ncmchat.py','w').write('"""NCM chat fine-tune: reuse released XSum NCM (frozen multiscale VQ tokenizer + prefix-conditioned generator),\nfine-tune the generator on (conversation prefix -> short reply). Needs HazyResearch/ncm on NCM_ROOT."""\nimport os, sys, random, torch, torch.nn.functional as F\nfrom omegaconf import OmegaConf\nNCM_ROOT = os.environ.get("NCM_ROOT", "/tmp/ncm"); sys.path.insert(0, NCM_ROOT)\nfrom models.next_concept import NextConceptTransformer\nfrom tokenization.vqvae_tokenizer import VQVAETokenizer\n\ndef load(ck, cfg_path=None, dev="cpu", load_gen=True):\n CFG = OmegaConf.load(cfg_path or os.path.join(ck, "config.yaml"))\n CFG.tokenizer.vqvae.checkpoint_path = os.path.join(ck, CFG.checkpoints.tokenizer)\n tk = VQVAETokenizer.from_config(CFG.tokenizer)\n tk.vqvae = tk.vqvae.to(dev).eval()\n for p in tk.vqvae.parameters(): p.requires_grad = False\n net = NextConceptTransformer(vqvae_local=tk.vqvae, config=CFG.model).to(dev)\n if load_gen:\n sd = torch.load(os.path.join(ck, CFG.checkpoints.generator), map_location="cpu", weights_only=False)\n net.load_state_dict(sd["model"], strict=True)\n return CFG, tk, net\n\ndef fmt_prompt(turns):\n """turns: list of (role, text) ending with a user turn -> prompt string"""\n s = "".join(f"{\'User\' if r==\'user\' else \'Assistant\'}: {t}\\n" for r, t in turns)\n return s + "Assistant:"\n\nclass Codec:\n def __init__(self, CFG, tk, net):\n self.tk, self.net, self.P = tk, net, int(CFG.model.prefix_len)\n self.pre = tk.pre_tokenizer\n self.EOS, self.PAD = self.pre.eos_token_id, self.pre.pad_token_id\n self.ppad = int(CFG.model.prefix_pad_token_id) # pad id used in the generator prefix\n self.T = tk.vqvae.context_length # target length (48)\n self.hf = self.pre.tokenizer if hasattr(self.pre, "tokenizer") else None\n def enc(self, text): return self.pre.encode(text)\n def prefix_ids(self, prompt):\n ids = self.enc(prompt)[-self.P:] # keep most recent context\n return ids + [self.ppad] * (self.P - len(ids))\n def target_ids(self, reply):\n ids = self.enc(" " + reply.strip())[: self.T - 1] + [self.EOS]\n return ids + [self.EOS] * (self.T - len(ids))\n\n@torch.no_grad()\ndef target_codes(codec, tgt): # tgt: LongTensor (B,T) -> idx (B, sum levels)\n return codec.tk.vqvae.encode(tgt)\n\ndef build_inputs(codec, prefix, idx, corrupt=0.1):\n net, vq = codec.net, codec.tk.vqvae.vq\n with torch.no_grad():\n idc = idx.clone()\n if corrupt > 0: # robustness: corrupt conditioning codes (as in paper)\n m = torch.rand_like(idc, dtype=torch.float) < corrupt\n idc[m] = torch.randint(0, 4096, (int(m.sum()),), device=idc.device)\n nc = vq.idx_to_ncm_input(idc) # cumulative recon fed to levels 1..K-1\n zp = codec.tk.vqvae.encoder.wte(prefix)[:, :codec.P, :]\n q0 = net.build_l0_query_input(prefix.shape[0])\n return torch.cat([zp, q0, nc], dim=1)\n\ndef loss_fn(codec, prefix, tgt, corrupt=0.1):\n idx = target_codes(codec, tgt)\n x = build_inputs(codec, prefix, idx, corrupt)\n pad_mask = prefix == codec.ppad\n logits = codec.net(x, prefix_pad_mask=None if not hasattr(codec.net, "_prefix_pad_bias") else pad_mask)\n return F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), idx.reshape(-1)), logits, idx\n\n@torch.no_grad()\ndef reply(codec, prompts, temperature=0.8, top_k=20, top_p=0.9, dev="cpu"):\n net, tk = codec.net, codec.tk\n pre = torch.tensor([codec.prefix_ids(p) for p in prompts], device=dev)\n g = net.generate(None, pre, temperature=temperature, top_k=top_k, top_p=top_p)\n _, zl = tk.vqvae.vq.idx_to_zhats(g, return_zlevels=True)\n ids = tk.vqvae.decoder(zl[-1]).argmax(-1).tolist()\n out = []\n for r in ids:\n if codec.EOS in r: r = r[: r.index(codec.EOS)]\n out.append(codec.pre.decode([t for t in r if t not in (codec.EOS, codec.PAD)], skip_special_tokens=True).strip())\n return out\n') |
| sys.path.insert(0,os.getcwd()) |
| |
| 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) |
|
|
