[PIR saveload] Add _suffix_varname_dict for jit and open tests (#64239)

changeyoung98 · web-flow · commit 90f0e80f0157 · 2024-05-13T15:55:49.000+08:00
* open more tests

* fix

* fix

* fix

* open load state_dict from old format
diff --git a/python/paddle/framework/io.py b/python/paddle/framework/io.py
@@ -154,12 +154,21 @@ def _load_state_dict_from_save_inference_model(model_path, config):
     # 1. load program desc & construct _ProgramHolder
     # TODO(GGBond8488):From a long-term perspective, it is inappropriate for the framework to
     # rely on jit. It is necessary to migrate the dependency from jit to the framework in the future
-    from paddle.jit.translated_layer import (
-        _construct_params_and_buffers,
-        _construct_program_holders,
-    )
+    if in_pir_mode():
+        from paddle.jit.pir_translated_layer import (
+            _construct_params_and_buffers,
+            _construct_program_holders,
+        )
 
-    programs = _construct_program_holders(model_path, config.model_filename)
+        programs = _construct_program_holders(model_path, config.model_filename)
+
+    else:
+        from paddle.jit.translated_layer import (
+            _construct_params_and_buffers,
+            _construct_program_holders,
+        )
+
+        programs = _construct_program_holders(model_path, config.model_filename)
 
     # 2. load layer parameters & buffers
     with base.dygraph.guard():
@@ -256,12 +265,18 @@ def _build_load_path_and_config(path, config):
     # raise error, avoid confusing behavior
     # TODO(GGBond8488):From a long-term perspective, it is inappropriate for the framework to
     # rely on jit. It is necessary to migrate the dependency from jit to the framework in the future
+    from paddle.jit.pir_translated_layer import (
+        PIR_INFER_MODEL_SUFFIX,
+    )
     from paddle.jit.translated_layer import (
         INFER_MODEL_SUFFIX,
         INFER_PARAMS_SUFFIX,
     )
 
-    prefix_format_path = path + INFER_MODEL_SUFFIX
+    if in_pir_mode():
+        prefix_format_path = path + PIR_INFER_MODEL_SUFFIX
+    else:
+        prefix_format_path = path + INFER_MODEL_SUFFIX
     prefix_format_exist = os.path.exists(prefix_format_path)
     directory_format_exist = os.path.isdir(path)
     if prefix_format_exist and directory_format_exist:
@@ -294,7 +309,10 @@ def _build_load_path_and_config(path, config):
                     "specified file prefix, the ``model_filename`` config does "
                     "not take effect."
                 )
-            config.model_filename = file_prefix + INFER_MODEL_SUFFIX
+            if in_pir_mode():
+                config.model_filename = file_prefix + PIR_INFER_MODEL_SUFFIX
+            else:
+                config.model_filename = file_prefix + INFER_MODEL_SUFFIX
             if config.params_filename is not None:
                 warnings.warn(
                     "When loading the result saved with the "
diff --git a/python/paddle/jit/pir_translated_layer.py b/python/paddle/jit/pir_translated_layer.py
@@ -56,13 +56,16 @@ def _generate_unique_var_name(prefix):
 def _get_pir_persistable_var_names(program):
     persistable_vars = []
     persistable_names = []
+    rename_new_old_dict = {}
     for block in program.blocks:
         for var in block.all_parameters():
             if var.persistable:
                 persistable_vars.append(var)
+                origin_name = var.name
                 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
+    return persistable_vars, persistable_names, rename_new_old_dict
 
 
 class _PirProgramHolder:
@@ -76,6 +79,8 @@ def __init__(self, program, trainable):
         self._persistable_names = []
 
         self.support_train = trainable
+        # append suffix var name dict
+        self._suffix_varname_dict = None
         self._infer_program = program
 
         self._preprocess()
@@ -84,6 +89,7 @@ def _preprocess(self):
         (
             self._persistable_vars,
             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:
@@ -214,6 +220,15 @@ def _load_pir_persistable_vars(model_path, program_holder, params_filename):
     return load_var_dict
 
 
+# NOTE(chenzhiyang): to adapt paddle.load to get state_dict
+def _remove_varname_suffix(var_dict, program_holder):
+    no_suffix_var_dict = {}
+    for var_name in var_dict:
+        no_suffix_name = program_holder._suffix_varname_dict[var_name]
+        no_suffix_var_dict[no_suffix_name] = var_dict[var_name]
+    return no_suffix_var_dict
+
+
 def _construct_program_holders(model_path, model_filename=None):
     # make sure the path has been checked
     program_holder_dict = {}
@@ -296,6 +311,10 @@ def _construct_params_and_buffers(
                     model_path, programs[func_name], file_name
                 )
             )
+
+        if not append_suffix:
+            var_dict = _remove_varname_suffix(var_dict, programs['forward'])
+
         return var_dict
 
 
diff --git a/test/deprecated/legacy_test/test_load_state_dict_from_old_format.py b/test/deprecated/legacy_test/test_load_state_dict_from_old_format.py
@@ -26,6 +26,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.framework import in_pir_mode
+from paddle.pir_utils import test_with_pir_api
 
 
 def convolutional_neural_network(img):
@@ -214,6 +215,7 @@ def check_load_state_dict(self, orig_dict, load_dict):
         for var_name, value in orig_dict.items():
             np.testing.assert_array_equal(value, load_dict[var_name])
 
+    @test_with_pir_api
     def test_load_default(self):
         self.save_dirname = os.path.join(
             self.temp_dir.name, "static_mnist.load_state_dict.default"
@@ -225,6 +227,7 @@ def test_load_default(self):
         new_load_param_dict = paddle.load(self.save_dirname)
         self.check_load_state_dict(orig_param_dict, new_load_param_dict)
 
+    @test_with_pir_api
     def test_load_with_model_filename(self):
         self.save_dirname = os.path.join(
             self.temp_dir.name, "static_mnist.load_state_dict.model_filename"
@@ -238,6 +241,7 @@ def test_load_with_model_filename(self):
         )
         self.check_load_state_dict(orig_param_dict, new_load_param_dict)
 
+    @test_with_pir_api
     def test_load_with_param_filename(self):
         self.save_dirname = os.path.join(
             self.temp_dir.name, "static_mnist.load_state_dict.param_filename"
@@ -251,6 +255,7 @@ def test_load_with_param_filename(self):
         )
         self.check_load_state_dict(orig_param_dict, new_load_param_dict)
 
+    @test_with_pir_api
     def test_load_with_model_and_param_filename(self):
         self.save_dirname = os.path.join(
             self.temp_dir.name,
