| """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 |
|
|
| 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() |
|
|
