File size: 6,979 Bytes
591cda4 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """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()
|