[PIR save/load]Fix pir_verison (#67563)

changeyoung98 · zhangbo9674 · web-flow · commit 32570b771ef6 · 2024-08-22T16:25:23.000+08:00
* fix

* fix pir_version

* fix

* fix pir_version

---------

Co-authored-by: zhangbo9674 &lt;zhangbo54@baidu.com&gt;
diff --git a/paddle/fluid/pir/serialize_deserialize/include/interface.h b/paddle/fluid/pir/serialize_deserialize/include/interface.h
@@ -63,7 +63,7 @@ void IR_API WriteModule(const pir::Program& program,
  */
 bool IR_API ReadModule(const std::string& file_path,
                        pir::Program* program,
-                       uint64_t pir_version);
+                       int64_t pir_version = -1);
 
 /**
  * @brief Save the given tensor into a single file at the specified file path
diff --git a/paddle/fluid/pir/serialize_deserialize/src/interface.cc b/paddle/fluid/pir/serialize_deserialize/src/interface.cc
@@ -69,17 +69,21 @@ void WriteModule(const pir::Program& program,
 
 bool ReadModule(const std::string& file_path,
                 pir::Program* program,
-                uint64_t pir_version) {
+                int64_t pir_version) {
   std::ifstream f(file_path);
   Json data = Json::parse(f);
-  pir_version = GetPirVersion();
+  if (pir_version < 0) {
+    pir_version = GetPirVersion();
+    VLOG(6) << "pir_version is null, get pir_version: " << pir_version;
+  }
+
   PatchBuilder builder(pir_version);
 
   if (data.contains(BASE_CODE) && data[BASE_CODE].contains(MAGIC) &&
       data[BASE_CODE][MAGIC] == PIR) {
     uint64_t file_version =
         data.at(BASE_CODE).at(PIRVERSION).template get<uint64_t>();
-    if (file_version != pir_version) {
+    if (file_version != (uint64_t)pir_version) {
       builder.SetFileVersion(file_version);
       const char* paddle_root = PADDLE_ROOT;
       VLOG(8) << "Paddle path: " << paddle_root;
diff --git a/paddle/fluid/pybind/io.cc b/paddle/fluid/pybind/io.cc
@@ -169,7 +169,11 @@ void BindIO(pybind11::module *m) {
          py::arg("overwrite") = true,
          py::arg("readable") = false,
          py::arg("trainable") = true);
-  m->def("deserialize_pir_program", &pir::ReadModule);
+  m->def("deserialize_pir_program",
+         &pir::ReadModule,
+         py::arg("file_path"),
+         py::arg("program"),
+         py::arg("pir_version") = -1);
 }
 }  // namespace pybind
 }  // namespace paddle
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_llama_save_load.py b/test/auto_parallel/hybrid_strategy/semi_auto_llama_save_load.py
@@ -221,19 +221,21 @@ def run_dy2static(self, tmp_ckpt_path):
                 break
 
         # check pir dist_model save&load
-        # paddle.enable_static()
-        # model_file_path = os.path.join(
-        #    tmp_ckpt_path,
-        #    "rank_" + str(paddle.distributed.get_rank()) + ".pd_dist_model",
-        # )
-        # paddle.save(
-        #    dist_model._engine._pir_dist_main_progs["train"], model_file_path
-        # )
-        # loaded_model = paddle.load(model_file_path)
-        # paddle.disable_static()
-        # self.check_program_equal(
-        #    dist_model._engine._pir_dist_main_progs["train"], loaded_model
-        # )
+        paddle.enable_static()
+        model_file_path = os.path.join(
+            tmp_ckpt_path,
+            "rank_" + str(paddle.distributed.get_rank()) + ".pd_dist_model",
+        )
+        paddle.save(
+            dist_model._engine._pir_dist_main_progs["train"], model_file_path
+        )
+        loaded_model = paddle.load(model_file_path)
+        self.check_program_equal(
+            dist_model._engine._pir_dist_main_progs["train"], loaded_model
+        )
+        paddle.disable_static()
+        paddle.distributed.barrier()
+
         time.sleep(10)
 
         loss_after_load = []
diff --git a/tools/auto_parallel/target_path_lists.sh b/tools/auto_parallel/target_path_lists.sh
@@ -29,6 +29,7 @@ target_lists_for_semi_auto_ci=(
     "paddle/fluid/ir_adaptor/"
     "paddle/fluid/pir/dialect"
     "paddle/fluid/pir/transforms"
+    "paddle/fluid/pir/serialize_deserialize"
 )
 
 target_lists_for_dygraph_ci=(

Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@ target_lists_for_semi_auto_ci=(`
`29`	`29`	`"paddle/fluid/ir_adaptor/"`
`30`	`30`	`"paddle/fluid/pir/dialect"`
`31`	`31`	`"paddle/fluid/pir/transforms"`
	`32`	`+ "paddle/fluid/pir/serialize_deserialize"`
`32`	`33`	`)`
`33`	`34`
`34`	`35`	`target_lists_for_dygraph_ci=(`