Fix PT2E Mul weight quantization for parameters

[ricardoasouz]

Summary

Fixes PT2E/TorchAO weight quantization for aten.mul.Tensor when a learned parameter can appear on either operand.

Previously, TorchAOQuantizerAdapter relied on metatype.weight_port_ids to determine whether a get_attr input represented a weight. PTMulMetatype has no fixed weight port because both:

• activation * weight
• weight * activation

are valid.

A torch.nn.Parameter feeding Mul could therefore be classified as an activation. Simply classifying it as a weight was not sufficient: WeightQuantizationInsertionPoint only preserved the target node, so the dynamically identified input port was lost before MinMax statistics collection.

This change:

• detects a torch.nn.Parameter on either Mul operand as a weight;
• preserves its exact input port in WeightQuantizationInsertionPoint;
• uses that explicit port during MinMax weight target-point creation;
• retains the existing backend weight-port discovery for legacy weight insertion points;
• leaves buffers and parameters used by other operations on their existing paths.

The optional input port is routing information only and does not change WeightQuantizerId identity semantics.

Concrete model case

I found a concrete case in torchvision's pretrained ConvNeXt-Tiny. Each CNBlock uses a learned LayerScale parameter in a Mul operation:

self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
...
result = self.layer_scale * self.block(input)

ConvNeXt-Tiny contains 18 such blocks.

I exported the blocks separately and used a custom TorchAO quantizer to isolate the LayerScale input of the Mul operation.

Without the fix, all 18 LayerScale parameters are classified as activation quantization points. With batch size 2, activation statistics treat axis 0 of the (dim, 1, 1) parameter as the sample axis.

For the first pretrained block:

• LayerScale absmax: 0.226632953
• activation-classified range bound: 0.092455307
• weight-classified range bound with this PR: 0.226632945

A regression test was also added using a learned parameter [0.25, -0.5, 1.0] with batch size 2 and covering both operand orders.

Without the fix, the resulting range bound is 0.5833333 instead of 1.0. With this PR, both operand orders use the correct weight range.

Tests

Coverage includes:

• activation * Parameter -> weight quantization point;
• Parameter * activation -> weight quantization point;
• activation * buffer -> remains activation quantization;
• end-to-end Q/DQ insertion for both parameter operand orders;
• end-to-end Q/DQ insertion for the buffer case;
• custom TorchAO quantizer with batch size 2;
• automatic batchwise_statistics=True;
• regression coverage for the resulting quantization range;
• verification that quantization is not silently skipped.

Validation:

• Python 3.12: tests/executorch/test_ptq.py -> 62 passed;
• Python 3.10.14: tests/executorch -> 74 passed, 6 expected xfails;
• make pre-commit -> all hooks passed;
• git diff --check -> passed.

A broader fallback for arbitrary torch.nn.Parameter inputs was also checked and rejected because it changed the existing XNNPACK ViT handling of positional embeddings feeding aten.add.Tensor. The final behavior is intentionally limited to PTMulMetatype.

Related to #3576.

AI-assisted development

AI tools were used during codebase exploration, hypothesis generation, and drafting. I personally reviewed, tested, and validated the final implementation and can explain and maintain the submitted changes.

[ricardoasouz]

Hi @anzr299, just rebased this onto the latest develop and ran tests/executorch/test_ptq.py again. All 57 tests pass. Let me know if there’s anything else you’d like me to check. Thanks!

[ricardoasouz]

Rebased onto the latest develop and revalidated using the repository-pinned ExecuTorch environment with Torch 2.13.0.
tests/executorch/test_ptq.py: 57 passed.
This includes the new Mul parameter detection cases for both operand orders, the buffer case, and the batchwise statistics regression test.

[ricardoasouz]

Updated this PR after re-checking it end-to-end. In the previous revision, the Mul parameter was classified as a weight, but WeightQuantizationInsertionPoint carried no port, and PTMulMetatype has no static weight ports, so MinMax created no target point, and the parameter ended up unquantized (on develop, it is still quantized as an activation). The weight insertion point now carries the input port found by the adapter, and MinMax uses it. End-to-end tests cover both operand orders and the buffer case; tests/executorch/test_ptq.py: 60 passed. The broader variant for other ops is described in the notes; happy to follow up if you prefer that direction.

[ricardoasouz]

Hi! It looks like the pre-commit workflow expired while waiting for approval.
Could someone please approve/re-run it when convenient? Thanks!

[ricardoasouz]

I @anzr299, Thanks very much by the observation.
I updated and included the example.