Qualcomm AI Engine Direct - Fold BatchNorm into the preceding conv - #22994
Conversation
In fp16 on HTP a standalone BatchNorm multiplies the conv's fp16 rounding error by its per-channel scale, which reaches ~100 after depthwise convs. EfficientNetV2-S fp16 on SM8850 drops to 0.1% top-1 (NaNs on some inputs); with the fold it is 74.25% (fp32: 74.40%) and 30% faster. Quantized graphs are unaffected, since PTQ already folds these pairs.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22994
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 1 PendingAs of commit 7d6b672 with merge base 8081eb8 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
nil-is-all
added
the
module: qnn
executorch-triage
Bot
added
the
community: contribution
qti-horodnic
left a comment
qti-horodnic
left a comment
There was a problem hiding this comment.
Hi @msluszniak, thank you for your contribution. A few comments:
-
This change breaks the CI,
test_qnn_conv1d_batch_normfails withExpect p_conv1d_fused_bn_weight to be parameter, buffer, get_attr, or lifted tensor constant. The fused placeholders are namedconv1d_fused_bn_weight, but export renames them top_...after decomposition, while the edge passes still look names up in the atenExportedProgramsignature soCanonicalizeConv'sget_parametermisses. I suggest registering the fused tensors the way the other QNN passes do (register_buffer + graph.get_attr, cf. CanonicalizeConv/ConvertLinearToConv2d), which would also drop theBUCKchange. -
Please add tests in the new rework pass test framework and the delegate tests. A
test_fuse_batch_norm_with_convinbackends/qualcomm/tests/rework/passes/test.pywould match the per-pass convention there, and an fp16 case intest_qnn_delegate.py(e.g.Conv1dBn / Conv2dBnHardtanhMean) asserting nobatch_normsurvives in the delegated graph would have caught thep_conv1d_fused_bn_weightfailure. The standalone test intest_passes.pyexercises the fold in isolation and passes regardless. -
This pass largely duplicates the generic
backends/transforms/fuse_batch_norm_with_conv.py, which other backends reuse. Could we reuse it or at least add a comment on why a QNN-specific copy is needed? I assume because that one matchesexir_ops.edge.aten.convolutionpost-decomposition, while this has to run pre-dispatch onconv1d/2d/3dinget_export_passes(). If a shared version isn't practical, note that the existing pass registers the fused tensors viaregister_parameter + graph.get_attr, which sidesteps the placeholder-renaming issue seen in the CI.
- Qualcomm copyright header, matching the other passes here. - Collect the erased BatchNorm's former args instead of an explicit list, so any leftover constant input is deleted rather than surviving as a dead one. - Name the fused placeholders "p_<key>" and register the state_dict key without the prefix. The to_edge passes resolve parameters against the pre-edge program by the node name they see in the edge graph, and to_edge renames a parameter placeholder to "p_" + its state_dict key; without the prefix the fused weight was renamed and CanonicalizeConv could no longer find it (test_qnn_conv1d_batch_norm).
c6371f8 to
c61df35
Compare
Addresses review: 1. Drop backends/qualcomm/_passes/fuse_batch_norm_with_conv.py and run the shared backends/transforms FuseBatchNormWithConvPass from the to_edge pipeline instead of the export pipeline. It matches the edge dialect, which is why the copy existed; at that stage it folds conv1d/2d/3d, with and without bias. This also removes the placeholder-renaming problem that broke test_qnn_conv1d_batch_norm, since the shared pass registers the fused tensors with register_parameter + graph.get_attr. BUCK now depends on :fuse_batch_norm_with_conv rather than :utils. CanonicalizeConv assumed a conv weight is a placeholder (fp) or a dequant node (QDQ), and raised IndexError on a get_attr weight. Only conv1d and dilated transpose convs dereference it, so nothing hit this before; accept get_attr there too. 2. Tests: test_fuse_batch_norm_with_conv in the rework per-pass framework (conv1d with and without bias, depthwise conv2d, conv3d), and a delegated graph assertion that no batch_norm survives, on the fp16 Conv1dBn and Conv2dBnHardtanhMean cases. Drops the isolated test_passes.py test, which passed regardless of what the QNN pipeline did.
Making CanonicalizeConv depend on FuseBatchNormWithConvPass moved it after LayoutTransform in the solved order. CanonicalizeConv rewrites conv1d into conv2d and copies the original node meta, so the 4D convolution inherited the 3-element QCOM_AXIS_ORDER LayoutTransform had already written, and the squeeze it inserts inherited it too. Lowering any conv1d followed by a batch norm then failed in op_squeeze with permute(): input.dim() = 4 is not equal to len(dims) = 3 The old pass list satisfied this ordering only by accident; state it in the dependency table.
assert_batch_norm_folded read the graph after partitioning, where nothing is left but the delegate call, so it passed on a tree without the fold at all. Run the same pipeline with no partitioner instead: on the parent commit the conv1d cases now report the surviving native_batch_norm, and they pass here.
| AnnotateStack: [RemoveRedundancy], | ||
| AnnotateUnbind: [RemoveRedundancy], | ||
| CanonicalizeConv: [FoldQDQ], | ||
| CanonicalizeConv: [FoldQDQ, FuseBatchNormWithConvPass], |
There was a problem hiding this comment.
FuseBatchNormWithConvPass doesn't look at transposed (conv.args[6]) and calls fuse_conv_bn_weights without transpose=, so ConvTranspose + BatchNorm folds against the wrong weight axis. On this branch ConvTranspose2d(8, 8, 3, bias=False) + BatchNorm2d(8) gives max abs diff 9.5 vs eager, and ConvTranspose2d(4, 8, ...) raises RuntimeError: The size of tensor a (4) must match the size of tensor b (8) at non-singleton dimension 0. This seems to actually be an existing issue in the shared pass, but QNN now hits it.
To keep this PR's scope, could we subclass it under backends/qualcomm/_passes/ and override can_fuse to return False when conv.args[ConvParamIdx.TRANSPOSED], so transposed convs keep their standalone BN as today? Something like:
class FuseBatchNormWithConv(FuseBatchNormWithConvPass):
"""Skips transposed convolutions.
fuse_conv_bn_weights() scales dim 0 of the weight, but a transposed weight
is [in, out/groups, *kernel], so folding there would be incorrect.
"""
@staticmethod
def can_fuse(conv, bn, program) -> bool:
if conv.args[ConvParamIdx.TRANSPOSED]:
return False
return FuseBatchNormWithConvPass.can_fuse(conv, bn, program)
...
Alos, a ConvTranspose2d + BN test case in the new pattern.py set would lock that in since none of the existing test_qnn_backend_conv_transpose models have a BN after them, which is why CI stays green.
There was a problem hiding this comment.
Ok, applied this one here and will create a separate PR for fix in shared passes.
qti-horodnic
left a comment
qti-horodnic
left a comment
There was a problem hiding this comment.
Thanks for making the changes, LGTM.
@psiddh Kindly take a look when you get the chance, change looks good to be merged to me.
fuse_conv_bn_weights() scales dim 0 of the convolution weight. That is the output channel for a regular conv, but a transposed weight is [in, out/groups, *kernel], so the fold hits the wrong axis: it raises when in != out, and silently returns wrong values when they match (ConvTranspose2d(8, 8, 3) + BatchNorm2d(8) drifts by 12.7 against eager). Subclass the shared pass under backends/qualcomm/_passes and decline transposed convs in can_fuse, so they keep their standalone BatchNorm. The existing conv_transpose models have no BatchNorm after them, which is why CI stayed green; the new conv_transpose2d_bn_is_skipped case covers it and fails with the original RuntimeError when the skip is removed.
8c49384 to
7d6b672
Compare
|
lgtm |
4 participants
Summary
Fold BatchNorm into the preceding conv for QNN, by running the shared
backends/transformspass in the to_edge pipeline. In HTP fp16 a standalone BatchNorm amplifies the conv's rounding error by its scale (~100 after depthwise convs): EfficientNetV2-S fp16 on SM8850 goes from 0.1% to 74.25% top-1 (fp32: 74.40%) and runs 30% faster. Quantized graphs are unaffected.Test plan
tests/rework/passes/test.py::test_fuse_batch_norm_with_conv,test_qnn_delegate.py::test_qnn_conv1d_batch_norm; EfficientNetV2-S fp16 on a Galaxy S26 Ultra (SM8850), 2000 ImageNetV2 images.cc @cbilgin @psiddh