code/ncm_export_kernel.py

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())
# Runs inside a Kaggle kernel (CPU ok). Inputs: /kaggle/input/<SRC>/generator_chat.pt ; outputs ncm_fp16w.onnx + report.
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)
# remove Identity nodes fed by initializers (they hid the weights from the int8 quantizer)
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"   # weights stored fp16 + Cast->fp32 at load (compute stays fp32)
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]
# tokenizer parity vectors for the browser
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)