Update on "[quant][pt2] Fix custom dtype per channel weight in QAT"

andrewor14 · andrewor14 · commit 7bccb52f43fe · 2023-11-06T15:39:17.000-08:00
Summary: Previously we only copied over q/dq args for the per
tensor case. This was because the qparams for `quantize_per_tensor`
are literals while the qparams for `quantize_per_channel` are
`get_attr` nodes (tensors), which disappear from the original
nodes in the graph after subgraph rewriting.

However, this is problematic because, in the per channel case,
not all q/dq args are tensors. In particular, the args after
the qparams (axis, qmin, qmax, dtype) are all literals. For
these literal args we simply used the hardcoded ones
(0, -127, 127, torch.int8 respectively), even if the user
explicitly specified to use a different weight dtype. This
commit fixes this by copying over these literal args for the
per channel case as well.

Test Plan:
python test/test_quantization.py TestQuantizePT2EQAT.test_qat_per_channel_weight_custom_dtype

Reviewers: jerryzh168, kimishpatel

Subscribers: jerryzh168, kimishpatel, supriyar

[ghstack-poisoned]
diff --git a/test/quantization/pt2e/test_quantize_pt2e_qat.py b/test/quantization/pt2e/test_quantize_pt2e_qat.py
@@ -738,7 +738,6 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
             dtype=torch.int32,
             quant_min=0,
             quant_max=2**31 - 1,
-
             qscheme=torch.per_channel_affine,
             observer_or_fake_quant_ctr=FusedMovingAvgObsFakeQuantize.with_args(
                 observer=MovingAveragePerChannelMinMaxObserver,
