120 lines
6.1 KiB
Python
120 lines
6.1 KiB
Python
"""CPU functional implementation tests; run on authorized remote server only."""
|
|
import json
|
|
from pathlib import Path
|
|
import struct
|
|
import tempfile
|
|
import unittest
|
|
import numpy as np
|
|
from functional.model import StreamingLM,causal_conv,silu,attention,norm
|
|
from functional.run import probability_metrics,error_rows
|
|
|
|
|
|
class FunctionalTests(unittest.TestCase):
|
|
def test_streamed_projection_embedding_and_output_head(self):
|
|
rng=np.random.default_rng(7)
|
|
arrays={'model.language_model.embed_tokens.weight':rng.normal(size=(19,4)).astype(np.float32),
|
|
'model.language_model.test.weight':rng.normal(size=(9,4)).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)]}
|
|
body+=raw;offset+=len(raw)
|
|
encoded=json.dumps(header).encode()
|
|
with tempfile.TemporaryDirectory() as d:
|
|
Path(d,'model.safetensors').write_bytes(struct.pack('<Q',len(encoded))+encoded+body)
|
|
lm=StreamingLM(d,{'rms_norm_eps':1e-6},enforce_memory=False)
|
|
ids=np.array([[2,4,1,18],[7,3,12,5]])
|
|
x=lm.embedding(ids)
|
|
np.testing.assert_array_equal(x,arrays['model.language_model.embed_tokens.weight'][ids])
|
|
actual=lm.linear(x,'test',rows=2)
|
|
np.testing.assert_allclose(actual,x@arrays['model.language_model.test.weight'].T,atol=1e-6)
|
|
logits=x[:,:-1]@arrays['model.language_model.embed_tokens.weight'].T
|
|
m=logits.max(-1)
|
|
expected=m+np.log(np.exp(logits-m[...,None]).sum(-1))-np.take_along_axis(logits,ids[:,1:,None],axis=-1)[...,0]
|
|
np.testing.assert_allclose(lm.nll(x,ids,rows=3),expected,atol=2e-6)
|
|
|
|
def test_causality_and_identical_function_metrics(self):
|
|
rng=np.random.default_rng(8)
|
|
q=rng.normal(size=(1,5,2,4)).astype(np.float32)
|
|
k=rng.normal(size=(1,5,1,4)).astype(np.float32)
|
|
v=rng.normal(size=(1,5,1,4)).astype(np.float32)
|
|
y,p=attention(q,k,v)
|
|
v2=v.copy();v2[:,3:]+=100
|
|
y2,_=attention(q,k,v2)
|
|
np.testing.assert_allclose(y[:,:3],y2[:,:3])
|
|
self.assertTrue(np.all(p[...,np.triu_indices(5,1)[0],np.triu_indices(5,1)[1]]==0))
|
|
m=probability_metrics(p,p)
|
|
self.assertAlmostEqual(m['js_mean'],0)
|
|
self.assertAlmostEqual(m['top5_overlap'],1)
|
|
rows=error_rows(y,y,[{'id':'one','split':'heldout'}],{})
|
|
self.assertEqual(rows[0]['relative_l2'],0)
|
|
|
|
def test_causal_convolution_and_norm_boundary(self):
|
|
x=np.arange(1,7,dtype=np.float32).reshape(1,3,2)
|
|
kernel=np.array([[1,2],[3,4]],dtype=np.float32)
|
|
expected=x*kernel[:,-1]
|
|
expected[:,1:]+=x[:,:-1]*kernel[:,0]
|
|
np.testing.assert_allclose(causal_conv(x,kernel),silu(expected))
|
|
weight=np.zeros(2,dtype=np.float32)
|
|
np.testing.assert_allclose(norm(x,weight),x/np.sqrt(np.mean(x*x,axis=-1,keepdims=True)+1e-6))
|
|
np.testing.assert_array_equal(norm(x,weight,centered=False),np.zeros_like(x))
|
|
|
|
|
|
def test_nonrotary_gain_identity_and_rotary_negative_control(self):
|
|
from functional.model import rope
|
|
from functional.gain_factorization import consolidate
|
|
rng=np.random.default_rng(19)
|
|
q=rng.standard_normal((2,7,8,256),dtype=np.float32)
|
|
k=rng.standard_normal((2,7,2,256),dtype=np.float32)
|
|
a=rng.uniform(.5,1.5,256).astype(np.float32)
|
|
b=rng.uniform(.5,1.5,256).astype(np.float32)
|
|
def scores(q,k):return np.einsum('bthd,bshd->bhts',rope(q),np.repeat(rope(k),4,axis=2))/16
|
|
reference=scores(norm(q,a,centered=False),norm(k,b,centered=False))
|
|
for direction in ('q_to_k','k_to_q'):
|
|
q1,k1=consolidate(q,k,a,b,64,direction)
|
|
np.testing.assert_allclose(scores(q1,k1),reference,atol=3e-6,rtol=3e-5)
|
|
# Moving gains through rotated coordinates is not a valid general identity.
|
|
q1,k1=consolidate(q,k,a,b,0,'q_to_k')
|
|
self.assertGreater(float(np.max(np.abs(scores(q1,k1)-reference))),1e-3)
|
|
|
|
|
|
def test_neuron_function_metrics_and_targeted_replacement(self):
|
|
from functional.neuron_reuse import pair_metrics,NeuronReuseLM
|
|
x=np.linspace(-2,2,128,dtype=np.float64)
|
|
m=pair_metrics(x,x*1.05)
|
|
self.assertAlmostEqual(m['relative_l2_a_to_b'],.05)
|
|
self.assertAlmostEqual(m['pearson'],1.)
|
|
self.assertGreater(m['effective_energy_tokens_a'],16)
|
|
dead=pair_metrics(np.zeros(128),np.zeros(128))
|
|
self.assertEqual(dead['rms_a'],0.)
|
|
class Tiny(NeuronReuseLM):
|
|
def __init__(self):
|
|
self.arrays={'gate_proj':np.array([[1,0],[0,2],[1,1]],dtype=np.float32),
|
|
'up_proj':np.array([[2,0],[0,1],[1,-1]],dtype=np.float32),
|
|
'down_proj':np.array([[1,2,3],[-1,1,0]],dtype=np.float32)}
|
|
def linear(self,h,name,rows=256):return h@self.arrays[name.rsplit('.',1)[-1]].T
|
|
lm=Tiny();lm.configure(5,0,1)
|
|
h=np.array([[[1,2],[-1,3]]],dtype=np.float32)
|
|
z=silu(h@lm.arrays['gate_proj'].T)*(h@lm.arrays['up_proj'].T)
|
|
z[...,0]=z[...,1]
|
|
np.testing.assert_allclose(lm.mlp(h,5),z@lm.arrays['down_proj'].T)
|
|
original=silu(h@lm.arrays['gate_proj'].T)*(h@lm.arrays['up_proj'].T)
|
|
np.testing.assert_allclose(lm.mlp(h,6),original@lm.arrays['down_proj'].T)
|
|
|
|
|
|
def test_uniform_rmsnorm_function_bound_is_tight(self):
|
|
from functional.norm_repetition import uniform_bounds
|
|
a=np.array([1.2,.8,1.5,1.1],dtype=np.float64)
|
|
b=np.array([1.22,.78,1.45,1.13],dtype=np.float64)
|
|
ab,ba=uniform_bounds(a,b)
|
|
basis=np.eye(4)
|
|
# Common RMS denominator cancels; basis inputs attain the exact maximum.
|
|
ya,yb=basis*a,basis*b
|
|
ea=np.linalg.norm(ya-yb,axis=-1)/np.linalg.norm(ya,axis=-1)
|
|
eb=np.linalg.norm(ya-yb,axis=-1)/np.linalg.norm(yb,axis=-1)
|
|
self.assertAlmostEqual(float(ea.max()),ab)
|
|
self.assertAlmostEqual(float(eb.max()),ba)
|
|
self.assertIsNone(uniform_bounds(np.array([0.,1.]),np.ones(2)))
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|