Repository navigation
XNNPACK: fix enable_bf16 load failures and BatchNorm folding into a 3-D linear output - #24477
Merged
Merged
Conversation
…out (#24380) ## What is wrong today XNNPACK supports tensors with at most 6 dimensions, but the XNNPACK partitioner, which decides which ops run on XNNPACK, does not check this. A permute or another op on a 7-D or 8-D tensor then lowers fine but fails to load: ``` Failed to define tensor 0 with code: xnn_status_unsupported_parameter ``` Some vision models hit this in a "patchify" step that cuts an image into patches: ```python y = x.reshape(2, 6, 8, 2, 16, 8, 2, 16).permute(0, 2, 5, 3, 6, 1, 4, 7) y = y.reshape(2, 64, -1) * 2 ``` With the 3-D convolution kernel from #24385, a statically quantized `Conv3d` is worse: it fails to load or, without a fused ReLU, returns one wrong value repeated everywhere, with no error. XNNPACK has no 3-D convolution, so the partitioner leaves the conv out, but only after marking the quantize and dequantize nodes around it as folded into it. Those nodes then land in delegates that skip them, so nothing writes the outputs. A statically quantized `Conv2d` or `Linear` whose output is not quantized hits the same bug, even without #24385. ## What this change does The partitioner now leaves out any op with an input or output above 6 dimensions, and it checks every node that a partition takes. A float linear on such an input still keeps its 2-D matrix multiply on XNNPACK. Quantize and dequantize nodes are now marked as folded only once the op they belong to is really partitioned, so a rejected op keeps them as separate ops. `force_non_static_weights_for_f32_linear` no longer changes how a quantized linear is partitioned. With it, a quantized linear decomposed into a matrix multiply (for example with `to_backend`) used to lose its quantize and dequantize nodes: a static one failed to lower or load, and a dynamic one ran as a float matrix multiply. Now it is partitioned as without the flag. The XNNPACK quantizer now leaves 3-D convolutions in float. The model file gets larger, but there are no extra quantize and dequantize steps and no quantization error. They run on the CPU kernels, which support 3-D convolution once #24385 lands. Without it, such a model fails with "Expect input tensor to be 3-D or 4-D". Still not covered: a linear on such an input, quantized with torchao's affine ops. Those ops have no portable kernels, so the model now fails in `to_executorch()` instead of at load. And a linear on such an input quantized with the XNNPACK quantizer runs outside XNNPACK, because its quantized input has too many dims. ## Testing - `test_permute.py` checks which ops are delegated for 6-D and 7-D permutes, and for a linear on an 8-D input: float (also with a ReLU), static and dynamic quantization, each in both partitioner modes. - New tests cover a statically quantized Conv2d and Linear with a float output, a statically and a dynamically quantized linear lowered through `to_backend` with `force_non_static_weights_for_f32_linear`, and 3-D convolutions that the quantizer leaves unquantized. - The per-channel transposed conv test now also runs the convs that XNNPACK rejects. Before this change they returned all zeros. - Each new test fails without the part of the change it covers. - On macOS arm64, the patchify example and other 7-D and 8-D cases load and match eager PyTorch. A Conv3d model quantized with the XNNPACK quantizer matches eager PyTorch with #24385, and fails without it. (cherry picked from commit 1549787)
…-D linear output (#24464) ### What is wrong today `XnnpackPartitioner(enable_bf16=True)` lets XNNPACK run bf16 nodes. Without it, bf16 matrix multiplies run on the single-threaded portable kernels, which is very slow. With it, a program can fail to load: 1. The partitioner hands XNNPACK bf16 nodes that XNNPACK cannot define in bf16: softmax, mean, average and max pooling, bilinear upsampling, and batch norm that does not fold into a linear. Loading fails with `xnn_status_invalid_parameter`. 2. Batch matrix multiply (`bmm`) is defined with a bf16 output, but XNNPACK only multiplies bf16 matrices into an fp32 result. Loading fails with `Failed to create BMM node`. A bf16 attention block written with `matmul` and `softmax` hits both. ### What this change does - A config can now say it does not support bf16 (`supports_bf16 = False`), so its bf16 nodes stay on the portable kernels and the partitioner logs why. Softmax, mean, both poolings, bilinear upsampling and convolution set it. The default is `True`: XNNPACK rewrites most bf16 ops to fp32 internally, and at the pinned commit only these reject bf16 when defined. A bf16 run of the other float ops found no other failure. - A bf16 batch norm is kept in XNNPACK only when `FuseBatchNormPass.can_fuse` says it folds into the linear before it. - `can_fuse` now also requires the batch norm to be the only user of the linear output, and the linear output to be 2-D. Without these checks, a bf16 batch norm could be kept but left unfused, so loading failed, and a batch norm after a 3-D linear output was folded along the wrong dim, which gave wrong results in fp32 and bf16. - A bf16 `bmm` now writes an fp32 result and converts it to bf16, as the linear layer already did. With this change, the SmolVLA robot policy (a vision-language-action model) in bf16 with `enable_bf16=True` loads, and then runs one action chunk in about 11 seconds on 16 x86 cores, instead of about 9 minutes with the default partitioner. ### Tests - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16`: each op, and a convolution followed by a batch norm, is delegated in fp32 but not in bf16, and the bf16 program matches eager. - `test_enable_bf16_fuses_batch_norm_into_a_linear`: a bf16 linear followed by `BatchNorm1d` runs fully in XNNPACK. When the linear output is also used elsewhere, the batch norm stays out. - `test_enable_bf16_runs_an_attention_block`: the linear layers and both `bmm`s run in XNNPACK, the softmax stays outside. - `test_bf16_bmm` in `test_bmm.py`, and `test_fp32_linear_batch_norm_on_3d_output` in `test_batch_norm_fusion.py`. All of these fail before this change, except the fused batch norm case, which guards the fusion that already worked. Ran the XNNPACK partitioner, op and pass tests on Linux x86_64. The only failures, in the per-channel int4 linear tests, also happen without this change. (cherry picked from commit 7bf0303)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24477
Note: Links to docs will display an error until the docs builds have been completed. ⏳ 3 Pending, 1 Unrelated FailureAs of commit 21014eb with merge base 4db778e ( BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This was referenced Nov 3, 2026
Base automatically changed from
cherry-pick-24380-by-pytorch_bot_bot_
to
release/1.6
November 3, 2026 04:02
JacobSzwejbka
approved these changes
Nov 3, 2026
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What is wrong today
XnnpackPartitioner(enable_bf16=True)lets XNNPACK run bf16 nodes. Without it, bf16 matrix multiplies run on the single-threaded portable kernels, which is very slow. With it, a program can fail to load:xnn_status_invalid_parameter.bmm) is defined with a bf16 output, but XNNPACK only multiplies bf16 matrices into an fp32 result. Loading fails withFailed to create BMM node.A bf16 attention block written with
matmulandsoftmaxhits both.What this change does
supports_bf16 = False), so its bf16 nodes stay on the portable kernels and the partitioner logs why. Softmax, mean, both poolings, bilinear upsampling and convolution set it. The default isTrue: XNNPACK rewrites most bf16 ops to fp32 internally, and at the pinned commit only these reject bf16 when defined. A bf16 run of the other float ops found no other failure.FuseBatchNormPass.can_fusesays it folds into the linear before it.can_fusenow also requires the batch norm to be the only user of the linear output, and the linear output to be 2-D. Without these checks, a bf16 batch norm could be kept but left unfused, so loading failed, and a batch norm after a 3-D linear output was folded along the wrong dim, which gave wrong results in fp32 and bf16.bmmnow writes an fp32 result and converts it to bf16, as the linear layer already did.With this change, the SmolVLA robot policy (a vision-language-action model) in bf16 with
enable_bf16=Trueloads, and then runs one action chunk in about 11 seconds on 16 x86 cores, instead of about 9 minutes with the default partitioner.Tests
test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16: each op, and a convolution followed by a batch norm, is delegated in fp32 but not in bf16, and the bf16 program matches eager.test_enable_bf16_fuses_batch_norm_into_a_linear: a bf16 linear followed byBatchNorm1druns fully in XNNPACK. When the linear output is also used elsewhere, the batch norm stays out.test_enable_bf16_runs_an_attention_block: the linear layers and bothbmms run in XNNPACK, the softmax stays outside.test_bf16_bmmintest_bmm.py, andtest_fp32_linear_batch_norm_on_3d_outputintest_batch_norm_fusion.py.All of these fail before this change, except the fused batch norm case, which guards the fusion that already worked. Ran the XNNPACK partitioner, op and pass tests on Linux x86_64. The only failures, in the per-channel int4 linear tests, also happen without this change.