Skip to content

Commit f95486d

Browse files
abeakkasabeakkas
authored
Fix int8 zero-point overflow and conv bias_scale in eager quantized ref kernels
Differential Revision: D110220645 Pull Request resolved: #20655
1 parent d54a0c0 commit f95486d

2 files changed

Lines changed: 83 additions & 36 deletions

File tree

‎backends/cadence/aot/ref_implementations.py‎

Lines changed: 33 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -318,22 +318,19 @@ def quantized_add_per_tensor(
318318
f"X and Y dtypes need to be in {supported_dtypes}. Got {dtype}"
319319
)
320320

321-
if dtype == torch.uint8:
322-
X = X.to(torch.int8)
323-
Y = Y.to(torch.int8)
321+
qmin = torch.iinfo(dtype).min
322+
qmax = torch.iinfo(dtype).max
324323

325-
# TODO(agrebenisan): This should be done in fixed point arithmetic, but to match the quantized_add_out.cpp
326-
# reference implementation, we'll do it in floating point.
327-
dequant_X = X_scale * (X - X_zero_point)
328-
dequant_Y = Y_scale * (Y - Y_zero_point)
324+
dequant_X = dequantize_per_tensor(X, X_scale, X_zero_point, qmin, qmax, dtype)
325+
dequant_Y = dequantize_per_tensor(Y, Y_scale, Y_zero_point, qmin, qmax, dtype)
329326

330327
# q_min/q_max are unused args
331328
return quantize_per_tensor(
332329
dequant_X + dequant_Y,
333330
out_scale,
334331
out_zero_point,
335-
torch.iinfo(dtype).min,
336-
torch.iinfo(dtype).max,
332+
qmin,
333+
qmax,
337334
dtype,
338335
)
339336

@@ -394,9 +391,9 @@ def quantized_add_asym8uxasym8u_asym8u_per_tensor(
394391
out_zero_point: int,
395392
) -> torch.Tensor:
396393
if X.dtype != torch.uint8:
397-
raise ValueError("X dtype must be torch.int8")
394+
raise ValueError("X dtype must be torch.uint8")
398395
if Y.dtype != torch.uint8:
399-
raise ValueError("Y dtype must be torch.int8")
396+
raise ValueError("Y dtype must be torch.uint8")
400397

