Fix precision tuning bug for ONNX CUDA EP (#1133)

Signed-off-by: yuwenzho <[email protected]>
intel · Aug 3, 2023 · d1f315f · d1f315f
1 parent 7fbcf54
commit d1f315f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/neural_compressor/strategy/utils/constant.py b/neural_compressor/strategy/utils/constant.py
@@ -25,7 +25,7 @@
                     ('weight','scheme'), ('weight','algorithm'), ('weight','granularity'),
                     ('weight','bits'), ('weight','group_size'), 'sampling_size']
 
-PRECISION_SET_V2_0 = {'fp32', 'bf16'}
+PRECISION_SET_V2_0 = {'fp32', 'bf16', 'fp16'}
 
 auto_query_order = ['static', 'dynamic', 'bf16', 'fp16', 'fp32']
 static_query_order = ['static', 'bf16', 'fp16', 'fp32']

diff --git a/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py b/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py
@@ -1231,6 +1231,12 @@ def eval(model):
             calib_dataloader=self.matmul_dataloader, eval_func=eval)
         self.assertTrue('QLinearMatMul' not in [i.op_type for i in q_model.nodes()])
 
+        config = PostTrainingQuantConfig(approach='static', backend='onnxrt_cuda_ep', device='gpu', quant_level=1)
+        q_model = quantization.fit(self.distilbert_model, config, 
+            calib_dataloader=DummyNLPDataloader_dict("distilbert-base-uncased-finetuned-sst-2-english"),
+            eval_func=eval)
+        self.assertTrue('QLinearMatMul' in [i.op_type for i in q_model.nodes()])
+
         config = PostTrainingQuantConfig(approach='static', recipes={'optypes_to_exclude_output_quant': ['MatMul']})
         q_model = quantization.fit(self.matmul_model, config,
             calib_dataloader=self.matmul_dataloader, eval_func=eval)