143 lines
7.7 KiB
Python
143 lines
7.7 KiB
Python
"""Remote-only numerical comparison with pinned official Torch definitions.
|
|
|
|
Reads only a known official Transformers source file; extracts named mathematical
|
|
operators/classes via AST, removing optional kernel decorators. Never imports the
|
|
complete model or executes checkpoint repository code. No compilation/build.
|
|
"""
|
|
import argparse
|
|
import ast
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import numpy as np
|
|
from .model import StreamingLM, delta_rule, rope, norm
|
|
|
|
EXPECTED_COMMIT = 'fd3cf0140de03f7619ac95b5cfed05b711ff9e22'
|
|
NAMES = {'apply_mask_to_padding_states', 'l2norm', 'torch_chunk_gated_delta_rule', 'torch_recurrent_gated_delta_rule',
|
|
'causal_conv1d_fn', 'Qwen3_5RMSNormGated', 'Qwen3_5GatedDeltaNet',
|
|
'rotate_half', 'apply_rotary_pos_emb', 'repeat_kv', 'eager_attention_forward',
|
|
'Qwen3_5Attention', 'Qwen3_5MLP', 'Qwen3_5RMSNorm', 'Qwen3_5DecoderLayer',
|
|
'Qwen3_5TextRotaryEmbedding'}
|
|
|
|
|
|
def official_namespace(source):
|
|
import torch
|
|
import torch.nn.functional as F
|
|
tree = ast.parse(source)
|
|
nodes = []
|
|
for node in tree.body:
|
|
if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in NAMES:
|
|
for child in ast.walk(node):
|
|
if isinstance(child, (ast.FunctionDef, ast.ClassDef)):
|
|
child.decorator_list = [d for d in child.decorator_list if isinstance(d, ast.Name) and d.id in ('staticmethod','classmethod')]
|
|
nodes.append(node)
|
|
assert {x.name for x in nodes} == NAMES
|
|
tree = ast.Module(body=[ast.ImportFrom(module='__future__',names=[ast.alias(name='annotations')],level=0),*nodes],type_ignores=[])
|
|
ns = {'torch':torch, 'nn':torch.nn, 'F':F, 'ACT2FN':{'silu':F.silu},
|
|
'GradientCheckpointingLayer':torch.nn.Module,
|
|
'is_torchdynamo_exporting':lambda:False,
|
|
'ALL_ATTENTION_FUNCTIONS':SimpleNamespace(get_interface=lambda name, fallback:fallback)}
|
|
exec(compile(ast.fix_missing_locations(tree),'pinned_transformers_reference','exec'),ns)
|
|
return ns
|
|
|
|
|
|
class ArrayLM(StreamingLM):
|
|
def __init__(self, arrays, config):
|
|
self.arrays, self.config, self.eps = arrays,config,config['rms_norm_eps']
|
|
def guard(self):
|
|
pass
|
|
def small(self,name):
|
|
return self.arrays[name]
|
|
def linear(self,x,name,rows=256):
|
|
return x@self.arrays[name+'.weight'].T
|
|
def embedding(self,ids):
|
|
return self.arrays['embed_tokens.weight'][ids]
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--output',type=Path,required=True)
|
|
a=p.parse_args()
|
|
import torch
|
|
torch.set_num_threads(1)
|
|
torch.manual_seed(42)
|
|
provenance=json.loads(Path(__file__).with_name('reference.json').read_text())
|
|
assert hashlib.sha256(a.source.read_bytes()).hexdigest()==provenance['sha256'], 'Reference source hash mismatch'
|
|
ns=official_namespace(a.source.read_text())
|
|
cfg={'hidden_size':16,'num_attention_heads':2,'num_key_value_heads':1,'head_dim':8,
|
|
'intermediate_size':24,'hidden_act':'silu','attention_bias':False,'attention_dropout':0.,
|
|
'linear_num_value_heads':2,'linear_num_key_heads':2,'linear_key_head_dim':4,
|
|
'linear_value_head_dim':4,'linear_conv_kernel_dim':4,'rms_norm_eps':1e-6,
|
|
'layer_types':['linear_attention','full_attention'],'max_position_embeddings':4096,
|
|
'rope_parameters':{'rope_type':'default','rope_theta':10000000.,'partial_rotary_factor':.5,'mrope_section':[1,0,0]},
|
|
'_attn_implementation':'eager'}
|
|
config=SimpleNamespace(**cfg)
|
|
layers=[ns['Qwen3_5DecoderLayer'](config,i).eval() for i in range(2)]
|
|
for layer in layers:
|
|
for name,t in layer.named_parameters():
|
|
with torch.no_grad():
|
|
t.uniform_(-.2,.2)
|
|
if name=='linear_attn.norm.weight':t.add_(1.)
|
|
arrays={f'layers.{i}.'+k:v.detach().numpy().copy() for i,l in enumerate(layers) for k,v in l.state_dict().items()}
|
|
lm=ArrayLM(arrays,cfg)
|
|
rng=np.random.default_rng(123)
|
|
x=rng.standard_normal((2,7,16),dtype=np.float32)
|
|
tx=torch.from_numpy(x)
|
|
positions=torch.arange(7)[None,:]
|
|
rotary=ns['Qwen3_5TextRotaryEmbedding'](config)
|
|
cos,sin=rotary(tx,positions)
|
|
mask=torch.triu(torch.full((7,7),-float('inf')),diagonal=1)[None,None]
|
|
evidence=[]
|
|
def check(name,actual,expected,atol=2e-5,rtol=2e-5):
|
|
actual=np.asarray(actual);expected=expected.detach().numpy() if hasattr(expected,'detach') else np.asarray(expected)
|
|
np.testing.assert_allclose(actual,expected,atol=atol,rtol=rtol)
|
|
evidence.append({'test':name,'shape':list(actual.shape),'max_abs_error':float(np.max(np.abs(actual-expected))),'atol':atol,'rtol':rtol,'passed':True})
|
|
with torch.no_grad():
|
|
for i,layer in enumerate(layers):
|
|
n=lm.normalization(x,f'layers.{i}.input_layernorm')
|
|
tn=layer.input_layernorm(torch.from_numpy(x))
|
|
check(f'layer{i}_rmsnorm',n,tn)
|
|
if i==0:
|
|
out=lm.linear_attention(n,i)
|
|
tout=layer.linear_attn(tn)
|
|
else:
|
|
out,pr=lm.full_attention(n,i)
|
|
tout,tpr=layer.self_attn(tn,(cos,sin),mask)
|
|
check('attention_probabilities',pr,tpr)
|
|
check(f'layer{i}_mixer',out,tout)
|
|
h=x+out
|
|
n=lm.normalization(h,f'layers.{i}.post_attention_layernorm')
|
|
check(f'layer{i}_mlp',lm.mlp(n,i),layer.mlp(layer.post_attention_layernorm(torch.from_numpy(h))))
|
|
actual=h+lm.mlp(n,i)
|
|
expected=layer(torch.from_numpy(x),(cos,sin),attention_mask=mask if i else None)
|
|
check(f'layer{i}_complete',actual,expected)
|
|
# Independent recurrence/chunk equality, including nonmultiple sequence length.
|
|
for length in (1,7,33):
|
|
q=rng.standard_normal((2,length,2,4),dtype=np.float32)
|
|
k=rng.standard_normal(q.shape,dtype=np.float32)
|
|
v=rng.standard_normal((2,length,2,5),dtype=np.float32)
|
|
g=-np.abs(rng.standard_normal((2,length,2),dtype=np.float32))
|
|
beta=rng.uniform(size=g.shape).astype(np.float32)
|
|
out,state=delta_rule(q,k,v,g,beta)
|
|
args=[torch.from_numpy(z) for z in (q,k,v,g,beta)]
|
|
for kind in ('torch_recurrent_gated_delta_rule','torch_chunk_gated_delta_rule'):
|
|
ref,refstate=ns[kind](*args,output_final_state=True,use_qk_l2norm_in_kernel=True)
|
|
check(f'{kind}_length{length}_output',out,ref)
|
|
check(f'{kind}_length{length}_state',state,refstate)
|
|
# Full two-layer residual forward, final norm and embedding lookup.
|
|
arrays['embed_tokens.weight']=rng.standard_normal((31,16),dtype=np.float32)
|
|
arrays['norm.weight']=rng.standard_normal((16,),dtype=np.float32)*.1
|
|
ids=rng.integers(0,31,size=(2,7))
|
|
h=torch.from_numpy(arrays['embed_tokens.weight'][ids])
|
|
for i,l in enumerate(layers):h=l(h,(cos,sin),attention_mask=mask if i else None)
|
|
expected=h*torch.rsqrt(h.square().mean(-1,keepdim=True)+1e-6)*(1+torch.from_numpy(arrays['norm.weight']))
|
|
actual,_=lm.forward(ids)
|
|
check('complete_two_layer_text_forward',actual,expected)
|
|
result={'passed':True,'reference_commit':EXPECTED_COMMIT,'source_sha256':hashlib.sha256(a.source.read_bytes()).hexdigest(),
|
|
'torch':torch.__version__,'numpy':np.__version__,'implementation_sha256':hashlib.sha256(Path(__file__).with_name('model.py').read_bytes()).hexdigest(),'tests':evidence,'limitations':'FP32 CPU semantics on tiny synthetic configurations; not full-checkpoint BF16 kernel equivalence'}
|
|
a.output.write_text(json.dumps(result,indent=2))
|
|
print(json.dumps({'passed':True,'checks':len(evidence),'maximum_absolute_discrepancy':max(x['max_abs_error'] for x in evidence)}))
|
|
|
|
|
|
if __name__=='__main__':main()
|