136 lines
8.4 KiB
Python
136 lines
8.4 KiB
Python
"""Architecture-derived Q/K gain consolidation on non-RoPE dimensions only.
|
|
|
|
Fixed layer3 interventions; no fitting, layer selection or learned alignment.
|
|
Run only after baseline/candidate pilot finishes, within the same cumulative budget.
|
|
"""
|
|
import argparse
|
|
import csv
|
|
import json
|
|
from pathlib import Path
|
|
import resource
|
|
import time
|
|
import numpy as np
|
|
from .model import StreamingLM,norm,rope
|
|
from .run import probability_metrics,write_json,save_csv
|
|
|
|
LAYERS=[3,7,11,15,19,23]
|
|
|
|
|
|
def normalized(x):
|
|
return x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+np.float32(1e-6))
|
|
|
|
|
|
def consolidate(q,k,a,b,r,move,joint=None):
|
|
q,k=normalized(q),normalized(k)
|
|
c=a[r:]*b[r:] if joint is None else joint
|
|
if move=='q_to_k':
|
|
q=np.concatenate([q[...,:r]*a[:r],q[...,r:]],axis=-1)
|
|
k=np.concatenate([k[...,:r]*b[:r],k[...,r:]*c],axis=-1)
|
|
else:
|
|
q=np.concatenate([q[...,:r]*a[:r],q[...,r:]*c],axis=-1)
|
|
k=np.concatenate([k[...,:r]*b[:r],k[...,r:]],axis=-1)
|
|
return q,k
|
|
|
|
|
|
class FactoredLM(StreamingLM):
|
|
def configure(self,layer,move):
|
|
self.factor_layer=layer;self.move=move
|
|
p=f'layers.{layer}.self_attn.'
|
|
self.a=1+self.small(p+'q_norm.weight');self.b=1+self.small(p+'k_norm.weight')
|
|
self.r=int(self.config['head_dim']*self.config['rope_parameters']['partial_rotary_factor'])
|
|
self.joint=self.a[self.r:]*self.b[self.r:]
|
|
|
|
def normalization(self,x,name):
|
|
p=f'layers.{self.factor_layer}.self_attn.'
|
|
if name not in (p+'q_norm',p+'k_norm'):return super().normalization(x,name)
|
|
# This path does not multiply the unit-gain tail, unlike a full-vector mask.
|
|
unit=x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+self.eps)
|
|
isq=name==p+'q_norm';head=self.a if isq else self.b
|
|
skip=(isq and self.move=='q_to_k') or (not isq and self.move=='k_to_q')
|
|
tail=unit[...,self.r:] if skip else unit[...,self.r:]*self.joint
|
|
return np.concatenate([unit[...,:self.r]*head[:self.r],tail],axis=-1)
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args();root=a.root
|
|
# Protocol is fixed before inspecting functional measurements.
|
|
protocol={'fixed_intervention_layer':3,'directions':['q_to_k','k_to_q'],'diagnostic_layers':LAYERS,'no_selection':True,'no_training':True,'precision':'FP32 product; no BF16 requantization','quality_gate':{'mean_nll_delta':.01,'sequence_p95_nll_delta':.05}}
|
|
protocol_path=root/'gain-factorization-protocol.json'
|
|
if protocol_path.exists():
|
|
old=json.loads(protocol_path.read_text());assert all(old[k]==v for k,v in protocol.items()), 'Protocol differs from preregistration'
|
|
else:write_json(protocol_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<60:raise TimeoutError('No remaining cumulative budget for deterministic identity verification')
|
|
start=time.monotonic()
|
|
report={'state':'running','records':[],'interventions':[]}
|
|
try:
|
|
cfg=json.loads((root/'model_config.json').read_text());c=cfg['text_config'];r=int(c['head_dim']*c['rope_parameters']['partial_rotary_factor'])
|
|
assert c['rms_norm_eps']==1e-6
|
|
assert c['head_dim']==256 and r==64 and c['num_attention_heads']==8 and c['num_key_value_heads']==2
|
|
lm=FactoredLM(root/'weights',cfg,seconds=remaining)
|
|
data=np.load(root/'baseline-activations.npz');samples=json.loads((root/'samples-tokenized.json').read_text())['samples']
|
|
rp=c['rope_parameters'];repeats=c['num_attention_heads']//c['num_key_value_heads']
|
|
records=[]
|
|
for layer in LAYERS:
|
|
lm.guard()
|
|
pfx=f'layers.{layer}.self_attn.'
|
|
aw=lm.small(pfx+'q_norm.weight');bw=lm.small(pfx+'k_norm.weight');ag,bg=1+aw,1+bw
|
|
q,k=data[f'qraw_{layer}'],data[f'kraw_{layer}']
|
|
qr,kr=rope(norm(q,aw),rp['rope_theta'],rp['partial_rotary_factor']),rope(norm(k,bw),rp['rope_theta'],rp['partial_rotary_factor'])
|
|
ref=np.einsum('bthd,bshd->bhts',qr,np.repeat(kr,repeats,axis=2))/np.float32(np.sqrt(c['head_dim']))
|
|
mask=np.triu(np.ones(ref.shape[-2:],dtype=bool),1)
|
|
def probs(s):
|
|
s=np.where(mask,-np.inf,s);e=np.exp(s-s.max(-1,keepdims=True));return e/e.sum(-1,keepdims=True)
|
|
pref=probs(ref)
|
|
for move in ('q_to_k','k_to_q'):
|
|
qq,kk=consolidate(q,k,ag,bg,r,move)
|
|
qq,kk=rope(qq,rp['rope_theta'],rp['partial_rotary_factor']),rope(kk,rp['rope_theta'],rp['partial_rotary_factor'])
|
|
changed=np.einsum('bthd,bshd->bhts',qq,np.repeat(kk,repeats,axis=2))/np.float32(np.sqrt(c['head_dim']))
|
|
pp=probs(changed)
|
|
for bi,s in enumerate(samples):
|
|
diff=changed[bi].astype(np.float64)-ref[bi]
|
|
records.append({'layer':layer,'direction':move,'sample_id':s['id'],'split':s['split'],
|
|
'score_max_abs_error':float(np.abs(diff).max()),'score_relative_l2':float(np.linalg.norm(diff)/max(np.linalg.norm(ref[bi].astype(np.float64)),1e-12)),
|
|
'probability_max_abs_error':float(np.abs(pp[bi]-pref[bi]).max()),**probability_metrics(pref[bi:bi+1],pp[bi:bi+1])})
|
|
save_csv(root/'gain-factorization-score-metrics.csv',records)
|
|
held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout']);ids=np.array([s['token_ids'] for s in samples],dtype=np.int64)
|
|
base=np.load(root/'baseline-heldout-token-nll.npy').mean(1)
|
|
for move in ('q_to_k','k_to_q'):
|
|
lm.configure(3,move)
|
|
h,_=lm.forward(ids[held])
|
|
nll=lm.nll(h,ids[held]);np.save(root/f'gain-factorization-{move}-token-nll.npy',nll)
|
|
delta=nll.mean(1)-base
|
|
report['interventions'].append({'layer':3,'direction':move,'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_delta':float(np.quantile(delta,.95)),'passed':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)})
|
|
write_json(root/'gain-factorization-summary.json',report)
|
|
# Local normalization timing uses precomputed gains for both implementations.
|
|
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')
|
|
def baseline():return normalized(q)*ag,normalized(k)*bg
|
|
joint=ag[r:]*bg[r:]
|
|
def factored():return consolidate(q,k,ag,bg,r,'q_to_k',joint=joint)
|
|
for _ in range(3):baseline();factored()
|
|
timing=[]
|
|
for repeat in range(10):
|
|
for label,fn in (('baseline',baseline),('factored',factored)) if repeat%2==0 else (('factored',factored),('baseline',baseline)):
|
|
lm.guard();t=time.perf_counter()
|
|
for _ in range(20):fn()
|
|
timing.append({'repeat':repeat,'implementation':label,'seconds_per_call':(time.perf_counter()-t)/20})
|
|
save_csv(root/'gain-factorization-kernel-timing.csv',timing)
|
|
report.update(state='completed',maximum_score_abs_error=max(x['score_max_abs_error'] for x in records),maximum_probability_abs_error=max(x['probability_max_abs_error'] for x in records),
|
|
bidirectional_heldout_pass=all(x['passed'] for x in report['interventions']),
|
|
elementwise_gain_multiplies_saved_per_token_layer={'q_to_k':1536,'k_to_q':384},
|
|
original_bf16_payload_bytes=1024,factored_lossless_mixed_payload_bytes=1024,
|
|
fp32_cached_gain_payload_before=2048,fp32_cached_gain_payload_after=1280,
|
|
claim_limit='Exact real-arithmetic identity only on non-RoPE coordinates; empirical FP32/NLL checks on tiny text pilot. Not evidence that two existing full layers interchange, not model-wide speedup.')
|
|
except Exception as exc:
|
|
report.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
|
|
finally:
|
|
elapsed=time.monotonic()-start
|
|
report.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.gain_factorization','seconds':elapsed});write_json(ledger_path,ledger)
|
|
write_json(root/'gain-factorization-summary.json',report)
|
|
|
|
|
|
if __name__=='__main__':main()
|