🐛 Describe the bug
Description:
When quantizing a transformer model with CortexMQuantizer and lowering to the Cortex-M backend, aten.bmm/aten.mul/aten._softmax are not quantized(These are mentioned in CMSIS-NN Supported Operators list).
import argparse
import logging
import torch
import torch.nn as nn
from torch.export import export,default_decompositions
from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge
from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
# Post-training static INT8 quantization (PT2E flow).
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("simple_transformer")
_SDPA_DECOMP_TABLE = {
op: decomp
for op, decomp in default_decompositions().items()
if "scaled_dot_product" or "_safe_softmax" in str(op) in str(op)
}
def make_calibration_inputs(example_inputs, num_batches: int = 8):
"""Yield representative inputs for static PTQ calibration.
TODO: Replace the synthetic tensors with REAL representative inference
inputs. Static quantization freezes activation scales from whatever is
observed here, so synthetic data yields meaningless ranges and poor
accuracy.
"""
for _ in range(num_batches):
yield tuple(torch.randn_like(t) for t in example_inputs)
def quantize_static_int8(model, example_inputs, calib_batches: int = 8):
"""Apply post-training static INT8 quantization via the PT2E flow."""
model.eval()
# prepared_graph = export(model, example_inputs).run_decompositions(_SDPA_DECOMP_TABLE)
prepared_graph = export(model, example_inputs).module()
from executorch.backends.arm._passes import ConstantFoldingPass
prepared_graph = ConstantFoldingPass().call(prepared_graph).graph_module # Pre-computes and removes the call_function op sitting between the linear nodes and their weights
quantizer = CortexMQuantizer()
prepared = prepare_pt2e(prepared_graph, quantizer)
with torch.no_grad():
for sample in make_calibration_inputs(example_inputs, calib_batches):
prepared(*sample)
return convert_pt2e(prepared)
def main(args):
S, B, D = args.seq_len, 1, args.d_model
# Minimal one-layer encoder-decoder transformer.
model = nn.Transformer(
d_model=D,
nhead=args.nhead,
num_encoder_layers=args.num_layers,
num_decoder_layers=args.num_layers,
dim_feedforward=args.dim_feedforward,
dropout=0.0,
batch_first=False,
)
model.eval()
example_input = (torch.randn(S, B, D), torch.randn(S, B, D))
print("Applying post-training static INT8 quantization")
quantized_model = quantize_static_int8(model, example_input, args.calib_batches)
quantized_exported_program = export(quantized_model, example_input)
config = EdgeCompileConfig(
preserve_ops=[torch.ops.aten.linear.default],
_check_ir_validity=False,
)
edge_program_manager = to_edge(quantized_exported_program, compile_config=config)
pass_manager = CortexMPassManager(edge_program_manager.exported_program())
edge_program_manager._edge_programs["forward"] = pass_manager.transform()
et_program = edge_program_manager.to_executorch(
config=ExecutorchBackendConfig(extract_delegate_segments=False)
)
for op in et_program.executorch_program.execution_plan[0].operators:
print(op.name)
pte_path = "simple_transformer_int8_CortexM.pte"
with open(pte_path, "wb") as f:
f.write(et_program.buffer)
print(f"Wrote {pte_path}")
print(model)
src = torch.rand((S, B, D))
tgt = torch.rand((S, B, D))
out = model(src, tgt)
print("input", src.size())
print("output", out.size())
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Export a minimal quantized encoder-decoder transformer to .pte for the Cortex-M backend"
)
parser.add_argument("--d_model", type=int, help="d_model", default=32)
parser.add_argument("--nhead", type=int, help="nhead", default=2)
parser.add_argument("--num_layers", type=int, help="number of encoder and decoder layers", default=1)
parser.add_argument("--dim_feedforward", type=int, help="dim_feedforward", default=64)
parser.add_argument("--seq_len", type=int, help="sequence length", default=8)
parser.add_argument("--calib_batches", type=int, help="number of calibration batches for static PTQ", default=8)
args = parser.parse_args()
main(args)Additional context:
Running decomposition before quantizing the model does quantize aten::_softmax and aten::bmm, aten::mul is the one that stays unquantized. (Separately, I also ran into the decomposed .pte not opening in Netron , unrelated to this issue, but flagging in case it's a known/expected
limitation for graphs with this many ops.)
Versions
Versions
PyTorch version: 2.13.0+cpu
OS: Ubuntu 24.04.4 LTS (x86_64)
Python version: 3.11.15 (main, Mar 11 2026, 17:20:07) [GCC 14.3.0] (64-bit runtime)
CPU:
Architecture: x86_64
Versions of relevant libraries:
[pip3] executorch==1.4.0+3dd7ccd
[pip3] flake8==6.1.0
[pip3] flake8-breakpoint==1.1.0
[pip3] flake8-bugbear==24.4.26
[pip3] flake8-comprehensions==3.14.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pyi==23.5.0
[pip3] mypy==1.14.1
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.4.4
[pip3] pytorch_tokenizers==1.4.1
[pip3] torch==2.13.0+cpu
[pip3] torchao==0.18.0
[pip3] torchaudio==2.11.0+cpu
[pip3] torchdata==0.11.0
[pip3] torchsr==1.0.4
[pip3] torchtune==0.0.0
[pip3] torchvision==0.28.0
[pip3] triton==3.7.1
[conda] executorch 1.4.0+3dd7ccd pypi_0 pypi
[conda] numpy 2.4.4 pypi_0 pypi
[conda] pytorch-tokenizers 1.4.1 pypi_0 pypi
[conda] torch 2.13.0+cpu pypi_0 pypi
[conda] torchao 0.18.0 pypi_0 pypi
[conda] torchaudio 2.11.0+cpu pypi_0 pypi
[conda] torchdata 0.11.0 pypi_0 pypi
[conda] torchfix 0.6.0 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchtune 0.0.0 pypi_0 pypi
[conda] torchvision 0.28.0 pypi_0 pypi
[conda] triton 3.7.1 pypi_0 pypi
🐛 Describe the bug
Versions
Versions
PyTorch version: 2.13.0+cpu
OS: Ubuntu 24.04.4 LTS (x86_64)
Python version: 3.11.15 (main, Mar 11 2026, 17:20:07) [GCC 14.3.0] (64-bit runtime)
CPU:
Architecture: x86_64
Versions of relevant libraries:
[pip3] executorch==1.4.0+3dd7ccd
[pip3] flake8==6.1.0
[pip3] flake8-breakpoint==1.1.0
[pip3] flake8-bugbear==24.4.26
[pip3] flake8-comprehensions==3.14.0
[pip3] flake8-plugin-utils==1.3.3
[pip3] flake8-pyi==23.5.0
[pip3] mypy==1.14.1
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.4.4
[pip3] pytorch_tokenizers==1.4.1
[pip3] torch==2.13.0+cpu
[pip3] torchao==0.18.0
[pip3] torchaudio==2.11.0+cpu
[pip3] torchdata==0.11.0
[pip3] torchsr==1.0.4
[pip3] torchtune==0.0.0
[pip3] torchvision==0.28.0
[pip3] triton==3.7.1
[conda] executorch 1.4.0+3dd7ccd pypi_0 pypi
[conda] numpy 2.4.4 pypi_0 pypi
[conda] pytorch-tokenizers 1.4.1 pypi_0 pypi
[conda] torch 2.13.0+cpu pypi_0 pypi
[conda] torchao 0.18.0 pypi_0 pypi
[conda] torchaudio 2.11.0+cpu pypi_0 pypi
[conda] torchdata 0.11.0 pypi_0 pypi
[conda] torchfix 0.6.0 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchtune 0.0.0 pypi_0 pypi
[conda] torchvision 0.28.0 pypi_0 pypi
[conda] triton 3.7.1 pypi_0 pypi