diff --git a/README.md b/README.md index fdd4420..33b3717 100644 --- a/README.md +++ b/README.md @@ -234,8 +234,10 @@ ml/purpose-classifier/venv/bin/python ml/purpose-classifier/quantize_coreml.py \ --overwrite-output ``` -Activation calibration is grouped to keep temporary Core ML packages bounded and prints -progress while it runs. The candidate uses per-tensor asymmetric uint8 activations, +Activation calibration writes its temporary packages under the candidate output directory +and removes each package immediately after prediction; this avoids Core ML Tools retaining +one full weight copy per calibration step until process exit. It prints progress while it +runs. The candidate uses per-tensor asymmetric uint8 activations, per-channel symmetric int8 linear weights, and per-tensor asymmetric uint8 embedding weights. It fails the command if the resulting package exceeds 25 MiB. diff --git a/quantize_coreml.py b/quantize_coreml.py index 3a6008e..628939c 100644 --- a/quantize_coreml.py +++ b/quantize_coreml.py @@ -4,10 +4,13 @@ from __future__ import annotations import argparse +import gc import hashlib import shutil import sys +import tempfile from collections import Counter +from contextlib import contextmanager from pathlib import Path from typing import Any, Sequence @@ -42,6 +45,38 @@ def _tree_sha256(package: Path) -> str: 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 + for package in temporary_root.glob("*.mlpackage"): + shutil.rmtree(package) + + def _optimization_configs(optimize: Any) -> tuple[Any, Any]: """Return the Core ML analogue of the accepted ONNX QDQ policy.""" @@ -101,6 +136,9 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]: 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( @@ -147,12 +185,18 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]: f"got {sorted(input_names)}" ) activation_config, weight_config = _optimization_configs(cto) - activation_quantized = cto.coreml.linear_quantize_activations( - source_model, - activation_config, - sample_data, - calibration_op_group_size=args.calibration_op_group_size, - ) + 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)): + activation_quantized = cto.coreml.linear_quantize_activations( + source_model, + activation_config, + sample_data, + calibration_op_group_size=args.calibration_op_group_size, + ) print("Core ML weight quantization: W8", flush=True) quantized = cto.coreml.linear_quantize_weights( activation_quantized, @@ -168,7 +212,6 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]: shutil.rmtree(args.output) else: args.output.unlink() - args.output.parent.mkdir(parents=True, exist_ok=True) quantized.save(str(args.output)) manifest = _package_manifest(args.output) @@ -226,7 +269,7 @@ def build_parser() -> argparse.ArgumentParser: 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=32) + parser.add_argument("--calibration-op-group-size", type=int, default=-1) parser.add_argument( "--shipping-budget-bytes", type=int, diff --git a/tests/test_quantize_coreml.py b/tests/test_quantize_coreml.py index aab057f..d01ebaa 100644 --- a/tests/test_quantize_coreml.py +++ b/tests/test_quantize_coreml.py @@ -74,6 +74,23 @@ class CoreMLQuantizationConfigTests(unittest.TestCase): with self.assertRaisesRegex(DataError, "must not overwrite"): quantize_coreml._validate_args(args) + def test_calibration_packages_are_removed_after_each_prediction(self): + class FakeDebugger: + def predict_intermediate_outputs(self): + package = Path(tempfile.mkdtemp(suffix=".mlpackage")) + (package / "weight.bin").write_bytes(b"weights") + return {"output": np.array([1.0])} + + with tempfile.TemporaryDirectory() as temp: + root = Path(temp) / "calibration" + with quantize_coreml._bounded_calibration_packages( + FakeDebugger, + root, + ): + output = FakeDebugger().predict_intermediate_outputs() + self.assertEqual([1.0], output["output"].tolist()) + self.assertEqual([], list(root.glob("*.mlpackage"))) + if __name__ == "__main__": unittest.main()