Skip to content

Commit 6d3c1c7

Browse files
JakeStevensJakeStevens
authored
Preserve Python scalar semantics in Vulkan and XNNPACK quantizers (#22804)
### Summary Preparing a model with the Vulkan or XNNPACK quantizer can change its results when a Python scalar is materialized using the operation's output dtype. For example, `torch.tensor([1e-4], dtype=torch.float16) * 100000.0` produces approximately 10 in PyTorch, but casting the scalar to FP16 first makes the result infinite. Smaller scalars can underflow, and other dtypes have their own promotion and rounding rules. Keep the original Python scalar operands for add and multiply operations whose output is not FP32. FP32 operations retain scalar lifting for quantization. Integer index arithmetic covered by #22372 and #22065 continues to work. ### Test plan All 10 focused tests and 208 scalar-promotion subtests pass across dtypes, add/multiply, scalar/vector inputs, and configured/unconfigured quantizers: ```bash OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 PYTHONPATH=src python -m pytest -o addopts= -q \ backends/vulkan/test/quantizer/test_vulkan_quantizer.py \ backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py \ -k 'scalar or int64 or add_mul_long or mul_float32_max or add_and_inplace_add or mul_and_inplace_mul' ``` Existing FP32 scalar quantization, tensor add/multiply, large-scalar handling, and integer-index regressions are included. `lintrunner --revision main` and `git diff --check` pass. Authored with OpenAI Codex. cc @SS-JIA @manuelcandales @digantdesai @cbilgin
1 parent 83d1c53 commit 6d3c1c7

4 files changed

Lines changed: 121 additions & 35 deletions

File tree

‎backends/vulkan/quantizer/vulkan_quantizer_utils.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,12 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu
201201
args = list(n.args)
202202
new_args = []
203203
for i in range(len(args)):
204-
if isinstance(args[i], torch.fx.Node):
204+
# Only FP32 binary ops need scalar lifting for quantization. Other
205+
# dtypes must retain Python scalar promotion and rounding semantics.
206+
if (
207+
isinstance(args[i], torch.fx.Node)
208+
or n.meta["val"].dtype != torch.float32
209+
):
205210
new_args.append(args[i])
206211
continue
207212
prefix = "_tensor_constant_"

‎backends/vulkan/test/quantizer/test_vulkan_quantizer.py‎

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7+
import itertools
78
import unittest
89

910
import torch
@@ -17,9 +18,6 @@
1718

1819
class TestVulkanQuantizer(unittest.TestCase):
1920
def test_int64_scalar_add_used_as_index(self):
20-
"""Scalars lifted to attrs must keep the op's output dtype; an int64
21-
add chain used as an index must not be promoted to float32."""
22-
2321
class M(torch.nn.Module):
2422
def forward(self, x):
2523
return x[:, torch.arange(4) + 0]
@@ -30,32 +28,56 @@ def forward(self, x):
3028
example_inputs = (torch.randn(1, 4, 5),)
3129
m = export(M(), example_inputs, strict=True).module()
3230
m = quantizer.transform_for_annotation(m)
33-
lifted_constants = [
34-
m.get_buffer(n.target)
35-
for n in m.graph.nodes
36-
if n.op == "get_attr" and n.target.startswith("_tensor_constant_")
37-
]
38-
self.assertEqual(len(lifted_constants), 1)
39-
self.assertEqual(lifted_constants[0].dtype, torch.int64)
4031
m = prepare_pt2e(m, quantizer)
41-
m(*example_inputs)
42-
43-
def test_int64_scalar_lifted_without_set_global(self):
44-
"""Passing through VulkanQuantizer without a config still lifts scalars
45-
and must preserve dtype."""
32+
torch.testing.assert_close(m(*example_inputs), M()(*example_inputs))
4633

34+
def test_int64_scalar_add_without_set_global(self):
4735
class M(torch.nn.Module):
4836
def forward(self, x):
4937
return x[:, torch.arange(4) + 0]
5038

5139
example_inputs = (torch.randn(1, 4, 5),)
5240
m = export(M(), example_inputs, strict=True).module()
5341
m = VulkanQuantizer().transform_for_annotation(m)
54-
lifted_constants = [
55-
m.get_buffer(n.target)
56-
for n in m.graph.nodes
57-
if n.op == "get_attr" and n.target.startswith("_tensor_constant_")
42+
torch.testing.assert_close(m(*example_inputs), M()(*example_inputs))
43+
44+
def test_scalar_type_promotion(self):
45+
class M(torch.nn.Module):
46+
def __init__(self, op, scalar):
47+
super().__init__()
48+
self.op = op
49+
self.scalar = scalar
50+
51+
def forward(self, x):
52+
return self.op(x, self.scalar)
53+
54+
cases = [
55+
(torch.float16, 1e-4, 100000.0),
56+
(torch.float16, 1e-4, 100000),
57+
(torch.float16, 10000.0, 1e-8),
58+
(torch.bfloat16, 100.0, 1.0039),
59+
(torch.float64, 1.0, 1.0 + 2**-30),
60+
(torch.int64, 2**54 + 1, 1),
61+
(torch.int32, 1, 1),
62+
(torch.int32, 1, 2**31),
63+
(torch.int8, 1, 256),
64+
(torch.bool, True, False),
65+
(torch.int32, 1, 0.5),
66+
(torch.float32, 1.0, 2),
67+
(torch.complex64, 1j, 1 + 2j),
5868
]
59-
self.assertEqual(len(lifted_constants), 1)
60-
self.assertEqual(lifted_constants[0].dtype, torch.int64)
61-
m(*example_inputs)
69+
for op, (dtype, value, scalar), shape, configured in itertools.product(
70+
(torch.add, torch.mul), cases, ((), (2,)), (False, True)
71+
):
72+
with self.subTest(
73+
op=op, dtype=dtype, scalar=scalar, shape=shape, configured=configured
74+
):
75+
model = M(op, scalar)
76+
example_inputs = (torch.full(shape, value, dtype=dtype),)
77+
expected = model(*example_inputs)
78+
quantizer = VulkanQuantizer()
79+
if configured:
80+
quantizer.set_global(get_symmetric_quantization_config())
81+
m = export(model, example_inputs, strict=True).module()
82+
m = prepare_pt2e(m, quantizer)
83+
torch.testing.assert_close(m(*example_inputs), expected, rtol=0, atol=0)

‎backends/xnnpack/quantizer/xnnpack_quantizer_utils.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -789,6 +789,9 @@ def _annotate_add_relu( # noqa: C901
789789
if filter_fn and any(not filter_fn(n) for n in partition):
790790
continue
791791

792+
if _is_input_non_float_tensor(add_node):
793+
continue
794+
792795
input_act_qspec = get_input_act_qspec(quantization_config)
793796
output_act_qspec = get_output_act_qspec(quantization_config)
794797

@@ -845,6 +848,9 @@ def _annotate_add(
845848
if filter_fn and any(not filter_fn(n) for n in partition):
846849
continue
847850

851+
if _is_input_non_float_tensor(add_node):
852+
continue
853+
848854
input_act_qspec = get_input_act_qspec(quantization_config)
849855
output_act_qspec = get_output_act_qspec(quantization_config)
850856

@@ -913,6 +919,9 @@ def _annotate_mul_relu( # noqa: C901
913919
if filter_fn and any(not filter_fn(n) for n in partition):
914920
continue
915921

922+
if _is_input_non_float_tensor(mul_node):
923+
continue
924+
916925
input_act_qspec = get_input_act_qspec(quantization_config)
917926
output_act_qspec = get_output_act_qspec(quantization_config)
918927

@@ -969,6 +978,9 @@ def _annotate_mul(
969978
if filter_fn and any(not filter_fn(n) for n in partition):
970979
continue
971980

981+
if _is_input_non_float_tensor(mul_node):
982+
continue
983+
972984
input_act_qspec = get_input_act_qspec(quantization_config)
973985
output_act_qspec = get_output_act_qspec(quantization_config)
974986

@@ -1154,7 +1166,12 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu
11541166
args = list(n.args)
11551167
new_args = []
11561168
for i in range(len(args)):
1157-
if isinstance(args[i], torch.fx.Node):
1169+
# Only FP32 binary ops need scalar lifting for quantization. Other
1170+
# dtypes must retain Python scalar promotion and rounding semantics.
1171+
if (
1172+
isinstance(args[i], torch.fx.Node)
1173+
or n.meta["val"].dtype != torch.float32
1174+
):
11581175
new_args.append(args[i])
11591176
continue
11601177
prefix = "_tensor_constant_"

‎backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py‎

Lines changed: 53 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# Owner(s): ["oncall: mobile"]
22
import copy
3+
import itertools
34
import operator
45

56
import torch
@@ -1123,9 +1124,6 @@ def forward(self, x):
11231124
)
11241125

11251126
def test_int64_scalar_add_used_as_index(self):
1126-
"""Scalars lifted to attrs must keep the op's output dtype; an int64
1127-
add chain used as an index must not be promoted to float32."""
1128-
11291127
class M(torch.nn.Module):
11301128
def forward(self, x):
11311129
return x[:, torch.arange(4) + 0]
@@ -1136,15 +1134,59 @@ def forward(self, x):
11361134
example_inputs = (torch.randn(1, 4, 5),)
11371135
m = export(M(), example_inputs, strict=True).module()
11381136
m = quantizer.transform_for_annotation(m)
1139-
lifted_constants = [
1140-
m.get_buffer(n.target)
1141-
for n in m.graph.nodes
1142-
if n.op == "get_attr" and n.target.startswith("_tensor_constant_")
1143-
]
1144-
self.assertEqual(len(lifted_constants), 1)
1145-
self.assertEqual(lifted_constants[0].dtype, torch.int64)
11461137
m = prepare_pt2e(m, quantizer)
1147-
m(*example_inputs)
1138+
torch.testing.assert_close(m(*example_inputs), M()(*example_inputs))
1139+
1140+
def test_scalar_type_promotion(self):
1141+
class M(torch.nn.Module):
1142+
def __init__(self, op, scalar):
1143+
super().__init__()
1144+
self.op = op
1145+
self.scalar = scalar
1146+
1147+
def forward(self, x):
1148+
return self.op(x, self.scalar)
1149+
1150+
cases = [
1151+
(torch.float16, 1e-4, 100000.0),
1152+
(torch.float16, 1e-4, 100000),
1153+
(torch.float16, 10000.0, 1e-8),
1154+
(torch.bfloat16, 100.0, 1.0039),
1155+
(torch.float64, 1.0, 1.0 + 2**-30),
1156+
(torch.int64, 2**54 + 1, 1),
1157+
(torch.int32, 1, 1),
1158+
(torch.int32, 1, 2**31),
1159+
(torch.int8, 1, 256),
1160+
(torch.bool, True, False),
1161+
(torch.int32, 1, 0.5),
1162+
(torch.float32, 1.0, 2),
1163+
(torch.float32, 1.0, 1 + 2j),
1164+
(torch.complex64, 1j, 1 + 2j),
1165+
]
1166+
for op, (dtype, value, scalar), shape, configured in itertools.product(
1167+
(torch.add, torch.mul), cases, ((), (2,)), (False, True)
1168+
):
1169+
with self.subTest(
1170+
op=op, dtype=dtype, scalar=scalar, shape=shape, configured=configured
1171+
):
1172+
model = M(op, scalar)
1173+
example_inputs = (torch.full(shape, value, dtype=dtype),)
1174+
expected = model(*example_inputs)
1175+
quantizer = XNNPACKQuantizer()
1176+
if configured:
1177+
quantizer.set_global(get_symmetric_quantization_config())
1178+
m = export(model, example_inputs, strict=True).module()
1179+
m = prepare_pt2e(m, quantizer)
1180+
torch.testing.assert_close(m(*example_inputs), expected, rtol=0, atol=0)
1181+
if (
1182+
configured
1183+
and dtype == torch.float32
1184+
and isinstance(scalar, complex)
1185+
):
1186+
m = convert_pt2e(m)
1187+
torch.testing.assert_close(
1188+
m(*example_inputs), expected, rtol=0, atol=0
1189+
)
11481190

11491191
def test_cat_same_node(self):
11501192
"""Ensure that concatenating the same node does not cause any unexpected behavior"""

0 commit comments

Comments
 (0)