Repository navigation
XNNPACKQuantizer: apply module filters to ops only, not weight nodes - #24056
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24056
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 05d9be3 with merge base dccea61 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Thanks! Overall this makes sense. One gap remains: conv+relu (and linear+relu) under set_module_type(nn.Conv2d) still never gets quantized. The trailing relu is called in the parent's forward, so its nn_module_stack doesn't include Conv2d, and the filter rejects the conv+relu group. Could we filter on the main op node only, as This would mean that you even if you exclude relu, it can get pulled into a preceding quantized conv. But, I dont think that is a very common case |
|
Done, filter now checks only the main op; added a conv+relu test. |
JakeStevens
left a comment
There was a problem hiding this comment.
try rebasing so we can check CI and should be good to go
|
These fails were infra specific, but rebasing either way. |
Annotators that put the weight/bias get_attr nodes into the partition (conv, conv_relu, conv_bn, linear_relu, ...) ran filter_fn over the whole partition. Those nodes carry no nn_module_stack, so set_module_type and set_module_name never matched a conv in an exported graph. Skip get_attr and placeholder nodes, as the linear annotator already does by checking only the op.
A relu called in the parent's forward has no Conv2d/Linear in its nn_module_stack, so set_module_type(nn.Conv2d) rejected every conv+relu partition. Check filter_fn on the conv/linear/add/mul node alone, as _annotate_linear does; a trailing relu is then fused under its op's config.
0ef33af to
05d9be3
Compare
Summary
Annotators that include weight/bias
get_attrnodes in the partition (conv, conv_relu, conv_bn, linear_relu, ...) ranfilter_fnover every node. Those nodes have nonn_module_stack, soset_module_type(torch.nn.Conv2d, ...)andset_module_namenever matched a conv in an exported graph. The filter now checks only the main op (conv/linear/add/mul), like the linear annotator, so a trailing relu is fused under its op's config.Test plan
python -m unittest backends.xnnpack.test.quantizer.test_xnnpack_quantizer backends.xnnpack.test.quantizer.test_pt2e_quantization(newtest_set_module_type_conv2dandtest_set_module_type_conv2d_relufail without the fix;test_grufails before and after on macOS with NoQEngine).cc @GregoryComer @digantdesai @cbilgin @JakeStevens @nil-is-all