Quantization with Composite Ops

Core AI recognizes certain well-known building blocks, such as SDPA or RMSNorm, as composite ops and applies optimized implementations for them at runtime. These composite ops are available via the coreai_torch.composite_ops API. When a PyTorch model uses one of these ops, it can be converted to Core AI using one of the following APIs:

• Using TorchConverter.add_pytorch_module() along with the externalize_modules arg: As described here, this approach takes as input an nn.Module torch model, along with the composite op torch module, specified via the externalize_modules arg. Any coreai-opt optimizer that yields a torch model of type nn.Module can be converted using this API.

• Using TorchConverter.add_exported_program() : This coreai-torch API operates on an already exported torch program. When using coreai-opt’s quantizer with graph execution mode, which results in a torch ExportedProgram, there are a few additional steps required to ensure correct conversion and externalization of the composite op modules. This process is explained below with an example.

Example

This example uses the same RMSNormComposite op from the Externalization guide, however, the same process applies for all composite ops with their respective ExternalizeSpecs.

First define the model that uses a composite op:

import torch
import torch.nn as nn


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


original_model = Model().eval()
example_inputs = (torch.randn(1, 32),)

Next, apply coreai-opt quantization with graph mode, and convert using coreai-torch:

import coreai_opt as opt
from coreai_opt.quantization import Quantizer, QuantizerConfig
from coreai_opt.quantization.config import ExecutionMode
import coreai_torch
from coreai_torch import (
    ExternalizeSpec,
    TorchConverter,
    _patch_model_for_externalization,
    _subexport_and_restore,
)

# Patch the model in-place
# to mark RMSNormComposite module for externalization
_patch_model_for_externalization(
    original_model,
    targets=[
        ExternalizeSpec(
            target_class=RMSNormComposite,
            composite_op_name="rms_norm",
            composite_attrs=["axes", "eps", "version"],
        )
    ],
)

# Usual flow using coreai-opt APIs for quantization
config = QuantizerConfig.presets.w8(execution_mode=ExecutionMode.GRAPH)
quantizer = Quantizer(original_model, config)
prepared_model = quantizer.prepare(example_inputs)
finalized_quantized_model = quantizer.finalize(backend=opt.ExportBackend.CoreAI)

# Export to Core AI
exported_program = torch.export.export(
    finalized_quantized_model, example_inputs
).run_decompositions(coreai_torch.get_decomp_table())
composite_module_exported_programs = _subexport_and_restore(
    original_model, exported_program
)
coreai_program = (
    TorchConverter()
    .add_exported_program(
        exported_program,
        _externalized_exported_programs=composite_module_exported_programs,
    )
    .to_coreai()
)

Compared to the usual flow, there are two key differences:

Patch original model with custom ops

In graph mode execution, coreai-opt’s quantizer.prepare invokes the torch.export.export API. By default, this causes the composite op body to be inlined in the graph and we lose the information about it’s boundary for externalization. In order to preserve this information, the _patch_model_for_externalization API modifies the original torch model in-place, before preparing the model for graph mode quantization using coreai-opt. Specifically, it replaces the forward method of the nn.modules specified in ExternalizeSpec(target_class) with a torch.library.custom_op that, in turn, invokes the original forward method of the composite op. This keeps the model functionally the same, however, now the sub-module (RMSNormComposite in this example) becomes opaque to the torch.export.export API, i.e., a single node instead of the composite op body being inlined in the exported graph.

Note that in addition to replacing the forward method of the composite ops with a torch.library.custom_op, the _patch_model_for_externalization API also registers some additional information/metadata as attributes on the composite op sub-modules. This metadata is then used by _subexport_and_restore as explained below, after which the model is restored by removing this metadata along with the torch.library.custom_op from the model.

Extract exported programs for submodules

Once a model is patched using the _patch_model_for_externalization API, it can be quantized as usual using coreai-opt’s graph mode (either for PTQ or QAT) and finalized to CoreAI backend.

Now in order to convert the finalized model into a coreai graph (aimodel), the composite op sub-module bodies need to be exported as torch exported programs in order to be included in the graph as externalized function calls. The _subexport_and_restore takes in the original torch model and does exactly that. It returns a list of torch exported programs corresponding to all the composite op sub-modules that were patched for externalization by _patch_model_for_externalization. These torch exported programs are then passed to the add_exported_program API via the _externalized_exported_programs argument to prepare the coreai graph (aimodel). After this process, the coreai graph contains both the main graph and the sub-graphs corresponding to the composite op submodules preserved as function calls.

Notes

• Quantization cannot be applied inside the composite op body because the _patch_model_for_externalization API replaces the composite op body with a torch custom op, thus making it opaque to coreai-opt’s graph mode quantizer.

• However, the boundaries (incoming and outgoing tensors) of the composite ops can be quantized as usual, using the module_input_spec and module_output_spec config kwargs as described in the documentation on coreai-opt configs here.

Warning

The externalization APIs used above, _patch_model_for_externalization and _subexport_and_restore in coreai-torch are currently experimental.