Merge nucleic/sleek-ember-seal-uady into dev
This commit is contained in:
+252
@@ -0,0 +1,252 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Verify MLX/PyTorch parity and checkpoint round-tripping before MLX training."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
from purpose_data import LABELS, DataError, load_jsonl
|
||||
from train import enable_quantization_aware_training, encode_fixed_shape
|
||||
from train_mlx import _checkpoint_config, encode_fixed_shape_numpy
|
||||
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_DATASET = SCRIPT_DIR / ".artifacts" / "dataset-v1" / "validation.jsonl"
|
||||
|
||||
|
||||
def verify(model_dir: Path, dataset: Path, records: int) -> None:
|
||||
try:
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
import torch
|
||||
from mlx.utils import tree_flatten
|
||||
from safetensors import safe_open
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
from mlx_model import (
|
||||
BertClassifierConfig,
|
||||
BertForSequenceClassification,
|
||||
fake_quantize_activation,
|
||||
fake_quantize_linear_weight,
|
||||
load_hugging_face_weights,
|
||||
save_hugging_face_weights,
|
||||
)
|
||||
except ImportError as exc:
|
||||
raise DataError(
|
||||
"verification requires requirements-mlx.txt on Apple Silicon"
|
||||
) from exc
|
||||
if not mx.metal.is_available():
|
||||
raise DataError("verification requires the MLX Metal backend")
|
||||
|
||||
# Check the fake-quantization contract independently of the full model. Tiny
|
||||
# backend-specific floating-point differences can cross later quantization
|
||||
# thresholds, especially on padded tokens, so full QAT logits are expected to
|
||||
# have more drift than the ordinary float graph.
|
||||
activation_values = np.asarray(
|
||||
[[-2.75, -0.125, 0.0, 0.625], [1.25, 3.5, -1.0, 0.25]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
torch_activation = torch.from_numpy(activation_values)
|
||||
activation_minimum = min(0.0, float(torch_activation.amin().item()))
|
||||
activation_maximum = max(0.0, float(torch_activation.amax().item()))
|
||||
activation_scale = max(
|
||||
(activation_maximum - activation_minimum) / 255.0,
|
||||
torch.finfo(torch.float32).eps,
|
||||
)
|
||||
activation_zero_point = max(
|
||||
0,
|
||||
min(255, round(-activation_minimum / activation_scale)),
|
||||
)
|
||||
torch_quantized_activation = torch.fake_quantize_per_tensor_affine(
|
||||
torch_activation,
|
||||
activation_scale,
|
||||
activation_zero_point,
|
||||
0,
|
||||
255,
|
||||
).numpy()
|
||||
mlx_quantized_activation = np.asarray(
|
||||
fake_quantize_activation(mx.array(activation_values))
|
||||
)
|
||||
np.testing.assert_array_equal(
|
||||
mlx_quantized_activation,
|
||||
torch_quantized_activation,
|
||||
)
|
||||
|
||||
weight_values = np.asarray(
|
||||
[[-1.5, -0.25, 0.75, 1.25], [0.125, -0.875, 2.0, -1.25]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
torch_weight = torch.from_numpy(weight_values)
|
||||
weight_scales = torch_weight.abs().amax(dim=1).div(127.0).clamp_min(
|
||||
torch.finfo(torch.float32).eps
|
||||
)
|
||||
torch_quantized_weight = torch.fake_quantize_per_channel_affine(
|
||||
torch_weight,
|
||||
weight_scales,
|
||||
torch.zeros_like(weight_scales, dtype=torch.int32),
|
||||
0,
|
||||
-127,
|
||||
127,
|
||||
).numpy()
|
||||
mlx_quantized_weight = np.asarray(
|
||||
fake_quantize_linear_weight(mx.array(weight_values))
|
||||
)
|
||||
np.testing.assert_array_equal(mlx_quantized_weight, torch_quantized_weight)
|
||||
print("fake-quant primitives: exact")
|
||||
|
||||
checkpoint = model_dir / "model.safetensors"
|
||||
if not checkpoint.is_file():
|
||||
raise DataError(f"{checkpoint}: checkpoint is missing")
|
||||
validation = load_jsonl(dataset)[:records]
|
||||
if not validation:
|
||||
raise DataError(f"{dataset}: no validation records")
|
||||
texts = [record["prompt"] for record in validation]
|
||||
labels = np.asarray(
|
||||
[LABELS.index(record["purpose"]) for record in validation],
|
||||
dtype=np.int32,
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
|
||||
numpy_tokens = encode_fixed_shape_numpy(tokenizer, texts)
|
||||
torch_tokens = encode_fixed_shape(tokenizer, texts, torch)
|
||||
|
||||
config_json = _checkpoint_config(model_dir)
|
||||
config = BertClassifierConfig.from_hugging_face(config_json)
|
||||
mlx_float = BertForSequenceClassification(
|
||||
config,
|
||||
quantization_aware=False,
|
||||
)
|
||||
load_hugging_face_weights(mlx_float, checkpoint)
|
||||
mlx_float.eval()
|
||||
mlx_float_logits = mlx_float(
|
||||
**{key: mx.array(value) for key, value in numpy_tokens.items()}
|
||||
)
|
||||
mx.eval(mlx_float_logits)
|
||||
mlx_float_logits = np.asarray(mlx_float_logits)
|
||||
|
||||
torch_model = AutoModelForSequenceClassification.from_pretrained(
|
||||
model_dir,
|
||||
local_files_only=True,
|
||||
)
|
||||
torch_model.eval()
|
||||
with torch.inference_mode():
|
||||
torch_float_logits = torch_model(**torch_tokens).logits.numpy()
|
||||
np.testing.assert_allclose(
|
||||
mlx_float_logits,
|
||||
torch_float_logits,
|
||||
rtol=1e-4,
|
||||
atol=2e-5,
|
||||
)
|
||||
if not np.array_equal(
|
||||
mlx_float_logits.argmax(axis=-1),
|
||||
torch_float_logits.argmax(axis=-1),
|
||||
):
|
||||
raise DataError("MLX and PyTorch float predictions differ")
|
||||
print(
|
||||
"float parity: "
|
||||
f"max_abs_error={np.max(np.abs(mlx_float_logits - torch_float_logits)):.3g}"
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="purpose-mlx-roundtrip-") as temp:
|
||||
roundtrip = Path(temp) / "model.safetensors"
|
||||
save_hugging_face_weights(mlx_float, roundtrip)
|
||||
with safe_open(checkpoint, framework="np") as original, safe_open(
|
||||
roundtrip,
|
||||
framework="np",
|
||||
) as converted:
|
||||
if set(original.keys()) != set(converted.keys()):
|
||||
raise DataError("MLX checkpoint round-trip changed parameter keys")
|
||||
for key in original.keys():
|
||||
np.testing.assert_array_equal(
|
||||
original.get_tensor(key),
|
||||
converted.get_tensor(key),
|
||||
)
|
||||
print("checkpoint round-trip: exact")
|
||||
|
||||
del mlx_float
|
||||
mx.clear_cache()
|
||||
mlx_qat = BertForSequenceClassification(
|
||||
config,
|
||||
quantization_aware=True,
|
||||
)
|
||||
load_hugging_face_weights(mlx_qat, checkpoint)
|
||||
mlx_qat.eval()
|
||||
mlx_qat_logits = mlx_qat(
|
||||
**{key: mx.array(value) for key, value in numpy_tokens.items()}
|
||||
)
|
||||
mx.eval(mlx_qat_logits)
|
||||
mlx_qat_logits = np.asarray(mlx_qat_logits)
|
||||
|
||||
enable_quantization_aware_training(torch, torch_model)
|
||||
torch_model.eval()
|
||||
with torch.inference_mode():
|
||||
torch_qat_logits = torch_model(**torch_tokens).logits.numpy()
|
||||
qat_max_abs_error = float(
|
||||
np.max(np.abs(mlx_qat_logits - torch_qat_logits))
|
||||
)
|
||||
if not np.isfinite(qat_max_abs_error) or qat_max_abs_error > 1.0:
|
||||
raise DataError(
|
||||
"MLX and PyTorch QAT logits have excessive backend drift: "
|
||||
f"{qat_max_abs_error:.3g}"
|
||||
)
|
||||
if not np.array_equal(
|
||||
mlx_qat_logits.argmax(axis=-1),
|
||||
torch_qat_logits.argmax(axis=-1),
|
||||
):
|
||||
raise DataError("MLX and PyTorch QAT predictions differ")
|
||||
print(
|
||||
"QAT parity: "
|
||||
f"predictions=exact max_abs_error={qat_max_abs_error:.3g}"
|
||||
)
|
||||
|
||||
mlx_qat.train()
|
||||
mlx_labels = mx.array(labels)
|
||||
|
||||
def loss_function() -> object:
|
||||
logits = mlx_qat(
|
||||
**{key: mx.array(value) for key, value in numpy_tokens.items()}
|
||||
)
|
||||
return nn.losses.cross_entropy(logits, mlx_labels, reduction="mean")
|
||||
|
||||
loss, gradients = nn.value_and_grad(mlx_qat, loss_function)()
|
||||
flat_gradients = [gradient for _, gradient in tree_flatten(gradients)]
|
||||
mx.eval(loss, gradients)
|
||||
if not all(
|
||||
bool(mx.all(mx.isfinite(gradient)).item())
|
||||
for gradient in flat_gradients
|
||||
):
|
||||
raise DataError("MLX QAT produced non-finite gradients")
|
||||
if not any(
|
||||
float(mx.max(mx.abs(gradient)).item()) > 0
|
||||
for gradient in flat_gradients
|
||||
):
|
||||
raise DataError("MLX QAT produced only zero gradients")
|
||||
print(f"QAT gradient smoke: loss={float(loss.item()):.6f}")
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--dataset", type=Path, default=DEFAULT_DATASET)
|
||||
parser.add_argument("--records", type=int, default=8)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
args = build_parser().parse_args(argv)
|
||||
if args.records <= 0:
|
||||
raise SystemExit("--records must be positive")
|
||||
try:
|
||||
verify(args.model.expanduser(), args.dataset.expanduser(), args.records)
|
||||
except (AssertionError, DataError) as exc:
|
||||
print(f"error: {exc}")
|
||||
return 2
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user