File size: 3,089 Bytes
3fd1a35 | 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 | #!/usr/bin/env python3
"""Locate representable-range clipping in exported KDA W4 input matrices."""
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from audit_absolute_weights import Package, TensorSource, source_linear_names, family, blank, add, measure, finish
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--source', type=Path, required=True)
p.add_argument('--original', type=Path, required=True)
p.add_argument('--calibrated', type=Path, required=True)
p.add_argument('--output', type=Path, required=True)
a = p.parse_args()
torch.set_num_threads(2)
src = TensorSource(a.source)
packs = {v:Package(getattr(a,v)) for v in ('original','calibrated')}
results = []
for name,e in packs['original'].entries.items():
if e[4] != 3: continue
base = name.removesuffix('.weight')
if family(base) != 'kda_in': continue
offset = 0
for key in source_linear_names(base):
t = src.tensor(key)
acc = {v:blank() for v in packs}
extra = {v:dict(out_of_range=0, clipped_error_sum_sq=0., worst=None) for v in packs}
for lo in range(0,len(t),256):
hi = min(lo+256,len(t))
ref = t[lo:hi].float().numpy()
for v,pack in packs.items():
q = pack.rows(base, offset+lo, offset+hi)
s = pack.array(base+'.scales','<f4')[offset+lo:offset+hi,None]
clipped = (ref > 7*s) | (ref < -8*s)
d = q.astype(np.float64)-ref.astype(np.float64)
m = measure(ref,q)
add(acc[v],m)
extra[v]['out_of_range'] += int(clipped.sum())
extra[v]['clipped_error_sum_sq'] += float(np.square(d[clipped]).sum())
i,j = np.unravel_index(np.abs(d).argmax(),d.shape)
worst = dict(row=int(lo+i), column=int(j), original=float(ref[i,j]),
reconstructed=float(q[i,j]), scale=float(s[i,0]),
abs_error=float(abs(d[i,j])), outside_range=bool(clipped[i,j]))
if extra[v]['worst'] is None or worst['abs_error'] > extra[v]['worst']['abs_error']:
extra[v]['worst'] = worst
results.append(dict(name=key, matrix=base, metrics={v:dict(**finish(acc[v]),**extra[v]) for v in packs}))
offset += len(t)
a.output.parent.mkdir(parents=True,exist_ok=True)
a.output.write_text(json.dumps(dict(scope='KDA input only; outside [-8s,7s] inclusive FP32 boundaries', tensors=results),indent=2)+'\n')
for v in packs:
count = sum(r['metrics'][v]['count'] for r in results)
clipped = sum(r['metrics'][v]['out_of_range'] for r in results)
print(json.dumps(dict(variant=v,count=count,clipped=clipped,clipped_fraction=clipped/count,
worst=max(results,key=lambda r:r['metrics'][v]['worst']['abs_error']))),flush=True)
if __name__ == '__main__': main()
|