"""Bounded 15-pair CPU study, calibration selection before any heldout feedback.""" import argparse from collections import defaultdict import hashlib import json from pathlib import Path import resource import time import traceback import numpy as np from functional.run import save_csv,write_json from .core import (QKVModel,InterventionLM,assignment,inverse,output_metrics, cross_head_output_metrics,distribution_metrics,value_gram_cka) from .protocol import LAYERS,PAIRS,PROTOCOL,TIERS PRIMARY=slice(7,None) TYPE_ORDER={'full_attention':0,'qkv_head':1,'qk_map':2,'qk_head':3,'v_group':4} def candidate_id(kind,i,j,a=None,b=None,suffix=''): return f'{kind}_{i}_{j}'+(f'_{a}_{b}' if a is not None else '')+suffix def rows_for(metrics,samples,meta): rows=[] for n,s in enumerate(samples): r={**meta,'sample':s['id']} for key,value in metrics.items(): v=value[n] r[key]=None if v is None else float(v) rows.append(r) return rows def outrows(a,b,samples,meta,floor_a=1e-6,floor_b=1e-6): m=output_metrics(a[:,PRIMARY],b[:,PRIMARY]);rows=rows_for(m,samples,meta) for n,r in enumerate(rows):r['active']=bool(m['rms_a'][n]>=np.broadcast_to(floor_a,m['rms_a'].shape)[n] and m['rms_b'][n]>=np.broadcast_to(floor_b,m['rms_b'].shape)[n]) return rows def head_floor(profile): rms=np.sqrt(np.mean(profile['heads'][:,7:].astype(np.float64)**2,axis=(1,3))) return np.maximum(1e-6,.1*np.median(rms,axis=1)) def candidate_pass(c,tier): t=TIERS[tier];out=c.get('output_rows',[]) if not out or not all(r['active'] and max(r['error_a'],r['error_b'])<=t['output_error'] and r['cosine']>=t['output_cosine'] for r in out):return False if c['kind'] in ('qk_head','qk_map'): p=c.get('probability_rows',[]) if not p or not all(r['js']<=t['js'] and r['js_p95']<=t['js_p95'] and r['top5']>=t['top5'] for r in p):return False if c['kind']=='v_group': v=c.get('value_rows',[]) if not v or not all(max(r['error_a'],r['error_b'])<=t['value_error'] and r['cosine']>=t['value_cosine'] and r['active'] for r in v):return False return True def preliminary_probability(c): t=TIERS['approximate'];r=c['probability_rows'] return bool(r) and all(x['js']<=t['js'] and x['js_p95']<=t['js_p95'] and x['top5']>=t['top5'] for x in r) def preliminary_value(c): t=TIERS['approximate'];r=c['value_rows'] return bool(r) and all(x['active'] and max(x['error_a'],x['error_b'])<=t['value_error'] and x['cosine']>=t['value_cosine'] for x in r) def load_profile(root,source,layer): with np.load(root/'cache'/f'calibration-source-{source}.npz') as data: prefix=f'layer{layer}_' return {key[len(prefix):]:data[key] for key in data.files if key.startswith(prefix)} def map_profile_values(p):return {k:v for k,v in p.items()} def new_bank():return {} def get_candidate(bank,kind,i,j,a=None,b=None,suffix='',perm=None): name=candidate_id(kind,i,j,a,b,suffix) if name not in bank: bank[name]={'id':name,'kind':kind,'i':i,'j':j,'a':a,'b':b,'perm':perm, 'output_rows':[],'probability_rows':[],'value_rows':[]} return bank[name] def add_output(c,a,b,samples,source,target=None,floor_a=1e-6,floor_b=1e-6): c['output_rows'].extend(outrows(a,b,samples,{'candidate':c['id'],'source_input_layer':source,'target':target},floor_a,floor_b)) def observe_pair(lm,bank,pi,pj,i,j,source,samples,diagnostics): fi,fj=head_floor(pi),head_floor(pj) whole=get_candidate(bank,'full_attention',i,j) add_output(whole,pi['output'],pj['output'],samples,source) hm=cross_head_output_metrics(pi['heads'][:,7:],pj['heads'][:,7:]) qm=cross_head_output_metrics(pi['q'][:,7:],pj['q'][:,7:]) repeats=lm.config['num_attention_heads']//lm.config['num_key_value_heads'] ki,kj=np.repeat(pi['k'],repeats,axis=2),np.repeat(pj['k'],repeats,axis=2) km=cross_head_output_metrics(ki[:,7:],kj[:,7:]) for a in range(pi['p'].shape[1]): for b in range(pj['p'].shape[1]): qk=get_candidate(bank,'qk_head',i,j,a,b);qkv=get_candidate(bank,'qkv_head',i,j,a,b) pm=distribution_metrics(pi['p'][:,a:a+1],pj['p'][:,b:b+1]) rr=rows_for(pm,samples,{'candidate':qk['id'],'source_input_layer':source}) qk['probability_rows'].extend(rr);qkv['probability_rows'].extend(rr) for n,s in enumerate(samples): rec={'candidate':qkv['id'],'source_input_layer':source,'sample':s['id'],**{key:float(v[n,a,b]) for key,v in hm.items()}} rec['active']=bool(rec['rms_a']>=fi[n] and rec['rms_b']>=fj[n]);qkv['output_rows'].append(rec) diagnostics.append({'i':i,'j':j,'head_i':a,'head_j':b,'source_input_layer':source,'sample':s['id'], 'q_error':float(max(qm['error_a'][n,a,b],qm['error_b'][n,a,b])), 'q_cosine':float(qm['cosine'][n,a,b]),'k_error':float(max(km['error_a'][n,a,b],km['error_b'][n,a,b])), 'k_cosine':float(km['cosine'][n,a,b]),'qkv_head_output_error':max(rec['error_a'],rec['error_b'])}) vm=cross_head_output_metrics(pi['v'][:,7:],pj['v'][:,7:]) vfi=np.maximum(1e-6,.1*np.median(np.sqrt(np.mean(pi['v'][:,7:].astype(np.float64)**2,axis=(1,3))),axis=1)) vfj=np.maximum(1e-6,.1*np.median(np.sqrt(np.mean(pj['v'][:,7:].astype(np.float64)**2,axis=(1,3))),axis=1)) for a in range(pi['v'].shape[2]): for b in range(pj['v'].shape[2]): c=get_candidate(bank,'v_group',i,j,a,b) gram=value_gram_cka(pi['v'][:,7:,a],pj['v'][:,7:,b]) for n,s in enumerate(samples): rec={'candidate':c['id'],'source_input_layer':source,'sample':s['id'],**{key:float(v[n,a,b]) for key,v in vm.items()},'centered_value_gram_cka':float(gram[n])} rec['active']=bool(rec['rms_a']>=vfi[n] and rec['rms_b']>=vfj[n]);c['value_rows'].append(rec) def probe_candidate(lm,c,pi,pj,source,samples): i,j=c['i'],c['j'];kind=c['kind'];fi,fj=head_floor(pi),head_floor(pj) if kind=='qk_map': perm=c['perm'];p=distribution_metrics(pi['p'],pj['p'][:,perm]) c['probability_rows'].extend(rows_for(p,samples,{'candidate':c['id'],'source_input_layer':source})) changed_i=lm.transplant(pi,pj,i,kind,perm=perm) changed_j=lm.transplant(pj,pi,j,kind,perm=inverse(perm)) add_output(c,pi['output'],changed_i,samples,source,i) add_output(c,pj['output'],changed_j,samples,source,j) elif kind=='qk_head': a,b=c['a'],c['b'] changed_i=lm.head_transplant_output(pi,pj,i,a,b) changed_j=lm.head_transplant_output(pj,pi,j,b,a) add_output(c,pi['heads'][:,:,a],changed_i,samples,source,i,fi,fi) add_output(c,pj['heads'][:,:,b],changed_j,samples,source,j,fj,fj) elif kind=='v_group': a,b=c['a'],c['b'];rep=lm.config['num_attention_heads']//lm.config['num_key_value_heads'] own_i=pi['heads'][:,:,a*rep:(a+1)*rep].sum(2);own_j=pj['heads'][:,:,b*rep:(b+1)*rep].sum(2) changed_i=lm.value_group_output(pi,pj,i,a,b);changed_j=lm.value_group_output(pj,pi,j,b,a) add_output(c,own_i,changed_i,samples,source,i,fi,fi);add_output(c,own_j,changed_j,samples,source,j,fj,fj) else:raise ValueError(kind) def summarize(c): out=c['output_rows'];p=c['probability_rows'];v=c['value_rows'] m={'id':c['id'],'kind':c['kind'],'i':c['i'],'j':c['j'],'a':c['a'],'b':c['b'],'perm':c['perm'], 'output_verified':bool(out),'output_observations':len(out),'probability_observations':len(p),'value_observations':len(v), 'strict_pass':candidate_pass(c,'strict'),'approximate_pass':candidate_pass(c,'approximate')} m['worst_output_error']=max((max(r['error_a'],r['error_b']) for r in out),default=None) m['minimum_output_cosine']=min((r['cosine'] for r in out),default=None) m['worst_mean_js']=max((r['js'] for r in p),default=None) m['minimum_top5']=min((r['top5'] for r in p),default=None) m['worst_value_error']=max((max(r['error_a'],r['error_b']) for r in v),default=None) return m def controls_for_profiles(profiles,source,samples,rows,h): rng=np.random.default_rng(20261109+source) ra=(h@rng.standard_normal((h.shape[-1],256),dtype=np.float32))/np.float32(np.sqrt(h.shape[-1])) rb=(h@rng.standard_normal((h.shape[-1],256),dtype=np.float32))/np.float32(np.sqrt(h.shape[-1])) random_cka=value_gram_cka(ra[:,7:],rb[:,7:]) for n,s in enumerate(samples):rows.append({'control':'value_gram_random_projection_pair','source_input_layer':source,'sample':s['id'],'gram_cka':float(random_cka[n])}) for layer,p in profiles.items(): for kv in range(p['v'].shape[2]): for label,other in [('value_gram_to_input',h),('value_gram_to_random_projection',ra)]: vcka=value_gram_cka(p['v'][:,7:,kv],other[:,7:]) for n,s in enumerate(samples):rows.append({'control':label,'source_input_layer':source,'projection_layer':layer,'kv_head':kv,'sample':s['id'],'gram_cka':float(vcka[n])}) probs=p['p'];uniform=np.zeros_like(probs);shuffled=probs.copy() for t in range(probs.shape[-1]): uniform[:,:,t,:t+1]=1/(t+1) for bi in range(probs.shape[0]): for head in range(probs.shape[1]):shuffled[bi,head,t,:t+1]=probs[bi,head,t,rng.permutation(t+1)] for name,other in [('identity',probs),('uniform_causal',uniform),('shuffled_legal_keys',shuffled)]: for head in range(probs.shape[1]): m=distribution_metrics(probs[:,head:head+1],other[:,head:head+1]) rows.extend(rows_for(m,samples,{'control':name,'source_input_layer':source,'projection_layer':layer,'head':head})) def validate_selected(lm,selected,inputs,samples,indices): i,j=selected['i'],selected['j'];c={k:v for k,v in selected.items() if k not in ('output_rows','probability_rows','value_rows')} c.update(output_rows=[],probability_rows=[],value_rows=[]) for source in (i,j): h=inputs[f'attn_{source}'][indices];pi,pj=lm.parts(h,i),lm.parts(h,j) temporary={};diag=[];observe_pair(lm,temporary,pi,pj,i,j,source,samples,diag) if c['kind'] in ('full_attention','qkv_head','qk_head','v_group'): original=temporary[candidate_id(c['kind'],i,j,c['a'],c['b'])] c['probability_rows'].extend(original['probability_rows']);c['value_rows'].extend(original['value_rows']) if c['kind'] in ('full_attention','qkv_head'):c['output_rows'].extend(original['output_rows']) if c['kind'] in ('qk_map','qk_head','v_group'):probe_candidate(lm,c,pi,pj,source,samples) return c def main(): p=argparse.ArgumentParser();p.add_argument('--source',required=True,type=Path);p.add_argument('--root',required=True,type=Path);p.add_argument('--seconds',type=int,default=7200) args=p.parse_args();source=args.source.resolve();root=args.root.resolve();root.mkdir(parents=True,exist_ok=True);(root/'cache').mkdir(exist_ok=True) assert not (root/'protocol.json').exists(),'Fresh experiment directory required; do not overwrite evidence' protocol={**PROTOCOL,'registered_unix':time.time(),'source_run':str(source),'seconds':args.seconds, 'fresh_sample_sha256':hashlib.sha256(Path(__file__).with_name('fresh_samples.json').read_bytes()).hexdigest()} write_json(root/'protocol.json',protocol) write_json(root/'retention-policy.json',{'automatic_cleanup':False,'weights_and_caches':'retain_until_explicit_user_deletion'}) start=time.monotonic();status={'state':'running','phase':'calibration_profiles','started_unix':time.time()} bank={};diagnostics=[];controls=[] try: manifest=json.loads((source/'manifest.json').read_text());assert manifest['revision']=='15852e8c16360a2fea060d615a32b45270f8a8fc' reference=json.loads((source/'reference-check.json').read_text());assert reference['passed'] assert reference['implementation_sha256']==hashlib.sha256(Path('functional/model.py').read_bytes()).hexdigest() config=json.loads((source/'model_config.json').read_text());lm=InterventionLM(source/'weights',config,seconds=args.seconds-180) assert [i for i,t in enumerate(lm.config['layer_types']) if t=='full_attention']==LAYERS samples=json.loads((source/'samples-tokenized.json').read_text())['samples'];cal=[s for s in samples if s['split']=='calibration'];valid=[s for s in samples if s['split']=='heldout'] ci=np.array([i for i,s in enumerate(samples) if s['split']=='calibration']);vi=np.array([i for i,s in enumerate(samples) if s['split']=='heldout']) write_json(root/'source-samples.json',{'samples':samples,'note':'old validation samples used in earlier studies; fresh test separately fixed'}) inputs=np.load(source/'baseline-activations.npz') write_json(root/'source-provenance.json',{'revision':manifest['revision'],'weights_manifest':str(source/'manifest.json'), 'baseline_activation_sha256':hashlib.sha256((source/'baseline-activations.npz').read_bytes()).hexdigest(), 'source_hashes':{str(f):hashlib.sha256(f.read_bytes()).hexdigest() for f in sorted(Path('crossqkv').glob('*')) if f.is_file()}}) for src in LAYERS: lm.guard();h=inputs[f'attn_{src}'][ci];profiles={layer:lm.parts(h,layer) for layer in LAYERS} np.savez(root/'cache'/f'calibration-source-{src}.npz',**{f'layer{layer}_{key}':value for layer,profile in profiles.items() for key,value in profile.items()}) controls_for_profiles(profiles,src,cal,controls,h) for i,j in PAIRS: if src in (i,j):observe_pair(lm,bank,profiles[i],profiles[j],i,j,src,cal,diagnostics) del profiles status.update(last_source_layer=src,elapsed_seconds=time.monotonic()-start);write_json(root/'status.json',status) save_csv(root/'raw-qk-diagnostics.csv',diagnostics);save_csv(root/'controls.csv',controls) print(json.dumps(status),flush=True) status['phase']='discrete_head_matching_and_output_probes';write_json(root/'status.json',status) mapping_rows=[] qk_pre=[c for c in bank.values() if c['kind']=='qk_head' and preliminary_probability(c)] qk_probe=sorted(qk_pre,key=lambda c:(max(r['js'] for r in c['probability_rows']),c['id']))[:32] v_pre=[c for c in bank.values() if c['kind']=='v_group' and preliminary_value(c)] v_probe=sorted(v_pre,key=lambda c:(max(max(r['error_a'],r['error_b']) for r in c['value_rows']),c['id']))[:16] for i,j in PAIRS: cost=np.empty((8,8),dtype=np.float64) for a in range(8): for b in range(8):cost[a,b]=np.mean([r['js'] for r in bank[candidate_id('qk_head',i,j,a,b)]['probability_rows']]) perm=assignment(cost) for a in range(8): for b in range(8):mapping_rows.append({'i':i,'j':j,'head_i':a,'head_j':b,'calibration_mean_js':float(cost[a,b]),'chosen':perm[a]==b}) maps=[get_candidate(bank,'qk_map',i,j,suffix='_identity',perm=list(range(8))),get_candidate(bank,'qk_map',i,j,suffix='_matched',perm=perm)] for src in (i,j): pi,pj=load_profile(root,src,i),load_profile(root,src,j) for c in [*maps,*[c for c in qk_probe+v_probe if c['i']==i and c['j']==j]]: lm.guard();probe_candidate(lm,c,pi,pj,src,cal) status.update(last_pair=[i,j],elapsed_seconds=time.monotonic()-start);write_json(root/'status.json',status) save_csv(root/'head-assignment-costs.csv',mapping_rows) summaries=[summarize(c) for c in bank.values()];write_json(root/'calibration-summary.json',summaries) save_csv(root/'calibration-output-metrics.csv',[r for c in bank.values() for r in c['output_rows']]) save_csv(root/'calibration-probability-metrics.csv',[r for c in bank.values() if c['kind']!='qkv_head' for r in c['probability_rows']]) save_csv(root/'calibration-value-metrics.csv',[r for c in bank.values() for r in c['value_rows']]) eligible=[c for c in summaries if c['strict_pass'] or c['approximate_pass']] eligible.sort(key=lambda c:(not c['strict_pass'],TYPE_ORDER[c['kind']],c['worst_output_error'],c['worst_mean_js'] or 0,c['id'])) selected=eligible[:1] write_json(root/'selection.json',{'selected':selected,'qk_head_prequalified':len(qk_pre),'qk_head_output_verified':len(qk_probe), 'v_group_prequalified':len(v_pre),'v_group_output_verified':len(v_probe),'selection_uses_validation':False}) status['selected']=selected if selected: c=selected[0];status['phase']='old_validation';write_json(root/'status.json',status) checked=validate_selected(lm,c,inputs,valid,vi);tier='strict' if c['strict_pass'] else 'approximate' write_json(root/'validation-summary.json',summarize(checked));save_csv(root/'validation-output-metrics.csv',checked['output_rows']);save_csv(root/'validation-probability-metrics.csv',checked['probability_rows']);save_csv(root/'validation-value-metrics.csv',checked['value_rows']) passed=candidate_pass(checked,tier);status['validation_pass']=passed if passed: from .fresh import run_fresh run_fresh(source,root,c,tier,lm,status) status.update(state='completed',phase='done',elapsed_seconds=time.monotonic()-start,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, candidate_counts={kind:sum(c['kind']==kind for c in summaries) for kind in TYPE_ORDER}, passing_counts={kind:sum(c['kind']==kind and c['approximate_pass'] for c in summaries) for kind in TYPE_ORDER}, original_11_and_norm_positive_not_used_as_qkv_evidence=True) except Exception as exc: status.update(state='stopped',error_type=type(exc).__name__,error=str(exc),elapsed_seconds=time.monotonic()-start,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss) (root/'error.txt').write_text(traceback.format_exc());raise finally:write_json(root/'status.json',status) if __name__=='__main__':main()