"""Bounded CPU pilot. One fixed model; calibration selects, heldout only validates.""" import argparse import csv import hashlib import json import os from pathlib import Path import resource import time import traceback import urllib.request import numpy as np from .model import StreamingLM, norm ATTN_PAIRS=[(3,7),(11,15),(19,23)] MLP_PAIRS=[(5,6),(20,21)] NORM_LAYERS=[3,7,11,15,19,23] REVISION='15852e8c16360a2fea060d615a32b45270f8a8fc' def write_json(path,obj): tmp=path.with_suffix('.tmp') tmp.write_text(json.dumps(obj,indent=2,allow_nan=False)+'\n') tmp.replace(path) def error_rows(a,b,samples,metadata): records=[] for x,y,sample in zip(a,b,samples): xf,yf=x.astype(np.float64).reshape(-1),y.astype(np.float64).reshape(-1) xn,yn=np.linalg.norm(xf),np.linalg.norm(yf) delta=np.linalg.norm(xf-yf) tokenerr=np.linalg.norm((x-y).reshape(x.shape[0],-1).astype(np.float64),axis=-1)/np.maximum(np.linalg.norm(x.reshape(x.shape[0],-1).astype(np.float64),axis=-1),1e-12) records.append({**metadata,'sample_id':sample['id'],'split':sample['split'], 'relative_l2':float(delta/max(xn,1e-12)), 'reverse_relative_l2':float(delta/max(yn,1e-12)), 'token_relative_p95':float(np.quantile(tokenerr,.95)), 'reverse_token_relative_p95':float(np.quantile(np.linalg.norm((x-y).reshape(x.shape[0],-1).astype(np.float64),axis=-1)/np.maximum(np.linalg.norm(y.reshape(x.shape[0],-1).astype(np.float64),axis=-1),1e-12),.95)), 'cosine':float(xf@yf/max(xn*yn,1e-30)), 'low_norm_tokens':int((np.linalg.norm(x.reshape(x.shape[0],-1),axis=-1)<1e-8).sum())}) return records def probability_metrics(a,b): m=(a+b)*.5 def term(p): return np.where(p>0,p*np.log(np.maximum(p,1e-30)/np.maximum(m,1e-30)),0) js=(term(a)+term(b)).sum(axis=-1)*.5 overlap=[] for row in range(a.shape[-2]): k=min(5,row+1) ai=np.argsort(a[:,:,row,:row+1],axis=-1)[...,-k:] bi=np.argsort(b[:,:,row,:row+1],axis=-1)[...,-k:] overlap.append(np.mean((ai[...,None]==bi[...,None,:]).any(axis=-1),axis=-1)) # Attention-sink control: remove key zero, renormalize valid remaining mass. aa,bb=a[...,1:],b[...,1:] mass_a,mass_b=aa.sum(-1),bb.sum(-1) valid=(mass_a>1e-12)&(mass_b>1e-12) aa=aa/np.maximum(mass_a[...,None],1e-30) bb=bb/np.maximum(mass_b[...,None],1e-30) mm=(aa+bb)*.5 conditional_js=.5*(np.where(aa>0,aa*np.log(np.maximum(aa,1e-30)/np.maximum(mm,1e-30)),0)+np.where(bb>0,bb*np.log(np.maximum(bb,1e-30)/np.maximum(mm,1e-30)),0)).sum(-1) return {'first_key_mass_a':float(a[...,0].mean()),'first_key_mass_b':float(b[...,0].mean()), 'js_without_first_key':float(conditional_js[valid].mean()) if valid.any() else None, 'without_first_key_valid_rows':int(valid.sum()),'js_mean':float(js.mean()),'js_p95':float(np.quantile(js,.95)), 'js_without_first_query':float(js[...,1:].mean()), 'top5_overlap':float(np.stack(overlap,axis=-1).mean())} def save_csv(path,rows): if not rows:return keys=list(dict.fromkeys(k for r in rows for k in r)) with path.open('w',newline='') as f: w=csv.DictWriter(f,fieldnames=keys);w.writeheader();w.writerows(rows) def tokenize(root): from tokenizers import Tokenizer path=root/'tokenizer.json' if not path.exists(): url=f'https://huggingface.co/Qwen/Qwen3.5-2B/resolve/{REVISION}/tokenizer.json' with urllib.request.urlopen(url,timeout=180) as r:path.write_bytes(r.read()) tokenizer=Tokenizer.from_file(str(path)) samples=json.loads(Path(__file__).with_name('samples.json').read_text())['samples'] ids=[] for s in samples: tokens=tokenizer.encode(s['text'],add_special_tokens=False).ids if len(tokens)<32:raise ValueError(f'Text too short: {s["id"]}') s['text_sha256']=hashlib.sha256(s['text'].encode()).hexdigest() s['full_token_count']=len(tokens);s['token_ids']=tokens[:32] ids.append(tokens[:32]) write_json(root/'samples-tokenized.json',{'revision':REVISION,'tokenizer_sha256':hashlib.sha256(path.read_bytes()).hexdigest(),'samples':samples}) return samples,np.array(ids,dtype=np.int64) def main(): p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);p.add_argument('--seconds',type=int,default=7200) args=p.parse_args();root=args.root;root.mkdir(exist_ok=True,parents=True) write_json(root/'retention-policy.json',{'weights':'retain_until_explicit_user_delete_request','automatic_cleanup':False,'instruction_date':'2026-11-09'}) samples,ids=tokenize(root) cal=np.array([i for i,s in enumerate(samples) if s['split']=='calibration']) held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout']) # Freeze all candidates/thresholds before opening weights or measuring functions. experiment={'revision':REVISION,'attention_pairs':ATTN_PAIRS,'mlp_pairs':MLP_PAIRS,'norm_layers':NORM_LAYERS, 'seconds':args.seconds,'output_relative_l2_limit':.05,'token_p95_limit':.1, 'attention_js_limit':.01,'attention_top5_min':.8,'heldout_nll_mean_delta_limit':.01, 'heldout_nll_sequence_p95_delta_limit':.05,'max_intervention_candidates':1, 'directions_per_selected_candidate':2,'training':False,'search_R':False,'automatic_weight_cleanup':False, 'dtype':'original weights decoded to FP32; not a BF16 backend benchmark', 'blas_threads':os.environ.get('OPENBLAS_NUM_THREADS')} write_json(root/'experiment.json',experiment) started=time.time();status={'state':'running','started_unix':started,'phase':'baseline'} write_json(root/'functional-status.json',status) rows=[];headrows=[];interventions=[] try: config=json.loads((root/'model_config.json').read_text()) lm=StreamingLM(root/'weights',config,seconds=args.seconds) overall_deadline=lm.deadline lm.deadline=min(overall_deadline,time.monotonic()+2400) lm.capture_qk=True capture=[('attn',i) for pair in ATTN_PAIRS for i in pair]+[('mlp',i) for pair in MLP_PAIRS for i in pair] def progress(layer): status.update(last_completed_layer=layer,elapsed_seconds=time.time()-started) write_json(root/'functional-status.json',status) print(json.dumps({'phase':status['phase'],'layer':layer,'elapsed':time.time()-started}),flush=True) h,saved=lm.forward(ids,capture=capture,progress=progress) np.savez(root/'baseline-activations.npz',**saved,**{f'qraw_{k}':v[0] for k,v in lm.qk_inputs.items()},**{f'kraw_{k}':v[1] for k,v in lm.qk_inputs.items()},final=h) # Baseline heldout NLL uses no selection information. nll=lm.nll(h[held],ids[held]);np.save(root/'baseline-heldout-token-nll.npy',nll) baseline_nll=nll.mean(axis=1) lm.capture_qk=False lm.deadline=min(overall_deadline,time.monotonic()+1800) status['phase']='calibration';write_json(root/'functional-status.json',status) candidates=[] for kind,pairs in [('attn',ATTN_PAIRS),('mlp',MLP_PAIRS)]: for i,j in pairs: pairid=f'{kind}_{i}_{j}';pairrows=[] for source in (i,j): x=saved[f'{kind}_{source}'][cal] if kind=='attn': a,pa=lm.full_attention(x,i);b,pb=lm.full_attention(x,j) pm={} for bi,si in enumerate(cal): for hi in range(pa.shape[1]): hpm=probability_metrics(pa[bi:bi+1,hi:hi+1],pb[bi:bi+1,hi:hi+1]) headrows.append({'candidate':pairid,'source':source,'sample_id':samples[si]['id'],'head':hi,**hpm}) save_csv(root/'attention-head-metrics.csv',headrows) else: a,b=lm.mlp(x,i),lm.mlp(x,j);pm={} rec=error_rows(a,b,[samples[t] for t in cal],{'candidate':pairid,'kind':kind,'source':source,**pm}) if kind=='attn': for bi,r in enumerate(rec):r.update(probability_metrics(pa[bi:bi+1],pb[bi:bi+1])) rows.extend(rec);pairrows.extend(rec) save_csv(root/'calibration-metrics.csv',rows) passing=all(r['relative_l2']<=.05 and r['reverse_relative_l2']<=.05 and r['token_relative_p95']<=.1 and r['reverse_token_relative_p95']<=.1 for r in pairrows) if kind=='attn':passing=passing and all(r['js_mean']<=.01 and r['top5_overlap']>=.8 for r in pairrows) prefix=f'model.language_model.layers.{i}.'+('self_attn.' if kind=='attn' else 'mlp.') size=sum(t.bytes for name,t in lm.tensors.items() if name.startswith(prefix)) candidates.append({'id':pairid,'kind':kind,'i':i,'j':j,'passed':passing,'maximum_symmetric_relative_error':max(max(r['relative_l2'],r['reverse_relative_l2']) for r in pairrows),'donor_component_bytes':size}) for layer in NORM_LAYERS: pairid=f'qk_norm_{layer}';pfx=f'layers.{layer}.self_attn.' qw,kw=lm.small(pfx+'q_norm.weight'),lm.small(pfx+'k_norm.weight') pairrows=[] for source,x in zip(('q','k'),lm.qk_inputs[layer]): a,b=norm(x[cal],qw,lm.eps),norm(x[cal],kw,lm.eps) rec=error_rows(a,b,[samples[t] for t in cal],{'candidate':pairid,'kind':'qk_norm','source':source}) pairrows.extend(rec);rows.extend(rec) passing=all(r['relative_l2']<=.05 and r['reverse_relative_l2']<=.05 and r['token_relative_p95']<=.1 and r['reverse_token_relative_p95']<=.1 for r in pairrows) candidates.append({'id':pairid,'kind':'qk_norm','i':layer,'j':layer,'passed':passing,'maximum_symmetric_relative_error':max(max(r['relative_l2'],r['reverse_relative_l2']) for r in pairrows),'donor_component_bytes':lm.tensors['model.language_model.'+pfx+'q_norm.weight'].bytes}) save_csv(root/'calibration-metrics.csv',rows) write_json(root/'candidate-summary.json',candidates) passed=[c for c in candidates if c['passed']] # Prefer larger potential savings, then smaller worst calibration error; never consult heldout. selected=sorted(passed,key=lambda c:(-c['donor_component_bytes'],c['maximum_symmetric_relative_error'],c['id']))[:1] write_json(root/'selection.json',{'selected':selected,'rule':'calibration pass; descending component bytes, ascending maximum symmetric error; maximum one candidate'}) if selected: c=selected[0] heldrows=[] if c['kind']=='qk_norm': pfx=f'layers.{c["i"]}.self_attn.' qw,kw=lm.small(pfx+'q_norm.weight'),lm.small(pfx+'k_norm.weight') for source,x in zip(('q','k'),lm.qk_inputs[c['i']]): heldrows.extend(error_rows(norm(x[held],qw,lm.eps),norm(x[held],kw,lm.eps),[samples[t] for t in held],{'candidate':c['id'],'kind':'qk_norm','source':source})) else: for source in (c['i'],c['j']): x=saved[f'{c["kind"]}_{source}'][held] if c['kind']=='attn': a,_=lm.full_attention(x,c['i']);b,_=lm.full_attention(x,c['j']) else:a,b=lm.mlp(x,c['i']),lm.mlp(x,c['j']) heldrows.extend(error_rows(a,b,[samples[t] for t in held],{'candidate':c['id'],'kind':c['kind'],'source':source})) save_csv(root/'heldout-function-metrics.csv',heldrows) held_function_pass=all(r['relative_l2']<=.05 and r['reverse_relative_l2']<=.05 and r['token_relative_p95']<=.1 and r['reverse_token_relative_p95']<=.1 for r in heldrows) status['heldout_function_pass']=held_function_pass lm.deadline=min(overall_deadline,time.monotonic()+1800) status['phase']='intervention';write_json(root/'functional-status.json',status) for reverse in (False,True): replacement=None;lm.norm_replacement=None if c['kind']=='qk_norm': pfx=f'layers.{c["i"]}.self_attn.' target,donor=('q_norm','k_norm') if not reverse else ('k_norm','q_norm') lm.norm_replacement=(pfx+target,pfx+donor) label=target+'_uses_'+donor else: target,donor=(c['i'],c['j']) if not reverse else (c['j'],c['i']) replacement=(c['kind'],target,donor);label=f'{target}_uses_{donor}' t=time.monotonic();modified,_=lm.forward(ids[held],replacement=replacement,progress=progress) changed=lm.nll(modified,ids[held]);np.save(root/f'intervention-{label}-token-nll.npy',changed) delta=changed.mean(axis=1)-baseline_nll rec={'candidate':c['id'],'direction':label,'elapsed_seconds':time.monotonic()-t, 'nll_mean_delta':float(delta.mean()),'nll_sequence_p95_delta':float(np.quantile(delta,.95)), 'sequence_nll_deltas':delta.tolist(),'passed_quality':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)} interventions.append(rec);write_json(root/'interventions.json',interventions) # No fallback candidate if heldout fails. status.update(state='completed',phase='done',selected_candidate=selected, positive_bidirectional_quality=bool(interventions and status.get('heldout_function_pass') and all(x['passed_quality'] for x in interventions)), elapsed_seconds=time.time()-started,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, claim_limit='Small text-only heldout pilot; no global interchangeability proof, no measured GPU speedup; component storage savings require explicit layout accounting.') except Exception as exc: status.update(state='stopped',error_type=type(exc).__name__,error=str(exc),elapsed_seconds=time.time()-started, peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) (root/'functional-error.txt').write_text(traceback.format_exc()) raise finally: write_json(root/'functional-status.json',status) if __name__=='__main__':main()