Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.

Commit d360f12

Browse files
authored
Fix ONNX export for MobileBERT (#1539)
* Account for the possibility of the quantized embedddings to be in int8 format (conversion to uint8 occurs later) * Set the padding value to match to the zero point accordingly. * Style and quality fixes
1 parent b282b80 commit d360f12

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

src/sparseml/pytorch/sparsification/quantization/quantize_qat_export.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1761,7 +1761,7 @@ def _propagate_mobilebert_embedding_quantization(model: ModelProto):
17611761
continue
17621762

17631763
embedding_array = numpy_helper.to_array(embedding_initializer)
1764-
if embedding_array.dtype != numpy.uint8:
1764+
if embedding_array.dtype not in [numpy.uint8, numpy.int8]:
17651765
continue
17661766

17671767
dequant_node = graph.get_node_single_child(gather_node)
@@ -1805,12 +1805,15 @@ def _propagate_mobilebert_embedding_quantization(model: ModelProto):
18051805
# switch position of dequantize node
18061806
for branch_node in graph.get_node_children(dequant_node):
18071807
if branch_node.op_type == "Slice":
1808+
zero_point = graph.get_init_by_name(dequant_node.input[2])
1809+
zero_point_array = numpy_helper.to_array(zero_point)
18081810
branch_node.input[0] = gather_node.output[0]
18091811
pad_node = graph.get_node_single_child(branch_node)
18101812
pad_value = graph.get_init_by_name(pad_node.input[2])
18111813
pad_value_array = numpy_helper.to_array(pad_value)
1812-
pad_value_array = pad_value_array + 128
1813-
pad_value_array = pad_value_array.astype(numpy.uint8)
1814+
pad_value_array = (
1815+
pad_value_array.astype(zero_point_array.dtype) + zero_point_array
1816+
)
18141817
model.graph.initializer.remove(pad_value)
18151818
pad_value = numpy_helper.from_array(
18161819
pad_value_array, name=pad_value.name

0 commit comments

Comments
 (0)