Repository navigation
Backend-neutral quantizer: core ATen op quantizers and PT2E composition - #24769
Open
mcremon-meta wants to merge 2 commits into
Open
mcremon-meta wants to merge 2 commits into
mcremon-meta wants to merge 2 commits into
Conversation
…ma-driven Summary: Pins down what a `fused_quant` op means, so that any backend can rely on it, and makes the output of `DecomposeFusedQuant`, which turns each fused op back into `dequantize -> aten -> quantize`, something any backend's partitioner can consume. Op contract (`exir/dialects/fused_quant/ops.py`): - `QuantParamsStruct.is_per_channel(tensor_shape)` / `is_per_group(tensor_shape)` can take the quantized tensor's shape. With it, a scale that tiles one axis counts as per-channel only when the block size is 1, and as per-group when it is larger (so a `[1, K // g]` group scale over `[1, K]` is no longer mistaken for per-channel). A scale size that does not divide the tensor size along that axis raises `ValueError`. Without it, they classify from the scale shape alone, as before. - `dequantized_dtype(tensor, qparams)` gives the float dtype of an operand after its input qparams are applied (`qparams.dtype`, or the tensor's own dtype when the operand is unquantized). The fake kernels use it, together with `torch.promote_types`, instead of assuming float32. - `sub.Scalar` is registered through the shared `_lib` (`define` / `impl` / `register_fake`) like `add.Scalar` and `mul.Scalar`, instead of through a standalone `custom_op`. `DecomposeFusedQuant` (`backends/transforms/fused_quant/`): - Granularity comes from the shape-aware `is_per_channel(tensor_shape)`. - Constant per-channel qparams become new 1-D constants (`add_constant`) instead of a view of the full-rank one, so a partitioner that needs per-channel qparams as placeholders works without a constant fold. - Kwargs left at their schema default are not emitted, so the rebuilt ATen node matches what `torch.export` would produce (otherwise, e.g., `alpha=1.0` reaches a two-operand `tosa.SUB`). - The ATen node's meta comes from the fused op's fake kernel when the ATen op has none (e.g. `_masked_softmax`). - Every node is validated before any rewrite, so an unsupported node fails the pass without leaving the graph half rewritten. - `is_fused_quant_node` and `AssertNoFusedQuantOps` are public. Nodes that carry a `delegation_tag` are still skipped, as before. `test_op_numerics` checks, for each op variant, that the reference kernel and its `DecomposeFusedQuant` lowering agree exactly. Reviewed By: rascani Differential Revision: D129194621
Summary: This diff makes the quantizer in `frontends/fused_quant/quantizer/` usable by any backend: - `get_core_aten_quantizers(config, weighted_config)` in `quantizer/quantizer.py`: the op quantizers every backend can share, with the contract below. `arg1_is_scalar` is the node filter its scalar add/sub/mul variants use. - `FusedQuantQuantizer` now composes with other PT2E quantizers on a first-claim-wins basis. Its annotations set `_annotated` (the flag other quantizers check), and a `fused_quant_owned` node flag lets `fuse()` skip nodes that another quantizer claimed. - `FusedQuantQuantizer.fuse` first repairs the `convert_pt2e` output: it canonicalizes torchao `quantize_affine` / `dequantize_affine` packets to `.default` and refreshes `meta["val"]`, both of which fusion depends on. The contract of `get_core_aten_quantizers`: - **Only core ATen ops.** Every op it quantizes is core ATen (`torch.Tag.core`): add, sub, mul, bmm, relu, max_pool2d_with_indices, avg_pool2d, convolution, sigmoid, tanh, hardtanh, gelu and native_layer_norm. It does not cover every core ATen op, and more can be added. - **Each quantizer always sees its op.** Export's decompositions keep core ATen ops, so a model's use of these ops reaches the quantizer without `preserved_ops()` keeping anything whole. - **Portable.** Whatever no backend claims decomposes to quantize/dequantize plus the same core ATen ops, all of which have portable kernels. - **Qparams come from the caller.** The package builds no `QuantizationConfig` of its own: callers build theirs with TorchAO or reuse their backend's. `weighted_config` applies to convolution, `config` to everything else. The only change it makes is moving the per-channel weight axis to 1 for transposed convolution. Grouped transposed convolution uses `config` instead. - **Non-core ops belong to the backend.** A backend that wants an op such as `linear` or `hardswish` kept whole adds its own `OpQuantizer`, and `preserved_ops()` keeps it out of decomposition. Non-core ops that many backends share, such as `linear`, can move into the shared list later. Differential Revision: D129194623
mcremon-meta
requested review from
JacobSzwejbka and
larryliu0820
as code owners
November 7, 2026 02:18
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24769
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated FailureAs of commit 389b986 with merge base 472953c ( NEW FAILURE - The following job has failed:
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. |
Contributor
|
@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D129194623. |
This PR needs a
|
This branch was previously 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.
Summary:
This diff makes the quantizer in
frontends/fused_quant/quantizer/usable byany backend:
get_core_aten_quantizers(config, weighted_config)inquantizer/quantizer.py: the op quantizers every backend can share, with thecontract below.
arg1_is_scalaris the node filter its scalar add/sub/mulvariants use.
FusedQuantQuantizernow composes with other PT2E quantizers on afirst-claim-wins basis. Its annotations set
_annotated(the flag otherquantizers check), and a
fused_quant_ownednode flag letsfuse()skipnodes that another quantizer claimed.
FusedQuantQuantizer.fusefirst repairs theconvert_pt2eoutput: itcanonicalizes torchao
quantize_affine/dequantize_affinepackets to.defaultand refreshesmeta["val"], both of which fusion depends on.The contract of
get_core_aten_quantizers:(
torch.Tag.core): add, sub, mul, bmm, relu, max_pool2d_with_indices,avg_pool2d, convolution, sigmoid, tanh, hardtanh, gelu and native_layer_norm.
It does not cover every core ATen op, and more can be added.
ATen ops, so a model's use of these ops reaches the quantizer without
preserved_ops()keeping anything whole.plus the same core ATen ops, all of which have portable kernels.
QuantizationConfigof its own: callers build theirs with TorchAO or reuse their backend's.
weighted_configapplies to convolution,configto everything else. Theonly change it makes is moving the per-channel weight axis to 1 for
transposed convolution. Grouped transposed convolution uses
configinstead.
linearorhardswishkept whole adds its ownOpQuantizer, andpreserved_ops()keeps it out of decomposition. Non-core ops that manybackends share, such as
linear, can move into the shared list later.Differential Revision: D129194623