modify op_type for set_local in 3.x API

examples/3.x_api/pytorch/nlp/huggingface_models/language-modeling/quantization/llm/run_clm_no_trainer.py

@@ -341,13 +341,13 @@ def run_fn_for_gptq(model, dataloader_for_calibration, *args):
                 quant_config = SmoothQuantConfig(alpha=args.alpha, folding=True)
 
             if re.search("gpt", user_model.config.model_type):
-                quant_config.set_local("add", SmoothQuantConfig(w_dtype="fp32", act_dtype="fp32"))
+                quant_config.set_local(torch.add, SmoothQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         else:
             from neural_compressor.torch.quantization import quantize, get_default_static_config, StaticQuantConfig
 
             quant_config =  get_default_static_config()
             if re.search("gpt", user_model.config.model_type):
-                quant_config.set_local("add", StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
+                quant_config.set_local(torch.add, StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
 
         from neural_compressor.torch.algorithms.smooth_quant import move_input_to_device
         from tqdm import tqdm

neural_compressor/torch/algorithms/static_quant/utility.py

@@ -291,15 +291,12 @@ def get_quantizable_ops_recursively(model, example_inputs):  # pragma: no cover
                     map_op_name_to_fqn[(tuple(name), ipex_op_type)] = module_fqn
                     if "class" in ipex_op_type:  # "<class 'torch.nn.modules.activation.ReLU'>"
                         op_type = ipex_op_type.split("'")[1]
-                        op_name_info.append((module_fqn, eval(op_type)))
+                        op_name_info.append((module_fqn, eval(op_type).__name__))
                     elif "method" in ipex_op_type:  # "<method 'add' of 'torch._C._TensorBase' objects>"
                         method = ipex_op_type.split("'")[1]
-                        op_type = getattr(
-                            torch._C._TensorBase if ipex_ver.release < Version("2.2") else torch._C.TensorBase, method
-                        )
-                        op_name_info.append((module_fqn, op_type))
-                    else:
-                        op_name_info.append((module_fqn, op_type))
+                        op_name_info.append((module_fqn, method))
+                    elif "Convolution" in ipex_op_type:  # "Convolution_Relu"
+                        op_name_info.append((module_fqn, "Conv2d"))
                 else:
                     re_flag = False
                     for pattern, unify_op_type in unify_op_type_mapping_ipex["re"].items():

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

@@ -55,6 +55,21 @@ def test_smooth_quant_auto(self):
         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!"
 
+    @pytest.mark.skipif(not is_ipex_available(), reason="Requires IPEX")
+    def test_smooth_quant_fallback(self):
+        fp32_model = copy.deepcopy(model)
+        quant_config = get_default_sq_config()
+        example_inputs = torch.randn([1, 3])
+        # fallback by op_type
+        quant_config.set_local(torch.nn.Linear, SmoothQuantConfig(w_dtype="fp32", act_dtype="fp32"))
+        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!"
+
+        for op, op_info in q_model.tune_cfg[" "]["q_op_infos"].items():
+            if op_info["op_type"] == "<class 'torch.nn.modules.linear.Linear'>":
+                dtype = q_model.tune_cfg[" "]["q_op_infos"][op]["input_tensor_infos"][0]["force_dtype"]
+                assert dtype == "torch.float32", "Failed to fallback linear op, please check!"
+
     @pytest.mark.skipif(not is_ipex_available(), reason="Requires IPEX")
     @pytest.mark.parametrize(
         "act_sym, act_algo, alpha, folding, scale_sharing",

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

@@ -61,15 +61,26 @@ def test_static_quant_fallback(self):
         quant_config = get_default_static_config()
         example_inputs = self.input
         # fallback by op_type
-        quant_config.set_local(torch.nn.modules.linear.Linear, StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
+        quant_config.set_local(torch.nn.Linear, StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         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!"
 
+        for op, op_info in q_model.tune_cfg[" "]["q_op_infos"].items():
+            if op_info["op_type"] == "<class 'torch.nn.modules.linear.Linear'>":
+                dtype = q_model.tune_cfg[" "]["q_op_infos"][op]["input_tensor_infos"][0]["force_dtype"]
+                assert dtype == "torch.float32", "Failed to fallback linear op, please check!"
+
         # fallback by op_name
+        quant_config = get_default_static_config()
         quant_config.set_local("fc1", StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         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!"
 
+        for op, op_info in q_model.tune_cfg[" "]["q_op_infos"].items():
+            if op_info["fqn"] == "fc1":
+                dtype = q_model.tune_cfg[" "]["q_op_infos"][op]["input_tensor_infos"][0]["force_dtype"]
+                assert dtype == "torch.float32", "Failed to fallback fc1 layer, please check!"
+
     @pytest.mark.skipif(not is_ipex_available(), reason="Requires IPEX")
     @pytest.mark.parametrize(
         "act_sym, act_algo",
@@ -184,19 +195,29 @@ def test_static_quant_fallback(self):
         quant_config = get_default_static_config()
         example_inputs = self.input
         # fallback by op_type
-        quant_config.set_local(torch.nn.modules.linear.Linear, StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
+        quant_config.set_local(torch.nn.Linear, StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         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!"
 
+        for op, op_info in q_model.tune_cfg[" "]["q_op_infos"].items():
+            if op_info["op_type"] == "<class 'torch.nn.modules.linear.Linear'>":
+                dtype = q_model.tune_cfg[" "]["q_op_infos"][op]["input_tensor_infos"][0]["force_dtype"]
+                assert dtype == "torch.float32", "Failed to fallback linear op, please check!"
+
         # fallback by op_name
         quant_config.set_local("fc1", StaticQuantConfig(w_dtype="fp32", act_dtype="fp32"))
         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!"
 
+        for op, op_info in q_model.tune_cfg[" "]["q_op_infos"].items():
+            if op_info["fqn"] == "fc1":
+                dtype = q_model.tune_cfg[" "]["q_op_infos"][op]["input_tensor_infos"][0]["force_dtype"]
+                assert dtype == "torch.float32", "Failed to fallback fc1 layer, please check!"
+
     @pytest.mark.skipif(not is_ipex_available(), reason="Requires IPEX")
     @pytest.mark.parametrize(
         "act_sym, act_algo",