Spaces:
Running
Running
File size: 3,729 Bytes
4ac3af3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """Compare cached ONNX inference against the official PyTorch trajectories."""
import json
from pathlib import Path
import time
import numpy as np
import onnxruntime as ort
out = Path('validation')
fixtures = json.loads((out / 'fixtures.json').read_text())
bf16 = np.load(out / 'reference-bf16.npz')
fp16 = np.load(out / 'reference-fp16.npz')
def metrics(reference, actual):
r = reference.astype(np.float64); a = actual.astype(np.float64)
p = np.exp(r - r.max(-1, keepdims=True)); p /= p.sum(-1, keepdims=True)
q = np.exp(a - a.max(-1, keepdims=True)); q /= q.sum(-1, keepdims=True)
return {'rmse': float(np.sqrt(np.mean((a-r)**2))),
'max_abs': float(np.max(np.abs(a-r))),
'cosine': float(np.mean(np.sum(a*r,-1)/(np.linalg.norm(a,axis=-1)*np.linalg.norm(r,axis=-1)))),
'argmax_agreement': float(np.mean(np.argmax(a,-1)==np.argmax(r,-1))),
'mean_kl_reference_to_actual': float(np.mean(np.sum(p*np.log((p+1e-30)/(q+1e-30)),axis=-1)))}
reports = {}
for name, path in [('fp16', 'models/own-fp16/model.onnx'),
('int4', 'models/own-int4/model.onnx'),
('packaged', 'models/minicpm5-webgpu/onnx/model_q4f16.onnx')]:
options = ort.SessionOptions(); options.intra_op_num_threads = 8; options.inter_op_num_threads = 1
start = time.time()
session = ort.InferenceSession(path, sess_options=options, providers=['CPUExecutionProvider'])
print('Loaded', name, time.time()-start, flush=True)
output_names = [o.name for o in session.get_outputs()]
samples = {}
for fixture in fixtures:
cache = {i.name: np.zeros((1,2,0,128), dtype=np.float16)
for i in session.get_inputs() if i.name.startswith('past_key_values')}
ids = np.array([fixture['input_ids']], dtype=np.int64)
total = ids.shape[1]
logits = []
for next_id in fixture['continuation']:
feeds = {'input_ids': ids, 'attention_mask': np.ones((1,total),dtype=np.int64), **cache}
results = dict(zip(output_names, session.run(None, feeds)))
assert results['logits'].shape == (1,1,130560)
assert np.isfinite(results['logits']).all()
logits.append(results['logits'][0,-1].astype(np.float32))
cache = {key.replace('present', 'past_key_values'): value
for key,value in results.items() if key.startswith('present')}
ids = np.array([[next_id]],dtype=np.int64); total += 1
samples[fixture['id']] = np.stack(logits)
print(name, fixture['id'], 'checked', len(logits), 'steps', flush=True)
np.savez(out / f'onnx-{name}.npz', **samples)
actual = np.concatenate(list(samples.values()))
reports[name] = {'vs_bf16': metrics(np.concatenate([bf16[f['id']] for f in fixtures]),actual),
'vs_fp16': metrics(np.concatenate([fp16[f['id']] for f in fixtures]),actual),
'positions': len(actual), 'elapsed_s': time.time()-start}
if name == 'packaged':
raw = np.load(out/'onnx-int4.npz')
reports[name]['vs_unmodified_export'] = metrics(np.concatenate([raw[f['id']] for f in fixtures]), actual)
assert all(np.array_equal(samples[f['id']],raw[f['id']]) for f in fixtures), 'Packaging changed logits'
print(json.dumps(reports[name]), flush=True)
(out/'onnx-validation.json').write_text(json.dumps(reports,indent=2))
del session, cache, results
assert reports['fp16']['vs_fp16']['cosine'] > 0.999, 'Unquantized export deviates from PyTorch'
assert reports['fp16']['vs_fp16']['mean_kl_reference_to_actual'] < 0.01, 'Unquantized export KL too high'
print('Export parity and lossless packaging checks passed', flush=True)
|