73 lines
7.3 KiB
Python
73 lines
7.3 KiB
Python
"""Paired CPU-only existing-adapter comparison for the published sqrt(JS) metric."""
|
|
import argparse,json,hashlib,time,resource,gc
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from tokenizers import Tokenizer
|
|
from dominance.core import InfluenceLM,concentration
|
|
from dominance.run import distances
|
|
from .lora import install,ExplicitLoRA,BASE_REVISION
|
|
SOURCES=(2,6,10,14,18,22)
|
|
|
|
def dump(p,d):p.write_text(json.dumps(d,indent=2,allow_nan=False)+'\n')
|
|
def single(lm,cache,layer,path):
|
|
base=lm.recipient(cache['next_raw'],layer+1);hs=cache['heads'].shape[2];t=cache['raw'].shape[1];assert hs==16
|
|
scores={k:[] for k in ('Q','K','V','QKV','JS')}
|
|
for start in range(0,hs,4):
|
|
masks=np.zeros((4,t,hs),dtype=np.float32)
|
|
for n in range(4):masks[n,:,start+n]=1
|
|
effects=distances(base,lm.intervene(cache,layer,masks))
|
|
for k,v in effects.items():scores[k].extend(v)
|
|
null=distances(base,lm.intervene(cache,layer,np.zeros((1,t,hs),dtype=np.float32)))
|
|
assert max(float(np.max(v)) for v in null.values())<2e-5
|
|
payload={'single_'+k:np.asarray(v,dtype=np.float32) for k,v in scores.items()};payload.update({'null_'+k:v for k,v in null.items()});np.savez_compressed(path,**payload)
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--parent',type=Path,required=True);p.add_argument('--assets',type=Path,required=True);p.add_argument('--out',type=Path,required=True);a=p.parse_args();a.out.mkdir(parents=True,exist_ok=False);(a.out/'cache').mkdir();(a.out/'raw').mkdir()
|
|
started=time.monotonic();meta=json.loads((a.assets/'adapter_meta.json').read_text());assert meta['base_revision']==BASE_REVISION
|
|
adapter_sha=hashlib.sha256((a.assets/'adapter.safetensors').read_bytes()).hexdigest();assert adapter_sha=='93a3f43a4356f84da0d13ecc622fee13e41ba15d237cdaad6094f243899c62d6'
|
|
tok=Tokenizer.from_file(str(a.source/'tokenizer.json'));legacy=json.loads((a.parent/'samples.json').read_text());gsm=[]
|
|
for s in legacy:assert s['token_ids']==tok.encode(s['text']).ids[:32]
|
|
for index,line in enumerate((a.assets/'gsm8k-test.jsonl').read_text().splitlines()):
|
|
q=json.loads(line)['question'];ids=tok.encode(q).ids
|
|
if len(ids)<32:continue
|
|
gsm.append({'id':f'gsm8k-test-{index:04d}','row_index':index,'text':q,'question_sha256':hashlib.sha256(q.encode()).hexdigest(),'token_ids':ids[:32]})
|
|
if len(gsm)==4:break
|
|
assert len(gsm)==4
|
|
cohorts={'legacy':legacy,'gsm8k':gsm};dump(a.out/'samples.json',cohorts);dump(a.out/'adapter_meta.json',meta);dump(a.out/'gsm8k-source.json',json.loads((a.assets/'gsm8k-source.json').read_text()))
|
|
protocol={'primary':'same six full-attention receivers and original4x32 token cohort as published8.931822/16','adapter_repo':'Avifenesh/gdn-beta-gate-qwen3.5-2b-grpo','adapter_revision':'038d8be9bd0edd3883fa019651b02bdf65609468','adapter_file':'lora/seed0/adapter.safetensors','adapter_sha256':adapter_sha,'base_revision':BASE_REVISION,'scale':2,'rank':16,'arm':'pure LoRA','training':'publisher-reported GSM8K GRPO; we do no training','sample_choice':'four original texts; first4 canonical GSM8K test questions with >=32tokens, length selection only, first32tokens; question only','source_layers':SOURCES,'primary_queries':list(range(7,32)),'main_metric':'concentration(sqrt(single_JS)); same FP32 tensors and sequence-wide donor-head deletion as prior run','cohort_pairing':'each comparison uses same question/tokenIDs before and after; base legacy arrays reused with elementwise checks','runtime':'read-only FP32 effective W+2BA; no BF16 re-rounding or quantization','controls':'head sum, null mask, explicit low-rank versus effective-weight full forward relL2<1e-4','budget_seconds':7200,'keep_all_weights_and_caches':True,'limitations':['one existing training seed','4 texts per cohort, 32 tokens each','dependent query/head observations, not iid textual samples','no generation or accuracy measurement','GDN targets only; ordinary attention weights stay unchanged','no pruning or functional interchangeability claim']}
|
|
dump(a.out/'protocol.json',protocol);status={'state':'running','links':[]};dump(a.out/'status.json',status);cfg=json.loads((a.source/'model_config.json').read_text());controls=[];references=[]
|
|
try:
|
|
for cohort,samples in cohorts.items():
|
|
for role in ('base','finetuned'):
|
|
reuse=cohort=='legacy' and role=='base';lm=None
|
|
if not reuse:
|
|
lm=InfluenceLM(a.source/'weights',cfg,seconds=max(1,7200-(time.monotonic()-started)))
|
|
if role=='finetuned':pairs,scale=install(lm,a.assets,meta)
|
|
for si,s in enumerate(samples):
|
|
if time.monotonic()-started>7200:raise TimeoutError('Paired CPU budget')
|
|
if not reuse:
|
|
cache=a.out/'cache'/f'{cohort}-{role}-{s["id"]}';cache.mkdir();ids=np.array([s['token_ids']],dtype=np.int64);final=lm.baseline(ids,cache);np.save(cache/'final.npy',final)
|
|
controls.extend({'cohort':cohort,'role':role,'sample':s['id'],**v} for v in lm.head_sum_errors)
|
|
if cohort=='legacy' and role=='finetuned' and si==0:
|
|
ref=ExplicitLoRA(a.source/'weights',cfg,seconds=max(1,7200-(time.monotonic()-started)));ref.pairs,ref.scale=pairs,scale;expected,_=ref.forward(ids);actual=lm.normalization(final,'norm');error=float(np.linalg.norm(actual.astype(np.float64)-expected)/np.linalg.norm(expected.astype(np.float64)));assert error<1e-4
|
|
status['explicit_lora_forward_relative_error']=error;del ref;gc.collect()
|
|
for layer in SOURCES:
|
|
out=a.out/'raw'/f'{cohort}-{role}-{s["id"]}-{layer:02d}.npz'
|
|
if reuse:
|
|
old=a.parent/'raw'/f'{s["id"]}-link-{layer:02d}.npz'
|
|
with np.load(old) as d:payload={prefix+k:d[prefix+k] for prefix in ('single_','null_') for k in ('Q','K','V','QKV','JS')}
|
|
np.savez_compressed(out,**payload)
|
|
with np.load(out) as d:
|
|
for k,v in payload.items():np.testing.assert_array_equal(d[k],v)
|
|
references.append({'source':str(old.resolve()),'sha256':hashlib.sha256(old.read_bytes()).hexdigest(),'derived':out.name})
|
|
else:
|
|
with np.load(cache/f'layer-{layer}.npz') as d:loaded={k:d[k] for k in d.files}
|
|
single(lm,loaded,layer,out);del loaded
|
|
status['links'].append({'cohort':cohort,'role':role,'sample':s['id'],'source':layer,'path':str(out.relative_to(a.out))});status['elapsed_seconds']=time.monotonic()-started;dump(a.out/'status.json',status);print('completed',cohort,role,s['id'],layer,flush=True)
|
|
dump(a.out/'head-sum-controls.json',controls);dump(a.out/'baseline-references.json',references)
|
|
if lm is not None:del lm;gc.collect()
|
|
status.update(state='completed',elapsed_seconds=time.monotonic()-started,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss);dump(a.out/'status.json',status)
|
|
except BaseException as e:
|
|
status.update(state='failed',error=repr(e),elapsed_seconds=time.monotonic()-started);dump(a.out/'status.json',status);raise
|
|
if __name__=='__main__':main()
|