228 lines
14 KiB
Python
228 lines
14 KiB
Python
"""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()
|