Add quantize/dequantize_per_channel_group support to QNN backend - #19629
Add quantize/dequantize_per_channel_group support to QNN backend#19629Hyungkeun-Park wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19629
Note: Links to docs will display an error until the docs builds have been completed.
|
|
@pytorchbot label "release notes: qualcomm" |
|
Thanks for the PR. @shewu-quic @haowhsu-quic @winskuo-quic @DannyYuyang-quic could any of you review this? |
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
|
@Hyungkeun-Park-Nota Thanks for this, sorry for the slow review turnaround. This pr is stale against
So the rebased change would be adding these to It's worth re-checking
Either way works and I don't want to hold you up further. If it's convenient, rebasing my pr merges would shrink this diff and avoid the |
quantized_decomposed.quantize_per_channel_group and its dequantize counterpart are what group-wise (int4) LLM weight quantization lowers to, but the QNN backend did not recognize them: the ops were decomposed away during export instead of reaching the backend, and InsertIOQDQ raised a KeyError when a pre-quantized weight annotated with a dequantize_per_channel_group encoding fed the graph output. Give them the same treatment torchao.quantize_affine / dequantize_affine already receive: - node_visitor: register both ops in q_ops / dq_ops and in q_dq_map, and route them through make_qnn_per_block_config via the new PER_CHANNEL_GROUP_ENCODING set (QNN models group-wise quantization as a per-block encoding). - partition/utils: keep both ops out of the decomposition table so they survive export and reach the backend unchanged. - _passes/utils: expose the block scale tensor as QCOM_SCALE, which make_qnn_per_block_config reads. The remap is gated on the encoding rather than on the presence of a "scales" key, because quantize_per_channel.default also has a "scales" argument but is consumed as a per-channel (not per-block) config. - qnn_pass_manager: register both ops in node_visitor.q_ops / dq_ops in get_to_edge_transform_passes, mirroring the torchao workaround. Test: test_insert_io_qdq_per_channel_group_resolves_through_q_dq_map annotates a parameter with a dequantize_per_channel_group encoding and wires it to the graph output, which is the branch of InsertIOQDQ that resolves the encoding through q_dq_map (the insert-quantize-after-input branch is skipped for parameters). Without the q_dq_map entries the test fails with KeyError at insert_io_qdq.py's q_dq_map lookup; with them it asserts a dequantize_per_channel_group node now feeds the output.
d2d5aab to
78c58b8
Compare
|
@qti-horodnic Thanks for the pointers. Rebuilt the branch on top of main (4b4df96) so it's no longer stale, and updated the PR description. You were right about On the test: the KeyError comes from the dequantize-before-output branch, which has no The test now picks the parameter placeholder via Good catch on if quant_node.target in PER_CHANNEL_GROUP_ENCODING:
quant_attrs[QCOM_SCALE] = quant_attrs[QCOM_SCALES]
That should leave the overlap with your PR at the |
| node_visitor.q_ops.add( | ||
| exir_ops.edge.quantized_decomposed.quantize_per_channel_group.default | ||
| ) |
There was a problem hiding this comment.
Minor: These lines are redundant since the ops are already in the literal sets in node_visitor.py. This block is a workaround that exists because the torchao namespace isn't resolvable at node_visitor import time. quantized_decomposed ops resolve fine which is why the node_visitor.py change works. Keeping both makes it look like these depend on the workaround, which will confuse whoever eventually removes the TODO.
There was a problem hiding this comment.
Removed. You're right — quantized_decomposed resolves at node_visitor import time, so the literal sets are sufficient and piggybacking on the torchao workaround was misleading.
| per_block_encoding = { | ||
| exir_ops.edge.torchao.quantize_affine.default, | ||
| exir_ops.edge.torchao.dequantize_affine.default, | ||
| *PER_CHANNEL_GROUP_ENCODING, |
There was a problem hiding this comment.
This routes to make_qnn_per_block_config, which hardcodes offset = 0. But dequantize_per_channel_group accepts zero_points (your test even sets it to None). An asymmetric group-quantized weight currently lowers to a symmetric encoding silently, which is a wrong-answer path rather than an error. Could we reject it explicitly? Somewhere near the top of make_qnn_per_block_config like:
zps = quant_attrs.get("zero_points")
if zps is not None and torch.any(zps != 0):
raise ValueError(
"per_channel_group with non-zero zero_points is not supported; "
"QNN blockwise expansion requires symmetric quantization"
)
Also, this lowers to QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION, the same encoding LPBQ uses, and validate_lpbq_support gates that at >= HtpArch.V69. Since pre-quantized weights bypass the quantizer, there's no arch check on this path at all. We've separately confirmed block-wise encodings hard-fail on V68 and have accuracy issues at V69 for some shapes, only becoming reliable around V79. Which archs have you validated this on? If it has the same V69+ constraint it needs an equivalent gate.
There was a problem hiding this comment.
Both points addressed in d67952a:
-
Asymmetric rejection: added at the top of
make_qnn_per_block_config, checking bothzero_pointsandzero_pointso the torchao affine path gets the same guard. Covered by a new test (test_make_qnn_per_block_config_rejects_asymmetric). -
Arch gate: I haven't validated this on any HTP device — this environment has no device attached, and the PR's scope was making the ops survive lowering. So per your point that the constraint applies structurally (same
BLOCKWISE_EXPANSIONencoding as LPBQ), I added a gate rather than claiming validation:QnnOperatorSupport.is_node_supportednow refuses to delegateper_channel_groupq/dq whensoc_info.htp_info.htp_arch < HtpArch.V69, mirroringvalidate_lpbq_support. The op then falls back to CPU with a warning instead of producing wrong numerics on V68. I put it in the partitioner because the pre-quantized path never reaches the quantizer's check, andNodeVisitorhas no access tosoc_info— happy to move it if you'd prefer it plumbed differently. If V69's shape-dependent accuracy issues warrant gating at V79 instead, I'd take your guidance on the threshold since I can't measure it here.
| # one quantize (input) and one dequantize (output) = +2 nodes. | ||
| self.assertEqual(node_count_after, node_count_before + 2) | ||
|
|
||
| def test_insert_io_qdq_per_channel_group_resolves_through_q_dq_map(self): |
There was a problem hiding this comment.
Good test. One gap: this covers the pass-level fix but nothing exercises make_qnn_per_block_config end-to-end with real group-quantized weights, so the scale-reshaping and the symmetry assumption are untested. Not asking you to add a device test here, but if you've validated numerics locally on an int4 model, could you mention it in the PR description?
There was a problem hiding this comment.
I have not validated numerics on a real int4 model on device — no HTP device is available in this environment. Updated the PR description to state that explicitly, and added the partitioner arch gate + asymmetric rejection so the untested paths fail loudly (or fall back) rather than lowering silently wrong.
| # on the presence of the key: quantize_per_channel.default also has a | ||
| # "scales" argument but is consumed as a per-channel (not per-block) config. | ||
| if quant_node.target in PER_CHANNEL_GROUP_ENCODING: | ||
| quant_attrs[QCOM_SCALE] = quant_attrs[QCOM_SCALES] |
There was a problem hiding this comment.
Minor: group_size gets read into quant_attrs here and then never used, QCOM_NUM_BLOCKS_PER_AXIS is inferred from q_scales.shape[1] in make_qnn_per_block_config instead. An assert that they match would catch a mismatched checkpoint early rather than producing a silently wrong encoding.
There was a problem hiding this comment.
Added in d67952a: get_quant_attrs now raises ValueError when scales.shape[-1] * group_size doesn't match the weight's input-feature count, so a mismatched checkpoint fails at annotation time instead of producing a silently wrong encoding.
@Hyungkeun-Park-Nota Thanks for the detailed writeup, you're right about
I've left a few comments, only one is blocking (on Regarding the merge order: yes, please don't block on my pr. If your change gets approved first it can be merged, I'll rebase accordingly, no worries. |
- Reject non-zero zero_points in make_qnn_per_block_config instead of silently lowering to a symmetric encoding - Skip delegating per_channel_group q/dq on HTP < V69 in QnnOperatorSupport, mirroring validate_lpbq_support for the pre-quantized path that bypasses the quantizer - Drop redundant dynamic q_ops/dq_ops registration in qnn_pass_manager (quantized_decomposed resolves at import time, unlike the torchao ops) - Validate group_size against the scale shape in get_quant_attrs
|
@qti-horodnic All four comments addressed in d67952a, replies inline. Summary:
|
qti-horodnic
left a comment
There was a problem hiding this comment.
Please rebase the changes, I think your main is stale again. I've left some comments mostly around keeping the scope contained, please take a look when you get a chance.
| exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, | ||
| exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor: exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, | ||
| exir_ops.edge.quantized_decomposed.quantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, | ||
| exir_ops.edge.quantized_decomposed.dequantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, |
There was a problem hiding this comment.
This key already exists on line 118 right? I think this is something you might have forgotten to remove, seems unrelated to the PR.
| is resolved, and it raised ``KeyError`` before per_channel_group was | ||
| added to the map. | ||
| """ | ||
| gm, ep = self._build_quantized_graph() |
There was a problem hiding this comment.
Minor: Let's give variables full meaningful names which we can derive their purpose from. I realize you probably followed the example of the other tests here, but I think we should improve upon those practices if we can. Applies in multiple places in this file, names like n, ep, u etc.
| exir_ops.edge.quantized_decomposed.quantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, | ||
| exir_ops.edge.quantized_decomposed.dequantize_per_channel.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, | ||
| exir_ops.edge.quantized_decomposed.quantize_per_channel_group.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default, | ||
| exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default, |
There was a problem hiding this comment.
Isn't this dead code? Since to_dq_op returns early for anything in dq_ops
| } | ||
|
|
||
|
|
||
| def to_q_op(target): |
There was a problem hiding this comment.
This returns the dq op, which then trips assert target in q_ops in insert_quant_node. We should mirror the per-channel/per-tensor pairs like:
exir_ops.edge.quantized_decomposed.quantize_per_channel_group.default: exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default,
exir_ops.edge.quantized_decomposed.dequantize_per_channel_group.default: exir_ops.edge.quantized_decomposed.quantize_per_channel_group.default,
| # one quantize (input) and one dequantize (output) = +2 nodes. | ||
| self.assertEqual(node_count_after, node_count_before + 2) | ||
|
|
||
| def test_insert_io_qdq_per_channel_group_resolves_through_q_dq_map(self): |
There was a problem hiding this comment.
I think you may need to rebase again, since this test no longer tests q_dq_map. insert_io_qdq is now to_dq_op, which short-circuits for anything in dq_ops and this PR adds dequantize_per_channel_group there, so the lookup never happens. The test passes with all the new q_dq_map entries reverted, and the KeyError repro in the description is against an old main.
The test is still valuable so please rename it accordingly and add a case that actually pins the mappingm. E.g. asserting to_q_op(dq_pcg) is q_pcg and to_dq_op(q_pcg) is dq_pcg.
| # Per-channel-group lowers to blockwise expansion (the encoding LPBQ | ||
| # uses), which requires HTP >= V69. Pre-quantized weights bypass the | ||
| # quantizer's validate_lpbq_support check, so gate the delegation here. | ||
| if ( |
There was a problem hiding this comment.
QnnOperatorSupport is used by other backends like GPU and LPAI, but this reads soc_info.htp_info.htp_arch unconditionally. Let's add a check here for backend_options.backend_type == QnnExecuTorchBackendType.kHtpBackend to keep the scope intentionally on HTP.
…test - q_dq_map: map dequantize_per_channel_group -> quantize_per_channel_group (it pointed back at itself, which trips `assert target in q_ops` in insert_quant_node) and drop a stray duplicate dequantize_per_channel entry. - QnnOperatorSupport: only apply the < V69 blockwise gate when the backend is HTP; the checker is reused by GPU/LPAI which have no htp_arch. - tests: rename the InsertIOQDQ test (to_dq_op short-circuits for dq_ops, so it no longer exercises q_dq_map), add a test pinning to_q_op/to_dq_op for the per_channel_group pair, use descriptive variable names. Claude-Session: https://claude.ai/code/session_01BjGb82xs4F5EDYGnYtP8Gp
|
@qti-horodnic Thanks for the second pass. All six comments are addressed in 79d51be, replies inline.
|
…annel-group-quantization
Motivation
quantized_decomposed.quantize_per_channel_groupanddequantize_per_channel_groupare what group-wise (int4) LLM weight quantization lowers to, but the QNN backend did not recognize them, causing two distinct failures:torch.exportinstead of being preserved for the backend.InsertIOQDQraised aKeyErrorwhen a pre-quantized weight annotated with adequantize_per_channel_groupencoding fed the graph output, because the op was absent fromq_dq_map.This change gives them the same treatment
torchao.quantize_affine/dequantize_affinealready receive.Changes
backends/qualcomm/builders/node_visitor.pyq_ops/dq_ops.q_dq_map(this is the dictInsertIOQDQimports and looks up, so this is what fixes theKeyError).PER_CHANNEL_GROUP_ENCODINGset and route it throughmake_qnn_per_block_config, since QNN models group-wise quantization as a per-block encoding.make_qnn_per_block_confignow rejects non-zerozero_points(zero_points/zero_point) with aValueError: QNN blockwise expansion hardcodes a per-channel offset of 0, so an asymmetric weight would otherwise silently lower to a symmetric encoding.backends/qualcomm/partition/utils.pyget_skip_decomp_tableso they survive export and reach the backend unchanged.backends/qualcomm/partition/qnn_partitioner.pyQnnOperatorSupportrefuses to delegateper_channel_groupq/dq whensoc_info.htp_info.htp_arch < HtpArch.V69, mirroringvalidate_lpbq_support. Pre-quantized weights bypass the quantizer, so without this there is no arch check on this path; blockwise expansion hard-fails on V68. The op falls back to CPU with a warning.backends/qualcomm/_passes/utils.pyQCOM_SCALEinget_quant_attrs, which is whatmake_qnn_per_block_configreads. The remap is gated on the encoding rather than on the presence of a"scales"key, becausequantize_per_channel.defaultalso has a"scales"argument but is consumed as a per-channel (not per-block) config.scales.shape[-1] * group_sizematches the weight's input-feature count, so a mismatched checkpoint fails at annotation time instead of producing a silently wrong encoding (QCOM_NUM_BLOCKS_PER_AXISis inferred from the scale shape downstream).Testing
test_q_dq_map_pins_per_channel_group_pairspins theq_dq_mapentries throughto_q_op/to_dq_op:to_q_op(dequantize_per_channel_group) is quantize_per_channel_group,to_dq_op(quantize_per_channel_group) is dequantize_per_channel_group, plus the identity cases.to_dq_opshort-circuits for targets already indq_ops, so this is the only place the cross-pair mapping is exercised; it fails with a self-referentialdequantize_per_channel_groupentry.test_insert_io_qdq_per_channel_group_dequantizes_outputannotates a parameter with adequantize_per_channel_groupencoding and wires it to the graph output. That is the dequantize-before-output branch ofInsertIOQDQ(the insert-quantize-after-input branch is skipped for parameters); it asserts adequantize_per_channel_groupnode now feeds the output.test_make_qnn_per_block_config_rejects_asymmetriccovers the non-zerozero_pointsrejection.The encoding gating in
get_quant_attrswas checked directly:quantize_per_channel.defaultno longer gets aQCOM_SCALEentry written, whiledequantize_per_channel_group.defaultdoes.Not validated on device: no HTP device is available in this environment, so end-to-end numerics of the per-block lowering (scale reshaping, symmetric encoding) have not been verified on any HTP arch. The partitioner arch gate (< V69 falls back to CPU) and the asymmetric/
group_sizevalidations are there so the unverified paths fail loudly or fall back instead of lowering silently wrong.cc @qti-horodnic @cbilgin @psiddh @cccclai @abhinaykukkadapu