Refactor int8 weight only quant to use quantize

Summary:
Similar to pytorch#294 we replaced the implementation
of int8 weight only quant to used the newly added `quantize` function, as a part of
the unification effort for affine quantization

Test Plan:
1. unit perf test:
python test/quantization/test_quant_api.py -k test_quantized_tensor_subclass_int8_wo_quant_perf

elapsed time: 0.23909856796264647, ref elapsed time: 0.25150911331176756
elapsed time: 0.24894208908081056, ref elapsed time: 0.2570047950744629
elapsed time: 0.21607391357421876, ref elapsed time: 0.22809568405151368

2. integration test:

TORCH_LOGS='output_code' python tutorials/quantize_vit/run_vit_b_quant.py

Reference: elapsed_time:  1.355208740234375  milliseconds
After refactor: elapsed_time:  1.32778857421875  milliseconds

code diff (gist): https://gist.github.com/jerryzh168/921a722cf20d476c8fc5888482e722dc
code diff (meta-only paste): https://www.internalfb.com/phabricator/paste/view/P1387333845

Reviewers:

Subscribers:

Tasks:

Tags:

test/quantization/test_quant_api.py

@@ -138,6 +138,40 @@ def _ref_change_linear_weights_to_int8_dqtensors(model, filter_fn=None, **kwargs
         model, _get_subclass_inserter(Int8DynamicallyQuantizedLinearWeight, enable_parametrization=False, **kwargs), filter_fn
     )
 
+def _ref_change_linear_weights_to_int8_woqtensors(model, filter_fn=None, **kwargs):
+    """
+    The deprecated implementation for int8 weight only quant API, used as a reference for
+    numerics and performance
+    """
+    from torchao.quantization.quant_api import _is_linear
+    from torchao.quantization.quant_api import _get_subclass_inserter
+    from torchao.quantization.subclass import Int8WeightOnlyQuantizedLinearWeight
+
+    filter_fn = kwargs.pop("filter_fn", _is_linear)
+
+    _replace_with_custom_fn_if_matches_filter(
+        model,
+        _get_subclass_inserter(Int8WeightOnlyQuantizedLinearWeight, enable_parametrization=True, **kwargs),
+        filter_fn,
+    )
+
+def _ref_change_linear_weights_to_int4_woqtensors(model, **kwargs):
+    """
+    The deprecated implementation for int4 weight only quant API, used as a reference for
+    numerics and performance
+    """
+    from torchao.quantization.quant_api import _is_linear
+    from torchao.quantization.quant_api import _get_subclass_inserter
+    from torchao.quantization.subclass import Int4WeightOnlyQuantizedLinearWeight
+
+    filter_fn = kwargs.pop("filter_fn", _is_linear)
+
+    _replace_with_custom_fn_if_matches_filter(
+        model,
+        _get_subclass_inserter(Int4WeightOnlyQuantizedLinearWeight, enable_parametrization=False, **kwargs),
+        filter_fn,
+    )
+
 class TestQuantFlow(unittest.TestCase):
     def test_dynamic_quant_gpu_singleline(self):
         m = ToyLinearModel().eval()
@@ -489,7 +523,7 @@ def test_quantized_tensor_subclass_int4(self):
 
     @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
     @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
-    def test_quantized_tensor_subclass_int8(self):
+    def test_quantized_tensor_subclass_int8_wo(self):
         m = ToyLinearModel().eval().to(torch.bfloat16)
         m_copy = copy.deepcopy(m)
         example_inputs = tuple(map(lambda x: x.to(torch.bfloat16), m.example_inputs()))
@@ -501,12 +535,12 @@ def test_quantized_tensor_subclass_int8(self):
 
         # reference
         from torchao.quantization.quant_api import change_linear_weights_to_int8_woqtensors
-        change_linear_weights_to_int8_woqtensors(m_copy)
+        _ref_change_linear_weights_to_int8_woqtensors(m_copy)
 
         res = m(*example_inputs)
         ref = m_copy(*example_inputs)
 
