63 lines
3.1 KiB
Python
63 lines
3.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))
|
|
|
|
|
|
if __name__=='__main__':unittest.main()
|