Repository navigation
Let a QuantizationRecipe decompose the capture before quantization - #25052
Merged
Merged
Conversation
Summary: X-link: pytorch/ao#5101 `AffineQuantizedObserverBase.convert` inserts `torchao.choose_qparams_affine`, `quantize_affine` and `dequantize_affine` as `OpOverloadPacket`s. The per-tensor and per-channel q/dq ops that `convert_pt2e` inserts are `OpOverload`s, and `torch.export` turns these packets into `.default` when it re-exports a converted model, so the same op shows up in two forms depending on whether a graph was re-exported. Each of these ops has a single overload, so convert now inserts `.default`. **The bug it fixes.** `convert_pt2e(..., use_reference_representation=True)` never applies the 4-bit groupwise dynamic linear rewrite to the output of torchao's own affine observers. The rewrite's patterns are traced with `torch.export`, so they carry `.default`, and the subgraph matcher compares targets exactly. `remove_tensor_overload_for_qdq_ops` lines the two forms up for the `quantized_decomposed` ops only, and the rewrite's filter checks for `.default` explicitly. The new `test_dynamic_linear_4bit_groupwise` (8-bit dynamic per-token activations, 4-bit per-group weights) fails without the observer change, because `reference_dqlinear_int4` is not in the graph, and passes with it. **Matchers that move with it.** Code that matches the packet on `convert_pt2e` output would silently stop matching, so it changes in the same diff: - torchao: `_QUANT_OPS` (which quantize nodes convert folds into constants), `constant_fold`'s `DEQUANT_OPS`, and `port_metadata_pass`'s op lists. `test_quantize_pt2e` now expects `.default`. - Helios `post_convert_transform`: Inductor's `add_dont_constant_fold` matches the exact target. With only the packet listed, `constant_fold` would fold the weight dequantize of `AffineQuantizedMinMaxObserver` weights into a float constant. - ExecuTorch QNN (`genai_pipeline/quant_utilities.py`, and the llama `llm_wrappers.py` and `mix_precision_analyzer.py`): the q/dq target sets accept both forms, because OSS ExecuTorch pins torchao (`third-party/ao`) and keeps getting the packet until that pin moves past this change. - Bolt `per_block_param_observer_test` compares against `.default`. Unchanged: eager calls such as `torch.ops.torchao.dequantize_affine(...)` resolve either way; substring matches on `str(node.target)` (XNNPACK `is_affine_qdq`, Arm, the native backend tests, Bolt's FileCheck) see both forms; and code that runs on exported or edge graphs already sees `.default`. SAS's `_normalize_affine_overloads`, which rewrites the packets to `.default` after convert, becomes a no-op. Removing it is a follow-up. Reviewed By: rascani, DrJessop Differential Revision: D130716026
Summary: `QuantizeStage` captures each model with a plain `torch.export.export(..., strict=True)` and hands its `.module()` to the `pre_prepare_passes` and to `prepare_pt2e`. A quantizer that matches core ATen ops, such as `aten.convolution` rather than `aten.conv2d`, cannot get that graph from a recipe: decomposition runs on the `ExportedProgram`, and the `pre_prepare_passes` only see the unlifted `GraphModule`. Add `QuantizationRecipe.preserve_ops`. When set, the stage decomposes the captured program with the default core ATen table, minus those ops, before taking its `.module()`. When it is None, the default, nothing is decomposed and the stage exports exactly as before. Graph rewrites the quantizer needs belong in its `transform_for_annotation`, which `prepare_pt2e` already runs. `ExportRecipe.combine` decomposes when any recipe sets `preserve_ops`, keeping every op any of them keeps. For post-training quantization the stage now also puts the model in eval mode before capturing it, as it already puts it in train mode for QAT. Calibration then runs batch norm on its running statistics and skips dropout, as the deployed model does, whatever mode the caller left the model in. As with QAT, the model stays in that mode afterwards. Reviewed By: digantdesai Differential Revision: D131123188
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/25052
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ You can merge normally! (1 Unrelated Failure)As of commit 596e900 with merge base 49fd4dd ( 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
Contributor
|
@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D131123188. |
digantdesai
approved these changes
Nov 13, 2026
digantdesai
left a comment
digantdesai
left a comment
Contributor
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
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:
QuantizeStagecaptures each model with a plaintorch.export.export(..., strict=True)and hands its.module()to thepre_prepare_passesand toprepare_pt2e. A quantizer that matches core ATenops, such as
aten.convolutionrather thanaten.conv2d, cannot get thatgraph from a recipe: decomposition runs on the
ExportedProgram, and thepre_prepare_passesonly see the unliftedGraphModule.Add
QuantizationRecipe.preserve_ops. When set, the stage decomposes thecaptured program with the default core ATen table, minus those ops, before
taking its
.module(). When it is None, the default, nothing is decomposed andthe stage exports exactly as before. Graph rewrites the quantizer needs belong
in its
transform_for_annotation, whichprepare_pt2ealready runs.ExportRecipe.combinedecomposes when any recipe setspreserve_ops, keepingevery op any of them keeps.
For post-training quantization the stage now also puts the model in eval mode
before capturing it, as it already puts it in train mode for QAT. Calibration
then runs batch norm on its running statistics and skips dropout, as the
deployed model does, whatever mode the caller left the model in. As with QAT,
the model stays in that mode afterwards.
Reviewed By: digantdesai
Differential Revision: D131123188