109 lines
7.3 KiB
Python
109 lines
7.3 KiB
Python
"""Existing learned RMSNorm modules: uniform function bound + independent probes.
|
|
|
|
Deterministic exhaustive selection by a mathematical all-input error bound, not
|
|
calibration/heldout fitting. Same role, shape and eps; one pair only.
|
|
"""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import resource
|
|
import time
|
|
import numpy as np
|
|
from .model import StreamingLM,norm
|
|
from .run import write_json,save_csv,error_rows
|
|
|
|
ROLES=['input_layernorm','post_attention_layernorm']
|
|
|
|
|
|
def uniform_bounds(a,b):
|
|
if np.any(a==0) or np.any(b==0):return None
|
|
diff=np.abs(a.astype(np.float64)-b.astype(np.float64))
|
|
return float(np.max(diff/np.abs(a))),float(np.max(diff/np.abs(b)))
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root
|
|
protocol={'roles':ROLES,'layers':list(range(24)),'same_role_only':True,'same_eps':1e-6,
|
|
'uniform_worst_error_limit':.20,'empirical_function_worst_relative_error_limit':.05,
|
|
'selection':'smallest maximum two-direction uniform bound, then role/layer; no input data used',
|
|
'minimum_gain_standard_deviation':.001,'max_selected_pairs':1,'heldout_nll_mean_delta_limit':.01,
|
|
'heldout_nll_sequence_p95_delta_limit':.05,'training':False,'cost_required':False}
|
|
path=root/'norm-repetition-protocol.json'
|
|
if path.exists():
|
|
existing=json.loads(path.read_text());assert all(existing[k]==v for k,v in protocol.items())
|
|
else:write_json(path,{**protocol,'registered_unix':time.time()})
|
|
ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-30)
|
|
if remaining<120:raise TimeoutError('Insufficient cumulative computation budget')
|
|
started=time.monotonic();result={'state':'running','selected':None,'interventions':[]}
|
|
try:
|
|
cfg=json.loads((root/'model_config.json').read_text());lm=StreamingLM(root/'weights',cfg,seconds=remaining)
|
|
assert lm.eps==1e-6
|
|
candidates=[];gains={}
|
|
for role in ROLES:
|
|
for layer in range(24):gains[(role,layer)]=1+lm.small(f'layers.{layer}.{role}.weight')
|
|
for i in range(24):
|
|
for j in range(i+1,24):
|
|
a,b=gains[(role,i)],gains[(role,j)]
|
|
bounds=uniform_bounds(a,b)
|
|
if bounds is None:continue
|
|
eligible=min(float(np.std(a)),float(np.std(b)))>=.001
|
|
candidates.append({'role':role,'a':i,'b':j,'a_to_b_uniform_bound':bounds[0],'b_to_a_uniform_bound':bounds[1],
|
|
'worst_uniform_bound':max(bounds),'eligible_nonconstant_gains':eligible,
|
|
'gain_std_a':float(np.std(a)),'gain_std_b':float(np.std(b)),
|
|
'gain_min_a':float(np.min(a)),'gain_min_b':float(np.min(b))})
|
|
save_csv(root/'norm-repetition-uniform-bounds.csv',candidates)
|
|
feasible=[c for c in candidates if c['eligible_nonconstant_gains'] and c['worst_uniform_bound']<=.20]
|
|
counts={str(t):sum(c['eligible_nonconstant_gains'] and c['worst_uniform_bound']<=t for c in candidates) for t in (.05,.10,.20)}
|
|
selected=sorted(feasible,key=lambda c:(c['worst_uniform_bound'],c['role'],c['a'],c['b']))[:1]
|
|
write_json(root/'norm-repetition-selection.json',{'selected':selected,'uniform_bound_counts':counts,'selection_reads_no_input_data':True})
|
|
result['uniform_bound_counts']=counts
|
|
if selected:
|
|
c=selected[0];result['selected']=c
|
|
# Selection is frozen before reading tokenized inputs or their activations.
|
|
samples=json.loads((root/'samples-tokenized.json').read_text())['samples'];ids=np.array([s['token_ids'] for s in samples],dtype=np.int64)
|
|
held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout'])
|
|
kind='attn' if c['role']=='input_layernorm' else 'mlp'
|
|
captures=[(kind,c['a']),(kind,c['b'])]
|
|
h,saved=lm.forward(ids,capture=captures)
|
|
# The untouched recomputation must reproduce the saved real baseline.
|
|
original=np.load(root/'baseline-activations.npz')['final']
|
|
np.testing.assert_array_equal(h,original)
|
|
np.savez(root/'norm-repetition-input-profiles.npz',**saved)
|
|
ga,gb=gains[(c['role'],c['a'])],gains[(c['role'],c['b'])]
|
|
rows=[]
|
|
for source in (c['a'],c['b']):
|
|
n=saved[f'{kind}_{source}'];g=gains[(c['role'],source)]
|
|
# n = u * gamma_source. Reconstruct common normalized input u;
|
|
# source gains are nonzero. Additional rounding is ~FP32 only.
|
|
unit=n/g
|
|
fa,fb=unit*ga,unit*gb
|
|
rec=error_rows(fa,fb,samples,{'source_layer':source,'role':c['role'],'a':c['a'],'b':c['b']})
|
|
rows.extend(rec)
|
|
save_csv(root/'norm-repetition-function-metrics.csv',rows)
|
|
heldrows=[r for r in rows if r['split']=='heldout']
|
|
function_pass=all(max(r['relative_l2'],r['reverse_relative_l2'])<=.05 for r in heldrows)
|
|
result.update(uniform_function_similarity_proven=True,heldout_function_pass=bool(function_pass),
|
|
worst_heldout_relative_l2=max(max(r['relative_l2'],r['reverse_relative_l2']) for r in heldrows),
|
|
worst_heldout_cosine=min(r['cosine'] for r in heldrows),untouched_full_forward_bitwise_matches_baseline=True)
|
|
if function_pass:
|
|
baseline=np.load(root/'baseline-heldout-token-nll.npy').mean(1)
|
|
for target,donor in ((c['a'],c['b']),(c['b'],c['a'])):
|
|
lm.norm_replacement=(f'layers.{target}.{c["role"]}',f'layers.{donor}.{c["role"]}')
|
|
changed,_=lm.forward(ids[held]);nll=lm.nll(changed,ids[held]);np.save(root/f'norm-repetition-{target}-uses-{donor}-token-nll.npy',nll)
|
|
delta=nll.mean(1)-baseline
|
|
result['interventions'].append({'target':target,'donor':donor,'sequence_nll_deltas':delta.tolist(),
|
|
'mean_nll_delta':float(delta.mean()),'sequence_p95_nll_delta':float(np.quantile(delta,.95)),
|
|
'passed_quality':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)})
|
|
write_json(root/'norm-repetition-summary.json',result)
|
|
result['learned_function_repetition_positive']=bool(function_pass)
|
|
result['bidirectional_network_quality_pass']=bool(len(result['interventions'])==2 and all(x['passed_quality'] for x in result['interventions']))
|
|
result.update(state='completed',claim_limit='Existing learned RMSNorm module functions with identical eps and shape; uniform real-arithmetic relative bound and tiny real-text probes. Does not establish attention/MLP-body duplication or broad task capability.')
|
|
except Exception as exc:result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
|
|
finally:
|
|
elapsed=time.monotonic()-started;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.norm_repetition','seconds':elapsed});write_json(ledger_path,ledger)
|
|
write_json(root/'norm-repetition-summary.json',result)
|
|
|
|
|
|
if __name__=='__main__':main()
|