Skip to content

Commit 05e8332

Browse files
Yuan, ChaoGitHub Enterprise
Yuan, Chao
authored and
GitHub Enterprise
committed
Skip QDQ pair scale/zp in duplicate_shared_initializers (#7097)
* Skip QDQ pair scale/zp in duplicate_shared_initializers Signed-off-by: Chao Yuan <chayuan@qti.qualcomm.com> * Fix asymmetric QDQ pair handling in _get_qdq_pair_initializers Signed-off-by: Chao Yuan <chayuan@qti.qualcomm.com> --------- Signed-off-by: Chao Yuan <chayuan@qti.qualcomm.com>
1 parent f830603 commit 05e8332

2 files changed

Lines changed: 288 additions & 3 deletions

File tree

‎TrainingExtensions/onnx/src/python/aimet_onnx/utils.py‎

Lines changed: 38 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
import platform
1111
import tempfile
1212
from pathlib import Path
13-
from typing import Dict, Iterable, List, Union, Tuple, Callable, Optional
13+
from typing import Dict, Iterable, List, Set, Union, Tuple, Callable, Optional
1414
from contextlib import contextmanager
1515
import os
1616
import pickle
@@ -984,6 +984,36 @@ def _add_value_info(model: onnx.ModelProto):
984984
model.graph.value_info.extend(initial_value_info)
985985

986986

987+
def _get_qdq_pair_initializers(graph: GraphProto) -> Set[str]:
988+
"""
989+
Return the set of initializer names that serve as scale or zero_point for
990+
at least one QuantizeLinear–DequantizeLinear pair.
991+
992+
A QDQ pair is identified by ``QuantizeLinear.output[0] == DequantizeLinear.input[0]``.
993+
Scale (input[1]) and zero_point (input[2]) of such a pair must remain shared
994+
between Q and DQ; duplicating them would decouple the encode/decode parameters
995+
and break quantization correctness.
996+
997+
:param graph: ONNX GraphProto to inspect
998+
:return: Set of initializer names that must not be duplicated
999+
"""
1000+
producers = {out: node for node in graph.node for out in node.output}
1001+
init_names = {init.name for init in graph.initializer}
1002+
protected: Set[str] = set()
1003+
for node in graph.node:
1004+
if node.op_type != "DequantizeLinear" or not node.input:
1005+
continue
1006+
q_node = producers.get(node.input[0])
1007+
if q_node is None or q_node.op_type != "QuantizeLinear":
1008+
continue
1009+
for i in range(1, max(len(q_node.input), len(node.input))):
1010+
if i < len(q_node.input) and q_node.input[i] in init_names:
1011+
protected.add(q_node.input[i])
1012+
if i < len(node.input) and node.input[i] in init_names:
1013+
protected.add(node.input[i])
1014+
return protected
1015+
1016+
9871017
def duplicate_shared_initializers(graph: GraphProto) -> int:
9881018
"""
9891019
Duplicate initializers that are shared across multiple nodes.
@@ -999,7 +1029,12 @@ def duplicate_shared_initializers(graph: GraphProto) -> int:
9991029
if not graph.node:
10001030
return 0
10011031

1002-
initializer_map = {init.name: init for init in graph.initializer}
1032+
# Exclude scale/zp initializers that belong to a QDQ pair — they must stay
1033+
# shared between QuantizeLinear and its paired DequantizeLinear.
1034+
qdq_protected = _get_qdq_pair_initializers(graph)
1035+
initializer_map = {
1036+
init.name: init for init in graph.initializer if init.name not in qdq_protected
1037+
}
10031038
if not initializer_map:
10041039
return 0
10051040

@@ -1029,7 +1064,7 @@ def duplicate_shared_initializers(graph: GraphProto) -> int:
10291064
new_initializers: List[TensorProto] = []
10301065
duplicate_count = 0
10311066
existing_names = (
1032-
set(initializer_map)
1067+
{init.name for init in graph.initializer}
10331068
| {out for node in graph.node for out in node.output}
10341069
| {inp.name for inp in graph.input}
10351070
)

‎TrainingExtensions/onnx/test/python/test_utils.py‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -942,6 +942,256 @@ def test_duplicate_shared_initializers_shape_consumer_is_skipped(self):
942942
"Shape consumer must not trigger duplication of the shared initializer"
943943
)
944944

945+
@staticmethod
946+
def _make_qdq_model(share_scale_across_pairs: bool = False) -> onnx.ModelProto:
947+
"""
948+
Build a minimal QDQ model.
949+
950+
When ``share_scale_across_pairs=False`` (default):
951+
A single QDQ pair wrapping a Conv weight:
952+
953+
weight → QuantizeLinear(scale, zp) → DequantizeLinear(scale, zp) → Conv
954+
955+
When ``share_scale_across_pairs=True``:
956+
Two independent QDQ pairs that happen to share the same scale/zp:
957+
958+
weight → QLinear1(scale, zp) → DQLinear1(scale, zp) → Conv1
959+
weight2 → QLinear2(scale, zp) → DQLinear2(scale, zp) → Conv2
960+
"""
961+
CHANNELS = 4
962+
weight = numpy_helper.from_array(
963+
np.random.randn(CHANNELS, CHANNELS, 1, 1).astype(np.float32),
964+
name="weight",
965+
)
966+
scale = numpy_helper.from_array(np.array(0.1, dtype=np.float32), name="scale")
967+
zp = numpy_helper.from_array(np.array(0, dtype=np.int8), name="zp")
968+
969+
nodes = [
970+
helper.make_node(
971+
"QuantizeLinear",
972+
inputs=["weight", "scale", "zp"],
973+
outputs=["weight_q"],
974+
name="q1",
975+
),
976+
helper.make_node(
977+
"DequantizeLinear",
978+
inputs=["weight_q", "scale", "zp"],
979+
outputs=["weight_dq"],
980+
name="dq1",
981+
),
982+
helper.make_node(
983+
"Conv",
984+
inputs=["model_input", "weight_dq"],
985+
outputs=["model_output"],
986+
name="conv1",
987+
),
988+
]
989+
initializers = [weight, scale, zp]
990+
991+
if share_scale_across_pairs:
992+
weight2 = numpy_helper.from_array(
993+
np.random.randn(CHANNELS, CHANNELS, 1, 1).astype(np.float32),
994+
name="weight2",
995+
)
996+
nodes += [
997+
helper.make_node(
998+
"QuantizeLinear",
999+
inputs=["weight2", "scale", "zp"],
1000+
outputs=["weight2_q"],
1001+
name="q2",
1002+
),
1003+
helper.make_node(
1004+
"DequantizeLinear",
1005+
inputs=["weight2_q", "scale", "zp"],
1006+
outputs=["weight2_dq"],
1007+
name="dq2",
1008+
),
1009+
helper.make_node(
1010+
"Conv",
1011+
inputs=["model_input", "weight2_dq"],
1012+
outputs=["branch2_output"],
1013+
name="conv2",
1014+
),
1015+
helper.make_node(
1016+
"Add",
1017+
inputs=["model_output", "branch2_output"],
1018+
outputs=["final_output"],
1019+
name="add",
1020+
),
1021+
]
1022+
initializers.append(weight2)
1023+
1024+
final_out = "final_output" if share_scale_across_pairs else "model_output"
1025+
graph = helper.make_graph(
1026+
nodes,
1027+
"QDQModel",
1028+
inputs=[
1029+
helper.make_tensor_value_info(
1030+
"model_input", TensorProto.FLOAT, [1, CHANNELS, 8, 8]
1031+
)
1032+
],
1033+
outputs=[
1034+
helper.make_tensor_value_info(
1035+
final_out, TensorProto.FLOAT, [1, CHANNELS, 8, 8]
1036+
)
1037+
],
1038+
initializer=initializers,
1039+
)
1040+
model = models_for_tests.make_model(graph)
1041+
onnx.checker.check_model(model)
1042+
return model
1043+
1044+
def test_duplicate_shared_initializers_qdq_scale_not_duplicated(self):
1045+
"""
1046+
Given a QDQ pair where scale and zero_point are shared between
1047+
QuantizeLinear and DequantizeLinear:
1048+
1049+
weight → QuantizeLinear(scale, zp) → DequantizeLinear(scale, zp) → Conv
1050+
1051+
When: duplicate_shared_initializers is called
1052+
Then: scale and zp must NOT be duplicated — Q and DQ of the same pair
1053+
must always refer to the same quantization parameters.
1054+
"""
1055+
model = self._make_qdq_model(share_scale_across_pairs=False)
1056+
1057+
n_duplicates = duplicate_shared_initializers(model.graph)
1058+
1059+
assert n_duplicates == 0, (
1060+
"scale/zp shared within a QDQ pair must not be duplicated"
1061+
)
1062+
init_names = {init.name for init in model.graph.initializer}
1063+
assert "scale" in init_names
1064+
assert "zp" in init_names
1065+
assert "scale_copy_1" not in init_names
1066+
assert "zp_copy_1" not in init_names
1067+
1068+
def test_duplicate_shared_initializers_qdq_weight_still_duplicated(self):
1069+
"""
1070+
In a QDQ model, a plain weight initializer (not scale/zp) that is
1071+
shared between two Conv nodes must still be duplicated, while the
1072+
QDQ pair's scale and zero_point remain untouched.
1073+
1074+
Graph:
1075+
shared_weight → Conv1
1076+
shared_weight → Conv2 ← ordinary shared weight; must be duplicated
1077+
1078+
qdq_weight → QuantizeLinear(scale, zp)
1079+
→ DequantizeLinear(scale, zp) → Conv3
1080+
↑ scale/zp shared within QDQ pair; must NOT be duplicated
1081+
"""
1082+
CHANNELS = 4
1083+
shared_weight = numpy_helper.from_array(
1084+
np.random.randn(CHANNELS, CHANNELS, 1, 1).astype(np.float32),
1085+
name="shared_weight",
1086+
)
1087+
qdq_weight = numpy_helper.from_array(
1088+
np.random.randn(CHANNELS, CHANNELS, 1, 1).astype(np.float32),
1089+
name="qdq_weight",
1090+
)
1091+
scale = numpy_helper.from_array(np.array(0.1, dtype=np.float32), name="scale")
1092+
zp = numpy_helper.from_array(np.array(0, dtype=np.int8), name="zp")
1093+
1094+
nodes = [
1095+
# Two Conv nodes directly sharing shared_weight
1096+
helper.make_node(
1097+
"Conv",
1098+
inputs=["model_input", "shared_weight"],
1099+
outputs=["conv1_out"],
1100+
name="conv1",
1101+
),
1102+
helper.make_node(
1103+
"Conv",
1104+
inputs=["model_input", "shared_weight"],
1105+
outputs=["conv2_out"],
1106+
name="conv2",
1107+
),
1108+
helper.make_node(
1109+
"Add",
1110+
inputs=["conv1_out", "conv2_out"],
1111+
outputs=["branch1_out"],
1112+
name="add1",
1113+
),
1114+
# QDQ pair with scale/zp shared between Q and DQ
1115+
helper.make_node(
1116+
"QuantizeLinear",
1117+
inputs=["qdq_weight", "scale", "zp"],
1118+
outputs=["qdq_weight_q"],
1119+
name="q1",
1120+
),
1121+
helper.make_node(
1122+
"DequantizeLinear",
1123+
inputs=["qdq_weight_q", "scale", "zp"],
1124+
outputs=["qdq_weight_dq"],
1125+
name="dq1",
1126+
),
1127+
helper.make_node(
1128+
"Conv",
1129+
inputs=["model_input", "qdq_weight_dq"],
1130+
outputs=["conv3_out"],
1131+
name="conv3",
1132+
),
1133+
helper.make_node(
1134+
"Add",
1135+
inputs=["branch1_out", "conv3_out"],
1136+
outputs=["model_output"],
1137+
name="add2",
1138+
),
1139+
]
1140+
graph = helper.make_graph(
1141+
nodes,
1142+
"QDQWithSharedWeight",
1143+
inputs=[
1144+
helper.make_tensor_value_info(
1145+
"model_input", TensorProto.FLOAT, [1, CHANNELS, 8, 8]
1146+
)
1147+
],
1148+
outputs=[
1149+
helper.make_tensor_value_info(
1150+
"model_output", TensorProto.FLOAT, [1, CHANNELS, 8, 8]
1151+
)
1152+
],
1153+
initializer=[shared_weight, qdq_weight, scale, zp],
1154+
)
1155+
model = models_for_tests.make_model(graph)
1156+
onnx.checker.check_model(model)
1157+
1158+
n_duplicates = duplicate_shared_initializers(model.graph)
1159+
1160+
# shared_weight is used by two Conv nodes → must be duplicated once
1161+
assert n_duplicates == 1
1162+
init_names = {init.name for init in model.graph.initializer}
1163+
assert "shared_weight_copy_1" in init_names
1164+
1165+
# Conv1 and Conv2 must now reference distinct initializers
1166+
conv_inputs = [
1167+
node.input[1]
1168+
for node in model.graph.node
1169+
if node.op_type == "Conv" and node.name in ("conv1", "conv2")
1170+
]
1171+
assert conv_inputs[0] != conv_inputs[1]
1172+
1173+
# scale and zp must NOT have been duplicated
1174+
assert "scale_copy_1" not in init_names
1175+
assert "zp_copy_1" not in init_names
1176+
1177+
def test_duplicate_shared_initializers_qdq_cross_pair_scale_excluded(self):
1178+
"""
1179+
When the same scale/zp initializer is shared across two different QDQ
1180+
pairs, it is currently excluded from duplication.
1181+
1182+
Graph:
1183+
weight → QLinear1(scale, zp) → DQLinear1(scale, zp) → Conv1
1184+
weight2 → QLinear2(scale, zp) → DQLinear2(scale, zp) → Conv2
1185+
"""
1186+
model = self._make_qdq_model(share_scale_across_pairs=True)
1187+
1188+
n_duplicates = duplicate_shared_initializers(model.graph)
1189+
1190+
assert n_duplicates == 0, "Cross-pair shared scale/zp is excluded"
1191+
init_names = {init.name for init in model.graph.initializer}
1192+
assert "scale_copy_1" not in init_names
1193+
assert "zp_copy_1" not in init_names
1194+
9451195

9461196
class TestORTInferenceSession:
9471197
"""

0 commit comments

Comments
 (0)