[mcremon-meta]

Summary:

X-link: pytorch/ao#5101

AffineQuantizedObserverBase.convert inserts torchao.choose_qparams_affine, quantize_affine and dequantize_affine as OpOverloadPackets. The per-tensor and per-channel q/dq ops that convert_pt2e inserts are OpOverloads, 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.

Differential Revision: D130716026

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24993

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 63b1381 with merge base 0b776dd ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[meta-codesync]

@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D130716026.