-        torch.testing.assert_close(res, ref, rtol=0.00001, atol=1e-2)
+        self.assertTrue(torch.equal(res, ref))
 
 
     @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
@@ -545,20 +579,20 @@ def test_quantized_tensor_subclass_int8_dyn_quant(self):
         # make sure it compiles
         torch._export.aot_compile(m_unwrapped, example_inputs)
 
-    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
-    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
-    @unittest.skip("This perf test is supposed to be run locally for sanity check performance when there is a change of int8 dynamic quant implementation")
-    def test_quantized_tensor_subclass_int8_dyn_quant_perf(self):
+
+    def _test_quantized_tensor_subclass_perf(self, api, ref_api, kwargs=None):
+        if kwargs is None:
+            kwargs = {}
+
         m = ToyLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")
         m_ref = copy.deepcopy(m)
         # setting batch_size to 20 to be compatible with the kernel
         example_inputs = m.example_inputs(batch_size=20, dtype=torch.bfloat16, device="cuda")
 
-        from torchao.quantization.quant_api import change_linear_weights_to_int8_dqtensors
-        change_linear_weights_to_int8_dqtensors(m)
+        api(m, **kwargs)
 
         # reference
-        _ref_change_linear_weights_to_int8_dqtensors(m_ref)
+        ref_api(m_ref, **kwargs)
 
         res = m(*example_inputs)
         ref = m_ref(*example_inputs)
@@ -583,7 +617,27 @@ def test_quantized_tensor_subclass_int8_dyn_quant_perf(self):
         print(f"elapsed time: {elapsed_time}, ref elapsed time: {ref_elapsed_time}")
         self.assertTrue(elapsed_time < 1.05 * ref_elapsed_time)
 
+    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    @unittest.skip("This perf test is supposed to be run locally for sanity check performance when there is a change of int8 dynamic quant implementation")
+    def test_quantized_tensor_subclass_int8_dyn_quant_perf(self):
+        from torchao.quantization.quant_api import change_linear_weights_to_int8_dqtensors
+        self._test_quantized_tensor_subclass_perf(change_linear_weights_to_int8_dqtensors, _ref_change_linear_weights_to_int8_dqtensors)
+
+    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    @unittest.skip("This perf test is supposed to be run locally for sanity check performance when there is a change of int8 weight only quant implementation")
+    def test_quantized_tensor_subclass_int8_wo_quant_perf(self):
+        from torchao.quantization.quant_api import change_linear_weights_to_int8_woqtensors
+        self._test_quantized_tensor_subclass_perf(change_linear_weights_to_int8_woqtensors, _ref_change_linear_weights_to_int8_woqtensors)
 
+    @unittest.skipIf(not TORCH_VERSION_AFTER_2_4, "Test only enabled for 2.4+")
+    @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
+    @unittest.skip("This perf test is supposed to be run locally for sanity check performance when there is a change of int4 weight only quant implementation")
+    def test_quantized_tensor_subclass_int4_wo_quant_perf(self):
+        kwargs = {"groupsize": 32}
+        from torchao.quantization.quant_api import change_linear_weights_to_int4_woqtensors
+        self._test_quantized_tensor_subclass_perf(change_linear_weights_to_int4_woqtensors, _ref_change_linear_weights_to_int4_woqtensors, kwargs)
 
 if __name__ == "__main__":
     unittest.main()

torchao/dtypes/aqt.py

@@ -574,29 +574,30 @@ def _quantized_linear_op(input_tensor, weight_qtensor, bias):
             scale_and_zero = weight_qtensor.layout_tensor.scale_and_zero
             return torch.ops.aten._weight_int4pack_mm(input_tensor.contiguous(), packed_weight, groupsize, scale_and_zero)
         elif (
-            is_cpu and
             weight_is_int8 and
             len(weight_qtensor.shape) == 2 and
             len(weight_qtensor.block_size) == 2 and
             weight_qtensor.block_size[0] == 1 and
             weight_qtensor.block_size[1] == weight_qtensor.shape[1] and
+            weight_qtensor.zero_point_domain == ZeroPointDomain.INT and
             weight_qtensor.layout == "plain"
         ):
             # TODO: enable cpu and mps efficient path
             # per channel int8 weight only quantizated mm
