122 lines
6.8 KiB
Python
122 lines
6.8 KiB
Python
"""Remote-only independent checks for matching, head decomposition and interventions."""
|
|
import itertools
|
|
import json
|
|
from pathlib import Path
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
import numpy as np
|
|
from crossqkv.core import (QKVModel,InterventionLM,assignment,inverse,distribution_metrics,
|
|
output_metrics,cross_head_output_metrics,value_gram_cka)
|
|
from functional.model import StreamingLM
|
|
|
|
|
|
def fixture(directory):
|
|
rng=np.random.default_rng(31);arrays={}
|
|
cfg={'hidden_size':16,'head_dim':4,'num_attention_heads':4,'num_key_value_heads':2,'rms_norm_eps':1e-6,
|
|
'rope_parameters':{'rope_theta':10000000.,'partial_rotary_factor':.5}}
|
|
for layer in (0,1):
|
|
p=f'model.language_model.layers.{layer}.self_attn.'
|
|
for role,shape in [('q_proj',(32,16)),('k_proj',(8,16)),('v_proj',(8,16)),('o_proj',(16,16)),('q_norm',(4,)),('k_norm',(4,))]:
|
|
arrays[p+role+'.weight']=(rng.standard_normal(shape)*.15).astype(np.float32)
|
|
header={};offset=0;body=b''
|
|
for name,arr in arrays.items():
|
|
raw=arr.tobytes();header[name]={'shape':list(arr.shape),'dtype':'F32','data_offsets':[offset,offset+len(raw)]};offset+=len(raw);body+=raw
|
|
h=json.dumps(header).encode();Path(directory,'model.safetensors').write_bytes(struct.pack('<Q',len(h))+h+body)
|
|
return cfg,rng.standard_normal((2,9,16),dtype=np.float32)
|
|
|
|
|
|
class CrossQKVTests(unittest.TestCase):
|
|
def test_assignment_against_exhaustive_and_inverse(self):
|
|
rng=np.random.default_rng(11);cost=rng.uniform(size=(4,4))
|
|
p=assignment(cost)
|
|
best=min((sum(cost[i,j] for i,j in enumerate(perm)),perm) for perm in itertools.permutations(range(4)))
|
|
self.assertEqual(tuple(p),best[1]);self.assertEqual([p[inverse(p)[i]] for i in range(4)],list(range(4)))
|
|
|
|
def test_probability_identity_and_nontrivial_queries(self):
|
|
rng=np.random.default_rng(12);p=np.zeros((2,3,9,9),dtype=np.float32)
|
|
for t in range(9):
|
|
x=rng.uniform(size=(2,3,t+1));p[:,:,t,:t+1]=x/x.sum(-1,keepdims=True)
|
|
metrics=distribution_metrics(p,p)
|
|
np.testing.assert_allclose(metrics['js'],0,atol=1e-14);np.testing.assert_array_equal(metrics['top5'],np.ones(2))
|
|
changed=p.copy();changed[:,:,:7]=0
|
|
# Primary query rows do not include the inevitable first row or t<7.
|
|
np.testing.assert_allclose(distribution_metrics(p,changed)['js'],0,atol=1e-14)
|
|
|
|
def test_profile_matches_canonical_and_head_sum(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
cfg,x=fixture(d);lm=QKVModel(d,cfg,enforce_memory=False)
|
|
for layer in (0,1):
|
|
parts=lm.parts(x,layer);expected,p=StreamingLM.full_attention(lm,x,layer)
|
|
np.testing.assert_allclose(parts['output'],expected,atol=2e-6,rtol=2e-6)
|
|
np.testing.assert_allclose(parts['p'],p,atol=2e-7)
|
|
np.testing.assert_allclose(parts['heads'].sum(2),expected,atol=2e-6,rtol=2e-6)
|
|
np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qk_map',perm=list(range(4))),expected,atol=2e-6)
|
|
np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qk_head',1,1),expected,atol=2e-6)
|
|
np.testing.assert_allclose(lm.transplant(parts,parts,layer,'v_group',1,1),expected,atol=2e-6)
|
|
np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qkv_head',1,1),expected,atol=2e-6)
|
|
|
|
def test_recipient_value_gate_output_and_gqa_group_stay_fixed(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
cfg,x=fixture(d);lm=QKVModel(d,cfg,enforce_memory=False);a,b=lm.parts(x,0),lm.parts(x,1)
|
|
v=a['v'].copy();v[:,:,1]=b['v'][:,:,0]
|
|
context=np.einsum('bhts,bshd->bthd',a['p'],np.repeat(v,2,axis=2))
|
|
expected=lm.project(context,a['gate'],0)
|
|
np.testing.assert_allclose(lm.transplant(a,b,0,'v_group',1,0),expected,atol=2e-6)
|
|
p=a['p'].copy();p[:,2]=b['p'][:,1]
|
|
context=np.einsum('bhts,bshd->bthd',p,np.repeat(a['v'],2,axis=2))
|
|
np.testing.assert_allclose(lm.transplant(a,b,0,'qk_head',2,1),lm.project(context,a['gate'],0),atol=2e-6)
|
|
|
|
def test_head_matrix_metrics_and_gram_are_diagnostic(self):
|
|
rng=np.random.default_rng(13);x=rng.standard_normal((2,9,3,6),dtype=np.float32)
|
|
y=rng.standard_normal((2,9,2,6),dtype=np.float32);m=cross_head_output_metrics(x,y)
|
|
for a in range(3):
|
|
for b in range(2):
|
|
ref=output_metrics(x[:,:,a],y[:,:,b])
|
|
for key in ref:np.testing.assert_allclose(m[key][:,a,b],ref[key],atol=1e-12)
|
|
v=rng.standard_normal((2,9,6),dtype=np.float32);r=np.eye(6,dtype=np.float32)[::-1]
|
|
np.testing.assert_allclose(value_gram_cka(v,v@r),np.ones(2),atol=1e-12)
|
|
self.assertGreater(float(np.linalg.norm(v-v@r)),1.)
|
|
|
|
def test_cached_mode_rejects_future_state_reuse(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
cfg,x=fixture(d);lm=InterventionLM(d,cfg,enforce_memory=False)
|
|
c={'i':0,'j':1,'kind':'qk_head','a':2,'b':1,'perm':None}
|
|
with self.assertRaises(AssertionError):lm.configure(c,reverse=False,cached=True)
|
|
lm.configure(c,reverse=True,cached=True)
|
|
lm.full_attention(x,0);out,_=lm.full_attention(x,1)
|
|
own=lm.parts(x,1);expected=lm.transplant(own,lm.donor_cache,1,'qk_head',1,2)
|
|
np.testing.assert_allclose(out,expected,atol=2e-6)
|
|
|
|
|
|
def test_controls_cover_multiple_projection_layers(self):
|
|
from crossqkv.run import controls_for_profiles
|
|
rng=np.random.default_rng(20)
|
|
h=rng.standard_normal((2,9,16),dtype=np.float32)
|
|
probs=np.zeros((2,4,9,9),dtype=np.float32)
|
|
for t in range(9):probs[:,:,t,:t+1]=1/(t+1)
|
|
profiles={layer:{'p':probs.copy(),'v':rng.standard_normal((2,9,2,4),dtype=np.float32)} for layer in (0,1)}
|
|
rows=[];controls_for_profiles(profiles,0,[{'id':'one'},{'id':'two'}],rows,h)
|
|
self.assertEqual({r['projection_layer'] for r in rows if 'projection_layer' in r},{0,1})
|
|
self.assertTrue(any(r['control']=='value_gram_random_projection_pair' for r in rows))
|
|
for r in rows:
|
|
for value in r.values():
|
|
if isinstance(value,float):self.assertTrue(np.isfinite(value))
|
|
|
|
|
|
def test_topk_tail_rank_can_change_with_tiny_probability_difference(self):
|
|
from crossqkv.conditional import tv_metrics
|
|
a=np.zeros((1,1,9,9),dtype=np.float64);b=a.copy()
|
|
a[0,0,0,0]=b[0,0,0,0]=1.
|
|
for t in range(1,9):
|
|
delta=np.linspace(-1e-6,1e-6,t)
|
|
a[0,0,t,:t]=.04/t+delta;b[0,0,t,:t]=.04/t-delta
|
|
a[0,0,t,t]=b[0,0,t,t]=.96
|
|
metric=distribution_metrics(a,b);tv=tv_metrics(a,b)
|
|
self.assertLess(float(metric['js'][0]),1e-6)
|
|
self.assertLess(float(tv['total_variation'][0]),1e-4)
|
|
self.assertLess(float(metric['top5'][0]),.8)
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|