bug fix for 3.x sq and static quant (#1823)

Signed-off-by: Cheng, Zixuan <[email protected]>

Author: violetch24

neural_compressor/torch/algorithms/smooth_quant/smooth_quant.py

@@ -389,8 +389,9 @@ def _ipex_post_quant_process(model, example_inputs, use_bf16, inplace=False):
                 else:
                     model = torch.jit.trace(model, example_inputs, strict=False)
                 model = torch.jit.freeze(model.eval())
-        # After freezing, run 1 time to warm up the profiling graph executor to insert prim::profile
-        # At the 2nd run, the llga pass will be triggered and the model is turned into
-        # an int8 model: prim::profile will be removed and will have LlgaFusionGroup in the graph
-        simple_inference(model, example_inputs, iterations=2)
-        return model
+
+    # After freezing, run 1 time to warm up the profiling graph executor to insert prim::profile
+    # At the 2nd run, the llga pass will be triggered and the model is turned into
+    # an int8 model: prim::profile will be removed and will have LlgaFusionGroup in the graph
+    simple_inference(model, example_inputs, iterations=2)
+    return model

neural_compressor/torch/algorithms/static_quant/static_quant.py

@@ -176,8 +176,9 @@ def _ipex_post_quant_process(model, example_inputs, use_bf16, inplace=False):
                 else:
                     model = torch.jit.trace(model, example_inputs, strict=False)
                 model = torch.jit.freeze(model.eval())
-        # After freezing, run 1 time to warm up the profiling graph executor to insert prim::profile
-        # At the 2nd run, the llga pass will be triggered and the model is turned into
-        # an int8 model: prim::profile will be removed and will have LlgaFusionGroup in the graph
-        simple_inference(model, example_inputs, iterations=2)
-        return model
+
+    # After freezing, run 1 time to warm up the profiling graph executor to insert prim::profile
+    # At the 2nd run, the llga pass will be triggered and the model is turned into
+    # an int8 model: prim::profile will be removed and will have LlgaFusionGroup in the graph
+    simple_inference(model, example_inputs, iterations=2)
+    return model

neural_compressor/torch/quantization/quantize.py

@@ -71,10 +71,10 @@ def quantize(
             from neural_compressor.torch.algorithms.smooth_quant import TorchSmoothQuant
 
             sq = TorchSmoothQuant(
-                model, dataloader=None, example_inputs=example_inputs, q_func=run_fn, record_max_info=True
+                q_model, dataloader=None, example_inputs=example_inputs, q_func=run_fn, record_max_info=True
             )
-            model.sq_info = sq
-            model = sq.transform(
+            q_model.sq_info = sq
+            q_model = sq.transform(
                 alpha=quant_config.alpha,
                 folding=quant_config.folding,
                 auto_alpha_args=quant_config.auto_alpha_args,

test/3x/torch/quantization/test_smooth_quant.py

@@ -193,6 +193,13 @@ def test_smooth_quant_mixed_precision(self):
         assert q_model is not None, "Quantization failed!"
 
         # quantize API
+        q_model = quantize(fp32_model, quant_config=quant_config, run_fn=run_fn, example_inputs=example_inputs)
+        assert q_model is not None, "Quantization failed!"
+
         quant_config.excluded_precisions = ["bf16"]
         q_model = quantize(fp32_model, quant_config=quant_config, run_fn=run_fn, example_inputs=example_inputs)
         assert q_model is not None, "Quantization failed!"
+
+        quant_config.folding = True
+        q_model = quantize(fp32_model, quant_config=quant_config, run_fn=run_fn, example_inputs=example_inputs)
+        assert q_model is not None, "Quantization failed!"

test/3x/torch/quantization/test_static_quant.py

@@ -195,9 +195,14 @@ def test_static_quant_with_quantize_API(self):
     @pytest.mark.skipif(not is_ipex_available(), reason="Requires IPEX")
     def test_static_quant_mixed_precision(self):
         fp32_model = copy.deepcopy(self.fp32_model)
+        example_inputs = self.input
         quant_config = get_default_static_config()
+        prepared_model = prepare(fp32_model, quant_config=quant_config, example_inputs=example_inputs)
+        run_fn(prepared_model)
+        q_model = convert(prepared_model)
+        assert q_model is not None, "Quantization failed!"
+
         quant_config.excluded_precisions = ["bf16"]
-        example_inputs = self.input
         prepared_model = prepare(fp32_model, quant_config=quant_config, example_inputs=example_inputs)
         run_fn(prepared_model)
         q_model = convert(prepared_model)