fix delete scale and zero_point var bug (#62225)

ming1753 · web-flow · commit 85ba93655e6e · 2024-03-01T14:33:27.000+08:00
* fix delete scale and zero_point var bug
diff --git a/paddle/fluid/framework/ir/delete_quant_dequant_linear_op_pass.cc b/paddle/fluid/framework/ir/delete_quant_dequant_linear_op_pass.cc
@@ -124,14 +124,18 @@ void DeleteQuantDequantLinearOpPass::ApplyImpl(ir::Graph* graph) const {
       return;
     }
     */
-    std::unordered_set<const Node*> nodes2rm = {};
-
-    // delete Scale and ZeroPoint tensor in scope
+    // Scale and ZeroPoint tensor should be removed in save_optimized_model_pass
     std::vector<std::string> vars2rm = {};
     vars2rm.emplace_back(quantize_linear_op->Op()->Input("Scale")[0]);
     vars2rm.emplace_back(quantize_linear_op->Op()->Input("ZeroPoint")[0]);
     vars2rm.emplace_back(dequantize_linear_op->Op()->Input("Scale")[0]);
     vars2rm.emplace_back(dequantize_linear_op->Op()->Input("ZeroPoint")[0]);
+    auto& scale_and_zero_point_param = g->GetOrInit<std::vector<std::string>>(
+        framework::ir::kScaleAndZeroPointParamAttr);
+    scale_and_zero_point_param.insert(
+        scale_and_zero_point_param.end(), vars2rm.begin(), vars2rm.end());
+
+    std::unordered_set<const Node*> nodes2rm = {};
 
     // Get input scale from tensor
     const phi::DenseTensor& input_scale_tensor =
@@ -182,13 +186,6 @@ void DeleteQuantDequantLinearOpPass::ApplyImpl(ir::Graph* graph) const {
     nodes2rm.insert(dequantize_linear_op);
     nodes2rm.insert(dequantize_linear_op_out);
     GraphSafeRemoveNodes(graph, nodes2rm);
-
-    for (auto& var_name : vars2rm) {
-      if (scope->FindVar(var_name)) {
-        scope->EraseVars({var_name});
-      }
-    }
-
     found_count++;
   };
   gpd(graph, handler);
diff --git a/paddle/fluid/framework/ir/fuse_pass_base.h b/paddle/fluid/framework/ir/fuse_pass_base.h
@@ -40,6 +40,11 @@ static const char kFuseStatisAttr[] = "__fuse_statis__";
 // allocation.
 static const char kRepetitiveParamAttr[] = "__repetitive_param__";
 
+// scale and zero point of the quantized/dequantized op should be removed in
+// save_optimized_model_pass.
+static const char kScaleAndZeroPointParamAttr[] =
+    "__scale_and_zero_point_param__";
+
 enum FuseOptions {
   DO_NOT_FUSE,  // fusing will not be done
   FUSE_NATIVE,  // fusing will be done without MKL-DNN
diff --git a/paddle/fluid/framework/ir/trt_delete_weight_dequant_linear_op_pass.cc b/paddle/fluid/framework/ir/trt_delete_weight_dequant_linear_op_pass.cc
@@ -231,13 +231,17 @@ void TrtDeleteWeightQuantDequantLinearOpPass::ApplyImpl(
       return;
     }
     */
-    std::unordered_set<const Node*> nodes2rm = {};
-
-    // delete Scale and ZeroPoint tensor in scope
+    // Scale and ZeroPoint tensor should be removed in save_optimized_model_pass
     std::vector<std::string> vars2rm = {};
     vars2rm.emplace_back(weight_dequantize_linear_op->Op()->Input("Scale")[0]);
     vars2rm.emplace_back(
         weight_dequantize_linear_op->Op()->Input("ZeroPoint")[0]);
+    auto& scale_and_zero_point_param = g->GetOrInit<std::vector<std::string>>(
+        framework::ir::kScaleAndZeroPointParamAttr);
+    scale_and_zero_point_param.insert(
+        scale_and_zero_point_param.end(), vars2rm.begin(), vars2rm.end());
+
+    std::unordered_set<const Node*> nodes2rm = {};
 
     int bit_length = PADDLE_GET_CONST(
         int, weight_dequantize_linear_op->Op()->GetAttr("bit_length"));
@@ -363,13 +367,6 @@ void TrtDeleteWeightQuantDequantLinearOpPass::ApplyImpl(
     }
 
     GraphSafeRemoveNodes(graph, nodes2rm);
-
-    for (auto& var_name : vars2rm) {
-      if (scope->FindVar(var_name)) {
-        scope->EraseVars({var_name});
-      }
-    }
-
     found_count++;
   };
   gpd(graph, handler);
diff --git a/paddle/fluid/inference/analysis/passes/save_optimized_model_pass.cc b/paddle/fluid/inference/analysis/passes/save_optimized_model_pass.cc
@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include <unordered_set>
 #include "paddle/fluid/framework/executor.h"
+#include "paddle/fluid/framework/ir/fuse_pass_base.h"
 #include "paddle/fluid/framework/ir/graph_helper.h"
 #include "paddle/fluid/framework/scope.h"
 
@@ -37,10 +38,17 @@ void SaveOptimizedModelPass::SaveOptimizedModel(Argument* argument) {
 
   framework::ir::GraphToProgram(*graph, &optimized_program_desc);
 
-  // Some vars may be deleted by pass, so we need to remove them in block
+  // Remove the scale and zero point parameters from optimized program.
+  auto scale_and_zero_point_param = graph->GetOrInit<std::vector<std::string>>(
+      framework::ir::kScaleAndZeroPointParamAttr);
   framework::BlockDesc* block = optimized_program_desc.MutableBlock(0);
   for (auto& var_desc : block->AllVars()) {
-    if (var_desc->Persistable() && !scope.FindVar(var_desc->Name())) {
+    auto var_name = var_desc->Name();
+    if (var_desc->Persistable() && scope.FindVar(var_name) &&
+        std::count(scale_and_zero_point_param.begin(),
+                   scale_and_zero_point_param.end(),
+                   var_name) > 0) {
+      scope.EraseVars({var_name});
       block->RemoveVar(var_desc->Name());
     }
   }
