[PIR save/load]Fix get_pir_parameters bug and open tests (#64625)

changeyoung98 · web-flow · commit 362c87729237 · 2024-05-29T10:45:14.000+08:00
* fix get_pir_parameters bug and open tests

* fix ci bug
diff --git a/python/paddle/static/io.py b/python/paddle/static/io.py
@@ -1735,6 +1735,7 @@ def set_program_state(program, state_dict):
     if in_pir_mode():
         params, opts = get_pir_parameters(program)
         parameter_list = params + opts
+        parameter_list = [var for var in parameter_list if var.persistable]
     else:
         parameter_list = list(filter(is_persistable, program.list_vars()))
 
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -567,8 +567,14 @@ def get_tensor(var):
 
     # get parameters and optimizer variables
     parameter_list, optimizer_param_list = get_pir_parameters(program)
-    param_dict = {var.name: get_tensor(var) for var in parameter_list}
-    opt_dict = {var.name: get_tensor(var) for var in optimizer_param_list}
+    param_dict = {
+        var.name: get_tensor(var) for var in parameter_list if var.persistable
+    }
+    opt_dict = {
+        var.name: get_tensor(var)
+        for var in optimizer_param_list
+        if var.persistable
+    }
 
     # save parameters
     param_dict = _unpack_saved_dict(param_dict, protocol)
@@ -643,10 +649,11 @@ def load_pir(program, model_path, executor=None, var_list=None):
             load_dict = _safe_load_pickle(f, encoding='latin1')
         load_dict = _pack_loaded_dict(load_dict)
     for var in parameter_list:
-        assert (
-            var.name in load_dict
-        ), f"Can not find [{var.name}] in model file [{parameter_file_name}]"
-        set_var(var.name, load_dict[var.name])
+        if var.persistable:
+            assert (
+                var.name in load_dict
+            ), f"Can not find [{var.name}] in model file [{parameter_file_name}]"
+            set_var(var.name, load_dict[var.name])
 
     if len(optimizer_param_list) > 0:
         opt_file_name = model_prefix + ".pdopt"
@@ -662,10 +669,11 @@ def load_pir(program, model_path, executor=None, var_list=None):
         with open(opt_file_name, 'rb') as f:
             load_dict = _safe_load_pickle(f, encoding='latin1')
         for var in optimizer_param_list:
-            assert (
-                var.name in load_dict
-            ), f"Can not find [{var.name}] in model file [{opt_file_name}]"
-            set_var(var.name, load_dict[var.name])
+            if var.persistable:
+                assert (
+                    var.name in load_dict
+                ), f"Can not find [{var.name}] in model file [{opt_file_name}]"
+                set_var(var.name, load_dict[var.name])
 
 
 @static_only
diff --git a/test/deprecated/legacy_test/test_executor_and_use_program_cache.py b/test/deprecated/legacy_test/test_executor_and_use_program_cache.py
@@ -22,6 +22,7 @@
 
 import paddle
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestExecutor(unittest.TestCase):
@@ -99,6 +100,7 @@ def train_and_save_inference_program(
                 program=self.main_program,
             )
 
+    @test_with_pir_api
     def test_inference_output(self):
         for rnn_model in ["static"]:
             # Set parallel to False to use the default executor.
diff --git a/test/legacy_test/test_eager_deletion_padding_rnn.py b/test/legacy_test/test_eager_deletion_padding_rnn.py
@@ -240,18 +240,32 @@ def encoder_static(
         init_cell, shape=[num_layers, -1, hidden_size]
     )
 
-    x_emb = paddle.static.nn.embedding(
-        input=x,
-        size=[vocab_size, hidden_size],
-        dtype='float32',
-        is_sparse=False,
-        param_attr=base.ParamAttr(
-            name='embedding_para',
-            initializer=paddle.nn.initializer.Uniform(
-                low=-init_scale, high=init_scale
+    if paddle.framework.in_pir_mode():
+        Emb = paddle.nn.Embedding(
+            vocab_size,
+            hidden_size,
+            sparse=False,
+            weight_attr=base.ParamAttr(
+                name='embedding_para',
+                initializer=paddle.nn.initializer.Uniform(
+                    low=-init_scale, high=init_scale
+                ),
             ),
-        ),
-    )
+        )
+        x_emb = Emb(x)
+    else:
+        x_emb = paddle.static.nn.embedding(
+            input=x,
+            size=[vocab_size, hidden_size],
+            dtype='float32',
+            is_sparse=False,
+            param_attr=base.ParamAttr(
+                name='embedding_para',
+                initializer=paddle.nn.initializer.Uniform(
+                    low=-init_scale, high=init_scale
+                ),
+            ),
+        )
 
     x_emb = paddle.reshape(x_emb, shape=[-1, num_steps, hidden_size])
     if dropout is not None and dropout > 0.0:
@@ -448,9 +462,9 @@ def _train_an_epoch(self, epoch_id, use_program_cache=True):
                 self.train_program,
                 feed=input_data_feed,
                 fetch_list=[
-                    self.loss.name,
-                    self.last_hidden.name,
-                    self.last_cell.name,
+                    self.loss,
+                    self.last_hidden,
+                    self.last_cell,
                 ],
                 use_program_cache=use_program_cache,
             )