-            w_vals_int8_t = weight_qtensor.layout_tensor.int_data.t().contiguous()
+            w_vals_int8_t = weight_qtensor.layout_tensor.int_data.t()
+            scale = weight_qtensor.layout_tensor.scale
             orig_dtype = input_tensor.dtype
             y = (
                 torch.mm(
                     input_tensor.reshape(-1, input_tensor.shape[-1]),
                     w_vals_int8_t.to(input_tensor.dtype),
                 )
-                * weight_qtensor.scale
+                * scale
             )
             y = y.reshape(*input_tensor.shape[:-1], y.shape[-1])
             if bias is not None:
                 y += bias
-                return y.to(orig_dtype)
+            return y.to(orig_dtype)
 
             # is_cpu and is_mps only, some issue with is_contiguous() currently
             # return torch.ops.aten._weight_int8pack_mm(input_tensor.contiguous(), w_vals_int8_t, weight_qtensor.layout_tensor.scale)

torchao/quantization/quant_api.py

@@ -206,13 +206,18 @@ def change_linear_weights_to_int8_woqtensors(model, filter_fn=None, **kwargs):
     Converts all linear weight tensors to the
     `Int8WeightOnlyQuantizedLinearWeight` tensor subclass,
     effectively applying the same form of quantization
-    as apply_dynamic_quant while not modifying the linear modules.
+    as apply_weight_only_int8_quant while not modifying the linear modules.
     """
-    _replace_with_custom_fn_if_matches_filter(
-        model,
-        _get_subclass_inserter(Int8WeightOnlyQuantizedLinearWeight, enable_parametrization=TORCH_VERSION_AFTER_2_4, **kwargs),
-        _is_linear if filter_fn is None else filter_fn,
-    )
+
+    if TORCH_VERSION_AFTER_2_4:
+        quantize(model, get_apply_int8wo_quant(), filter_fn)
+        unwrap_tensor_subclass(model, filter_fn)
+    else:
+        _replace_with_custom_fn_if_matches_filter(
+            model,
+            _get_subclass_inserter(Int8WeightOnlyQuantizedLinearWeight, enable_parametrization=False, **kwargs),
+            _is_linear if filter_fn is None else filter_fn,
+        )
 
 
 def change_linear_weights_to_int4_woqtensors(model, **kwargs):

tutorials/quantize_vit/run_vit_b_quant.py

@@ -14,18 +14,27 @@
 input_tensor = torch.randn(1, 3, 224, 224, dtype=torch.bfloat16, device='cuda')
 
 ## Quantization code - start
+# int8 act, int8 weight dynamic quantization
 torchao.apply_dynamic_quant(model)
-from torch._inductor import config as inductorconfig
-inductorconfig.force_fuse_int_mm_with_mul = True
+
+# int8 weight only quantization
+# torchao.quantization.change_linear_weights_to_int8_woqtensors(model)
 ## Quantization code - end
 
+
+## compilation configs
+torch._dynamo.config.automatic_dynamic_shapes = False
+torch._inductor.config.force_fuse_int_mm_with_mul = True
+torch._inductor.config.use_mixed_mm = True
+## compilation configs end
+
 model = torch.compile(model, mode='max-autotune', fullgraph=True)
 
 # Must run with no_grad when optimizing for inference
 with torch.no_grad():
     # warmup
-    benchmark_model(model, 5, input_tensor)
+    benchmark_model(model, 20, input_tensor)
     # benchmark
-    print("elapsed_time: ", benchmark_model(model, 100, input_tensor), " milliseconds")
+    print("elapsed_time: ", benchmark_model(model, 1000, input_tensor), " milliseconds")
     # Create a trace
     profiler_runner("quant.json.gz", benchmark_model, model, 5, input_tensor)