diff --git a/exir/passes/_quant_patterns_and_replacements.py b/exir/passes/_quant_patterns_and_replacements.py index 463abc37c5e..fceedbe0274 100644 --- a/exir/passes/_quant_patterns_and_replacements.py +++ b/exir/passes/_quant_patterns_and_replacements.py @@ -298,7 +298,8 @@ def _( num_embeddings, packed_embedding_dim = weight.shape embedding_dim = packed_embedding_dim * 4 embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device) - return embedding(indices) + # As in the reference op, the variant without dtype returns the scales' dtype. + return embedding(indices).to(weight_scales.dtype) @register_fake("quantized_decomposed::embedding_2bit.out") @@ -473,7 +474,8 @@ def _( num_embeddings, packed_embedding_dim = weight.shape embedding_dim = packed_embedding_dim * 2 embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device) - return embedding(indices) + # As in the reference op, the variant without dtype returns the scales' dtype. + return embedding(indices).to(weight_scales.dtype) @register_fake("quantized_decomposed::embedding_4bit.out") diff --git a/exir/tests/test_quant_fusion_pass.py b/exir/tests/test_quant_fusion_pass.py index 499419a1e10..80c2637adfb 100644 --- a/exir/tests/test_quant_fusion_pass.py +++ b/exir/tests/test_quant_fusion_pass.py @@ -19,6 +19,7 @@ QuantFusionPass, ) from executorch.exir.tests.common import register_additional_test_aten_ops +from torch._subclasses.fake_tensor import FakeTensorMode from torch.ao.quantization import ( # @manual float_qparams_weight_only_qconfig, get_default_qconfig_mapping, @@ -381,6 +382,34 @@ def forward(self, indices): # m.dump_graph_module().code # ) + def test_embedding_xbit_fake_dtype(self) -> None: + # The variants without a dtype argument return the scales' dtype; their + # fake kernels must report the same dtype, or export records the wrong one. + indices = torch.tensor([1, 2, 3]) + for op, values_per_byte, quant_min, quant_max in ( + (torch.ops.quantized_decomposed.embedding_2bit.default, 4, -2, 1), + (torch.ops.quantized_decomposed.embedding_4bit.default, 2, -8, 7), + ): + for scales_dtype in (torch.float16, torch.bfloat16, torch.float32): + with self.subTest(op=str(op), scales_dtype=scales_dtype): + weight = torch.randint( + 0, 256, (10, 32 // values_per_byte), dtype=torch.uint8 + ) + scales = torch.rand(10, dtype=scales_dtype) + 0.5 + expected = op(weight, scales, None, quant_min, quant_max, indices) + with FakeTensorMode() as mode: + actual = op( + mode.from_tensor(weight), + mode.from_tensor(scales), + None, + quant_min, + quant_max, + mode.from_tensor(indices), + ) + self.assertEqual(expected.dtype, scales_dtype) + self.assertEqual(actual.dtype, expected.dtype) + self.assertEqual(actual.shape, expected.shape) + def test_embedding_torchao(self) -> None: for bit_width, use_dtype_variant, test_per_group, mapping_type in zip( [2, 4, 8], diff --git a/kernels/quantized/cpu/embeddingxb.cpp b/kernels/quantized/cpu/embeddingxb.cpp index ed5bc3f4397..e3269777fb7 100644 --- a/kernels/quantized/cpu/embeddingxb.cpp +++ b/kernels/quantized/cpu/embeddingxb.cpp @@ -125,9 +125,10 @@ void check_embedding_xbit_args( static_cast(weight_scales.dim())); ET_CHECK_MSG( - opt_weight_zero_points.value().scalar_type() == out.scalar_type(), + opt_weight_zero_points.value().scalar_type() == + weight_scales.scalar_type(), "weight zero points scalar type %" PRId8 - " does not match out.scalar_type()", + " does not match weight_scales.scalar_type()", static_cast(opt_weight_zero_points.value().scalar_type())); for (int32_t i = 0; i < weight_scales.dim(); ++i) { @@ -290,6 +291,15 @@ Tensor& quantized_embedding_xbit_out( out, weight_nbit); + // Without a dtype argument the output takes the scales' dtype, and the + // kernel reads the scales as the output type. + ET_CHECK_MSG( + weight_scales.scalar_type() == out_type, + "weight_scales.scalar_type() %" PRId8 + " does not match out.scalar_type() %" PRId8, + static_cast(weight_scales.scalar_type()), + static_cast(out_type)); + constexpr auto name = "quantized_decomposed::embedding_xbit.out"; ScalarType indices_type = indices.scalar_type(); ET_SWITCH_THREE_TYPES( diff --git a/kernels/quantized/cpu/op_embedding.cpp b/kernels/quantized/cpu/op_embedding.cpp index 682e2046c2e..b57d6612d30 100644 --- a/kernels/quantized/cpu/op_embedding.cpp +++ b/kernels/quantized/cpu/op_embedding.cpp @@ -87,9 +87,10 @@ void check_embedding_byte_args( static_cast(weight_scales.dim())); ET_CHECK_MSG( - opt_weight_zero_points.value().scalar_type() == out.scalar_type(), + opt_weight_zero_points.value().scalar_type() == + weight_scales.scalar_type(), "weight zero points scalar type %" PRId8 - " does not match out.scalar_type()", + " does not match weight_scales.scalar_type()", static_cast(opt_weight_zero_points.value().scalar_type())); for (int32_t i = 0; i < weight_scales.dim(); ++i) { @@ -259,6 +260,15 @@ Tensor& quantized_embedding_byte_out( out_type, out); + // Without a dtype argument the output takes the scales' dtype, and the + // kernel reads the scales as the output type. + ET_CHECK_MSG( + weight_scales.scalar_type() == out_type, + "weight_scales.scalar_type() %" PRId8 + " does not match out.scalar_type() %" PRId8, + static_cast(weight_scales.scalar_type()), + static_cast(out_type)); + constexpr auto name = "quantized_decomposed::embedding_byte.out"; ET_SWITCH_TWO_TYPES(Byte, Char, w_type, ctx, name, CTYPE_W, [&]() { ET_SWITCH_THREE_TYPES( diff --git a/kernels/quantized/test/op_embedding4b_test.cpp b/kernels/quantized/test/op_embedding4b_test.cpp index f9636e14376..95ffb2b7f2d 100644 --- a/kernels/quantized/test/op_embedding4b_test.cpp +++ b/kernels/quantized/test/op_embedding4b_test.cpp @@ -216,6 +216,79 @@ TEST(OpQuantizedEmbedding4bTest, TestGroupWiseQuantizedEmbeddingBFloat16) { EXPECT_TENSOR_CLOSE(out, expected); } +TEST(OpQuantizedEmbedding4bTest, DtypeOutHalfQParamsFloatOut) { + et_pal_init(); + TensorFactory tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfl; + + Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5}); + Tensor weight_zero_points = tfh.make({3}, {1, -5, 0}); + Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126}); + Tensor indices = tfl.make({3}, {0, 2, 1}); + + Tensor out = tf.zeros({3, 4}); + Tensor expected = tf.make( + {3, 4}, {-2.0, 0.0, 2.5, 3.0, -12.0, 4.5, -1.5, 9.0, 7.0, 0.0, 1.0, 5.0}); + + quantized_embedding_4bit_dtype_out( + qweight, + weight_scales, + weight_zero_points, + -8, + 7, + indices, + ScalarType::Float, + out); + + EXPECT_TENSOR_CLOSE(out, expected); +} + +TEST(OpQuantizedEmbedding4bTest, DtypeOutZeroPointsScalesMismatchDies) { + et_pal_init(); + TensorFactory tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfl; + + Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5}); + Tensor weight_zero_points = tf.make({3}, {1, -5, 0}); + Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126}); + Tensor indices = tfl.make({3}, {0, 2, 1}); + + Tensor out = tf.zeros({3, 4}); + ET_EXPECT_DEATH( + quantized_embedding_4bit_dtype_out( + qweight, + weight_scales, + weight_zero_points, + -8, + 7, + indices, + ScalarType::Float, + out), + "does not match weight_scales"); +} + +TEST(OpQuantizedEmbedding4bTest, OutScalesOutMismatchDies) { + et_pal_init(); + TensorFactory tfb; + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfl; + + Tensor weight_scales = tfh.make({3}, {0.5, 1.0, 1.5}); + Tensor qweight = tfb.make({3, 2}, {89, 239, 163, 72, 11, 126}); + Tensor indices = tfl.make({3}, {0, 2, 1}); + + Tensor out = tf.zeros({3, 4}); + ET_EXPECT_DEATH( + quantized_embedding_4bit_out( + qweight, weight_scales, std::nullopt, -8, 7, indices, out), + "does not match out.scalar_type"); +} + TEST(OpQuantizedEmbedding4bTest, TestGroupWiseQuantizedEmbeddingDeath2) { et_pal_init(); TensorFactory tfb; diff --git a/kernels/quantized/test/op_embedding_test.cpp b/kernels/quantized/test/op_embedding_test.cpp index ad27daf866e..c499ebfface 100644 --- a/kernels/quantized/test/op_embedding_test.cpp +++ b/kernels/quantized/test/op_embedding_test.cpp @@ -510,3 +510,78 @@ TEST(OpQuantizedEmbeddingTest, BFloat16Rounding) { EXPECT_TENSOR_CLOSE_WITH_TOL(out, expected, 1e-2, 1e-2); } + +// Zero points share the scales' dtype, which may differ from the output dtype. +TEST(OpQuantizedEmbeddingTest, DtypeOutHalfQParamsFloatOut) { + et_pal_init(); + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory tf_l; + + Tensor weight_scales = tfh.full({3}, 0.5); + Tensor weight_zero_points = tfh.full({3}, 1); + Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27}); + Tensor indices = tf_l.make({2}, {0, 2}); + + Tensor out = tf.zeros({2, 2}); + Tensor expected = tf.make({2, 2}, {3.5, 2, 5.5, 13}); + + quantized_embedding_byte_dtype_out( + qweight, + weight_scales, + weight_zero_points, + 0, + 255, + indices, + ScalarType::Float, + out); + + EXPECT_TENSOR_CLOSE(out, expected); +} + +TEST(OpQuantizedEmbeddingTest, DtypeOutZeroPointsScalesMismatchDies) { + et_pal_init(); + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory tf_l; + + Tensor weight_scales = tfh.full({3}, 0.5); + Tensor weight_zero_points = tf.full({3}, 1); + Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27}); + Tensor indices = tf_l.make({2}, {0, 2}); + + Tensor out = tf.zeros({2, 2}); + ET_EXPECT_DEATH( + quantized_embedding_byte_dtype_out( + qweight, + weight_scales, + weight_zero_points, + 0, + 255, + indices, + ScalarType::Float, + out), + "does not match weight_scales"); +} + +// Without a dtype argument the scales are read as the output type, so a +// narrower scales tensor would be read past its end. +TEST(OpQuantizedEmbeddingTest, OutScalesOutMismatchDies) { + et_pal_init(); + TensorFactory tfh; + TensorFactory tf; + TensorFactory tfb; + TensorFactory tf_l; + + Tensor weight_scales = tfh.full({3}, 0.5); + Tensor qweight = tfb.make({3, 2}, {8, 5, 9, 3, 12, 27}); + Tensor indices = tf_l.make({2}, {0, 2}); + + Tensor out = tf.zeros({2, 2}); + ET_EXPECT_DEATH( + quantized_embedding_byte_out( + qweight, weight_scales, std::nullopt, 0, 255, indices, out), + "does not match out.scalar_type"); +}