"""Audit FP32 conversion and compare native MLX outputs with PyTorch.""" from __future__ import annotations import argparse import hashlib import json import sys from pathlib import Path import mlx.core as mx import numpy as np import torch from safetensors import safe_open from transformers import AutoModelForMaskedLM, AutoTokenizer def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def metrics(reference: np.ndarray, candidate: np.ndarray) -> dict: error = candidate.astype(np.float64) - reference.astype(np.float64) ref64 = reference.astype(np.float64) return { "max_abs_error": float(np.max(np.abs(error))), "mean_abs_error": float(np.mean(np.abs(error))), "rmse": float(np.sqrt(np.mean(error**2))), "relative_l2_error": float(np.linalg.norm(error.ravel()) / np.linalg.norm(ref64.ravel())), "cosine_similarity": float( np.dot(ref64.ravel(), candidate.astype(np.float64).ravel()) / (np.linalg.norm(ref64.ravel()) * np.linalg.norm(candidate.astype(np.float64).ravel())) ), } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("source", type=Path) parser.add_argument("mlx_model", type=Path) args = parser.parse_args() sys.path.insert(0, str(args.mlx_model)) from model import load_model # noqa: PLC0415 source_weights = args.source / "model.safetensors" mlx_weights = args.mlx_model / "model.safetensors" tensor_count = 0 parameter_count = 0 exact_tensor_count = 0 with safe_open(source_weights, framework="np") as source, safe_open( mlx_weights, framework="np" ) as converted: if set(source.keys()) != set(converted.keys()): raise AssertionError("Source and MLX tensor key sets differ") for name in source.keys(): before = source.get_tensor(name) after = converted.get_tensor(name) expected = before.transpose(0, 2, 1) if name.endswith("conv.conv.weight") else before if before.dtype != np.float32 or after.dtype != np.float32: raise AssertionError(f"Non-FP32 tensor: {name}: {before.dtype} -> {after.dtype}") if np.array_equal(expected, after): exact_tensor_count += 1 else: raise AssertionError(f"Tensor values changed during conversion: {name}") tensor_count += 1 parameter_count += before.size tokenizer = AutoTokenizer.from_pretrained(args.source, trust_remote_code=True) texts = [ f"The capital of France is {tokenizer.mask_token}.", f"Liquid neural networks run efficiently on {tokenizer.mask_token} devices.", f"Paris is the {tokenizer.mask_token} of France.", ] encoded = tokenizer(texts, padding=True, return_tensors="pt") torch.set_grad_enabled(False) reference_model = AutoModelForMaskedLM.from_pretrained( args.source, trust_remote_code=True, dtype=torch.float32, attn_implementation="eager", ).eval() reference_output = reference_model( **encoded, output_hidden_states=True, return_dict=True ) reference_logits = reference_output.logits.cpu().numpy() reference_hidden = reference_output.hidden_states[-1].cpu().numpy() mlx_model = load_model(args.mlx_model) mlx_logits, mlx_hidden = mlx_model( mx.array(encoded["input_ids"].numpy()), mx.array(encoded["attention_mask"].numpy()), return_hidden=True, ) mx.eval(mlx_logits, mlx_hidden) candidate_logits = np.array(mlx_logits) candidate_hidden = np.array(mlx_hidden) valid = encoded["attention_mask"].numpy().astype(bool) mask_positions = encoded["input_ids"].numpy() == tokenizer.mask_token_id top1_reference = np.argmax(reference_logits, axis=-1) top1_candidate = np.argmax(candidate_logits, axis=-1) hidden_metrics = metrics(reference_hidden[valid], candidate_hidden[valid]) logits_metrics = metrics(reference_logits[valid], candidate_logits[valid]) valid_top1_agreement = float( np.mean(top1_reference[valid] == top1_candidate[valid]) ) masked_top1_agreement = float( np.mean(top1_reference[mask_positions] == top1_candidate[mask_positions]) ) predicted = [] for batch_index, position in np.argwhere(mask_positions): ref_id = int(top1_reference[batch_index, position]) mlx_id = int(top1_candidate[batch_index, position]) predicted.append( { "text": texts[batch_index], "reference_token_id": ref_id, "mlx_token_id": mlx_id, "reference_token": tokenizer.decode([ref_id]), "mlx_token": tokenizer.decode([mlx_id]), "match": ref_id == mlx_id, } ) report = { "source": "LiquidAI/LFM2.5-Encoder-350M", "source_revision": "b886781f7c6f10ca9b7096e21b83e30a073c2f39", "runtime": { "mlx_device": str(mx.default_device()), "torch_device": "cpu", "dtype": "float32", }, "tensor_audit": { "tensor_count": tensor_count, "exact_tensor_count": exact_tensor_count, "parameter_count": parameter_count, "all_source_tensors_fp32": True, "all_mlx_tensors_fp32": True, "all_values_bit_exact_after_layout_transform": exact_tensor_count == tensor_count, "source_weights_sha256": sha256(source_weights), "mlx_weights_sha256": sha256(mlx_weights), }, "parity": { "input_shape": list(encoded["input_ids"].shape), "valid_token_count": int(valid.sum()), "hidden": hidden_metrics, "logits": logits_metrics, "top1_agreement_valid_tokens": valid_top1_agreement, "top1_agreement_masked_tokens": masked_top1_agreement, "masked_predictions": predicted, }, } checks = { "bit_exact_fp32_weights": exact_tensor_count == tensor_count, "hidden_relative_l2_below_1e-4": hidden_metrics["relative_l2_error"] < 1e-4, "logits_relative_l2_below_1e-4": logits_metrics["relative_l2_error"] < 1e-4, "valid_token_top1_agreement_100_percent": valid_top1_agreement == 1.0, "masked_token_top1_agreement_100_percent": masked_top1_agreement == 1.0, } report["acceptance"] = {"passed": all(checks.values()), "checks": checks} output = args.mlx_model / "validation.json" output.write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(report, indent=2)) if not report["acceptance"]["passed"]: raise SystemExit("FP32 parity acceptance failed") if __name__ == "__main__": main()