84 lines
4.5 KiB
Python
84 lines
4.5 KiB
Python
"""One allocation-reduction attempt, with a similarly optimized baseline.
|
|
|
|
Does not change weights or run the network again. Checks bitwise equality to the
|
|
network-tested factorization on every captured real Q/K input before timing.
|
|
"""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import resource
|
|
import time
|
|
import numpy as np
|
|
from .model import StreamingLM
|
|
from .gain_factorization import consolidate,normalized,LAYERS
|
|
from .run import save_csv,write_json
|
|
|
|
|
|
def baseline_inplace(q,k,a,b):
|
|
q=normalized(q);q*=a
|
|
k=normalized(k);k*=b
|
|
return q,k
|
|
|
|
|
|
def factored_inplace(q,k,a,b,r,combined,move='q_to_k'):
|
|
q,k=normalized(q),normalized(k)
|
|
if move=='q_to_k':
|
|
q[...,:r]*=a[:r];k*=combined
|
|
else:
|
|
q*=combined;k[...,:r]*=b[:r]
|
|
return q,k
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args();root=a.root
|
|
ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-20)
|
|
if remaining<60:raise TimeoutError('Cumulative budget insufficient')
|
|
start=time.monotonic();result={'state':'running'}
|
|
try:
|
|
cfg=json.loads((root/'model_config.json').read_text());lm=StreamingLM(root/'weights',cfg,seconds=min(remaining,300))
|
|
r=64;data=np.load(root/'baseline-activations.npz');checks=[]
|
|
for layer in LAYERS:
|
|
lm.guard();pfx=f'layers.{layer}.self_attn.'
|
|
ag,bg=1+lm.small(pfx+'q_norm.weight'),1+lm.small(pfx+'k_norm.weight')
|
|
q,k=data[f'qraw_{layer}'],data[f'kraw_{layer}']
|
|
for move in ('q_to_k','k_to_q'):
|
|
combined=np.concatenate([(bg if move=='q_to_k' else ag)[:r],ag[r:]*bg[r:]])
|
|
expected=consolidate(q,k,ag,bg,r,move)
|
|
actual=factored_inplace(q,k,ag,bg,r,combined,move)
|
|
for x,y in zip(expected,actual):np.testing.assert_array_equal(x,y)
|
|
checks.append({'layer':layer,'direction':move,'query_elements':int(q.size),'key_elements':int(k.size),'bitwise_equal':True})
|
|
q,k=data['qraw_3'][:1],data['kraw_3'][:1];pfx='layers.3.self_attn.'
|
|
ag,bg=1+lm.small(pfx+'q_norm.weight'),1+lm.small(pfx+'k_norm.weight')
|
|
combined=np.concatenate([bg[:r],ag[r:]*bg[r:]])
|
|
funcs={'baseline_legacy':lambda:(normalized(q)*ag,normalized(k)*bg),
|
|
'baseline_inplace':lambda:baseline_inplace(q,k,ag,bg),
|
|
'factored_inplace':lambda:factored_inplace(q,k,ag,bg,r,combined)}
|
|
for _ in range(3):
|
|
for fn in funcs.values():fn()
|
|
timing=[];names=list(funcs)
|
|
for repeat in range(30):
|
|
order=names[repeat%3:]+names[:repeat%3]
|
|
for name in order:
|
|
lm.guard();t=time.perf_counter()
|
|
for _ in range(50):funcs[name]()
|
|
timing.append({'repeat':repeat,'implementation':name,'seconds_per_call':(time.perf_counter()-t)/50,'calls_per_block':50})
|
|
save_csv(root/'gain-factorization-optimized-timing.csv',timing)
|
|
values={name:np.array([row['seconds_per_call'] for row in timing if row['implementation']==name]) for name in names}
|
|
ratio=values['baseline_inplace']/values['factored_inplace']
|
|
medians={name:float(np.median(v)) for name,v in values.items()}
|
|
result.update(state='completed',bitwise_checks=checks,medians_seconds=medians,
|
|
paired_baseline_over_factored_median=float(np.median(ratio)),
|
|
paired_ratio_p10=float(np.quantile(ratio,.1)),paired_ratio_p90=float(np.quantile(ratio,.9)),
|
|
fraction_blocks_factored_faster=float(np.mean(ratio>1)),
|
|
descriptive_local_gain=bool(np.median(ratio)>1.05 and np.mean(ratio>1)>=.8 and np.quantile(ratio,.1)>1),
|
|
claim_limit='Hot-array CPU normalization microbenchmark against equally allocation-optimized baseline; no full-model speedup or independent statistical significance claim.')
|
|
except Exception as exc:
|
|
result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
|
|
finally:
|
|
elapsed=time.monotonic()-start;result.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
|
|
ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.benchmark_factorization','seconds':elapsed});write_json(ledger_path,ledger)
|
|
write_json(root/'gain-factorization-optimized-summary.json',result)
|
|
|
|
|
|
if __name__=='__main__':main()
|