401398
return quantized_add_per_tensor(
402399
X, X_scale, X_zero_point, Y, Y_scale, Y_zero_point, out_scale, out_zero_point
@@ -447,19 +444,18 @@ def quantized_mul_per_tensor(
447444
f"X and Y dtypes need to be in {supported_dtypes}. Got {dtype}"
448445
)
449446

450-
if dtype == torch.uint8:
451-
X = X.to(torch.int8)
452-
Y = Y.to(torch.int8)
447+
qmin = torch.iinfo(dtype).min
448+
qmax = torch.iinfo(dtype).max
453449

454-
dequant_X = X_scale * (X - X_zero_point)
455-
dequant_Y = Y_scale * (Y - Y_zero_point)
450+
dequant_X = dequantize_per_tensor(X, X_scale, X_zero_point, qmin, qmax, dtype)
451+
dequant_Y = dequantize_per_tensor(Y, Y_scale, Y_zero_point, qmin, qmax, dtype)
456452

457453
return quantize_per_tensor(
458454
dequant_X * dequant_Y,
459455
out_scale,
460456
out_zero_point,
461-
torch.iinfo(dtype).min,
462-
torch.iinfo(dtype).max,
457+
qmin,
458+
qmax,
463459
dtype,
464460
)
465461

@@ -503,8 +499,8 @@ def quantized_linear_common(
503499
)
504500

505501
out = torch.nn.functional.linear(
506-
(src - in_zero_point).float(),
507-
(weight - weight_zero_point).float(),
502+
src.float() - in_zero_point,
503+
weight.float() - weight_zero_point,
508504
bias.float(),
509505
)
510506
return quantize_per_tensor(
@@ -673,8 +669,8 @@ def quantized_matmul(
673669
out_scale = 1.0 / (-out_multiplier * (1 / (1 << 31)) * (2**out_shift))
674670

675671
out = torch.matmul(
676-
(X - X_zero_point).float(),
677-
(Y - Y_zero_point).float(),
672+
X.float() - X_zero_point,
673+
Y.float() - Y_zero_point,
678674
)
679675
return quantize_per_tensor(
680676
out,
@@ -857,21 +853,21 @@ def quantized_conv_per_tensor(
857853
- out_shift (int): Unused
858854
"""
859855
if len(input_tensor.shape) == 3:
860-
float_out = torch.nn.functional.conv1d(
861-
(input_tensor - in_zero_point).float(),
862-
(weight - weight_zero_point).float(),
863-
(bias * bias_scale).float(),
856+
acc = torch.nn.functional.conv1d(
857+
input_tensor.float() - in_zero_point,
858+
weight.float() - weight_zero_point,
859+
bias.float(),
864860
stride[-1],
865861
padding[-1],
866862
dilation[-1],
867863
groups,
868864
)
869865

870866
elif len(input_tensor.shape) == 4:
871-
float_out = torch.nn.functional.conv2d(
872-
(input_tensor - in_zero_point).float(),
873-
(weight - weight_zero_point).float(),
874-
(bias * bias_scale).float(),
867+
acc = torch.nn.functional.conv2d(
868+
input_tensor.float() - in_zero_point,
869+
weight.float() - weight_zero_point,
870+
bias.float(),
875871
stride,
876872
padding,
877873
dilation,
@@ -880,6 +876,11 @@ def quantized_conv_per_tensor(
880876
else:
881877
raise ValueError("Input tensor must be 3D or 4D")
882878

879+
# conv accumulates in the integer domain (scale = in_scale * weight_scale =
880+
# bias_scale) with the integer bias added pre-scale; dequantize the whole
881+
# accumulation by bias_scale to get the floating-point result.
882+
float_out = acc * bias_scale
883+
883884
return quantize_per_tensor(
884885
float_out,
885886
output_scale,
@@ -1944,8 +1945,8 @@ def quantized_relu_common(
19441945

19451946
out_scale = 1.0 / (-out_multiplier * (1 / (1 << 31)) * (2**out_shift))
19461947
dequantized_X = torch.where(
1947-
X > X_zero_point, X - X_zero_point, torch.zeros_like(X)
1948-
).to(torch.float32)
1948+
X > X_zero_point, X.float() - X_zero_point, torch.zeros_like(X)
1949+
)
19491950
out = quantize_per_tensor(
19501951
dequantized_X,
19511952
out_scale,

‎backends/cadence/aot/tests/test_ref_implementations.py‎

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -670,10 +670,11 @@ def test_quantized_layer_norm_per_tensor(
670670
0, # unused out_shift
671671
torch.uint8, # dtype
672672
torch.tensor(
673-
[[[[238]]]], dtype=torch.uint8
674-
), # (130 - 128) + (134 - 128) = 10
675-
# + bias -> 10 + 1 = 11
676-
# round(11 / 0.1 + 128) = 238,
673+
[[[[148]]]], dtype=torch.uint8
674+
), # conv_acc = sum((input - 128) * (weight - 128))
675+
# = 2*1 + 4*0 + 6*0 + 8*1 = 10
676+
# float_out = bias_scale * (conv_acc + bias) = 0.1 * (10 + 10) = 2.0
677+
# round(2.0 / 0.1 + 128) = 148
677678
memory_format,
678679
)
679680
for memory_format in [torch.contiguous_format, torch.channels_last]
@@ -918,6 +919,34 @@ def test_quantized_layer_norm_per_tensor(
918919
)
919920
for memory_format in [torch.contiguous_format, torch.channels_last]
920921
],
922+
# Zero-point overflow: int8 input minus a negative zero point exceeds
923+
# the int8 range and wraps unless the subtraction is upcast first.
924+
*[
925+
(
926+
torch.tensor(
927+
[[[[120, 120], [120, 120]]]], dtype=torch.int8
928+
), # input
929+
torch.tensor([[[[1, 0], [0, 1]]]], dtype=torch.int8), # weight
930+
torch.tensor([0], dtype=torch.int32), # bias
931+
(1, 1), # stride
932+
(0, 0), # padding
933+
(1, 1), # dilation
934+
1, # groups
935+
-20, # in_zero_point (120 - (-20) = 140 wraps to -116 in int8)
936+
0, # weight_zero_point
937+
1.0, # bias_scale
938+
4.0, # output_scale
939+
0, # output_zero_point
940+
0, # unused out_multiplier
941+
0, # unused out_shift
942+
torch.int8, # dtype
943+
torch.tensor(
944+
[[[[70]]]], dtype=torch.int8
945+
), # (120 + 20) * 2 / 4 = 70
946+
memory_format,
947+
)
948+
for memory_format in [torch.contiguous_format, torch.channels_last]
949+
],
921950
]
922951
)
923952
def test_quantized_conv_per_tensor(
@@ -1407,6 +1436,23 @@ def test_quantized_w8a32_linear(
14071436
)
14081437
for dtype in [torch.uint8]
14091438
],
1439+
# Zero-point overflow: int8 X minus a negative zero point exceeds the
1440+
# int8 range and wraps unless the subtraction is upcast first.
1441+
*[
1442+
(
1443+
"int8_negative_zp_overflow",
1444+
torch.tensor([120], dtype=dtype), # input
1445+
-20, # X_zero_point (120 - (-20) = 140 wraps to -116 in int8)
1446+
0, # out_zero_point
1447+
1073741824, # out_multiplier (0.5 * 2^31)
1448+
0, # out_shift
1449+
dtype, # dtype
1450+
torch.tensor(
1451+
[-70], dtype=dtype
1452+
), # shifted = 140; 140 * (-0.5) = -70
1453+
)
1454+
for dtype in [torch.int8]
1455+
],
14101456
]
14111457
)
14121458
def test_quantized_relu_per_tensor(

0 commit comments

Comments
 (0)