diff --git a/README.md b/README.md index f155e90..ba88e06 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,12 @@ weights were deleted after archival; the functional pilot executes streaming FP3 The current CPU functional checkpoint is retained until the user explicitly requests deletion. No model weights are committed here. +The latest cross-layer QKV study covers all 15 full-attention layer pairs. Its primary +protocol had no passing candidate; an exploratory H23-conditioned self-token retrieval +pattern is reported separately with its fresh-test failures. + +- [Cross-layer QKV findings and evidence limits](docs/CROSS_LAYER_QKV_FINDINGS.md) +- [Cross-layer QKV fixed protocol](docs/CROSS_LAYER_QKV_PROTOCOL.md) - [CPU learned-function repetition findings](docs/CPU_FUNCTIONAL_FINDINGS.md) - [Current learned-function repetition criterion](docs/LEARNED_FUNCTION_REPETITION.md) - [CPU-first feasibility and validation](docs/CPU_FIRST_VALIDATION.md) diff --git a/crossqkv/README.md b/crossqkv/README.md new file mode 100644 index 0000000..5141a61 --- /dev/null +++ b/crossqkv/README.md @@ -0,0 +1,42 @@ +# Cross-layer QKV CPU study + +This phase targets learned QKV functions across different layers, independent of cost. +The six full-attention layers (indices 3/7/11/15/19/23) form 15 unordered pairs. +No within-layer Q/K/V substitution, training, continuous R fitting, GPU, or new model download. + +Read [fixed primary protocol](../docs/CROSS_LAYER_QKV_PROTOCOL.md) and +[findings with evidence limits](../docs/CROSS_LAYER_QKV_FINDINGS.md). +The original main protocol had no passing candidate. A later fixed one-source audit +showed near-self-token retrieval and small readout errors, but failed strict fresh +mean-JS/row-JS-P95/row-TV-P95 criteria; it is not relabeled as a primary pass. + +- `core.py`: exact discrete head assignment, actual attention parts, common-space head + decomposition, fixed recipient V/gate/O transplants and temporal-cache controls. +- `protocol.py`, `fresh_samples.json`: predeclared layer/unit/gate scope and untouched fresh texts. +- `run.py`, `fresh.py`: calibration, frozen selection, old validation and gated fresh/network steps. +- `conditional.py`: separately registered post-screen metric audit on one fixed H23 candidate; + no fallback selection or later threshold changes. +- `audit_case.py`: recompute selected fresh probabilities, self mass, shared top5 mass and + actual-H3 versus same-H23 difference; no full-network intervention after failure. +- `report.py`, `audit.py`: saved-data scientific plots, raw-count/matching/probability audits, + SHA256 and explicit small Git delivery with external large-cache metadata. + +Use the authorized remote CPU environment with the retained source checkpoint/cache: + +```sh +OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 venv/bin/python -B -m unittest tests.test_crossqkv -v +OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 timeout 7200 venv/bin/python -B -m crossqkv.run --source runs/cpu-qwen35-2b --root runs/NEW_QKV_DIRECTORY +``` + +The main run requires a fresh directory and never overwrites prior protocol evidence. +The recorded execution used `runs/cross-layer-qkv-v2` after fixing a controls-variable +shadowing error in the first attempt; old logs and its profile cache are retained. +The first attempt took about 10.3 seconds, then a 60-second reserve was deducted from +the next 7140-second run limit. The main study finished in about 74.8 seconds. +The fixed conditional audit used only remaining shared budget and took about 105.3 seconds. +Final tests/plots/audits fit comfortably within the two-hour computation budget. + +All weights and caches stay on the server until explicit user deletion. Git contains +all figure tables and the selected small probability matrices; large caches have exact +paths/sizes/hashes in `external-caches.json`. Old validation texts were used in earlier +experiments; the four fresh texts were fixed before the new cross-layer measurements. diff --git a/crossqkv/__init__.py b/crossqkv/__init__.py new file mode 100644 index 0000000..973d897 --- /dev/null +++ b/crossqkv/__init__.py @@ -0,0 +1 @@ +"""Cross-layer learned QKV function comparisons, CPU only.""" diff --git a/crossqkv/audit.py b/crossqkv/audit.py new file mode 100644 index 0000000..0825a80 --- /dev/null +++ b/crossqkv/audit.py @@ -0,0 +1,95 @@ +"""Audit and package only science outputs; retain all weights/profile caches.""" +import argparse +import csv +import hashlib +import json +from pathlib import Path +import shutil +import time +import numpy as np +from functional.run import write_json +from .protocol import LAYERS,PAIRS +from .core import assignment + + +def digest(path): + h=hashlib.sha256() + with path.open('rb') as f: + while b:=f.read(4*1024*1024):h.update(b) + return h.hexdigest() + + +def main(): + p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);p.add_argument('--source',type=Path,required=True);p.add_argument('--output',type=Path,required=True);a=p.parse_args();root=a.root.resolve();source=a.source.resolve();out=a.output.resolve();assert root!=out + status=json.loads((root/'status.json').read_text());assert status['state']=='completed' and status['last_source_layer']==23 and status['last_pair']==[19,23] + expected={'full_attention':15,'qkv_head':960,'qk_map':30,'qk_head':960,'v_group':60};assert status['candidate_counts']==expected + summary=json.loads((root/'calibration-summary.json').read_text());assert len(summary)==2025 + assert not status['selected'] and all(not x['strict_pass'] and not x['approximate_pass'] for x in summary) + selection=json.loads((root/'selection.json').read_text());assert selection['qk_head_prequalified']==selection['v_group_prequalified']==0 + from collections import defaultdict + rawout=defaultdict(list) + for r in csv.DictReader((root/'calibration-output-metrics.csv').open()):rawout[r['candidate']].append(r) + for c in summary: + if c['output_verified']: + rr=rawout[c['id']];assert len(rr)==c['output_observations'] + worst=max(max(float(r['error_a']),float(r['error_b'])) for r in rr) + assert abs(worst-c['worst_output_error'])<1e-12 and worst>.20 + assert min(c['worst_mean_js'] for c in summary if c['kind']=='qk_head')>.05 + assert min(c['worst_value_error'] for c in summary if c['kind']=='v_group')>.20 + + rows=list(csv.DictReader((root/'head-assignment-costs.csv').open()));assert len(rows)==960 + for i,j in PAIRS: + r=[x for x in rows if int(x['i'])==i and int(x['j'])==j];cost=np.empty((8,8));chosen=[] + for x in r: + ha,hb=int(x['head_i']),int(x['head_j']);cost[ha,hb]=float(x['calibration_mean_js']) + if x['chosen']=='True':chosen.append((ha,hb)) + assert len(chosen)==8;assert [b for _,b in sorted(chosen)]==assignment(cost) + tests=(root/'final-tests.txt').read_text();assert 'Ran 20 tests' in tests and '\nOK\n' in tests + cond=json.loads((root/'conditional-status.json').read_text());assert cond['state']=='completed' and cond['fresh_function_pass'] is False + case=json.loads((root/'conditional-case-audit.json').read_text());assert case['passed_recomputation'] and case['no_success_relabeling'] + assert len(case['strict_failures'])==3 and case['network_intervention_not_run_due_to_fresh_gate_failure'] + assert cond['shared_budget_seconds_including_60s_early_reserve']+600<7200 + prob=np.load(root/'conditional-fresh-probabilities.npz');target=prob['target_layer23_head6'];donor=prob['donor_layer3_head7_on_H23'] + from .core import distribution_metrics + from .conditional import tv_metrics + metric={**distribution_metrics(target,donor),**tv_metrics(target,donor)} + fresh=[r for r in csv.DictReader((root/'conditional-function-metrics.csv').open()) if r['split']=='fresh_test'] + assert len(fresh)==4 + for k,values in metric.items(): + if k in fresh[0] and values[0] is not None:np.testing.assert_allclose(values,[float(r[k]) for r in fresh],atol=1e-12,rtol=1e-12) + for file in root.glob('*.csv'): + with file.open() as f: + for r in csv.DictReader(f): + for value in r.values(): + if value in ('',None):continue + try:x=float(value) + except ValueError:continue + assert np.isfinite(x),(file,value) + weights=json.loads((source/'manifest.json').read_text());assert weights['revision']=='15852e8c16360a2fea060d615a32b45270f8a8fc' + for entry in weights['files']: + if entry['name'].endswith('.safetensors'):assert (source/'weights'/entry['name']).stat().st_size==entry['bytes'] and entry['sha256']==entry['upstream_sha256'] + assert json.loads((source/'retention-policy.json').read_text())['automatic_cleanup'] is False + reference=json.loads((source/'reference-check.json').read_text());assert reference['implementation_sha256']==digest(Path('functional/model.py')) + write_json(root/'final-source-hashes.json',{str(file):digest(file) for file in [*sorted(Path('crossqkv').glob('*.py')),Path('crossqkv/fresh_samples.json'),Path('tests/test_crossqkv.py'),Path('functional/model.py')]}) + caches=[] + for file in sorted((root/'cache').glob('*')): + if file.is_file():caches.append({'server_path':str(file),'bytes':file.stat().st_size,'sha256':digest(file),'retained':True}) + old_cache=root.parent/'cross-layer-qkv'/'cache'/'calibration-source-3.npz' + if old_cache.exists():caches.append({'server_path':str(old_cache),'bytes':old_cache.stat().st_size,'sha256':digest(old_cache),'retained':True,'failed_attempt_cache':True}) + write_json(root/'external-caches.json',{'host':'8.211.154.134','source_weights_directory':str(source/'weights'),'source_baseline_cache':str(source/'baseline-activations.npz'), + 'files':caches,'automatic_deletion':False,'all_figure_source_tables_and_selected_probabilities_in_git':True}) + (root/'DELIVERY.md').write_text('# Cross-layer QKV scientific delivery\n\nAll plot CSV/JSON, the selected fresh probability matrices, protocols, reference tests and failure logs are included.\nLarge intermediate profiles remain on 8.211.154.134 with locations/sizes/SHA256 in external-caches.json.\nThe fixed checkpoint and all caches are retained until explicit user deletion. No weights are committed.\nThe primary protocol has no passing candidate. The later fixed conditional audit is exploratory and its strict fresh checks failed; observed readout overlap is reported without relabeling it a complete pass.\n') + out.mkdir(exist_ok=True,parents=True) + files=[f for f in root.iterdir() if f.is_file() and f.name not in ('checksums.json','audit-validation.json','archive-files.json')] + files += list((root/'figures').glob('*')) + for file in files: + dest=out/file.relative_to(root);dest.parent.mkdir(exist_ok=True,parents=True);shutil.copy2(file,dest) + checks={str(f.relative_to(out)):digest(f) for f in sorted(out.rglob('*')) if f.is_file() and f.name not in ('checksums.json','audit-validation.json','archive-files.json')} + write_json(out/'checksums.json',checks) + write_json(out/'archive-files.json',[*checks,'checksums.json','audit-validation.json','archive-files.json']) + write_json(out/'audit-validation.json',{'passed':True,'validated_unix':time.time(),'checksummed_files':len(checks),'primary_counts_rechecked':expected, + 'fresh_probabilities_independently_recomputed':True,'no_unrun_network_quality_claims':True,'all_weights_and_caches_retained':True}) + print(json.dumps({'passed':True,'checksummed_files':len(checks),'external_cache_files':len(caches),'weights_retained':True})) + + +if __name__=='__main__':main() diff --git a/crossqkv/audit_case.py b/crossqkv/audit_case.py new file mode 100644 index 0000000..7ecde78 --- /dev/null +++ b/crossqkv/audit_case.py @@ -0,0 +1,62 @@ +"""Recompute the frozen exploratory case for transparent plotted probabilities. + +No new selection, no threshold changes, no full-network intervention after failure. +""" +import argparse +import json +from pathlib import Path +import numpy as np +from functional.run import save_csv,write_json +from .core import QKVModel,distribution_metrics +from .conditional import tv_metrics + + +def ranks_and_mass(a,b): + result=[] + for bi in range(a.shape[0]): + d={'target_self_mass':[],'donor_self_mass':[],'argmax_agreement':[],'shared_top5_mass_target':[],'shared_top5_mass_donor':[]} + for t in range(7,32): + x,y=a[bi,0,t,:t+1],b[bi,0,t,:t+1] + common=np.intersect1d(np.argsort(x)[-5:],np.argsort(y)[-5:]) + d['target_self_mass'].append(float(x[t]));d['donor_self_mass'].append(float(y[t])) + d['argmax_agreement'].append(float(np.argmax(x)==np.argmax(y))) + d['shared_top5_mass_target'].append(float(x[common].sum()));d['shared_top5_mass_donor'].append(float(y[common].sum())) + result.append({k:float(np.mean(v)) for k,v in d.items()}) + return result + + +def main(): + p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root;source=args.source + status=json.loads((root/'conditional-status.json').read_text());assert status['state']=='completed' + data=np.load(root/'cache'/'conditional-fresh-inputs.npz') + samples=json.loads((root/'conditional-fresh-samples.json').read_text())['samples'] + lm=QKVModel(source/'weights',json.loads((source/'model_config.json').read_text()),seconds=300) + own=lm.parts(data['attn_23'],23);same=lm.parts(data['attn_23'],3);actual=lm.parts(data['attn_3'],3) + target=own['p'][:,6:7];donor=same['p'][:,7:8];cached=actual['p'][:,7:8] + np.savez(root/'conditional-fresh-probabilities.npz',target_layer23_head6=target,donor_layer3_head7_on_H23=donor,donor_layer3_head7_on_H3=cached) + rows=[] + from .run import outrows + for mode,profile,prob in [('same_H23',same,donor),('actual_H3_to_H23',actual,cached)]: + stats={**distribution_metrics(target,prob),**tv_metrics(target,prob)} + mass=ranks_and_mass(target,prob) + changed=lm.head_transplant_output(own,profile,23,6,7) + out=outrows(own['heads'][:,:,6],changed,samples,{},1e-6,1e-6) + for n,s in enumerate(samples):rows.append({'mode':mode,'sample':s['id'],**mass[n],**{k:(None if v[n] is None else float(v[n])) for k,v in stats.items()},'head_output_error':max(out[n]['error_a'],out[n]['error_b'])}) + save_csv(root/'conditional-fresh-source-diagnostic.csv',rows) + with (root/'conditional-function-metrics.csv').open() as f:import csv;allrows=list(csv.DictReader(f)) + failures=[] + for r in allrows: + if r['split']!='fresh_test':continue + for field,limit in [('js',.01),('js_p95',.05),('total_variation',.10),('total_variation_p95',.20)]: + if float(r[field])>limit:failures.append({'sample':r['sample'],'field':field,'value':float(r[field]),'limit':limit}) + same_rows=[r for r in rows if r['mode']=='same_H23'] + classification='self_token_dominated' if min(min(r['target_self_mass'],r['donor_self_mass']) for r in same_rows)>.85 else 'inspect detailed source diagnostic' + write_json(root/'conditional-case-audit.json',{'passed_recomputation':True,'strict_failures':failures,'pattern_classification':classification, + 'minimum_same_input_target_self_mass':min(r['target_self_mass'] for r in same_rows),'minimum_same_input_donor_self_mass':min(r['donor_self_mass'] for r in same_rows), + 'minimum_same_input_argmax_agreement':min(r['argmax_agreement'] for r in same_rows), + 'minimum_shared_top5_probability_mass':min(min(r['shared_top5_mass_target'],r['shared_top5_mass_donor']) for r in same_rows), + 'maximum_same_input_head_output_error':max(r['head_output_error'] for r in same_rows), + 'no_success_relabeling':True,'network_intervention_not_run_due_to_fresh_gate_failure':True}) + + +if __name__=='__main__':main() diff --git a/crossqkv/conditional.py b/crossqkv/conditional.py new file mode 100644 index 0000000..c8a4921 --- /dev/null +++ b/crossqkv/conditional.py @@ -0,0 +1,113 @@ +"""One explicitly exploratory, fixed-candidate audit of unweighted top-k rejection. + +Primary protocol remains unchanged/negative. A new confirmation uses untouched +fresh texts; no further candidate selection, threshold tuning, weights or R fitting. +""" +import argparse +import hashlib +import json +from pathlib import Path +import resource +import time +import traceback +import numpy as np +from tokenizers import Tokenizer +from functional.run import write_json,save_csv +from .core import InterventionLM,distribution_metrics,output_metrics +from .run import head_floor,rows_for + +CANDIDATE={'id':'conditional_qk_head_3_23_7_6','kind':'qk_head','i':3,'j':23,'a':7,'b':6,'perm':None} +PROTOCOL={'classification':'post-screen exploratory metric audit; not a primary-protocol positive', + 'candidate':CANDIDATE,'conditioning_input_layer':23,'donor_projection_layer':3,'target_head':6,'donor_head':7, + 'selection_basis':'one minimum worst-text JS among the already measured source-conditioned pairs; no raw readout/fresh data used', + 'js_max':.01,'js_row_p95_max':.05,'mean_total_variation_max':.10,'row_total_variation_p95_max':.20, + 'target_head_output_relative_l2_max':.20,'target_head_output_cosine_min':.98, + 'primary_query_positions':list(range(7,32)),'top5_is_diagnostic_only':True, + 'active_head_required':True,'no_candidate_fallback':True,'fresh_text_file':'crossqkv/fresh_samples.json', + 'same_input_network_direction':'layer23 head6 uses QK probability of layer3 head7 recomputed on H23', + 'actual_cache_separate':'layer3 probability on H3 reused at23; different-input diagnostic, not function-equivalence evidence', + 'quality_gate':{'mean_nll_delta_max':.01,'per_text_p95_nll_delta_max':.05},'training':False,'R_fitting':False,'cost_required':False} + + +def tv_metrics(a,b): + tv=.5*np.sum(np.abs(a.astype(np.float64)-b.astype(np.float64)),axis=-1)[:,:,7:] + return {'total_variation':tv.mean((1,2)),'total_variation_p95':np.quantile(tv.reshape(tv.shape[0],-1),.95,axis=1)} + + +def evaluate(lm,h,samples,split,donor_h=None): + target=lm.parts(h,23);donor=lm.parts(h if donor_h is None else donor_h,3) + a,b=target['p'][:,6:7],donor['p'][:,7:8] + pm={**distribution_metrics(a,b),**tv_metrics(a,b)} + changed=lm.head_transplant_output(target,donor,23,6,7) + om=output_metrics(target['heads'][:,7:,6],changed[:,7:]);floor=head_floor(target) + uniform=np.zeros_like(a) + for t in range(a.shape[-1]):uniform[:,:,t,:t+1]=1/(t+1) + null=distribution_metrics(a,uniform) + rows=rows_for({**pm,**{f'output_{k}':v for k,v in om.items()},'target_uniform_js':null['js']},samples,{'split':split,'mode':'same_H23' if donor_h is None else 'actual_H3_to_H23'}) + for n,r in enumerate(rows):r['active']=bool(om['rms_a'][n]>=floor[n] and om['rms_b'][n]>=floor[n]) + passed=all(r['js']<=.01 and r['js_p95']<=.05 and r['total_variation']<=.10 and r['total_variation_p95']<=.20 and + max(r['output_error_a'],r['output_error_b'])<=.20 and r['output_cosine']>=.98 and r['active'] for r in rows) + return rows,passed + + +def main(): + p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--root',type=Path,required=True);a=p.parse_args();source=a.source.resolve();root=a.root.resolve() + assert not (root/'conditional-protocol.json').exists(),'Fresh conditional audit protocol required' + write_json(root/'conditional-protocol.json',{**PROTOCOL,'registered_unix':time.time(),'fresh_text_sha256':hashlib.sha256(Path(__file__).with_name('fresh_samples.json').read_bytes()).hexdigest()}) + prior=json.loads((root/'status.json').read_text());assert prior['state']=='completed' and not prior['selected'] + reserve=prior['elapsed_seconds']+60 + remaining=int(7200-reserve-180) + if remaining<120:raise TimeoutError('Insufficient shared two-hour budget') + started=time.monotonic();status={'state':'running','phase':'calibration_readout','primary_protocol_unchanged_and_negative':True} + rows=[] + try: + lm=InterventionLM(source/'weights',json.loads((source/'model_config.json').read_text()),seconds=remaining) + old=json.loads((source/'samples-tokenized.json').read_text())['samples'];data=np.load(source/'baseline-activations.npz') + for split in ('calibration','heldout'): + index=np.array([i for i,s in enumerate(old) if s['split']==split]);samples=[old[i] for i in index] + rr,passed=evaluate(lm,data['attn_23'][index],samples,split) + rows.extend(rr);save_csv(root/'conditional-function-metrics.csv',rows) + status[split+'_pass']=passed + write_json(root/'conditional-status.json',status) + if not passed: + status['stop_reason']=f'{split} probability-mass/readout/activity screen failed';break + else: + fresh=json.loads(Path(__file__).with_name('fresh_samples.json').read_text())['samples'];tokenizer=Tokenizer.from_file(str(source/'tokenizer.json'));ids=[] + for s in fresh: + tokens=tokenizer.encode(s['text'],add_special_tokens=False).ids;assert len(tokens)>=32 + s.update(token_ids=tokens[:32],full_token_count=len(tokens),text_sha256=hashlib.sha256(s['text'].encode()).hexdigest());ids.append(tokens[:32]) + ids=np.array(ids,dtype=np.int64);write_json(root/'conditional-fresh-samples.json',{'samples':fresh,'candidate_already_fixed':True}) + status['phase']='fresh_baseline';write_json(root/'conditional-status.json',status) + h,inputs=lm.forward(ids,capture=[('attn',3),('attn',23)]) + np.savez(root/'cache'/'conditional-fresh-inputs.npz',**inputs,final=h) + base=lm.nll(h,ids);np.save(root/'conditional-fresh-baseline-token-nll.npy',base) + rr,passed=evaluate(lm,inputs['attn_23'],fresh,'fresh_test');rows.extend(rr);save_csv(root/'conditional-function-metrics.csv',rows) + status['fresh_function_pass']=passed + if passed: + status['phase']='one_direction_same_input_network';write_json(root/'conditional-status.json',status) + lm.configure(CANDIDATE,reverse=True,cached=False);changed,_=lm.forward(ids);nll=lm.nll(changed,ids);np.save(root/'conditional-same-input-network-token-nll.npy',nll) + delta=nll.mean(1)-base.mean(1) + network={'mode':'layer23_head6_uses_layer3_head7_on_H23','sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)), + 'quality_pass':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05),'one_direction_only':True} + write_json(root/'conditional-network.json',network);status['same_input_network_quality_pass']=network['quality_pass'] + # Check the actual-cache probability/readout separately, then run its network experiment. + cached_rows,cached_function_pass=evaluate(lm,inputs['attn_23'],fresh,'fresh_test_actual_cache',donor_h=inputs['attn_3']) + save_csv(root/'conditional-actual-cache-metrics.csv',cached_rows) + status['actual_cache_function_screen_pass']=cached_function_pass + lm.configure(CANDIDATE,reverse=True,cached=True);changed,_=lm.forward(ids);nll=lm.nll(changed,ids);np.save(root/'conditional-actual-cache-network-token-nll.npy',nll) + delta=nll.mean(1)-base.mean(1) + cache={'mode':'actual_layer3_H3_probability_reused_at23_H23','sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)), + 'quality_pass':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05),'different_hidden_inputs':True} + write_json(root/'conditional-actual-cache-network.json',cache);status['actual_cache_network_quality_pass']=cache['quality_pass'] + status.update(state='completed',phase='done',conditional_function_positive=bool(status.get('fresh_function_pass')), + position_pattern_classification='near_uniform' if rows and max(r['target_uniform_js'] for r in rows)<=.01 else 'inspect first-key/self/recent3 masses', + claim_limit='Single-source H23-conditioned, single-head QK probability/readout similarity; not bidirectional, not V similarity, not full-QKV-head equivalence; exploratory addendum independently fresh-tested.') + except Exception as exc: + status.update(state='stopped',error_type=type(exc).__name__,error=str(exc));(root/'conditional-error.txt').write_text(traceback.format_exc());raise + finally: + status.update(elapsed_seconds=time.monotonic()-started,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, + shared_budget_seconds_including_60s_early_reserve=reserve+time.monotonic()-started) + write_json(root/'conditional-status.json',status) + + +if __name__=='__main__':main() diff --git a/crossqkv/core.py b/crossqkv/core.py new file mode 100644 index 0000000..1101421 --- /dev/null +++ b/crossqkv/core.py @@ -0,0 +1,171 @@ +"""Attention functions and discrete diagnostics, no continuous alignment/fitting.""" +import numpy as np +from functional.model import StreamingLM,PREFIX,attention,rope,sigmoid + + +def assignment(cost): + """Exact minimum-cost bijection via O(n^2 2^n) dynamic programming.""" + cost=np.asarray(cost,dtype=np.float64);n=cost.shape[0] + assert cost.shape==(n,n) and n<=12 and np.isfinite(cost).all() + states={0:(0.,())} + for i in range(n): + nxt={} + for mask,(value,perm) in states.items(): + for j in range(n): + if mask&(1<=7).""" + a,b=a.astype(np.float64),b.astype(np.float64) + assert a.shape==b.shape and a.shape[-1]>=8 + m=(a+b)*.5 + def term(x):return np.where(x>0,x*np.log(np.maximum(x,1e-300)/np.maximum(m,1e-300)),0) + js=np.maximum(.5*np.sum(term(a)+term(b),axis=-1),0) + eligible=js[...,7:] + overlaps=[] + for t in range(7,a.shape[-2]): + aa=np.argsort(a[:,:,t,:t+1],axis=-1)[...,-5:] + bb=np.argsort(b[:,:,t,:t+1],axis=-1)[...,-5:] + overlaps.append(np.mean((aa[...,None]==bb[...,None,:]).any(-1),-1)) + # First-key ablation is diagnostic, not an unreported filter of positional repetitions. + ar,br=a[...,1:],b[...,1:];ma,mb=ar.sum(-1),br.sum(-1) + valid=(ma>1e-8)&(mb>1e-8) + ar=ar/np.maximum(ma[...,None],1e-300);br=br/np.maximum(mb[...,None],1e-300) + mr=(ar+br)*.5 + cond=.5*np.sum(np.where(ar>0,ar*np.log(np.maximum(ar,1e-300)/np.maximum(mr,1e-300)),0)+np.where(br>0,br*np.log(np.maximum(br,1e-300)/np.maximum(mr,1e-300)),0),-1) + sink=[] + for bi in range(a.shape[0]): + which=valid[bi,:,7:];values=cond[bi,:,7:] + sink.append(float(np.maximum(values[which],0).mean()) if which.any() else None) + diag=np.stack([a[:,:,t,t] for t in range(7,a.shape[-2])],-1).mean((1,2)) + local=np.stack([a[:,:,t,max(0,t-2):t+1].sum(-1) for t in range(7,a.shape[-2])],-1).mean((1,2)) + return {'js':eligible.mean((1,2)),'js_p95':np.quantile(eligible.reshape(a.shape[0],-1),.95,axis=1), + 'top5':np.stack(overlaps,-1).mean((1,2)),'js_all_queries':js.mean((1,2)), + 'first_key_mass_a':a[:,:,7:,0].mean((1,2)),'first_key_mass_b':b[:,:,7:,0].mean((1,2)), + 'self_mass_a':diag,'recent3_mass_a':local,'js_without_first_key':sink} + + +def cross_head_output_metrics(a,b): + """All head pairs, per text. a/b: [B,T,H,Dhidden].""" + aa=a.transpose(0,2,1,3).reshape(a.shape[0],a.shape[2],-1).astype(np.float64) + bb=b.transpose(0,2,1,3).reshape(b.shape[0],b.shape[2],-1).astype(np.float64) + xx=np.sum(aa*aa,-1);yy=np.sum(bb*bb,-1) + dot=np.einsum('bif,bjf->bij',aa,bb,optimize=True) + diff=np.maximum(xx[:,:,None]+yy[:,None,:]-2*dot,0) + return {'error_a':np.sqrt(diff/np.maximum(xx[:,:,None],1e-30)), + 'error_b':np.sqrt(diff/np.maximum(yy[:,None,:],1e-30)), + 'cosine':dot/np.maximum(np.sqrt(xx[:,:,None]*yy[:,None,:]),1e-30), + 'rms_a':np.broadcast_to(np.sqrt(xx/aa.shape[-1])[:,:,None],dot.shape), + 'rms_b':np.broadcast_to(np.sqrt(yy/bb.shape[-1])[:,None,:],dot.shape)} + + +def value_gram_cka(a,b): + """Centered token Gram agreement, diagnostic only; not interchangeability.""" + a=a.astype(np.float64)-a.mean(axis=1,keepdims=True,dtype=np.float64) + b=b.astype(np.float64)-b.mean(axis=1,keepdims=True,dtype=np.float64) + ga=np.einsum('btd,bsd->bts',a,a);gb=np.einsum('btd,bsd->bts',b,b) + return np.sum(ga*gb,(1,2))/np.maximum(np.sqrt(np.sum(ga*ga,(1,2))*np.sum(gb*gb,(1,2))),1e-30) + + +class QKVModel(StreamingLM): + def project(self,context,gate,layer): + return self.linear((context*sigmoid(gate)).reshape(*context.shape[:2],-1),f'layers.{layer}.self_attn.o_proj') + + def project_heads(self,context,gate,layer): + """Each head's contribution in the common residual output coordinates.""" + c=context*sigmoid(gate);b,t,h,d=c.shape + weight=self.tensors[PREFIX+f'layers.{layer}.self_attn.o_proj.weight'] + outdim,indim=weight.shape;assert indim==h*d + out=np.empty((b,t,h,outdim),dtype=np.float32) + for start in range(0,outdim,256): + self.guard();end=min(start+256,outdim) + w=weight.read(slice(start*indim,end*indim)).reshape(end-start,indim) + for head in range(h):out[:,:,head,start:end]=c[:,:,head,:]@w[:,head*d:(head+1)*d].T + del w;weight.evict() + return out + + def parts(self,h,layer,heads=True): + p=f'layers.{layer}.self_attn.';c=self.config;d=c['head_dim'];nh=c['num_attention_heads'];nk=c['num_key_value_heads'] + mixed=self.linear(h,p+'q_proj').reshape(*h.shape[:2],nh,2*d) + q,gate=mixed[...,:d],mixed[...,d:].copy() + k=self.linear(h,p+'k_proj').reshape(*h.shape[:2],nk,d) + v=self.linear(h,p+'v_proj').reshape(*h.shape[:2],nk,d) + q=self.normalization(q,p+'q_norm');k=self.normalization(k,p+'k_norm') + r=c['rope_parameters'];q,k=(rope(x,r['rope_theta'],r['partial_rotary_factor']) for x in (q,k)) + context,probs=attention(q,k,v) + out={'q':q,'k':k,'p':probs,'v':v,'gate':gate,'context':context,'output':self.project(context,gate,layer)} + if heads:out['heads']=self.project_heads(context,gate,layer) + return out + + def transplant(self,own,donor,layer,kind,target_head=None,donor_head=None,perm=None): + repeats=self.config['num_attention_heads']//self.config['num_key_value_heads'] + if kind=='full_attention':return donor['output'] + if kind=='qkv_head':return own['output']-own['heads'][:,:,target_head]+donor['heads'][:,:,donor_head] + probs=own['p'].copy();v=own['v'].copy() + if kind=='qk_map':probs=donor['p'][:,perm] + elif kind=='qk_head':probs[:,target_head]=donor['p'][:,donor_head] + elif kind=='v_group':v[:,:,target_head]=donor['v'][:,:,donor_head] + else:raise ValueError(kind) + context=np.einsum('bhts,bshd->bthd',probs,np.repeat(v,repeats,axis=2)) + return self.project(context,own['gate'],layer) + + def head_transplant_output(self,own,donor,layer,target_head,donor_head): + """QK-only transplanted target head contribution; all other heads omitted.""" + repeats=self.config['num_attention_heads']//self.config['num_key_value_heads'] + c=own['context'].copy() + c[:,:,target_head]=np.einsum('bts,bsd->btd',donor['p'][:,donor_head],own['v'][:,:,target_head//repeats]) + # Zero remaining head contexts before projection to isolate this learned subfunction. + for head in range(c.shape[2]): + if head!=target_head:c[:,:,head]=0 + return self.project(c,own['gate'],layer) + + def value_group_output(self,own,donor,layer,target_group,donor_group): + repeats=self.config['num_attention_heads']//self.config['num_key_value_heads'];c=np.zeros_like(own['context']) + for head in range(target_group*repeats,(target_group+1)*repeats): + c[:,:,head]=np.einsum('bts,bsd->btd',own['p'][:,head],donor['v'][:,:,donor_group]) + return self.project(c,own['gate'],layer) + + +class InterventionLM(QKVModel): + def configure(self,candidate,reverse=False,cached=False): + self.candidate=candidate;self.reverse=reverse;self.use_cache=cached;self.donor_cache=None + self.target_layer=candidate['j'] if reverse else candidate['i'] + self.donor_layer=candidate['i'] if reverse else candidate['j'] + self.target_head=candidate.get('b') if reverse else candidate.get('a') + self.donor_head=candidate.get('a') if reverse else candidate.get('b') + self.perm=inverse(candidate['perm']) if reverse and candidate.get('perm') else candidate.get('perm') + if cached:assert self.donor_layer=32,s['id'] + s.update(token_ids=encoded[:32],full_token_count=len(encoded),text_sha256=hashlib.sha256(s['text'].encode()).hexdigest()) + ids.append(encoded[:32]) + ids=np.array(ids,dtype=np.int64) + write_json(root/'fresh-samples-tokenized.json',{'samples':samples,'selection_already_frozen':True}) + status['phase']='fresh_test_baseline';write_json(root/'status.json',status) + h,inputs=lm.forward(ids,capture=[('attn',i) for i in LAYERS]) + np.savez(root/'cache'/'fresh-baseline-activations.npz',final=h,**inputs) + nll=lm.nll(h,ids);np.save(root/'fresh-baseline-token-nll.npy',nll) + checked=validate_selected(lm,c,inputs,samples,np.arange(len(samples))) + passed=candidate_pass(checked,tier) + write_json(root/'fresh-function-summary.json',summarize(checked)) + save_csv(root/'fresh-output-metrics.csv',checked['output_rows']);save_csv(root/'fresh-probability-metrics.csv',checked['probability_rows']);save_csv(root/'fresh-value-metrics.csv',checked['value_rows']) + status.update(fresh_function_pass=passed,fresh_confirmation_tier=tier) + if not passed:return + status['phase']='same_input_network_substitution';write_json(root/'status.json',status) + interventions=[] + for reverse in (False,True): + lm.configure(c,reverse=reverse,cached=False) + t=time.monotonic();changed,_=lm.forward(ids);loss=lm.nll(changed,ids) + label=f'{lm.target_layer}_uses_{lm.donor_layer}' + np.save(root/f'network-{label}-token-nll.npy',loss) + delta=loss.mean(1)-nll.mean(1) + rec={'mode':'same_input_donor_function','target_layer':lm.target_layer,'donor_layer':lm.donor_layer, + 'target_head_or_group':lm.target_head,'donor_head_or_group':lm.donor_head, + 'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)), + 'quality_pass':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05),'elapsed_seconds':time.monotonic()-t} + interventions.append(rec);write_json(root/'network-interventions.json',interventions) + # Temporal reuse is a separate intervention and only goes from earlier to later. + status['phase']='earlier_to_later_actual_cache';write_json(root/'status.json',status) + lm.configure(c,reverse=True,cached=True) + changed,_=lm.forward(ids);loss=lm.nll(changed,ids);np.save(root/'network-earlier-actual-cache-token-nll.npy',loss) + delta=loss.mean(1)-nll.mean(1) + cache={'mode':'earlier_actual_state_reused_later','donor_layer':lm.donor_layer,'target_layer':lm.target_layer, + 'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)), + 'quality_pass':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05), + 'not_same_input_comparison':True,'no_reverse_future_state_claim':True} + write_json(root/'actual-cache-intervention.json',cache) + status.update(bidirectional_same_input_network_pass=all(r['quality_pass'] for r in interventions),actual_cache_quality_pass=cache['quality_pass']) diff --git a/crossqkv/fresh_samples.json b/crossqkv/fresh_samples.json new file mode 100644 index 0000000..bcdf309 --- /dev/null +++ b/crossqkv/fresh_samples.json @@ -0,0 +1,9 @@ +{ + "origin":"Four experiment-authored fresh plain texts, fixed for this cross-layer QKV goal before measurements; no chat template, first 32 tokenizer tokens.", + "samples":[ + {"id":"fresh-zh-1","language":"zh","text":"植物学家在不同海拔采集了叶片样本，并为每一份标本记录温度、光照和土壤湿度。回到实验室后，她按地点整理标签，先核对缺失记录，再比较不同环境中的叶片形状。"}, + {"id":"fresh-zh-2","language":"zh","text":"如果一个盒子里有十二枚蓝色棋子和八枚白色棋子，随机取出一枚后不放回，第二次取到蓝色的概率需要根据第一次的结果分别计算。请先列出各个事件，再检查它们的概率之和。"}, + {"id":"fresh-en-1","language":"en","text":"A technician recorded the vibration of a small motor before and after replacing its bearings. She used the same mounting position and sampling interval for both measurements. The final notebook included the raw traces, instrument settings, and a brief explanation of unusual observations."}, + {"id":"fresh-en-2","language":"en","text":"The function walks through a list of integers and stores each running total in a new list. An empty input should return an empty result. Before changing the implementation, compare a short example by hand and make sure negative numbers preserve the intended behavior."} + ] +} diff --git a/crossqkv/protocol.py b/crossqkv/protocol.py new file mode 100644 index 0000000..990b8b7 --- /dev/null +++ b/crossqkv/protocol.py @@ -0,0 +1,30 @@ +"""Fixed protocol, constants frozen before any new cross-layer measurements.""" +from itertools import combinations + +LAYERS=[3,7,11,15,19,23] +PAIRS=list(combinations(LAYERS,2)) +TIERS={ + 'strict':{'output_error':.05,'output_cosine':.99,'js':.01,'js_p95':.05,'top5':.8,'value_error':.05,'value_cosine':.99}, + 'approximate':{'output_error':.20,'output_cosine':.98,'js':.05,'js_p95':.15,'top5':.6,'value_error':.20,'value_cosine':.98}} +PROTOCOL={'layers':LAYERS,'layer_pairs':PAIRS,'heads':8,'kv_heads':2,'head_dim':256, + 'calibration_samples':['cal-zh-1','cal-zh-2','cal-en-1','cal-en-2'], + 'validation_samples':['test-zh-1','test-zh-2','test-en-1','test-en-2'], + 'fresh_test_file':'crossqkv/fresh_samples.json','fresh_test_count':4,'tokens_per_text':32, + 'query_positions_for_primary_metrics':list(range(7,32)), + 'candidate_types':['full_attention','qkv_head','qk_map','qk_head','v_group'], + 'discrete_head_matching':'Exact 8-head minimum mean-JS bijection, two source-layer calibration input distributions; frozen before validation; P injection only', + 'continuous_alignment':False,'training':False,'scale_fitting':False,'gpu':False,'cost_required':False, + 'tiers':TIERS,'max_qk_head_output_candidates':32,'max_v_group_output_candidates':16, + 'all_qk_map_output_candidates':30,'max_final_candidate':1, + 'selection':'strict before approximate; type priority full_attention/qkv_head/qk_map/qk_head/v_group; then worst output error, JS, ID; no validation-driven fallback', + 'head_output_activity_floor':'max(1e-6, 0.1 * within-layer median head RMS) for every calibration/validation/test text', + 'fresh_test_run_only_if_frozen_candidate_passes_validation':True, + 'network_interventions':'two same-input donor-function substitution directions; plus one earlier-to-later actual cached-state reuse test reported separately', + 'network_quality':{'mean_token_nll_delta_max':.01,'per_text_p95_nll_delta_max':.05}, + 'retention':'all weights and caches retained until explicit user deletion request', + 'seconds':7200,'controls':['identity','uniform causal P','seeded within-causal-row key shuffle','first-key ablation','self/recent3 mass','value Gram to input','value Gram of two seeded Gaussian projections'], + 'positional_repeat_labels':{'sink_dominated':.90,'self_dominated':.90,'recent3_dominated':.90}, + 'limitations':['Old validation texts were used in prior experiments; four fresh texts are reserved for this new goal', + 'Direct V-coordinate disagreement does not disprove V/O functional equivalence; centered value Gram CKA is diagnostic only', + 'Arbitrary P-head assignment does not authorize arbitrary GQA KV-weight permutation', + 'Short text-only pilot does not establish broad task equivalence or linear-attention reuse']} diff --git a/crossqkv/report.py b/crossqkv/report.py new file mode 100644 index 0000000..78eb776 --- /dev/null +++ b/crossqkv/report.py @@ -0,0 +1,126 @@ +"""Scientific plots from saved tables/probabilities only; no model access.""" +import argparse +import csv +import json +from pathlib import Path +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from matplotlib.patches import Rectangle +from .protocol import LAYERS,PAIRS + + +def read_csv(path): + with path.open() as f:return list(csv.DictReader(f)) + + +def main(): + p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root;out=root/'figures';out.mkdir(exist_ok=True) + plt.rcParams.update({'font.size':9,'pdf.fonttype':42,'svg.fonttype':'none','savefig.bbox':'tight'}) + def save(fig,name): + for ext in ('png','pdf','svg'):fig.savefig(out/f'{name}.{ext}',dpi=180) + plt.close(fig) + status=json.loads((root/'status.json').read_text());summary=json.loads((root/'calibration-summary.json').read_text()) + matrices=[] + for mode in ('qk_head','identity','matched'): + matrix=np.full((6,6),np.nan) + for i,j in PAIRS: + rows=[r for r in summary if r['i']==i and r['j']==j and ((mode=='qk_head' and r['kind']=='qk_head') or (mode!='qk_head' and r['id']==f'qk_map_{i}_{j}_{mode}'))] + value=min(r['worst_mean_js'] for r in rows) + a,b=LAYERS.index(i),LAYERS.index(j);matrix[a,b]=matrix[b,a]=value + matrices.append(matrix) + fig,axes=plt.subplots(1,3,figsize=(13,4.7));vmax=max(np.nanmax(m) for m in matrices) + for ax,m,title in zip(axes,matrices,['Best single head pair','All heads: original order','All heads: calibration bijection']): + im=ax.imshow(m,cmap='YlOrRd',vmin=0,vmax=vmax) + for a in range(6): + ax.add_patch(Rectangle((a-.5,a-.5),1,1,facecolor='#eeeeee',edgecolor='#aaaaaa',hatch='///',lw=.3)) + for b in range(6): + if a!=b:ax.text(b,a,f'{m[a,b]:.3f}',ha='center',va='center',fontsize=7,color='black' if m[a,b]