AtenToCortexMPass crashes with "unsupported param type, call_function" when aten.linear.default is in preserve_ops

[Tharanga01]

🐛 Describe the bug

Description:
When lowering a CortexMQuantizer-quantized model to the Cortex-M backend using the standard EdgeCompileConfig(preserve_ops=[torch.ops.aten.linear.default, ...], _check_ir_validity=False) pattern (as used in the backend's own CortexMToEdge test stage), CortexMPassManager.transform() crashes during AtenToCortexMPass with:
RuntimeError: unsupported param type, call_function.

import argparse
import logging

import torch
import torch.nn as nn
from torch.export import export

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


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).module()

    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)

    # preserve_ops=[aten.linear.default] keeps `linear` nodes intact so
    # AtenToCortexMPass can substitute them with cortex_m::quantized_linear.
    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)

RuntimeError: unsupported param type, call_function.

Workaround found
Removing torch.ops.aten.linear.default from preserve_ops avoids the crash

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

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @psiddh @AdrianLundell

[mansnils]

Thanks for reporting this. We will have a look.

[AdrianLundell]

Hi, the reason for this crash is that the Cortex-M backend currently requires constant weights since they are repacked AOT to the format required for the CMSIS-NN kernels used under the hood. As a workaround, try applying the ConstantFoldingPass from the Arm backend before quantization, i.e update your example to:

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

With this I managed to run your script and got

...
cortex_m::quantize_per_tensor
cortex_m::quantized_linear
aten::unsqueeze_copy
aten::permute_copy
aten::squeeze_copy
dim_order_ops::_clone_dim_order
aten::select_copy
cortex_m::dequantize_per_tensor
aten::full
aten::mul
aten::expand_copy
aten::bmm
aten::_softmax
aten::eq
aten::logical_not
aten::any
aten::full_like
aten::where
cortex_m::transpose
cortex_m::quantized_add
aten::native_layer_norm
Wrote simple_transformer_int8_CortexM.pte
Transformer(
  (encoder): TransformerEncoder(
    (layers): ModuleList(
      (0): TransformerEncoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=32, out_features=32, bias=True)
        )
        (linear1): Linear(in_features=32, out_features=64, bias=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (linear2): Linear(in_features=64, out_features=32, bias=True)
        (norm1): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
        (norm2): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
        (dropout1): Dropout(p=0.0, inplace=False)
        (dropout2): Dropout(p=0.0, inplace=False)
      )
    )
    (norm): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
  )
  (decoder): TransformerDecoder(
    (layers): ModuleList(
      (0): TransformerDecoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=32, out_features=32, bias=True)
        )
        (multihead_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=32, out_features=32, bias=True)
        )
        (linear1): Linear(in_features=32, out_features=64, bias=True)
        (dropout): Dropout(p=0.0, inplace=False)
        (linear2): Linear(in_features=64, out_features=32, bias=True)
        (norm1): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
        (norm2): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
        (norm3): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
        (dropout1): Dropout(p=0.0, inplace=False)
        (dropout2): Dropout(p=0.0, inplace=False)
        (dropout3): Dropout(p=0.0, inplace=False)
      )
    )
    (norm): LayerNorm((32,), eps=1e-05, elementwise_affine=True, bias=True)
  )
)

input torch.Size([8, 1, 32])
output torch.Size([8, 1, 32])

Running transformers on the Cortex-M backend is really pushing new grounds so I am very interested in following your development here!

[Tharanga01]

Thanks @AdrianLundell ! Confirmed the ConstantFoldingPass workaround fixes it on my end too.