heterogeneous quantization for different backends

[goka-wu]

My team aims to develop a heterogeneous inference framework using ExecuTorch, but we are currently grappling with challenges in heterogeneous quantization.

Consider this scenario: For a single model, we plan to simultaneously delegate computations to NPU-A, NPU-B, and CPU backends. The CPU will utilize XNNPACKQuantizer, while NPU-A and NPU-B require custom quantization algorithms. How can we apply these three distinct quantization methods to their respective partitions before the graph is partitioned?

Based on my understanding of ExecuTorch's workflow: Quantization at Aten IR → Partitioning at EdgeIR,
if the graph is partitioned into:

P1 (executed on NPU-A)
P2 (executed on NPU-B)
P3 (executed on CPU via XNNPACK)

How can we ensure that the NPU-A-specific quantization algorithm is applied to P1 during the Aten IR quantization stage, given that partitioning occurs later at the Edge IR stage?

[yujiaoliang]

I also encountered the same issue. After quantization, the partition fails.

[GregoryComer]

@kimishpatel @metascroy Do you have any suggestions on how to handle this?

[kimishpatel]

have you looked at https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/quantizer/composable_quantizer.py#L17

[Bunleab25]

Hi @goka-wu,

The core challenge with heterogeneous quantization in ExecuTorch is phase ordering: Quantization happens on the ATen graph (ExportedProgram) via PT2E / TorchAO, whereas backend delegation and partitioning occur downstream at Edge IR (to_edge_transform_and_lower).

If you apply a single global quantizer, it will annotate ops in ways incompatible with other backends, causing subsequent partitioners to fail pattern matching at partition boundaries.

Here is how you can structure and solve this:

---

Approach 1: Use ComposableQuantizer with Submodule/Node Scoping (Recommended)

TorchAO / PyTorch 2.x provides torch.ao.quantization.quantizer.composable_quantizer.ComposableQuantizer, which allows chaining multiple quantizers that only target specific submodules or nodes.

1. Scope Quantizers by Submodule Stack

When designing custom quantizers for NPU-A and NPU-B, override annotate() to inspect each node's nn_module_stack metadata. Only annotate nodes belonging to that backend's target submodule:

import torch
from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
from executorch.exir import to_edge_transform_and_lower
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner

# Custom Quantizer for NPU-A that only annotates its assigned submodules
class NPUAQuantizer(CustomBaseQuantizer):
    def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
        for node in model.graph.nodes:
            module_stack = node.meta.get("nn_module_stack", {})
            if any("submodule_npu_a" in str(mod_path) for mod_path in module_stack.values()):
                # Apply NPU-A specific quantization annotation (e.g. INT8 symmetric per-tensor)
                self._annotate_node(node)
        return model

# Similarly for NPU-B
class NPUBQuantizer(CustomBaseQuantizer):
    def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
        for node in model.graph.nodes:
            module_stack = node.meta.get("nn_module_stack", {})
            if any("submodule_npu_b" in str(mod_path) for mod_path in module_stack.values()):
                self._annotate_node(node)
        return model

2. Compose and Quantize

Compose the quantizers in a single pass. Order matters if there are fallback nodes:

# 1. Export unquantized model to ATen IR
exported_program = torch.export.export(model, example_inputs)

# 2. Compose the quantizers
composed_quantizer = ComposableQuantizer([
    npu_a_quantizer,
    npu_b_quantizer,
    xnnpack_quantizer  # CPU fallback
])

# 3. Prepare, calibrate, and convert
prepared_graph = prepare_pt2e(exported_program.graph_module, composed_quantizer)
# Calibrate with representative inputs...
converted_graph = convert_pt2e(prepared_graph)

3. Handling Partition Boundaries at Edge IR

The common reason partitions fail after quantization is dangling dequantize/quantize (q/dq) pairs at partition boundaries:

• If NPU-A outputs an INT8 tensor with scale $S_A$, and NPU-B expects INT8 with scale $S_B$ (or CPU expects Float32), a standard partitioner will fail to absorb the boundary conversion.
• Rule: Keep boundary activations in Float32 (dequantize at the output of P1, re-quantize at the input of P2). Each backend's partitioner will then cleanly match its internal quantize -> op -> dequantize pattern without attempting to ingest the other backend's quantization nodes.

4. Multi-backend Lowering

Pass the partitioners in sequence to to_edge_transform_and_lower:

# Lower to Edge IR with chained partitioners
edge_program = to_edge_transform_and_lower(
    torch.export.ExportedProgram(converted_graph, ...),
    partitioner=[
        NPUAPartitioner(),
        NPUBPartitioner(),
        XnnpackPartitioner()  # Remaining unlowered ops stay on CPU
    ]
)

# Compile to final .pte
pte_program = edge_program.to_executorch()
with open("heterogeneous_model.pte", "wb") as f:
    f.write(pte_program.buffer)

---

Approach 2: Modular Submodule Export & Lowering (Alternative)

If your model has distinct functional blocks (e.g. Feature Extractor on NPU-A, Backbone on NPU-B, Classifier on CPU):

1. Split the model into submodules: model.p1, model.p2, model.p3.
2. Export and quantize each submodule independently with its native quantizer.
3. Lower each independently to its respective backend delegate using to_edge_transform_and_lower().
4. Wrap them in a parent coordinator module, or execute them sequentially in your C++ ExecuTorch runtime passing the output buffer of method 1 as input to method 2. This completely avoids cross-quantizer graph contamination.