Update docs and add tests to validate externalization workflow with graph mode

docs/src/index.md

@@ -42,6 +42,7 @@ utils/mixed_precision
 utils/activation_comparison
 utils/casting
 utils/coreai_compression
+utils/composite_op_quantization
 ```
 
 ```{toctree}

docs/src/quantization/overview.md

@@ -204,7 +204,7 @@ The two modes are expected to produce very similar models for weight-only quanti
 A few scenarios where `eager` mode may need to be used instead of `graph`:
 
 - If you run into any errors during the `prepare` call which, under the hood, invokes the `torch.export.export` and `torchao`'s `prepare_qat_pt2e`/`convert_pt2e` APIs. See [Graph Mode Troubleshooting](../debugging/graph_mode_troubleshooting.md) for common export errors and workarounds before falling back to eager mode.
-- When `torch.nn.Module` needs to be provided as an input, instead of `ExportedProgram` to the conversion API of [coreai-torch](https://github.com/apple/coreai-torch). This happens when the `coreai-torch` conversion needs to "externalize" certain sub-modules to map them to _composite ops_ for better runtime performance.
+- When `torch.nn.Module` needs to be provided as an input, instead of `ExportedProgram` to the conversion API of [coreai-torch](https://github.com/apple/coreai-torch). Note that models whose submodules must be "externalized" to map them to _composite ops_ for better runtime performance can still be quantized in graph mode. See [Quantizing Models with Core AI Composite Ops in Graph Mode](../utils/composite_op_quantization.md).
 
 #### Weights and activations quantization
 

docs/src/utils/composite_op_quantization.md

@@ -0,0 +1,251 @@
+# Quantizing Models with Core AI Composite Ops in Graph Mode
+
+Core AI recognizes certain well-known building blocks, such as SDPA or RMSNorm, as _composite ops_ and applies optimized implementations for them.
+`coreai-torch` establishes those boundaries through _externalization_.
+Refer to the [Externalization](https://apple.github.io/coreai-torch/main/guides/externalization.html) guide for details.
+Here, we will discuss the steps required to quantize a model in `graph` mode using `coreai-opt`.
+
+`graph`-mode quantization invokes `torch.export.export` under the hood, which decomposes a submodule's `forward` into aten ops.
+In order to preserve the composite op structure during this process for externalization, the following APIs are provided:
+
+- `_patch_model_for_externalization`: Patch the model **before** `quantizer.prepare`, so that the composite op call sites survive export and all subsequent quantization passes as opaque nodes.
+- `_subexport_and_restore`: The submodule bodies of the composite ops themselves are then exported and restored before lowering to `CoreAI`.
+
+Quantization treats each composite op as opaque, i.e., no fake-quantize op is placed inside the composite body.
+The composite's input and output boundary can still be quantized, see [Quantizing the composite op boundary](#quantizing-the-composite-op-boundary) below for details.
+
+:::{warning}
+The externalization APIs used below, `_patch_model_for_externalization` and `_subexport_and_restore` in `coreai-torch` are currently experimental.
+:::
+
+```mermaid
+---
+title: "Graph mode Quantization Workflow with Externalization"
+---
+flowchart LR
+    model["Full Precision<br>Model"] --> patch["Patch Model for<br>Externalization"]
+    patch --> prepare["Prepare and<br>Calibrate"]
+    prepare --> qfin["Finalize and<br>Export"]
+    qfin --> sub["Sub-export<br>and Restore"]
+    sub --> convert["Convert to<br>Core AI"]
+    style model fill:#f9f9f9,stroke:#999
+    style patch fill:#e8f0fe,stroke:#4285f4
+    style sub fill:#e8f0fe,stroke:#4285f4
+```
+
+## Step 1: Patch the model before prepare
+
+`_patch_model_for_externalization` replaces the `forward` of every matching submodule in the model with a `torch.library.custom_op`, in place.
+Call it before constructing the `Quantizer`.
+The example below demonstrates this using the same `RMSNormComposite` op from the [Externalization](https://apple.github.io/coreai-torch/main/guides/externalization.html) guide, however, the same process applies for all composite ops with their respective `ExternalizeSpec`s.
+
+```python
+import torch
+import torch.nn as nn
+from coreai_torch import ExternalizeSpec, _patch_model_for_externalization
+
+
+# The composite op
+class RMSNormComposite(nn.Module):
+    def __init__(self, axes=-1, eps=1e-5, version=1):
+        super().__init__()
+        self.axes = axes
+        self.eps = eps
+        self.version = version
+
+    def forward(self, input: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+        x_f32 = input.to(torch.float32)
+        inv_rms = torch.rsqrt((x_f32 * x_f32).mean(self.axes, keepdim=True) + self.eps)
+        return (input * inv_rms).to(input.dtype) * scale
+
+
+# A model that uses the composite op
+class Model(nn.Module):
+    def __init__(self, dim=32):
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+        self.norm = RMSNormComposite()
+        self.norm_weight = nn.Parameter(torch.ones(dim))
+        self.out = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        return self.out(self.norm(self.proj(x), self.norm_weight))
+
+
+model = Model().eval()
+example_inputs = (torch.randn(1, 32),)
+
+# Patch the model in-place
+# to externalize the RMSNormComposite
+_patch_model_for_externalization(
+    model,
+    targets=[
+        ExternalizeSpec(
+            target_class=RMSNormComposite,
+            composite_op_name="rms_norm",
+            composite_attrs=["axes", "eps", "version"],
+        )
+    ],
+)
+```
+
+## Step 2: Prepare, calibrate and finalize
+
+Nothing about the quantizer configuration or the calibration workflow changes.
+The composite op holds no weights of its own here, so weight quantization applies to the surrounding `Linear` layers only.
+
+```python
+import coreai_opt as opt
+from coreai_opt.quantization import ModuleQuantizerConfig, Quantizer, QuantizerConfig
+from coreai_opt.quantization.spec import (
+    default_activation_quantization_spec,
+    default_weight_quantization_spec,
+)
+
+global_config = ModuleQuantizerConfig(
+    op_state_spec={"weight": default_weight_quantization_spec()},
+    op_input_spec={"*": default_activation_quantization_spec()},
+    op_output_spec={"*": default_activation_quantization_spec()},
+)
+quant_config = QuantizerConfig(global_config=global_config)
+
+quantizer = Quantizer(model, quant_config)
+prepared_model = quantizer.prepare(example_inputs)
+
+with quantizer.calibration_mode():
+    for batch in calibration_dataloader:
+        prepared_model(batch)
+
+final_model = quantizer.finalize(backend=opt.ExportBackend.CoreAI)
+```
+
+## Step 3: Export and convert to Core AI
+
+After quantization is complete and the model is finalized, `_subexport_and_restore` API exports each patched composite op and restores the original `forward` method in the model.
+Note that the first argument to `_subexport_and_restore` is the original module that was patched in Step 1, not the finalized `GraphModule`.
+
+```python
+import coreai_torch
+from coreai_torch import TorchConverter, _subexport_and_restore
+
+exported_program = torch.export.export(final_model, example_inputs).run_decompositions(
+    coreai_torch.get_decomp_table()
+)
+externalized = _subexport_and_restore(model, exported_program)
+
+coreai_program = (
+    TorchConverter()
+    .add_exported_program(
+        exported_program, _externalized_exported_programs=externalized
+    )
+    .to_coreai()
+)
+```
+
+In the Core AI graph, the composite op is emitted as a separate private graph that `@main` reaches through `coreai.invoke`:
+
+```text
+// composite op body
+coreai.graph private noinline @norm_57e2d4a8(%arg0: tensor<1x32xf32> {coreai.name = "input"}, %arg1: tensor<32xf32> {coreai.name = "scale"}) -> (tensor<1x32xf32>) attributes {composite_decl = ...} {
+  %2 = coreai.decomposable.broadcasting_mul %0, %1 : (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
+  %4 = coreai.reduce_mean %2, %3 : (tensor<1x32xf32>, tensor<1xsi32>) -> tensor<1x1xf32>
+  %8 = coreai.decomposable.broadcasting_add %6, %7 : (tensor<1x1xf32>, tensor<f32>) -> tensor<1x1xf32>
+  %9 = coreai.rsqrt %8 : tensor<1x1xf32> -> tensor<1x1xf32>
+  %12 = coreai.decomposable.broadcasting_mul %10, %11 : (tensor<1x32xf32>, tensor<1x1xf32>) -> tensor<1x32xf32>
+  %15 = coreai.decomposable.broadcasting_mul %13, %14 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  coreai.output %15 : tensor<1x32xf32>
+}
+
+coreai.graph @main(%arg0: tensor<1x32xf32> {coreai.name = "x"}) -> (tensor<1x32xf32>) {
+  %44 = coreai.decomposable.broadcasting_add %43, %2 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  %53 = coreai.quantize %44, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %62 = coreai.dequantize %53, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+
+  // externalized composite op invocation
+  %63 = coreai.invoke @norm_57e2d4a8(%62, %0) : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  %72 = coreai.quantize %63, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %81 = coreai.dequantize %72, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+  %84 = coreai.decomposable.broadcasting_batch_matmul %81, %83 : (tensor<1x32xf32>, tensor<32x32xf32>) -> tensor<1x32xf32>
+}
+```
+
+(`coreai.cast`, `coreai.constant` and `coreai.reshape` ops omitted above for brevity.)
+
+The composite body carries no `coreai.quantize` or `coreai.dequantize` op and stays in full precision.
+
+## Quantizing the composite op boundary
+
+The `coreai.quantize` pairs surrounding the `coreai.invoke` above come from the global config. They are the output quantizer of the preceding `Linear` and the input quantizer of the following one.
+The composite op boundary itself is not targeted by the global config.
+
+To target the boundary specifically, use `module_input_spec` and `module_output_spec` on a {class}`~coreai_opt.quantization.config.ModuleQuantizerConfig` scoped by `module_type_configs` or `module_name_configs`.
+
+To see this in isolation, the following example uses a model with the composite op alone and specifies a module level spec to quantize it's boundary.
+
+```python
+from coreai_opt.quantization.spec import (
+    PerTensorGranularity,
+    QuantizationScheme,
+    QuantizationSpec,
+)
+
+
+class RMSNormOnly(nn.Module):
+    def __init__(self, dim=32):
+        super().__init__()
+        self.norm = RMSNormComposite()
+        self.norm_weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return self.norm(x, self.norm_weight)
+
+
+boundary_spec = QuantizationSpec(
+    dtype=torch.int8,
+    qscheme=QuantizationScheme.SYMMETRIC,
+    granularity=PerTensorGranularity(),
+)
+quant_config = QuantizerConfig(
+    module_type_configs={
+        RMSNormComposite: ModuleQuantizerConfig(
+            module_input_spec={"*": boundary_spec},
+            module_output_spec={"*": boundary_spec},
+        )
+    },
+)
+```
+
+Running the same patch, prepare, calibrate, finalize and convert steps as above, gives a `@main` graph containing just the boundary quantization and the composite call.
+
+```text
+coreai.graph private noinline @norm_20ea9665(%arg0: tensor<1x32xf32> {coreai.name = "input"}, %arg1: tensor<32xf32> {coreai.name = "scale"}) -> (tensor<1x32xf32>) attributes {composite_decl = ...} {
+  %2 = coreai.decomposable.broadcasting_mul %0, %1 : (tensor<1x32xf32>, tensor<1x32xf32>) -> tensor<1x32xf32>
+  %4 = coreai.reduce_mean %2, %3 : (tensor<1x32xf32>, tensor<1xsi32>) -> tensor<1x1xf32>
+  %8 = coreai.decomposable.broadcasting_add %6, %7 : (tensor<1x1xf32>, tensor<f32>) -> tensor<1x1xf32>
+  %9 = coreai.rsqrt %8 : tensor<1x1xf32> -> tensor<1x1xf32>
+  %12 = coreai.decomposable.broadcasting_mul %10, %11 : (tensor<1x32xf32>, tensor<1x1xf32>) -> tensor<1x32xf32>
+  %15 = coreai.decomposable.broadcasting_mul %13, %14 : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+  coreai.output %15 : tensor<1x32xf32>
+}
+
+coreai.graph @main(%arg0: tensor<1x32xf32> {coreai.name = "x"}) -> (tensor<1x32xf32>) {
+
+  // Input boundary quantizers for the composite op
+  %13 = coreai.quantize %arg0, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %22 = coreai.dequantize %13, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+
+  // externalized composite op invocation
+  %23 = coreai.invoke @norm_20ea9665(%22, %0) : (tensor<1x32xf32>, tensor<32xf32>) -> tensor<1x32xf32>
+
+  // Output boundary quantizers for the composite op
+  %32 = coreai.quantize %23, ... : (tensor<1x32xf32>, ...) -> tensor<1x32xsi8>
+  %41 = coreai.dequantize %32, ... : (tensor<1x32xsi8>, ...) -> tensor<1x32xf32>
+  coreai.output %41 : tensor<1x32xf32>
+}
+```
+
+(`coreai.cast`, `coreai.constant` and `coreai.reshape` ops omitted above for brevity.)
+
+## Notes
+
+- The same set of APIs and steps apply for Quantization Aware Training in `graph` mode as well.

pyproject.toml

@@ -57,7 +57,7 @@ name = "Apple Core AI Optimization Team"
 [project.optional-dependencies]
 coreai = [
     "coreai-core==1.0.0b2",
-    "coreai-torch==0.4.1",
+    "coreai-torch>=0.4.1",
     "scikit-learn>=1.7.2",
 ]
 coreml = [
@@ -205,6 +205,12 @@ conflicts = [
     ],
 ]
 [tool.uv.sources]
+# TEMPORARY: resolve coreai-torch from the module externalization API branch
+# instead of the PyPI release. That branch adds the `_patch_model_for_externalization`
+# and `_subexport_and_restore` entry points used by the externalization tests.
+# It lives on a fork; apple/coreai-torch does not carry it yet. Drop this entry
+# (and re-pin the versions above) once the work lands upstream and is released.
+coreai-torch = { git = "https://github.com/gokulkrishna98/coreai-torch.git", branch = "dev/gokul/module-externalization-api" }
 torch = [
     { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
     { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },

tests/conftest.py

@@ -32,6 +32,7 @@
     "tests.models.mnist",
     "tests.models.resnet",
     "tests.models.simple",
+    "tests.models.composite",
 ]
 
 _DEFAULT_SEED: int = 42

tests/export/export_utils.py

@@ -436,10 +436,14 @@ def convert(
         self,
         traced_model: torch.export.ExportedProgram,
         input_data: torch.Tensor,
+        externalized_model: Any = None,
         **kwargs: Any,
     ) -> AIProgram:
         _, _ = input_data, kwargs
-        coreai_program = self._lower_to_coreai(traced_model)
+        coreai_program = self._lower_to_coreai(
+            traced_model,
+            externalized_model=externalized_model,
+        )
         assert type(coreai_program) is AIProgram
 
         return coreai_program
@@ -526,10 +530,25 @@ def _verify_custom_ops_in_torch_program(
     @staticmethod
     def _lower_to_coreai(
         exported_program: torch.export.ExportedProgram,
+        externalized_model: Any = None,
     ) -> AIProgram:
-        """Lower exported program to Core AI."""
+        """Lower exported program to Core AI.
+
+        Args:
+            exported_program: The exported program to lower.
+            externalized_model: Optional ``torch.nn.Module`` that was marked in
+                place by ``coreai_torch._patch_model_for_externalization``.
+        """
         converter = coreai_torch.TorchConverter()
-        converter.add_exported_program(exported_program)
+        externalized_exported_programs = (
+            coreai_torch._subexport_and_restore(externalized_model, exported_program)
+            if externalized_model is not None
+            else None
+        )
+        converter.add_exported_program(
+            exported_program,
+            _externalized_exported_programs=externalized_exported_programs,
+        )
         return converter.to_coreai()
 
 
@@ -554,6 +573,7 @@ def convert_and_verify(
     expected_ops: Mapping[str, int],
     export_backend: ExportBackend,
     prepared_model_output: torch.Tensor | tuple[torch.Tensor, ...],
+    externalized_model: torch.nn.Module | None = None,
     snr_thresh: float = 20.0,
     psnr_thresh: float = 22.0,
     skip_finalized_model_verify: bool = False,
@@ -568,6 +588,9 @@ def convert_and_verify(
         export_backend: Target inference stack (CoreML or CoreAI)
         prepared_model_output: Pre-computed reference output from the prepared
             PyTorch model (single tensor or tuple).
+        externalized_model: Optional ``torch.nn.Module`` that was patched in place by
+            ``coreai_torch._patch_model_for_externalization``. Only supported by the
+            CoreAI backend.
         snr_thresh: Minimum acceptable SNR value
         psnr_thresh: Minimum acceptable PSNR value
         skip_finalized_model_verify: If True, skip forward pass verification on
@@ -578,9 +601,20 @@ def convert_and_verify(
     Returns:
         The converted model in the specified format
 
+    Raises:
+        ValueError: If externalized_model is given for a non-CoreAI backend.
+
     """
     converter = create_converter(export_backend)
 
+    if externalized_model is not None:
+        if export_backend is not ExportBackend.CoreAI:
+            msg = (
+                f"externalized_model is only supported by the CoreAI backend, got {export_backend}"
+            )
+            raise ValueError(msg)
+        converter_kwargs["externalized_model"] = externalized_model
+
     # Run finalized model forward pass BEFORE tracing. torch.export.export()
     # (called in trace) may mutate the model (e.g., strip parametrizations on
     # older PyTorch versions), so the forward pass must happen first.