Merge nucleic/sleek-ember-seal-uady into dev

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,