pass_ci_fix_test_ptq (#68049)

xiaoluomi · web-flow · commit 94589c3ebafc · 2024-09-06T18:20:06.000+08:00
diff --git a/paddle/fluid/pir/serialize_deserialize/include/schema.h b/paddle/fluid/pir/serialize_deserialize/include/schema.h
@@ -61,6 +61,7 @@ namespace pir {
 #define ATTRS "A"
 #define OPRESULTS_ATTRS "OA"
 #define DIST_ATTRS "DA"
+#define QUANT_ATTRS "QA"
 
 // value's key:
 //  value's type which should be pir::Type's json object(ID or ID and DATA).
@@ -88,6 +89,7 @@ std::pair<std::string, std::string> GetContentSplitByDot(
     const std::string& str);
 
 std::vector<std::string> GetOpDistAttr();
+std::vector<std::string> GetOpQuantAttr();
 
 void GetCompressOpName(std::string* op_name);
 
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc
@@ -159,6 +159,17 @@ pir::Operation* ProgramReader::ReadParameterOp(Json* op_json) {
       }
     }
   }
+  if (op_json->contains(QUANT_ATTRS)) {
+    Json& quant_attrs_json = op_json->at(QUANT_ATTRS);
+    for (auto& attr_json : quant_attrs_json) {
+      auto attr_name = attr_json.at(NAME).template get<std::string>();
+      if (attr_json.contains(ATTR_TYPE)) {
+        attributes.insert({attr_name, ReadAttribute(&attr_json)});
+      } else {
+        VLOG(6) << "Attribute " << attr_name << " Deleted.";
+      }
+    }
+  }
 
   if (op_json->contains(OPRESULTS_ATTRS)) {
     Json& other_attrs_json = op_json->at(OPRESULTS_ATTRS);
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_serialize.cc
@@ -187,8 +187,12 @@ Json ProgramWriter::WriteParameterOP(const pir::Operation& op) {
                                             "trainable",
                                             "op_callstack" /*no need*/};
   std::vector<std::string> DistAttrsNameList = GetOpDistAttr();
+  std::vector<std::string> QuantAttrsNameList = GetOpQuantAttr();
   AttrsNameList.insert(
       AttrsNameList.end(), DistAttrsNameList.begin(), DistAttrsNameList.end());
+  AttrsNameList.insert(AttrsNameList.end(),
+                       QuantAttrsNameList.begin(),
+                       QuantAttrsNameList.end());
   for (auto attr : op.attributes()) {
     auto attr_name = attr.first;
     auto it = std::find(AttrsNameList.begin(), AttrsNameList.end(), attr_name);
@@ -241,6 +245,15 @@ Json ProgramWriter::WriteParameterOP(const pir::Operation& op) {
   }
   op_json[DIST_ATTRS] = dist_attrs_json;
 
+  Json quant_attrs_json = Json::array();
+  for (auto key : GetOpQuantAttr()) {
+    if (op.attributes().count(key) > 0) {
+      quant_attrs_json.emplace_back(
+          WriteAttribute(key, op.attributes().at(key)));
+    }
+  }
+  op_json[QUANT_ATTRS] = quant_attrs_json;
+
   Json other_attrs_json = Json::array();
   OPTIONAL_CHECK(other_attrs_json, "persistable", 1)
   OPTIONAL_CHECK(other_attrs_json, "stop_gradient", 1)
@@ -333,6 +346,11 @@ Json ProgramWriter::WriteAttributesMapOpinfo(pir::Operation* op,
         attrs_json.emplace_back(WriteAttribute(key, attr_map.at(key)));
       }
     }
+    for (auto key : GetOpQuantAttr()) {
+      if (attr_map.count(key) > 0) {
+        attrs_json.emplace_back(WriteAttribute(key, attr_map.at(key)));
+      }
+    }
   } else {
     for (auto& attr : attr_map) {
       if (attr.first != "stop_gradient" && attr.first != "persistable" &&
diff --git a/paddle/fluid/pir/serialize_deserialize/src/schema.cc b/paddle/fluid/pir/serialize_deserialize/src/schema.cc
@@ -28,7 +28,7 @@ std::pair<std::string, std::string> GetContentSplitByDot(
 }
 
 std::vector<std::string> GetOpDistAttr() { return {"op_dist_attr", "op_role"}; }
-
+std::vector<std::string> GetOpQuantAttr() { return {"struct_name"}; }
 void GetCompressOpName(std::string* op_name) {
   std::pair<std::string, std::string> name = GetContentSplitByDot(*op_name);
   *op_name = pir::DialectIdMap::Instance()->GetCompressDialectId(name.first) +
diff --git a/python/paddle/nn/quant/format.py b/python/paddle/nn/quant/format.py
@@ -17,8 +17,9 @@
 import abc
 
 import paddle
-from paddle import _legacy_C_ops as _C_ops
-from paddle.framework import in_dynamic_mode
+from paddle import _C_ops, _legacy_C_ops
+from paddle.base import unique_name
+from paddle.framework import in_dynamic_mode, in_pir_mode
 
 from ..layer.layers import Layer
 
@@ -104,6 +105,8 @@ def __init__(
             shape=scales.shape, attr=scale_attr, dtype="float32"
         )
         self._scales.set_value(scales)
+        self.in_accum = paddle.to_tensor(0.0, dtype="float32")
+        self.in_state = paddle.to_tensor(0.0, dtype="float32")
         zero_point = zero_point if zero_point is not None else zero_point
         zero_point = paddle.to_tensor(zero_point, dtype="float32")
         zp_attr = paddle.framework.ParamAttr(
@@ -177,7 +180,7 @@ def forward(self, input):
                     input, self._scales, self._quant_axis, type='e5m2'
                 )
 
-            return _C_ops.quantize_linear(
+            return _legacy_C_ops.quantize_linear(
                 input.cast('float32'),
                 self._scales,
                 self._zero_point,
@@ -190,6 +193,31 @@ def forward(self, input):
                 "qmax",
                 self._qmax,
             ).cast(input.dtype)
+        if in_pir_mode():
+            input.stop_gradient = True
+            quant_out = paddle.pir.core.create_persistable_value(
+                dtype='float32',
+                shape=input.shape,
+                name=unique_name.generate("quant_out"),
+                initializer=paddle.nn.initializer.Constant(0.0),
+                stop_gradient=True,
+            )
+            # TODO(xiaoluomi): need to add only observer pass for quantize_linear
+            quant_out, out_state, out_accum, out_scale = _C_ops.quantize_linear(
+                input,
+                self._scales,
+                self._zero_point,
+                self.in_accum,
+                self.in_state,
+                self._quant_axis,
+                self._bit_length,
+                self._qmin,
+                self._qmax,
+                0,
+                True,
+                False,
+            )
+            return quant_out
         else:
             out = self._helper.create_variable_for_type_inference(input.dtype)
             self._helper.append_op(
@@ -239,6 +267,8 @@ def __init__(
             shape=scales.shape, attr=scale_attr, dtype="float32"
         )
         self._scales.set_value(scales)
+        self.in_accum = paddle.to_tensor(0.0, dtype="float32")
+        self.in_state = paddle.to_tensor(0.0, dtype="float32")
         zero_point = zero_point if zero_point is not None else zero_point
         zero_point = paddle.to_tensor(zero_point, dtype="float32")
         zp_attr = paddle.framework.ParamAttr(
@@ -306,7 +336,7 @@ def forward(self, input):
                     input, self._scales, self._quant_axis, type='e5m2'
                 )
 
-            return _C_ops.dequantize_linear(
+            return _legacy_C_ops.dequantize_linear(
                 input.cast('float32'),
                 self._scales,
                 self._zero_point,
@@ -319,6 +349,33 @@ def forward(self, input):
                 "qmax",
                 self._qmax,
             ).cast(input.dtype)
+        if in_pir_mode():
+            input.stop_gradient = True
+            dequant_out = paddle.pir.core.create_persistable_value(
+                dtype='float32',
+                shape=input.shape,
+                name=unique_name.generate("quant_out"),
+                initializer=paddle.nn.initializer.Constant(0.0),
+                stop_gradient=True,
+            )
+            # TODO(xiaoluomi): need to add only observer pass for dequantize_linear
+            dequant_out, out_state, out_accum, out_scale = (
+                _C_ops.dequantize_linear(
+                    input,
+                    self._scales,
+                    self._zero_point,
+                    self.in_accum,
+                    self.in_state,
+                    self._quant_axis,
+                    self._bit_length,
+                    self._qmin,
+                    self._qmax,
+                    0,
+                    True,
+                    False,
+                )
+            )
+            return dequant_out
         else:
             out = self._helper.create_variable_for_type_inference(input.dtype)
             self._helper.append_op(
diff --git a/test/quantization/CMakeLists.txt b/test/quantization/CMakeLists.txt
@@ -218,6 +218,7 @@ endfunction()
 if(WIN32)
   list(REMOVE_ITEM TEST_OPS test_light_nas)
   list(REMOVE_ITEM TEST_OPS test_post_training_quantization_mnist)
+  list(REMOVE_ITEM TEST_OPS test_ptq)
   list(REMOVE_ITEM TEST_OPS test_post_training_quantization_mobilenetv1)
   list(REMOVE_ITEM TEST_OPS test_post_training_quantization_resnet50)
   list(REMOVE_ITEM TEST_OPS test_post_training_quantization_program_resnet50)
@@ -487,6 +488,7 @@ if(NOT WIN32)
   set_tests_properties(test_post_training_quantization_mnist PROPERTIES TIMEOUT
                                                                         150)
   set_tests_properties(test_imperative_ptq PROPERTIES TIMEOUT 120)
+  set_tests_properties(test_ptq PROPERTIES TIMEOUT 200)
   set_tests_properties(test_quant_aware_config PROPERTIES TIMEOUT 200)
 endif()
 
@@ -522,4 +524,5 @@ if(APPLE)
                                                                         300)
   set_tests_properties(test_imperative_ptq PROPERTIES TIMEOUT 300)
   set_tests_properties(test_imperative_skip_op PROPERTIES TIMEOUT 300)
+  set_tests_properties(test_ptq PROPERTIES TIMEOUT 300)
 endif()
diff --git a/test/quantization/test_ptq.py b/test/quantization/test_ptq.py

Original file line number	Diff line number	Diff line change
`@@ -159,6 +159,17 @@ pir::Operation* ProgramReader::ReadParameterOp(Json* op_json) {`
`159`	`159`	`}`
`160`	`160`	`}`
`161`	`161`	`}`
	`162`	`+ if (op_json->contains(QUANT_ATTRS)) {`
	`163`	`+ Json& quant_attrs_json = op_json->at(QUANT_ATTRS);`
	`164`	`+ for (auto& attr_json : quant_attrs_json) {`
	`165`	`+ auto attr_name = attr_json.at(NAME).template get<std::string>();`
	`166`	`+ if (attr_json.contains(ATTR_TYPE)) {`
	`167`	`+ attributes.insert({attr_name, ReadAttribute(&attr_json)});`
	`168`	`+ } else {`
	`169`	`+ VLOG(6) << "Attribute " << attr_name << " Deleted.";`
	`170`	`+ }`
	`171`	`+ }`
	`172`	`+ }`
`162`	`173`
`163`	`174`	`if (op_json->contains(OPRESULTS_ATTRS)) {`
`164`	`175`	`Json& other_attrs_json = op_json->at(OPRESULTS_ATTRS);`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ std::pair<std::string, std::string> GetContentSplitByDot(`
`28`	`28`	`}`
`29`	`29`
`30`	`30`	`std::vector<std::string> GetOpDistAttr() { return {"op_dist_attr", "op_role"}; }`
`31`		`-`
	`31`	`+std::vector<std::string> GetOpQuantAttr() { return {"struct_name"}; }`
`32`	`32`	`void GetCompressOpName(std::string* op_name) {`
`33`	`33`	`std::pair<std::string, std::string> name = GetContentSplitByDot(*op_name);`
`34`	`34`	`*op_name = pir::DialectIdMap::Instance()->GetCompressDialectId(name.first) +`