Summary
to_edge_transform_and_lower does not constant-fold subgraphs whose inputs are all parameters, so a convolution or a linear whose weight is computed — anything wrapped in torch.nn.utils.parametrize, which is what weight_norm and spectral_norm are — never reaches the delegate. It falls to the portable kernels silently, and the cost is not small.
ExecuTorch already ships the pass that fixes it. exir.passes.constant_prop_pass is described as "constant propagation for Exported Program with lifted parameters", which is exactly this case. The Qualcomm and Samsung backends call it; XNNPACK's own test/ops/test_conv1d.py runs it as a test stage. The generic lowering path does not.
executorch 1.4.0, torch 2.13.0, Python 3.12, macOS arm64, pip wheel.
Measurement
wav2vec2-large's positional convolution, Conv1d(1024, 1024, kernel_size=128, groups=16) under weight_norm, at sequence length 499 — the module on its own, nothing else in the graph:
delegated median latency
as exported 70.0% 3293.7 ms
+ constant_prop_pass 100.0% 3.5 ms
941x, with the output unchanged (max_abs_diff against eager: 4.578e-05 before, 2.003e-05 after).
What stays outside the delegate without the pass is the weight computation and the convolution it feeds:
{'aten.sum.dim_IntList': 1, 'aten.pow.Tensor_Scalar': 1, 'aten.convolution.default': 1}
XNNPACK's partitioner configs require the weight to be a static parameter (is_param_node in partition/config/gemm_configs.py), and a parametrized weight is not one, so ConvolutionConfig declines and takes the convolution with it. The refusal is logged through WhyNoPartition at DEBUG, so at default log level nothing is printed.
Repro
import torch
from transformers import Wav2Vec2ForCTC
from executorch.exir import to_edge_transform_and_lower
from executorch.exir.passes.constant_prop_pass import constant_prop_pass
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
model = Wav2Vec2ForCTC.from_pretrained(
"jonatasgrosman/wav2vec2-large-xlsr-53-japanese", dtype=torch.float32).eval()
pos = model.wav2vec2.encoder.pos_conv_embed
x = (torch.randn(1, 499, 1024),)
for fold in (False, True):
ep = torch.export.export(pos, x)
if fold:
ep = constant_prop_pass(ep)
edge = to_edge_transform_and_lower(ep, partitioner=[XnnpackPartitioner()])
# count executorch_call_delegate vs call_function at top level, then time
# edge.to_executorch() through the runtimeWhole models pay it in one lump: the same .pte for that model goes from 3485.4 ms to 267.8 ms once the parametrization is out of the graph, on a build whose correlation against eager is 1.000000 either way.
Why it is easy to miss
Nothing looks wrong. The export succeeds, the delegation report shows a healthy-looking 64%, correlation against eager is 1.000000, and the model simply runs slow. The three ops left outside are a pow, a sum and one convolution in a list of hundreds. I spent an afternoon ruling out the attention mask, the portable layer norms, XNNPACK's GEMM path and the attention itself before subtracting the measured parts from the total and going looking for what was missing.
Suggestions
- Run
constant_prop_pass inside to_edge_transform_and_lower before partitioning, or expose it as a flag there. It is already in the tree and already used by two backends.
- Failing that, a warning when a partitioner config declines an op because its weight is not a parameter would turn this into a one-line fix for the user (
torch.nn.utils.parametrize.remove_parametrizations).
Happy to send a PR for either.
Summary
to_edge_transform_and_lowerdoes not constant-fold subgraphs whose inputs are all parameters, so a convolution or a linear whose weight is computed — anything wrapped intorch.nn.utils.parametrize, which is whatweight_normandspectral_normare — never reaches the delegate. It falls to the portable kernels silently, and the cost is not small.ExecuTorch already ships the pass that fixes it.
exir.passes.constant_prop_passis described as "constant propagation for Exported Program with lifted parameters", which is exactly this case. The Qualcomm and Samsung backends call it; XNNPACK's owntest/ops/test_conv1d.pyruns it as a test stage. The generic lowering path does not.executorch 1.4.0, torch 2.13.0, Python 3.12, macOS arm64, pip wheel.
Measurement
wav2vec2-large's positional convolution,
Conv1d(1024, 1024, kernel_size=128, groups=16)underweight_norm, at sequence length 499 — the module on its own, nothing else in the graph:941x, with the output unchanged (max_abs_diff against eager: 4.578e-05 before, 2.003e-05 after).
What stays outside the delegate without the pass is the weight computation and the convolution it feeds:
XNNPACK's partitioner configs require the weight to be a static parameter (
is_param_nodeinpartition/config/gemm_configs.py), and a parametrized weight is not one, soConvolutionConfigdeclines and takes the convolution with it. The refusal is logged throughWhyNoPartitionat DEBUG, so at default log level nothing is printed.Repro
Whole models pay it in one lump: the same
.ptefor that model goes from 3485.4 ms to 267.8 ms once the parametrization is out of the graph, on a build whose correlation against eager is 1.000000 either way.Why it is easy to miss
Nothing looks wrong. The export succeeds, the delegation report shows a healthy-looking 64%, correlation against eager is 1.000000, and the model simply runs slow. The three ops left outside are a
pow, asumand oneconvolutionin a list of hundreds. I spent an afternoon ruling out the attention mask, the portable layer norms, XNNPACK's GEMM path and the attention itself before subtracting the measured parts from the total and going looking for what was missing.Suggestions
constant_prop_passinsideto_edge_transform_and_lowerbefore partitioning, or expose it as a flag there. It is already in the tree and already used by two backends.torch.nn.utils.parametrize.remove_parametrizations).Happy to send a PR for either.