280 lines
18 KiB
Python
280 lines
18 KiB
Python
"""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()
|