79 lines
5.8 KiB
Python
79 lines
5.8 KiB
Python
"""Independent algebra, mixed-layer canonical, and joint-mask intervention checks."""
|
|
import json,struct,tempfile,unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
import numpy as np
|
|
from dominance.core import InfluenceLM,concentration,masks_from_influences,js_rows
|
|
from dominance.run import distances,link
|
|
from functional.model import StreamingLM,delta_rule,attention
|
|
from crossqkv.core import QKVModel
|
|
|
|
|
|
def fixture(path,kinds):
|
|
rng=np.random.default_rng(922);h=16;nh=4;dim=4
|
|
c={'rms_norm_eps':1e-6,'hidden_size':h,'head_dim':dim,'num_attention_heads':nh,'num_key_value_heads':2,'rope_parameters':{'rope_theta':10000000.,'partial_rotary_factor':.5},'linear_num_key_heads':2,'linear_num_value_heads':4,'linear_key_head_dim':4,'linear_value_head_dim':4,'layer_types':kinds}
|
|
a={}
|
|
def add(name,shape):a['model.language_model.'+name]=(rng.standard_normal(shape)*.15).astype(np.float32)
|
|
add('embed_tokens.weight',(19,h));add('norm.weight',(h,))
|
|
for l,kind in enumerate(kinds):
|
|
p=f'layers.{l}.'
|
|
for k in ('input_layernorm','post_attention_layernorm'):add(p+k+'.weight',(h,))
|
|
for k,shape in [('gate_proj',(24,h)),('up_proj',(24,h)),('down_proj',(h,24))]:add(p+'mlp.'+k+'.weight',shape)
|
|
if kind=='full_attention':
|
|
for k,shape in [('q_proj',(32,h)),('k_proj',(8,h)),('v_proj',(8,h)),('o_proj',(h,16)),('q_norm',(4,)),('k_norm',(4,))]:add(p+'self_attn.'+k+'.weight',shape)
|
|
else:
|
|
for k,shape in [('in_proj_qkv',(32,h)),('in_proj_z',(16,h)),('in_proj_a',(4,h)),('in_proj_b',(4,h)),('out_proj',(h,16)),('conv1d',(32,1,4)),('norm',(4,))]:add(p+'linear_attn.'+k+'.weight',shape)
|
|
add(p+'linear_attn.A_log',(4,));add(p+'linear_attn.dt_bias',(4,))
|
|
header={};body=b''
|
|
for k,v in a.items():
|
|
raw=v.tobytes();header[k]={'shape':list(v.shape),'dtype':'F32','data_offsets':[len(body),len(body)+len(raw)]};body+=raw
|
|
j=json.dumps(header).encode();Path(path,'model.safetensors').write_bytes(struct.pack('<Q',len(j))+j+body)
|
|
return c
|
|
|
|
|
|
class DominanceTests(unittest.TestCase):
|
|
def test_continuous_limits_and_undefined_zero(self):
|
|
x=concentration([[1,1,1,1],[4,0,0,0],[0,0,0,0]])
|
|
np.testing.assert_allclose(x['effective_heads'][:2],[4,1]);np.testing.assert_allclose(x['concentration'][:2],[0,1]);self.assertFalse(x['valid'][2]);self.assertTrue(np.isnan(x['effective_heads'][2]))
|
|
self.assertAlmostEqual(float(concentration([3.,1.])['effective_heads']),1.6)
|
|
self.assertAlmostEqual(float(concentration([30.,10.])['effective_heads']),1.6)
|
|
with self.assertRaises(ValueError):concentration([-1,2])
|
|
def test_mask_ranks_vary_by_token_and_recipient(self):
|
|
d=np.arange(4*3*2).reshape(4,3,2).astype(float);d[:,1]=d[::-1,1]
|
|
w=masks_from_influences(d);s=masks_from_influences(d,False)
|
|
np.testing.assert_array_equal(w.sum(-1),2);np.testing.assert_array_equal(w+s,1)
|
|
self.assertFalse(np.array_equal(w[:,0],w[:,1]))
|
|
def test_js_identity_and_disjoint(self):
|
|
np.testing.assert_allclose(js_rows([[1,0]],[[1,0]]),0)
|
|
np.testing.assert_allclose(js_rows([[1,0]],[[0,1]]),np.log(2))
|
|
def test_mixed_layer_forward_and_effective_qkv(self):
|
|
for kinds in [('linear_attention','full_attention'),('full_attention','linear_attention')]:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
cfg=fixture(d,list(kinds));lm=InfluenceLM(d,cfg,enforce_memory=False);ids=np.array([[1,2,3,4,5,6,7,8,9]])
|
|
cache=Path(d)/'cache';cache.mkdir();final=lm.baseline(ids,cache);ref,_=StreamingLM.forward(lm,ids)
|
|
# baseline returns pre-final-norm residual state, canonical returns normalized state.
|
|
np.testing.assert_allclose(lm.normalization(final,'norm'),ref,atol=2e-6,rtol=2e-6)
|
|
for layer in (0,1):
|
|
raw=np.load(cache/f'layer-{layer}.npz')['raw'];obs=lm.recipient(raw,layer);h=lm.normalization(raw,f'layers.{layer}.input_layernorm')
|
|
if kinds[layer]=='full_attention':
|
|
expected=QKVModel(d,cfg,enforce_memory=False).parts(h,layer,heads=False)
|
|
np.testing.assert_allclose(obs['P'],expected['p'],atol=2e-6)
|
|
np.testing.assert_allclose(obs['Q'],expected['q'],atol=2e-6)
|
|
else:
|
|
with patch('functional.model.delta_rule',wraps=delta_rule) as fn:lm.linear_attention(h,layer)
|
|
q,k,v=fn.call_args.args[:3]
|
|
q=q/np.sqrt(np.sum(q*q,-1,keepdims=True)+1e-6)/np.sqrt(cfg['linear_key_head_dim']);k=k/np.sqrt(np.sum(k*k,-1,keepdims=True)+1e-6)
|
|
for key,val in [('Q',q),('K',k),('V',v)]:np.testing.assert_allclose(obs[key],val,atol=2e-6)
|
|
def test_joint_recomputation_null_and_saved_raw(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
cfg=fixture(d,['full_attention','linear_attention']);lm=InfluenceLM(d,cfg,enforce_memory=False);ids=np.array([[1,2,3,4,5,6,7,8,9]])
|
|
cache=Path(d)/'cache';cache.mkdir();lm.baseline(ids,cache)
|
|
data=dict(np.load(cache/'layer-0.npz'));masks=np.ones((1,9,4),dtype=np.float32)
|
|
actual=lm.intervene(data,0,masks);mid=data['raw']+data['out']-data['heads'].sum(2)
|
|
next_raw=mid+lm.mlp(lm.normalization(mid,'layers.0.post_attention_layernorm'),0)
|
|
expected=lm.recipient(next_raw,1)
|
|
for k in expected:np.testing.assert_allclose(actual[k],expected[k],atol=2e-6)
|
|
info=link(lm,data,0,Path(d)/'effects.npz',99);self.assertEqual(info['source_heads'],4);self.assertLess(info['null_max'],1e-6)
|
|
raw=dict(np.load(Path(d)/'effects.npz'));self.assertEqual(raw['single_Q'].shape,(4,9,4));self.assertEqual(raw['joint_weak_half_Q'].shape,(9,4))
|
|
np.testing.assert_array_equal(raw['mask_weak_half']+raw['mask_strong_half'],1)
|