91 lines
8.6 KiB
Python
91 lines
8.6 KiB
Python
"""Recompute saved concentration, masks, scale and scope checks; preserve caches."""
|
|
import json,hashlib,csv,sys
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from .core import concentration
|
|
|
|
|
|
def main(root,source):
|
|
root=Path(root);source=Path(source);status=json.loads((root/'status.json').read_text());assert status['state']=='completed'
|
|
samples=json.loads((root/'samples.json').read_text());assert len(samples)==4 and all(len(s['token_ids'])==32 for s in samples)
|
|
expected={(s['id'],i) for s in samples for i in range(23)};actual={(x['sample'],x['source']) for x in status['completed_links']};assert actual==expected and len(status['completed_links'])==92
|
|
cfg=json.loads((source/'model_config.json').read_text())['text_config'];obs=dict(np.load(root/'observations.npz'));lookup=sorted(s['id'] for s in samples);counts={};undefined={};full_js_links=0;joint_values={}
|
|
for sample,layer in sorted(actual):
|
|
f=root/'raw'/f'{sample}-link-{layer:02d}.npz';d=dict(np.load(f));hs=8 if cfg['layer_types'][layer]=='full_attention' else 16;hr=8 if cfg['layer_types'][layer+1]=='full_attention' else 16
|
|
if hr==8:
|
|
full_js_links+=1;d['single_sqrtJS']=np.sqrt(d['single_JS'])
|
|
for label in ('weak_half','strong_half','random_half','all_heads'):
|
|
d['joint_'+label+'_sqrtJS']=np.sqrt(d['joint_'+label+'_JS'])
|
|
mask=d.get('mask_'+label,np.ones((hr,32,hs),dtype=np.float32))
|
|
d['individual_sum_'+label+'_sqrtJS']=np.sum(d['single_sqrtJS'].transpose(2,1,0)*mask,-1).T
|
|
d['null_sqrtJS']=np.sqrt(d['null_JS'])
|
|
masks={label:d['mask_'+label] for label in ('weak_half','strong_half','random_half')}
|
|
for label,m in masks.items():
|
|
assert m.shape==(hr,32,hs);assert np.isin(m,[0,1]).all();np.testing.assert_array_equal(m.sum(-1),hs//2)
|
|
np.testing.assert_array_equal(masks['weak_half']+masks['strong_half'],1)
|
|
ranked=np.argsort(d['single_QKV'].transpose(2,1,0),axis=-1,kind='stable')
|
|
np.testing.assert_array_equal(np.take_along_axis(masks['weak_half'],ranked[...,:hs//2],-1),1)
|
|
for key in ('Q','K','V','QKV')+(('JS','sqrtJS') if hr==8 else ()):
|
|
single=d['single_'+key];assert single.shape==(hs,32,hr) and np.isfinite(single).all() and np.min(single)>=0
|
|
effect=single[:,7:,:].transpose(1,2,0);stats=concentration(effect);counts[key]=counts.get(key,0)+effect.shape[0]*hr;undefined[key]=undefined.get(key,0)+int((~stats['valid']).sum())
|
|
choose=(obs[key+'_source']==layer)&(obs[key+'_sample_index']==lookup.index(sample))
|
|
assert int(choose.sum())==25*hr
|
|
for metric,values in stats.items():
|
|
if metric=='valid':continue
|
|
np.testing.assert_allclose(obs[key+'_'+metric][choose],values.ravel(),rtol=1e-6,atol=1e-7,equal_nan=True)
|
|
for label in ('weak_half','strong_half','random_half','all_heads'):
|
|
mask=masks.get(label,np.ones((hr,32,hs),dtype=np.float32));expected_sum=np.sum(single.transpose(2,1,0)*mask,-1).T
|
|
np.testing.assert_allclose(d['individual_sum_'+label+'_'+key],expected_sum,rtol=1e-6,atol=1e-8)
|
|
assert d['joint_'+label+'_'+key].shape==(32,hr) and np.isfinite(d['joint_'+label+'_'+key]).all()
|
|
joint_values.setdefault((key,label),[]).extend(d['joint_'+label+'_'+key][7:].ravel())
|
|
assert float(np.max(d['null_'+key]))<=2e-5
|
|
assert full_js_links==24
|
|
summary=json.loads((root/'analysis-summary.json').read_text())
|
|
def check_stats(values,record):
|
|
v=np.asarray(values,dtype=np.float64);v=v[np.isfinite(v)];assert len(v)==record['count']
|
|
if len(v):
|
|
for name,x in zip(['p10','median','p90','mean'],[*np.quantile(v,[.1,.5,.9]),v.mean()]):np.testing.assert_allclose(record[name],x,rtol=2e-6,atol=1e-8)
|
|
for key in counts:
|
|
for metric,record in summary['concentration'][key].items():check_stats(obs[key+'_'+metric],record)
|
|
primary=json.loads((root/'primary-analysis-summary.json').read_text())
|
|
for section in ('concentration','joint','paired_joint','by_recipient_type'):
|
|
for key,value in primary[section].items():assert summary[section][key]==value,(section,key)
|
|
full_sources=sorted({x['source'] for x in status['completed_links'] if x['full_attention_recipient']})
|
|
assert summary['full_attention_source_indices']==full_sources
|
|
for key,groups in summary['by_recipient_type'].items():
|
|
src=obs[key+'_source']
|
|
for label,metrics in groups.items():
|
|
choose=np.isin(src,full_sources) if label=='full_attention' else ~np.isin(src,full_sources)
|
|
for metric,record in metrics.items():check_stats(obs[key+'_'+metric][choose],record)
|
|
for (key,label),values in joint_values.items():check_stats(values,summary['joint'][key][label])
|
|
for key in counts:
|
|
weak=np.asarray(joint_values[(key,'weak_half')]);strong=np.asarray(joint_values[(key,'strong_half')]);random=np.asarray(joint_values[(key,'random_half')]);valid=weak>0;record=summary['paired_joint'][key]
|
|
assert record['zero_weak_effect_count']==int((~valid).sum())
|
|
np.testing.assert_allclose(record['weak_less_than_strong_fraction'],np.mean(weak<strong))
|
|
np.testing.assert_allclose(record['weak_less_than_random_fraction'],np.mean(weak<random))
|
|
check_stats(strong[valid]/weak[valid],record['strong_over_weak']);check_stats(random[valid]/weak[valid],record['random_over_weak'])
|
|
controls=json.loads((root/'head-sum-controls.json').read_text());assert len(controls)==96 and max(x['relative_l2'] for x in controls)<2e-5
|
|
assert status['canonical_forward_relative_error']<2e-5
|
|
caches=[]
|
|
for f in sorted((root/'cache').rglob('*')):
|
|
if f.is_file():caches.append({'path':str(f.resolve()),'bytes':f.stat().st_size,'sha256':hashlib.sha256(f.read_bytes()).hexdigest()})
|
|
weights=list((source/'weights').rglob('*.safetensors'));assert weights and sum(f.stat().st_size for f in weights)==4548221488
|
|
manifest=json.loads((source/'manifest.json').read_text());weight_records=[]
|
|
for f in weights:
|
|
with f.open('rb') as stream:sha=hashlib.file_digest(stream,'sha256').hexdigest()
|
|
expected=next(x['upstream_sha256'] for x in manifest['files'] if x['name']==f.name);assert sha==expected
|
|
weight_records.append({'path':str(f.resolve()),'bytes':f.stat().st_size,'sha256':sha,'matches_upstream':True})
|
|
external={'host':'8.211.154.134','checkpoint_directory':str((source/'weights').resolve()),'checkpoint_bytes':sum(f.stat().st_size for f in weights),'checkpoint_files':weight_records,'revision':'15852e8c16360a2fea060d615a32b45270f8a8fc','baseline_caches':caches,'automatic_deletion':False,'prior_caches_untouched':True}
|
|
(root/'external-caches.json').write_text(json.dumps(external,indent=2)+'\n')
|
|
reference=json.loads((source/'reference-check.json').read_text())
|
|
assert reference['passed'] and all(x['passed'] for x in reference['tests'])
|
|
assert hashlib.sha256(Path('functional/model.py').read_bytes()).hexdigest()==reference['implementation_sha256']
|
|
(root/'prior-official-reference-check.json').write_bytes((source/'reference-check.json').read_bytes())
|
|
files=list(Path('dominance').glob('*.py'))+[Path('dominance/PROTOCOL.md'),Path('dominance/README.md'),Path('tests/test_dominance.py'),Path('functional/model.py'),Path('crossqkv/fresh_samples.json')]
|
|
(root/'source-hashes.json').write_text(json.dumps({str(f):hashlib.sha256(f.read_bytes()).hexdigest() for f in files if not f.name.startswith('._')},indent=2)+'\n')
|
|
audit={'passed':True,'transitions':23,'texts':4,'link_observations':92,'primary_observations_per_metric':counts,'undefined_zero_total_per_metric':undefined,'full_attention_js_links':full_js_links,'head_sum_controls':96,'canonical_forward_relative_error':status['canonical_forward_relative_error'],'all_raw_concentration_and_masks_recomputed':True,'pooled_concentration_and_joint_summary_recomputed':True,'original_primary_summaries_unchanged':True,'sqrt_JS_scale_sensitivity_recomputed':True,'retained_baseline_cache_files':len(caches),'checkpoint_sha_matches_upstream':True,'prior_official_reference_check_passed_same_implementation':True,'prior_official_reference_tests':len(reference['tests'])}
|
|
(root/'audit.json').write_text(json.dumps(audit,indent=2)+'\n');print(json.dumps(audit,indent=2),flush=True)
|
|
included=[f for f in root.rglob('*') if f.is_file() and 'cache' not in f.relative_to(root).parts and f.name!='checksums.json']
|
|
(root/'checksums.json').write_text(json.dumps({str(f.relative_to(root)):hashlib.sha256(f.read_bytes()).hexdigest() for f in sorted(included)},indent=2)+'\n')
|
|
if __name__=='__main__':main(sys.argv[1],sys.argv[2])
|