| """Compare converted checkpoints against pinned Mistral CPU FP32 inference.""" |
| import argparse |
| import gc |
| import importlib.metadata |
| import json |
| import platform |
| import sys |
| import subprocess |
| import time |
| from pathlib import Path |
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / "upstream" / "src")) |
| import numpy as np |
| import torch |
| import mlx.core as mx |
| from timestral import TimestralForecaster as SourceForecaster |
| from timestral_mlx import TimestralForecaster |
| from timestral_mlx.loading import sha256 |
|
|
| VARIANTS = ("5M", "28M", "123M", "209M") |
|
|
|
|
| def fixtures(): |
| rng = np.random.default_rng(20261020) |
| def signal(b, v, t): |
| x = (3*np.sin(np.arange(t)*2*np.pi/24) + rng.normal(0,.2,(b,v,t)).cumsum(-1) + 10).astype(np.float32) |
| return x |
| x = signal(1, 1, 257) |
| x[..., 10:16] = np.nan |
| yield "univariate_missing_96", x, 96, {} |
| x = signal(2, 2, 129); x[0,0,:40] = np.nan; x[1,1,60:70] = np.nan |
| yield "multivariate_batched_65", x, 65, {} |
| x = signal(1, 1, 193) |
| past,known = signal(1,2,193),signal(1,2,193+97) |
| past[0,0,20:28] = np.nan |
| known[:,1] = np.arange(193+97)%2 |
| known[0,0,210:215] = np.nan |
| yield "numeric_covariates_97", x, 97, dict(past_covariates=past,known_covariates=known) |
| codes = (np.arange(193+97)%4 + 2**26).astype(np.float64)[None,None] |
| codes[...,5:9] = np.nan |
| yield "categorical_covariates_97", x, 97, dict(categorical_past_covariates=codes[...,:193],categorical_known_covariates=codes) |
| x = signal(1,1,257) |
| yield "autoregressive_1057", x, 1057, dict(symmetric_averaging=False) |
| yield "constant_65", np.full(65, 5, np.float32), 65, {} |
| yield "all_missing_65", np.full(65, np.nan, np.float32), 65, {} |
|
|
|
|
| def validate(size): |
| print(f"Validating {size} against upstream...", flush=True) |
| src = SourceForecaster.from_pretrained(ROOT / "models" / "source" / size) |
| dst = TimestralForecaster.from_pretrained(ROOT / "models" / size) |
| records=[] |
| for name,x,h,kw in fixtures(): |
| start=time.perf_counter(); expected=src.forecast(x,h,**kw).quantiles.numpy(); source_s=time.perf_counter()-start |
| start=time.perf_counter(); actual=dst.forecast(x,h,**kw).quantiles; mlx_s=time.perf_counter()-start |
| delta=actual.astype(np.float64)-expected |
| rms=float(np.sqrt(np.mean(expected.astype(np.float64)**2))) |
| err=float(np.sqrt(np.mean(delta**2))) |
| normalized_rmse=err/max(rms,1e-6) |
| normalized_max=float(np.max(np.abs(delta)))/max(rms,1e-6) |
| |
| passed=bool(np.isfinite(actual).all() and (np.diff(actual,axis=-1)>=0).all() |
| and normalized_rmse <= 1e-4 and normalized_max <= 2e-3) |
| record=dict(fixture=name,shape=list(actual.shape),source_seconds=source_s,mlx_seconds=mlx_s, |
| max_absolute_error=float(np.max(np.abs(delta))),rmse=err,reference_rms=rms, |
| normalized_rmse=normalized_rmse,normalized_max_error=normalized_max, |
| finite=bool(np.isfinite(actual).all()),monotone=bool((np.diff(actual,axis=-1)>=0).all()),passed=passed) |
| records.append(record) |
| print(f" {name}: normalized RMSE={normalized_rmse:.3g}, max error={record['max_absolute_error']:.3g}, pass={passed}",flush=True) |
| report=dict(variant=size,passed=all(r['passed'] for r in records), |
| acceptance=dict(normalized_rmse_max=1e-4,normalized_max_error_max=2e-3), |
| bundle_manifest_sha256=sha256(ROOT/'models'/size/'manifest.json'), |
| runtime_sha256={p.name:sha256(p) for p in sorted((ROOT/'timestral_mlx').glob('*.py')) if not p.name.startswith('._')}, |
| environment=dict(python=platform.python_version(),platform=platform.platform(), |
| torch=torch.__version__,mlx=importlib.metadata.version('mlx'),numpy=np.__version__, |
| source_device='cpu',mlx_device=str(mx.default_device()),torch_threads=torch.get_num_threads()), |
| notes='Deterministic synthetic parity fixtures; timings include first-call overhead and are not a controlled benchmark. No GIFT-Eval/TIME/fev-bench evaluation was run.',fixtures=records) |
| (ROOT/'reports'/f'{size}.json').write_text(json.dumps(report,indent=2)+'\n') |
| del src,dst; gc.collect(); mx.clear_cache() |
| if not report['passed']: |
| raise RuntimeError(f'{size} did not pass numerical validation') |
|
|
|
|
| if __name__=='__main__': |
| p=argparse.ArgumentParser(description=__doc__) |
| p.add_argument('variants',nargs='*',choices=VARIANTS,default=list(VARIANTS)) |
| revision=subprocess.check_output(['git','-C',str(ROOT/'upstream'),'rev-parse','HEAD'],text=True).strip() |
| if revision!='a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff': |
| raise ValueError('upstream source revision mismatch') |
| subprocess.run(['git','-C',str(ROOT/'upstream'),'diff','--exit-code','HEAD','--','src'],check=True) |
| torch.set_num_threads(1) |
| for size in p.parse_args().variants: |
| validate(size) |
|
|
