380 lines
14 KiB
Python
380 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Calibrate a Core ML W8A8 candidate from the selected float16 ML Program."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from collections import Counter
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
|
|
from convert_coreml import COREMLTOOLS_VERSION, _package_manifest
|
|
from export import stratified_calibration_sample
|
|
from purpose_data import DataError, load_jsonl, prompt_hash, write_json
|
|
from train import encode_fixed_shape
|
|
|
|
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
CANDIDATE_DIR = (
|
|
SCRIPT_DIR / "outputs" / "purpose-lite-v1-distilled-qat-mlx-4e"
|
|
)
|
|
DEFAULT_MODEL = CANDIDATE_DIR / "coreml" / "purpose-lite-v1-fp16.mlpackage"
|
|
DEFAULT_MODEL_DIR = CANDIDATE_DIR / "model"
|
|
DEFAULT_VALIDATION = SCRIPT_DIR / ".artifacts" / "dataset-v1" / "validation.jsonl"
|
|
DEFAULT_OUTPUT = CANDIDATE_DIR / "coreml" / "purpose-lite-v1-w8a8.mlpackage"
|
|
SHIPPING_BUDGET_BYTES = 25 * 1024 * 1024
|
|
|
|
|
|
def _tree_sha256(package: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
for path in sorted(item for item in package.rglob("*") if item.is_file()):
|
|
relative = str(path.relative_to(package)).encode("utf-8")
|
|
digest.update(relative)
|
|
digest.update(b"\0")
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
@contextmanager
|
|
def _bounded_calibration_packages(debugger_type: Any, temporary_root: Path):
|
|
"""Eagerly remove Core ML Tools' per-prediction temporary packages.
|
|
|
|
Core ML Tools registers these packages for process-exit cleanup. Activation
|
|
calibration creates one package per intermediate-output group per record, so retaining
|
|
all of them can consume tens of gigabytes before the process exits.
|
|
"""
|
|
|
|
original_predict = debugger_type.predict_intermediate_outputs
|
|
previous_tempdir = tempfile.tempdir
|
|
|
|
def predict_and_cleanup(*args: Any, **kwargs: Any) -> Any:
|
|
try:
|
|
return original_predict(*args, **kwargs)
|
|
finally:
|
|
gc.collect()
|
|
for package in temporary_root.glob("*.mlpackage"):
|
|
shutil.rmtree(package)
|
|
|
|
temporary_root.mkdir(parents=True, exist_ok=True)
|
|
tempfile.tempdir = str(temporary_root)
|
|
debugger_type.predict_intermediate_outputs = predict_and_cleanup
|
|
try:
|
|
yield
|
|
finally:
|
|
debugger_type.predict_intermediate_outputs = original_predict
|
|
tempfile.tempdir = previous_tempdir
|
|
|
|
|
|
def _activation_cache_contract(
|
|
source_sha256: str,
|
|
calibration: Sequence[dict[str, Any]],
|
|
*,
|
|
calibration_seed: int,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"schemaVersion": 1,
|
|
"sourcePackageSha256": source_sha256,
|
|
"coremltoolsVersion": COREMLTOOLS_VERSION,
|
|
"activationQuantization": "linear:per-tensor-asymmetric-uint8",
|
|
"activationOpTypes": ["linear"],
|
|
"calibrationRecords": len(calibration),
|
|
"calibrationSeed": calibration_seed,
|
|
"calibrationPromptHashes": sorted(
|
|
prompt_hash(item["prompt"]) for item in calibration
|
|
),
|
|
}
|
|
|
|
|
|
def _activation_cache_is_valid(
|
|
package: Path,
|
|
manifest: Path,
|
|
expected: dict[str, Any],
|
|
) -> bool:
|
|
if not package.is_dir() or not manifest.is_file():
|
|
return False
|
|
try:
|
|
return json.loads(manifest.read_text(encoding="utf-8")) == expected
|
|
except (OSError, UnicodeError, json.JSONDecodeError):
|
|
return False
|
|
|
|
|
|
def _optimization_configs(optimize: Any) -> tuple[Any, Any]:
|
|
"""Return the Core ML analogue of the accepted ONNX QDQ policy."""
|
|
|
|
activation = optimize.coreml.OpLinearQuantizerConfig(
|
|
mode="linear",
|
|
dtype=np.uint8,
|
|
granularity="per_tensor",
|
|
)
|
|
activation_config = optimize.coreml.OptimizationConfig(
|
|
# A global activation policy also selects integer `add` operations in the
|
|
# embedding/index path. Core ML's quantize op requires its floating-point scale
|
|
# to match a floating-point input, so quantize only accelerator-supported linear
|
|
# activations. Attention matmul activation quantization is not supported by this
|
|
# Core ML graph pass.
|
|
op_type_configs={"linear": activation},
|
|
)
|
|
|
|
linear_weight = optimize.coreml.OpLinearQuantizerConfig(
|
|
mode="linear_symmetric",
|
|
dtype=np.int8,
|
|
granularity="per_channel",
|
|
weight_threshold=2048,
|
|
)
|
|
embedding_weight = optimize.coreml.OpLinearQuantizerConfig(
|
|
mode="linear",
|
|
dtype=np.uint8,
|
|
granularity="per_tensor",
|
|
weight_threshold=2048,
|
|
)
|
|
weight_config = optimize.coreml.OptimizationConfig(
|
|
op_type_configs={
|
|
"gather": embedding_weight,
|
|
"linear": linear_weight,
|
|
"matmul": linear_weight,
|
|
}
|
|
)
|
|
return activation_config, weight_config
|
|
|
|
|
|
def _validate_args(args: argparse.Namespace) -> None:
|
|
if not args.model.is_dir():
|
|
raise DataError(f"{args.model}: source Core ML package is missing")
|
|
if not args.model_dir.is_dir():
|
|
raise DataError(f"{args.model_dir}: tokenizer directory is missing")
|
|
if args.output.suffix != ".mlpackage":
|
|
raise DataError("Core ML output must end in .mlpackage")
|
|
try:
|
|
same_output = args.model.resolve() == args.output.resolve()
|
|
except OSError as exc:
|
|
raise DataError(f"cannot resolve Core ML package paths: {exc}") from exc
|
|
if same_output:
|
|
raise DataError("W8A8 output must not overwrite its float16 source package")
|
|
if args.output.exists() and not args.overwrite_output:
|
|
raise DataError(
|
|
f"{args.output}: output exists; pass --overwrite-output intentionally"
|
|
)
|
|
|
|
|
|
def quantize(args: argparse.Namespace) -> dict[str, Any]:
|
|
_validate_args(args)
|
|
try:
|
|
import coremltools as ct
|
|
import coremltools.optimize as cto
|
|
import torch
|
|
from coremltools.optimize.coreml.experimental._model_debugger import (
|
|
ModelDebugger,
|
|
)
|
|
from transformers import AutoTokenizer
|
|
except ImportError as exc:
|
|
raise DataError(
|
|
"Core ML quantization requires requirements-coreml.txt on macOS"
|
|
) from exc
|
|
if ct.__version__ != COREMLTOOLS_VERSION:
|
|
raise DataError(
|
|
f"expected coremltools {COREMLTOOLS_VERSION}, found {ct.__version__}"
|
|
)
|
|
|
|
validation = load_jsonl(args.validation)
|
|
calibration = stratified_calibration_sample(
|
|
validation,
|
|
args.calibration_records,
|
|
seed=args.calibration_seed,
|
|
)
|
|
source_sha256 = _tree_sha256(args.model)
|
|
activation_cache = args.output.with_name(
|
|
f"{args.output.stem}-a8-cache.mlpackage"
|
|
)
|
|
activation_cache_manifest = activation_cache.with_name(
|
|
f"{activation_cache.stem}-manifest.json"
|
|
)
|
|
cache_contract = _activation_cache_contract(
|
|
source_sha256,
|
|
calibration,
|
|
calibration_seed=args.calibration_seed,
|
|
)
|
|
reuse_activation_cache = _activation_cache_is_valid(
|
|
activation_cache,
|
|
activation_cache_manifest,
|
|
cache_contract,
|
|
)
|
|
sample_data = []
|
|
if not reuse_activation_cache:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.model_dir,
|
|
local_files_only=True,
|
|
)
|
|
for record in calibration:
|
|
encoded = encode_fixed_shape(tokenizer, [record["prompt"]], torch)
|
|
sample_data.append(
|
|
{
|
|
name: value.numpy().astype(np.int32, copy=False)
|
|
for name, value in encoded.items()
|
|
}
|
|
)
|
|
|
|
source_model = ct.models.MLModel(
|
|
str(args.model),
|
|
compute_units=ct.ComputeUnit.CPU_ONLY,
|
|
)
|
|
input_names = {item.name for item in source_model.get_spec().description.input}
|
|
expected_inputs = {"input_ids", "attention_mask", "token_type_ids"}
|
|
if input_names != expected_inputs:
|
|
raise DataError(
|
|
f"Core ML inputs changed: expected {sorted(expected_inputs)}, "
|
|
f"got {sorted(input_names)}"
|
|
)
|
|
activation_config, weight_config = _optimization_configs(cto)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(
|
|
prefix=".purpose-coreml-calibration-",
|
|
dir=args.output.parent,
|
|
) as temporary:
|
|
with _bounded_calibration_packages(ModelDebugger, Path(temporary)):
|
|
if reuse_activation_cache:
|
|
print(
|
|
f"Reusing Core ML A8 cache: {activation_cache}",
|
|
flush=True,
|
|
)
|
|
activation_quantized = ct.models.MLModel(
|
|
str(activation_cache),
|
|
compute_units=ct.ComputeUnit.CPU_ONLY,
|
|
)
|
|
else:
|
|
print(
|
|
f"Core ML activation calibration: {len(sample_data)} records",
|
|
flush=True,
|
|
)
|
|
activation_quantized = cto.coreml.linear_quantize_activations(
|
|
source_model,
|
|
activation_config,
|
|
sample_data,
|
|
calibration_op_group_size=args.calibration_op_group_size,
|
|
)
|
|
activation_quantized.save(str(activation_cache))
|
|
write_json(activation_cache_manifest, cache_contract)
|
|
print(
|
|
f"Core ML A8 cache: {activation_cache}",
|
|
flush=True,
|
|
)
|
|
print("Core ML weight quantization: W8", flush=True)
|
|
quantized = cto.coreml.linear_quantize_weights(
|
|
activation_quantized,
|
|
weight_config,
|
|
)
|
|
quantized.user_defined_metadata["com.nucleic.model.quantization"] = (
|
|
"W8A8"
|
|
)
|
|
quantized.user_defined_metadata[
|
|
"com.nucleic.model.quantizationCalibration"
|
|
] = f"stratified:{len(calibration)}:seed={args.calibration_seed}"
|
|
|
|
if args.output.exists():
|
|
if args.output.is_dir():
|
|
shutil.rmtree(args.output)
|
|
else:
|
|
args.output.unlink()
|
|
# Save while the Core ML Tools result package and its source weights are
|
|
# still alive inside the dedicated temporary directory.
|
|
quantized.save(str(args.output))
|
|
|
|
manifest = _package_manifest(args.output)
|
|
manifest.update(
|
|
{
|
|
"sourcePackage": str(args.model),
|
|
"sourcePackageSha256": source_sha256,
|
|
"coremltoolsVersion": ct.__version__,
|
|
"quantization": {
|
|
"name": "W8A8",
|
|
"activations": "per-tensor asymmetric uint8",
|
|
"linearWeights": "per-channel symmetric int8",
|
|
"embeddingWeights": "per-tensor asymmetric uint8",
|
|
},
|
|
"calibrationRecords": len(calibration),
|
|
"calibrationSeed": args.calibration_seed,
|
|
"calibrationOpGroupSize": args.calibration_op_group_size,
|
|
"calibrationSample": {
|
|
"strategy": "stratified by purpose, slice, and primary language",
|
|
"purposeCounts": dict(
|
|
sorted(Counter(item["purpose"] for item in calibration).items())
|
|
),
|
|
"sliceCounts": dict(
|
|
sorted(Counter(item["slice"] for item in calibration).items())
|
|
),
|
|
"promptHashes": sorted(
|
|
prompt_hash(item["prompt"]) for item in calibration
|
|
),
|
|
},
|
|
"shippingBudgetBytes": args.shipping_budget_bytes,
|
|
"shippingBudgetPassed": (
|
|
manifest["bytes"] <= args.shipping_budget_bytes
|
|
),
|
|
}
|
|
)
|
|
manifest_path = args.output.with_name(f"{args.output.stem}-manifest.json")
|
|
write_json(manifest_path, manifest)
|
|
print(
|
|
f"Core ML W8A8 package: {args.output} ({manifest['bytes']} bytes)",
|
|
flush=True,
|
|
)
|
|
if not manifest["shippingBudgetPassed"]:
|
|
raise DataError(
|
|
f"{args.output}: {manifest['bytes']} bytes exceeds the "
|
|
f"{args.shipping_budget_bytes}-byte shipping budget"
|
|
)
|
|
return manifest
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
|
|
parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
|
|
parser.add_argument("--validation", type=Path, default=DEFAULT_VALIDATION)
|
|
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
|
parser.add_argument("--calibration-records", type=int, default=256)
|
|
parser.add_argument("--calibration-seed", type=int, default=20260730)
|
|
parser.add_argument("--calibration-op-group-size", type=int, default=-1)
|
|
parser.add_argument(
|
|
"--shipping-budget-bytes",
|
|
type=int,
|
|
default=SHIPPING_BUDGET_BYTES,
|
|
)
|
|
parser.add_argument("--overwrite-output", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
if (
|
|
args.calibration_records <= 0
|
|
or args.calibration_op_group_size == 0
|
|
or args.calibration_op_group_size < -1
|
|
or args.shipping_budget_bytes <= 0
|
|
):
|
|
parser.error(
|
|
"calibration records/budget must be positive and op group size must be "
|
|
"-1 or positive"
|
|
)
|
|
try:
|
|
quantize(args)
|
|
except (DataError, OSError, RuntimeError, ValueError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|