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

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.
 

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,6 +185,12 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]:
             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)):
             activation_quantized = cto.coreml.linear_quantize_activations(
                 source_model,
                 activation_config,
@@ -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,

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()