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

README.md

@@ -237,7 +237,10 @@ ml/purpose-classifier/venv/bin/python ml/purpose-classifier/quantize_coreml.py \
 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,
+runs. The successfully rewritten A8 package is cached beside the W8A8 output and reused
+only when its source hash, Core ML Tools version, activation policy, calibration seed, and
+prompt hashes match exactly. This prevents a later weight-stage failure from forcing
+another calibration. The candidate uses per-tensor asymmetric uint8 activations,
 per-channel symmetric int8 linear weights, and per-tensor asymmetric uint8 embedding
 weights. Activation quantization is limited to floating-point linear operations; applying
 Core ML Tools' global policy also selects integer embedding-index additions and produces

quantize_coreml.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import gc
 import hashlib
+import json
 import shutil
 import sys
 import tempfile
@@ -73,8 +74,39 @@ def _bounded_calibration_packages(debugger_type: Any, temporary_root: Path):
     finally:
         debugger_type.predict_intermediate_outputs = original_predict
         tempfile.tempdir = previous_tempdir
-        for package in temporary_root.glob("*.mlpackage"):
-            shutil.rmtree(package)
+
+
+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]:
@@ -160,11 +192,29 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]:
         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,
         )
-    sample_data = []
         for record in calibration:
             encoded = encode_fixed_shape(tokenizer, [record["prompt"]], torch)
             sample_data.append(
@@ -174,10 +224,6 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]:
                 }
             )
 
-    print(
-        f"Core ML activation calibration: {len(sample_data)} records",
-        flush=True,
-    )
     source_model = ct.models.MLModel(
         str(args.model),
         compute_units=ct.ComputeUnit.CPU_ONLY,
@@ -196,34 +242,58 @@ def quantize(args: argparse.Namespace) -> dict[str, Any]:
         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:{args.calibration_records}:seed={args.calibration_seed}"
+            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": _tree_sha256(args.model),
+            "sourcePackageSha256": source_sha256,
             "coremltoolsVersion": ct.__version__,
             "quantization": {
                 "name": "W8A8",

tests/test_quantize_coreml.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import sys
 import tempfile
 import unittest
@@ -94,6 +95,36 @@ class CoreMLQuantizationConfigTests(unittest.TestCase):
                 self.assertEqual([1.0], output["output"].tolist())
                 self.assertEqual([], list(root.glob("*.mlpackage")))
 
+                result_package = Path(tempfile.mkdtemp(suffix=".mlpackage"))
+            self.assertTrue(result_package.exists())
+
+    def test_activation_cache_requires_exact_contract(self):
+        expected = {
+            "schemaVersion": 1,
+            "sourcePackageSha256": "abc",
+        }
+        with tempfile.TemporaryDirectory() as temp:
+            root = Path(temp)
+            package = root / "a8.mlpackage"
+            package.mkdir()
+            manifest = root / "a8-manifest.json"
+            manifest.write_text(json.dumps(expected), encoding="utf-8")
+            self.assertTrue(
+                quantize_coreml._activation_cache_is_valid(
+                    package,
+                    manifest,
+                    expected,
+                )
+            )
+            manifest.write_text("{}", encoding="utf-8")
+            self.assertFalse(
+                quantize_coreml._activation_cache_is_valid(
+                    package,
+                    manifest,
+                    expected,
+                )
+            )
+
 
 if __name__ == "__main__":
     unittest.main()