【pir save load】 modify multi jit program parmeter name. test_jit_layer.py (#64281)

xiaoguoguo626807 · web-flow · commit e8bf283d99ab · 2024-05-15T15:32:43.000+08:00
* modify name

* modify mkloutput

* modify ci
diff --git a/paddle/fluid/pir/serialize_deserialize/include/interface.h b/paddle/fluid/pir/serialize_deserialize/include/interface.h
@@ -128,7 +128,8 @@ void LoadFunction(const std::string& file_path,
                   int64_t seek,
                   const std::vector<int64_t>& shape,
                   bool load_as_fp16,
-                  phi::DenseTensor* out);
+                  phi::DenseTensor* out,
+                  phi::Place place = phi::Place());
 
 /**
  * @brief Save the given tensor into a single file at the specified file path
@@ -146,5 +147,6 @@ void LoadFunction(const std::string& file_path,
 void LoadCombineFunction(const std::string& file_path,
                          const std::vector<std::string>& names,
                          std::vector<phi::DenseTensor*>* out,
-                         bool load_as_fp16);
+                         bool load_as_fp16,
+                         phi::Place place = phi::Place());
 }  // namespace pir
diff --git a/paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc b/paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc
@@ -21,20 +21,29 @@ limitations under the License. */
 
 namespace pir {
 
-const phi::DeviceContext* GetDeviceContext(const phi::DenseTensor& x) {
+const phi::DeviceContext* GetDeviceContext(
+    const phi::DenseTensor& x, const phi::Place& place = phi::Place()) {
   phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
   const phi::DeviceContext* dev_ctx = nullptr;
-  auto place = x.place();
-  if (place.GetType() == phi::AllocationType::UNDEFINED) {
+  auto x_place = x.place();
+  if (x_place.GetType() != phi::AllocationType::UNDEFINED) {
+    dev_ctx = pool.Get(x_place);
+    return dev_ctx;
+  } else if (place.GetType() != phi::AllocationType::UNDEFINED) {
+    dev_ctx = pool.Get(place);
+    return dev_ctx;
+  } else {
+    phi::Place compile_place;
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-    place = phi::GPUPlace();
+    compile_place = phi::GPUPlace();
 #elif defined(PADDLE_WITH_XPU)
-    place = phi::XPUPlace();
+    compile_place = phi::XPUPlace();
 #else
-    place = phi::CPUPlace();
+    compile_place = phi::CPUPlace();
 #endif
+    dev_ctx = pool.Get(compile_place);
+    return dev_ctx;
   }
-  dev_ctx = pool.Get(place);
   return dev_ctx;
 }
 
@@ -145,7 +154,8 @@ void LoadFunction(const std::string& file_path,
                   int64_t seek,
                   const std::vector<int64_t>& shape,
                   bool load_as_fp16,
-                  phi::DenseTensor* out) {
+                  phi::DenseTensor* out,
+                  phi::Place place) {
   std::ifstream fin(file_path, std::ios::binary);
   PADDLE_ENFORCE_EQ(static_cast<bool>(fin),
                     true,
@@ -156,7 +166,7 @@ void LoadFunction(const std::string& file_path,
   PADDLE_ENFORCE_NOT_NULL(out,
                           phi::errors::InvalidArgument(
                               "The variable to be loaded cannot be found."));
-  const phi::DeviceContext* dev_ctx = GetDeviceContext(*out);
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(*out, place);
 
   if (seek != -1) {
     PADDLE_ENFORCE_GE(seek,
@@ -179,7 +189,8 @@ void LoadFunction(const std::string& file_path,
 void LoadCombineFunction(const std::string& file_path,
                          const std::vector<std::string>& names,
                          std::vector<phi::DenseTensor*>* out,
-                         bool load_as_fp16) {
+                         bool load_as_fp16,
+                         phi::Place place) {
   std::ifstream fin(file_path, std::ios::binary);
   PADDLE_ENFORCE_EQ(static_cast<bool>(fin),
                     true,
@@ -194,7 +205,7 @@ void LoadCombineFunction(const std::string& file_path,
                         "The number of variables to be saved is %d, expect "
                         "it to be greater than 0.",
                         out->size()));
-  const phi::DeviceContext* dev_ctx = GetDeviceContext(*(out->at(0)));
+  const phi::DeviceContext* dev_ctx = GetDeviceContext(*(out->at(0)), place);
   for (size_t i = 0; i < names.size(); i++) {
     auto tensor = out->at(i);
     paddle::framework::DeserializeFromStream(fin, tensor, *dev_ctx);
diff --git a/paddle/fluid/pybind/io.cc b/paddle/fluid/pybind/io.cc
@@ -19,13 +19,31 @@ limitations under the License. */
 #include "paddle/fluid/framework/selected_rows_utils.h"
 #include "paddle/fluid/pir/serialize_deserialize/include/interface.h"
 #include "paddle/fluid/platform/enforce.h"
+#include "paddle/fluid/pybind/eager_utils.h"
 #include "paddle/fluid/pybind/pybind_variant_caster.h"
 #include "paddle/utils/pybind.h"
 
 namespace py = pybind11;
 namespace paddle {
 namespace pybind {
+template <typename PlaceType>
+void LoadCombine(const std::string &file_path,
+                 const std::vector<std::string> &names,
+                 std::vector<phi::DenseTensor *> *out,
+                 bool load_as_fp16,
+                 const PlaceType place) {
+  pir::LoadCombineFunction(file_path, names, out, load_as_fp16, place);
+}
 
+template <typename PlaceType>
+void Load(const std::string &file_path,
+          int64_t seek,
+          const std::vector<int64_t> &shape,
+          bool load_as_fp16,
+          phi::DenseTensor *out,
+          const PlaceType place) {
+  pir::LoadFunction(file_path, seek, shape, load_as_fp16, out, place);
+}
 void BindIO(pybind11::module *m) {
   m->def("save_lod_tensor",
          [](const phi::DenseTensor &tensor, const std::string &str_file_name) {
@@ -128,9 +146,20 @@ void BindIO(pybind11::module *m) {
 
   m->def("save_combine_func", &pir::SaveCombineFunction);
 
-  m->def("load_func", &pir::LoadFunction);
-
-  m->def("load_combine_func", &pir::LoadCombineFunction);
+  m->def("load_func", &Load<paddle::platform::CPUPlace>);
+  m->def("load_func", &Load<paddle::platform::CustomPlace>);
+  m->def("load_func", &Load<paddle::platform::XPUPlace>);
+  m->def("load_func", &Load<paddle::platform::CUDAPinnedPlace>);
+  m->def("load_func", &Load<paddle::platform::CUDAPlace>);
+  m->def("load_func", &Load<paddle::platform::IPUPlace>);
+  m->def("load_func", &Load<paddle::platform::Place>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::CPUPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::CustomPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::XPUPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::CUDAPinnedPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::CUDAPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::IPUPlace>);
+  m->def("load_combine_func", &LoadCombine<paddle::platform::Place>);
 
   m->def("serialize_pir_program",
          &pir::WriteModule,
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -2165,6 +2165,7 @@ All parameter, weight, gradient are variables in Paddle.
   py::class_<framework::Executor>(m, "Executor")
       .def(py::init<const platform::Place &>())
       .def("close", &Executor::Close)
+      .def("get_place", &Executor::GetPlace)
       .def("run_from_dataset",
            &Executor::RunFromDataset,
            py::call_guard<py::gil_scoped_release>())
diff --git a/python/paddle/jit/pir_dy2static/parameter_recorder.py b/python/paddle/jit/pir_dy2static/parameter_recorder.py
@@ -45,6 +45,7 @@ def get(self, program, tensor):
                 dtype=dtype,
                 shape=tensor.shape,
                 type=tensor.type,
+                name=tensor.name,
                 initializer=non_used_initializer,
                 trainable=(not tensor.stop_gradient),
             )
diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py
@@ -65,7 +65,11 @@ def _get_pir_persistable_var_names(program):
                 var.name = _generate_unique_var_name(var.name)
                 rename_new_old_dict[var.name] = origin_name
                 persistable_names.append(var.name)
-    return persistable_vars, persistable_names, rename_new_old_dict
+    return (
+        persistable_vars,
+        persistable_names,
+        rename_new_old_dict,
+    )
 
 
 class _PirProgramHolder:
@@ -91,6 +95,7 @@ def _preprocess(self):
             self._persistable_names,
             self._suffix_varname_dict,
         ) = _get_pir_persistable_var_names(self._infer_program)
+
         block = self._infer_program.global_block()
         for op in block.ops:
             if op.name() == 'pd_op.data':
@@ -174,9 +179,8 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
     load_var_list = []
     load_densetensor_list = []
     persistable_var = program_holder.persistable_vars
-    persistable_var_name = program_holder.persistable_names
-
-    for name, var in sorted(zip(persistable_var_name, persistable_var)):
+    persistable_name = program_holder.persistable_names
+    for name, var in sorted(zip(persistable_name, persistable_var)):
         if var.persistable:
             # use default shape and dtype
             new_var = framework.EagerParamBase(
@@ -210,6 +214,7 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
             list(load_var_dict.keys()),
             load_densetensor_list,
             False,
+            framework._current_expected_place(),
         )
     else:
         raise ValueError(
@@ -312,10 +317,10 @@ def _construct_params_and_buffers(
                 )
             )
 
-        if not append_suffix:
-            var_dict = _remove_varname_suffix(var_dict, programs['forward'])
+    if not append_suffix:
+        var_dict = _remove_varname_suffix(var_dict, programs['forward'])
 
-        return var_dict
+    return var_dict
 
 
 def _run_dygraph(instance, input, program_holder):
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -483,7 +483,14 @@ def load_vars_pir(
                         "The directory path and params cannot be None at the same time."
                     )
                 file_path = os.path.join(dirname, v.name)
-                core.load_func(file_path, -1, [], False, var.get_tensor())
+                core.load_func(
+                    file_path,
+                    -1,
+                    [],
+                    False,
+                    var.get_tensor(),
+                    executor._default_executor.get_place(),
+                )
             else:
                 load_var_map[v.name] = var
 
@@ -498,7 +505,11 @@ def load_vars_pir(
                 filename = os.path.join(dirname, filename)
 
             core.load_combine_func(
-                filename, load_var_names, load_var_list, False
+                filename,
+                load_var_names,
+                load_var_list,
+                False,
+                executor._default_executor.get_place(),
             )
             for name, var in zip(load_var_names, load_var_list):
                 set_var(name, np.array(var))
diff --git a/test/deprecated/legacy_test/test_jit_layer.py b/test/deprecated/legacy_test/test_jit_layer.py
@@ -22,6 +22,7 @@
 import paddle
 from paddle.base.framework import _dygraph_place_guard
 from paddle.jit.layer import Layer
+from paddle.pir_utils import test_with_dygraph_pir
 from paddle.static import InputSpec
 
 sys.path.append("../../dygraph_to_static")
@@ -82,6 +83,22 @@ def test_multi_load(self):
         np.testing.assert_allclose(forward_out1, forward_out2[0], rtol=1e-05)
         np.testing.assert_allclose(infer_out1, infer_out2[0], rtol=1e-05)
 
+    @test_with_dygraph_pir
+    def test_multi_jit_load(self):
+        x = paddle.full([2, 4], 2)
+        model = Net()
+        with enable_to_static_guard(False):
+            forward_out1 = model.forward(x)
+            infer_out1 = model.infer(x)
+        model_path = os.path.join(self.temp_dir.name, 'multi_program')
+        paddle.jit.save(model, model_path, combine_params=True)
+
+        jit_layer = paddle.jit.load(model_path)
+        forward_out2 = jit_layer.forward(x)
+        infer_out2 = jit_layer.infer(x)
+        np.testing.assert_allclose(forward_out1, forward_out2, rtol=1e-05)
+        np.testing.assert_allclose(infer_out1, infer_out2, rtol=1e-05)
+
 
 class SaveLinear(paddle.nn.Layer):
     def __init__(self):
@@ -117,6 +134,24 @@ def test_mkl_output(self):
             out = paddle.unsqueeze(out[0], 0)
             np.testing.assert_equal(out.shape, [1, 498, 80])
 
+    @test_with_dygraph_pir
+    def test_mkl_jit_output(self):
+        with _dygraph_place_guard(place=paddle.CPUPlace()):
+            net = SaveLinear()
+            x = paddle.ones([498, 80])
+            orig_out = net.forward(x)
+            model_path = os.path.join(self.temp_dir.name, 'save_linear')
+            paddle.jit.save(net, model_path, combine_params=True)
+
+            layer = paddle.jit.load(model_path)
+
+            out = layer.forward(x)
+            np.testing.assert_equal(
+                np.mean(orig_out.numpy()), np.mean(out.numpy())
+            )
+            out = paddle.unsqueeze(out, 0)
+            np.testing.assert_equal(out.shape, [1, 498, 80])
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/ir/pir/test_save_load_params.py b/test/ir/pir/test_save_load_params.py
@@ -132,15 +132,29 @@ def test_params_cpp(self):
                     paddle.base.core.save_func(v, k, path, True, True)
                     tensor = param_dict[k]
                     tensor.set(np.zeros_like(np.array(tensor)), self.place)
-                    paddle.base.core.load_func(path, -1, [], False, tensor)
+                    paddle.base.core.load_func(
+                        path,
+                        -1,
+                        [],
+                        False,
+                        tensor,
+                        paddle.framework._current_expected_place_(),
+                    )
                     np.testing.assert_array_equal(tensor, v)
 
                 for k, v in opt_dict.items():
                     path = os.path.join(save_dir, k, '.pdopt')
                     paddle.base.core.save_func(v, k, path, True, False)
                     tensor = opt_dict[k]
                     tensor.set(np.zeros_like(np.array(tensor)), self.place)
-                    paddle.base.core.load_func(path, -1, [], False, tensor)
+                    paddle.base.core.load_func(
+                        path,
+                        -1,
+                        [],
+                        False,
+                        tensor,
+                        paddle.framework._current_expected_place_(),
+                    )
                     np.testing.assert_array_equal(tensor, v)
 
                 # test save_combine_func and load_combine_func
@@ -157,7 +171,11 @@ def test_params_cpp(self):
                     tensor.set(np.zeros_like(np.array(tensor)), self.place)
                     param_new.append(tensor)
                 paddle.base.core.load_combine_func(
-                    path, list(param_dict.keys()), param_new, False
+                    path,
+                    list(param_dict.keys()),
+                    param_new,
+                    False,
+                    paddle.framework._current_expected_place_(),
                 )
                 np.testing.assert_equal(param_new, param_vec)
                 # save to memory
@@ -170,7 +188,11 @@ def test_params_cpp(self):
                 )
                 # load as fp16
                 paddle.base.core.load_combine_func(
-                    path, list(param_dict.keys()), param_new, True
+                    path,
+                    list(param_dict.keys()),
+                    param_new,
+                    True,
+                    paddle.framework._current_expected_place_(),
                 )
 
                 # test save_vars

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@ def get(self, program, tensor):`
`45`	`45`	`dtype=dtype,`
`46`	`46`	`shape=tensor.shape,`
`47`	`47`	`type=tensor.type,`
	`48`	`+ name=tensor.name,`
`48`	`49`	`initializer=non_used_initializer,`
`49`	`50`	`trainable=(not tensor.stop_gradient),`
`50`	`51`	`)`