| |
| """Apple Silicon AuroraASR qualification on a pinned acoustic tensor. |
| |
| Run each backend in an independent fresh process, interleaving benchmark arms. |
| Produces a JSON record and fails numeric/transcription gates. Prints requested |
| Core ML compute unit only; ANE placement requires MLComputePlan or Instruments. |
| """ |
| import argparse |
| import json |
| import os |
| import platform |
| import resource |
| import statistics |
| import time |
| from collections import Counter |
| from pathlib import Path |
| import numpy as np |
|
|
| p=argparse.ArgumentParser() |
| p.add_argument('--backend',choices=['mlx','onnx_cpu','coreml'],required=True) |
| p.add_argument('--model-dir',type=Path,required=True) |
| p.add_argument('--runs',type=int,default=40) |
| p.add_argument('--precision',choices=['fp32','fp16'],default='fp32') |
| p.add_argument('--compute-unit',choices=['all','cpu','gpu'],default='all') |
| p.add_argument('--dynamic',action='store_true') |
| a=p.parse_args() |
| if a.dynamic and a.backend=='coreml': |
| raise SystemExit('--dynamic applies to ONNX, not the static Core ML package') |
| if a.backend!='coreml' and a.precision!='fp32': |
| raise SystemExit('The ONNX exports are FP32; --precision fp16 applies only to Core ML') |
| if platform.system() != 'Darwin' or platform.machine() != 'arm64': |
| raise SystemExit('Requires macOS Apple Silicon (arm64).') |
| fix=np.load(a.model_dir/'reference_fixture.npz') |
| x=fix['acoustic_input'];expected=fix['reference_logits'] |
| assert x.dtype==np.float32 and expected.ndim==3 |
| mac_profile=None |
| provider_info={} |
| start=time.perf_counter() |
| if a.backend in ('mlx','onnx_cpu'): |
| import onnxruntime as ort |
| providers=['CPUExecutionProvider'] |
| if a.backend=='mlx': |
| import onnxruntime_ep_mlx |
| onnxruntime_ep_mlx.register_execution_provider_library() |
| providers=['MLXExecutionProvider','CPUExecutionProvider'] |
| opts=ort.SessionOptions() |
| opts.enable_profiling=a.backend=='mlx' |
| fname='aurora_dynamic.onnx' if a.dynamic else 'aurora_asr.onnx' |
| sess=ort.InferenceSession(str(a.model_dir/fname),opts,providers=providers) |
| if a.backend=='mlx' and 'MLXExecutionProvider' not in sess.get_providers(): |
| raise RuntimeError('MLX provider not active. Do not call this an MLX benchmark.') |
| name=sess.get_inputs()[0].name |
| def run(): |
| return sess.run(None,{name:x})[0] |
| provider_info['registered_providers']=sess.get_providers() |
| else: |
| import coremltools as ct |
| package=a.model_dir/f'aurora_{a.precision}.mlpackage' |
| if not package.is_dir() or not (package/'Manifest.json').exists(): |
| raise SystemExit( |
| f'Core ML package missing: {package.resolve()}\n' |
| 'Use the sibling AuroraASR-1M-CoreML directory; AuroraASR-1M-MLX ' |
| 'contains ONNX weights only. Do not mix --backend coreml with an MLX --model-dir.') |
| unit={'all':ct.ComputeUnit.ALL,'cpu':ct.ComputeUnit.CPU_ONLY, |
| 'gpu':ct.ComputeUnit.CPU_AND_GPU}[a.compute_unit] |
| model=ct.models.MLModel(str(package),compute_units=unit) |
| inname=model.get_spec().description.input[0].name |
| outname=model.get_spec().description.output[0].name |
| def run(): |
| return np.asarray(model.predict({inname:x})[outname]) |
| provider_info.update({'coreml_package':package.name,'coreml_requested_compute_unit':a.compute_unit, |
| 'actual_ane_placement':'NOT_VERIFIED'}) |
| load_ms=(time.perf_counter()-start)*1000 |
| for _ in range(8): |
| out=np.asarray(run()) |
| if out.shape!=expected.shape: |
| raise RuntimeError(f'Output shape mismatch: got {out.shape}, expected {expected.shape}') |
| times=[] |
| for _ in range(max(1,a.runs)): |
| start=time.perf_counter();out=np.asarray(run());times.append((time.perf_counter()-start)*1000) |
| absdiff=np.abs(out-expected) |
| max_abs=float(absdiff.max());mean_abs=float(absdiff.mean()) |
| agreement=float(np.mean(out.argmax(axis=-1)==expected.argmax(axis=-1))) |
| |
| atol=0.05 if (a.backend=='coreml' and a.precision=='fp16') else 0.005 |
| rtol=0.05 if (a.backend=='coreml' and a.precision=='fp16') else 0.005 |
| numerical=bool(np.allclose(out,expected,atol=atol,rtol=rtol) and agreement>=0.995) |
| try: |
| from aurora_runtime import ctc_text |
| config=json.loads((a.model_dir/'upstream_reference'/'config.json').read_text()) |
| spm=a.model_dir/'upstream_reference'/'bpe256.model' |
| transcript_ref=ctc_text(expected,spm,config.get('blank_token_id',0)) |
| transcript_got=ctc_text(out,spm,config.get('blank_token_id',0)) |
| transcript_exact=bool(transcript_ref==transcript_got) |
| except Exception as exc: |
| transcript_ref=None;transcript_got=None;transcript_exact=False |
| provider_info['transcript_error']=f'{type(exc).__name__}: {exc}' |
| node_providers={};node_durations_us={};mlx_claimed=False |
| if a.backend=='mlx': |
| profile=Path(sess.end_profiling()) |
| try: |
| events=json.loads(profile.read_text()) |
| counts=Counter(); durations=Counter() |
| for e in events: |
| if e.get('cat')!='Node':continue |
| provider=e.get('args',{}).get('provider') |
| if not provider: continue |
| counts[str(provider)]+=1 |
| durations[str(provider)]+=float(e.get('dur',0)) |
| node_providers=dict(counts) |
| node_durations_us=dict(durations) |
| mlx_claimed=counts.get('MLXExecutionProvider',0)>0 |
| finally: |
| profile.unlink(missing_ok=True) |
| |
| peak_rss_mib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1048576 |
| import hashlib |
| model_input_path=(a.model_dir/('aurora_dynamic.onnx' if a.dynamic else 'aurora_asr.onnx')) if a.backend!='coreml' else package |
| if model_input_path.is_file(): |
| model_sha256=hashlib.sha256(model_input_path.read_bytes()).hexdigest() |
| else: |
| model_sha256=hashlib.sha256(b''.join( |
| str(p.relative_to(model_input_path)).encode('utf-8')+b'\x00'+hashlib.sha256(p.read_bytes()).digest() |
| for p in sorted(model_input_path.rglob('*')) if p.is_file())).hexdigest() |
| summary={'backend':a.backend,'precision':a.precision,'dynamic':bool(a.dynamic), |
| 'model_file':model_input_path.name,'model_sha256':model_sha256,'shape':list(x.shape), |
| 'cold_load_ms':load_ms,'p50_warm_ms':float(np.median(times)), |
| 'p95_warm_ms':float(np.percentile(times,95)), |
| 'max_abs_logit_error':max_abs,'mean_abs_logit_error':mean_abs, |
| 'frame_argmax_agreement':agreement,'allclose_atol_rtol':[atol,rtol], |
| 'reference_transcript':transcript_ref,'actual_transcript':transcript_got, |
| 'transcript_exact':transcript_exact,'peak_process_rss_mib':peak_rss_mib, |
| 'ep_node_events':node_providers,'ep_node_duration_us':node_durations_us, |
| 'mlx_claimed_any_subgraph':mlx_claimed if a.backend=='mlx' else None, |
| 'provider_details':provider_info, 'numeric_pass':numerical} |
| filename=f"mac_{a.backend}_{a.precision}_{a.compute_unit}_{'dynamic' if a.dynamic else 'static'}_v4_qualification.json" |
| (a.model_dir/filename).write_text(json.dumps(summary,indent=2)) |
| print(json.dumps(summary,indent=2)) |
| if not numerical:raise SystemExit('FAIL numeric gate') |
| if not transcript_exact:raise SystemExit('FAIL CTC transcript gate / missing tokenizer') |
| if a.backend=='mlx' and not mlx_claimed:raise SystemExit('FAIL no MLX subgraph proved') |
| print('PASS: numeric & decoded CTC parity. ANE device placement remains unverified.') |
|
|
