| """Stratified weight reconstruction and GPU matmul probe, not an image-quality score.""" |
| import argparse |
| import gc |
| import json |
| import time |
| from pathlib import Path |
| import mlx.core as mx |
| from scripts.common import SOURCE, eligible, tensors, write_json |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--source', type=Path, default=SOURCE) |
| ap.add_argument('--output', default='artifacts/precision-probe.json') |
| args = ap.parse_args() |
| results = [] |
| for component, layers in [('transformer', ('0', '15', '31')), ('text_encoder', ('0','17','35'))]: |
| for path, name, info in tensors(args.source / component): |
| if not eligible(component, name, info['shape']): |
| continue |
| if not any(f'.{i}.' in name for i in layers): |
| continue |
| if not any(x in name for x in ('to_q.weight', 'img_mlp.out.weight', 'q_proj.weight', 'down_proj.weight')): |
| continue |
| all_weights = mx.load(str(path)) |
| |
| full = all_weights[name] |
| weight = full[mx.linspace(0, full.shape[0]-1, 256).astype(mx.int32)] |
| mx.eval(weight) |
| del full, all_weights |
| wf = weight.astype(mx.float32) |
| mx.random.seed(2026) |
| x = mx.random.normal((1,256,weight.shape[-1])).astype(mx.bfloat16) |
| expected = x @ weight.T |
| mx.eval(expected) |
| for bits in (4,6,8): |
| q,s,b = mx.quantize(weight, bits=bits, group_size=64) |
| rec = mx.dequantize(q,s,b,bits=bits,group_size=64).astype(mx.float32) |
| err = float(mx.sqrt(mx.mean((rec-wf)**2) / mx.mean(wf**2)).item()) |
| fn = lambda: mx.quantized_matmul(x,q,s,b,bits=bits,group_size=64,transpose=True) |
| for _ in range(3): mx.eval(fn()) |
| start=time.perf_counter() |
| for _ in range(10): mx.eval(fn()) |
| ms=(time.perf_counter()-start)*100 |
| y=fn().astype(mx.float32); e=expected.astype(mx.float32) |
| out_err=float(mx.sqrt(mx.mean((y-e)**2)/mx.mean(e**2)).item()) |
| row=dict(component=component,tensor=name,source_shape=info['shape'],sample_rows=256, |
| bits=bits,weight_relative_rmse=err,synthetic_matmul_relative_rmse=out_err,matmul_ms=ms) |
| results.append(row) |
| print(json.dumps(row),flush=True) |
| del weight,wf,x,expected,q,s,b,rec,y,e |
| gc.collect(); mx.clear_cache() |
| write_json(args.output,dict(group_size=64,mode='affine',device=mx.metal.device_info(),results=results, |
| limitations='Sampled rows and synthetic inputs; not calibrated activations, image quality or end-to-end speed.')) |
|
|
| if __name__ == '__main__': main() |
|
|
