File size: 5,082 Bytes
d35e174 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """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)
# Fixed, declared acceptance criteria; not adjusted per model.
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)
|