[PIR] Fix 4 unittests in pir mode (#65886)

huangjiyi · web-flow · commit 4e4c113cbd75 · 2024-07-10T11:10:52.000+08:00
diff --git a/python/paddle/utils/install_check.py b/python/paddle/utils/install_check.py
@@ -136,7 +136,7 @@ def _run_static_single(use_cuda, use_xpu, use_custom, custom_device_name):
         with paddle.static.program_guard(train_prog, startup_prog):
             input, out, weight = _simple_network()
             param_grads = paddle.static.append_backward(
-                out, parameter_list=[weight.name]
+                out, parameter_list=[weight]
             )[0]
 
         if use_cuda:
@@ -153,7 +153,7 @@ def _run_static_single(use_cuda, use_xpu, use_custom, custom_device_name):
         exe.run(
             train_prog,
             feed={input.name: _prepare_data()},
-            fetch_list=[out.name, param_grads[1].name],
+            fetch_list=[out, param_grads[1]],
         )
     paddle.disable_static()
 
diff --git a/test/deprecated/legacy_test/CMakeLists.txt b/test/deprecated/legacy_test/CMakeLists.txt
@@ -396,7 +396,6 @@ list(REMOVE_ITEM TEST_OPS test_fetch_lod_tensor_array)
 list(REMOVE_ITEM TEST_OPS test_data_norm_op_deprecated)
 list(REMOVE_ITEM TEST_OPS test_layers_deprecated)
 list(REMOVE_ITEM TEST_OPS test_imperative_ocr_attention_model)
-list(REMOVE_ITEM TEST_OPS test_install_check)
 list(REMOVE_ITEM TEST_OPS test_basic_gru_api)
 list(REMOVE_ITEM TEST_OPS test_basic_gru_unit_op)
 list(REMOVE_ITEM TEST_OPS test_basic_lstm_api)
@@ -475,12 +474,6 @@ set_tests_properties(test_logcumsumexp_op PROPERTIES TIMEOUT 30)
 py_test_modules(
   test_imperative_ocr_attention_model MODULES
   test_imperative_ocr_attention_model ENVS FLAGS_cudnn_deterministic=1)
-py_test_modules(test_install_check MODULES test_install_check ENVS
-                FLAGS_cudnn_deterministic=1)
-set_tests_properties(test_install_check PROPERTIES LABELS "RUN_TYPE=DIST")
-py_test_modules(test_install_check_pir MODULES test_install_check ENVS
-                FLAGS_cudnn_deterministic=1 FLAGS_enable_pir_in_executor=1)
-set_tests_properties(test_install_check_pir PROPERTIES LABELS "RUN_TYPE=DIST")
 
 if((WITH_GPU) AND (WITH_CUDNN_FRONTEND))
   py_test_modules(test_fuse_dot_product_attention_pass MODULES
@@ -653,9 +646,6 @@ set_tests_properties(test_imperative_selected_rows_to_lod_tensor
                      PROPERTIES TIMEOUT 200)
 set_tests_properties(test_argsort_op_deprecated PROPERTIES TIMEOUT 120)
 set_tests_properties(test_gather_nd_op PROPERTIES TIMEOUT 120)
-set_tests_properties(test_imperative_ptb_rnn_sorted_gradient PROPERTIES TIMEOUT
-                                                                        120)
-set_tests_properties(test_imperative_ptb_rnn PROPERTIES TIMEOUT 120)
 set_tests_properties(test_masked_select_op PROPERTIES TIMEOUT 120)
 set_tests_properties(test_sigmoid_cross_entropy_with_logits_op
                      PROPERTIES TIMEOUT 120)
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -441,6 +441,7 @@ list(REMOVE_ITEM TEST_OPS test_basic_lstm_unit_op)
 list(REMOVE_ITEM TEST_OPS test_fuse_bn_add_act_pass)
 list(REMOVE_ITEM TEST_OPS test_conv3d_transpose_op)
 list(REMOVE_ITEM TEST_OPS test_layers)
+list(REMOVE_ITEM TEST_OPS test_install_check)
 
 # disable this unittest temporarily
 list(REMOVE_ITEM TEST_OPS test_imperative_data_loader_exception)
@@ -1210,3 +1211,12 @@ set_tests_properties(test_squeeze2_op_rename PROPERTIES TIMEOUT 120)
 set_tests_properties(test_reduce_as_op PROPERTIES TIMEOUT 30)
 set_tests_properties(test_reduce_op PROPERTIES TIMEOUT 500)
 set_tests_properties(test_reduce_op_static_build PROPERTIES TIMEOUT 500)
+set_tests_properties(test_imperative_ptb_rnn_sorted_gradient PROPERTIES TIMEOUT
+                                                                        120)
+set_tests_properties(test_imperative_ptb_rnn PROPERTIES TIMEOUT 120)
+py_test_modules(test_install_check MODULES test_install_check ENVS
+                FLAGS_cudnn_deterministic=1)
+set_tests_properties(test_install_check PROPERTIES LABELS "RUN_TYPE=DIST")
+py_test_modules(test_install_check_pir MODULES test_install_check ENVS
+                FLAGS_cudnn_deterministic=1 FLAGS_enable_pir_in_executor=1)
+set_tests_properties(test_install_check_pir PROPERTIES LABELS "RUN_TYPE=DIST")
diff --git a/test/legacy_test/test_imperative_ptb_rnn.py b/test/legacy_test/test_imperative_ptb_rnn.py
@@ -20,10 +20,35 @@
 
 import paddle
 from paddle import base
-from paddle.base import core, framework
+from paddle.autograd.backward_utils import ValueDict
+from paddle.base import core
 from paddle.nn import Embedding
 
 
+def create_parameter_mapping(startup_program, main_program):
+    startup_params = {}
+    main_params = {}
+    parameter_mapping = ValueDict()
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.operand(0).source()
+            startup_params[name] = param
+
+    for op in main_program.global_block().ops:
+        if op.name() == "builtin.parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.result(0)
+            main_params[name] = param
+
+    assert len(startup_params) == len(main_params)
+    for name, startup_param in startup_params.items():
+        assert name in main_params
+        main_param = main_params[name]
+        parameter_mapping[main_param] = startup_param
+    return parameter_mapping
+
+
 class SimpleLSTMRNN(paddle.nn.Layer):
     def __init__(
         self, hidden_size, num_steps, num_layers=2, init_scale=0.1, dropout=None
@@ -251,7 +276,14 @@ def ptb_rnn_cpu_float32(self, is_sparse):
 
         with base.dygraph.guard():
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
+
             # TODO: marsyang1993 Change seed to
             ptb_model = PtbModel(
                 hidden_size=hidden_size,
@@ -309,7 +341,14 @@ def ptb_rnn_cpu_float32(self, is_sparse):
 
         with new_program_scope():
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
+
             ptb_model = PtbModel(
                 hidden_size=hidden_size,
                 vocab_size=vocab_size,
@@ -328,17 +367,18 @@ def ptb_rnn_cpu_float32(self, is_sparse):
             x = paddle.static.data(
                 name="x", shape=[-1, num_steps], dtype='int64'
             )
-            x.desc.set_need_check_feed(False)
             y = paddle.static.data(name="y", shape=[-1, 1], dtype='float32')
-            y.desc.set_need_check_feed(False)
             init_hidden = paddle.static.data(
                 name="init_hidden", shape=[-1, 1], dtype='float32'
             )
-            init_hidden.desc.set_need_check_feed(False)
             init_cell = paddle.static.data(
                 name="init_cell", shape=[-1, 1], dtype='float32'
             )
-            init_cell.desc.set_need_check_feed(False)
+            if not paddle.framework.use_pir_api():
+                x.desc.set_need_check_feed(False)
+                y.desc.set_need_check_feed(False)
+                init_hidden.desc.set_need_check_feed(False)
+                init_cell.desc.set_need_check_feed(False)
 
             static_loss, static_last_hidden, static_last_cell = ptb_model(
                 x, y, init_hidden, init_cell
@@ -347,15 +387,31 @@ def ptb_rnn_cpu_float32(self, is_sparse):
             static_param_updated = {}
             static_param_init = {}
             static_param_name_list = []
+            static_params = []
             for param in ptb_model.parameters():
                 static_param_name_list.append(param.name)
+                static_params.append(param)
+
+            if paddle.framework.use_pir_api():
+                parameter_mapping = create_parameter_mapping(
+                    paddle.static.default_startup_program(),
+                    paddle.static.default_main_program(),
+                )
+                startup_params = [
+                    parameter_mapping[param] for param in static_params
+                ]
+            else:
+                startup_params = static_params
 
             out = exe.run(
-                framework.default_startup_program(),
-                fetch_list=static_param_name_list,
+                paddle.static.default_startup_program(),
+                fetch_list=startup_params,
             )
-            for i in range(len(static_param_name_list)):
-                static_param_init[static_param_name_list[i]] = out[i]
+
+            for i in range(len(static_params)):
+                param_name = static_param_name_list[i]
+                static_param_init[param_name] = out[i]
+
             static_loss_value = None
             static_last_cell_value = None
             static_last_hidden_value = None
@@ -371,7 +427,7 @@ def ptb_rnn_cpu_float32(self, is_sparse):
                     (num_layers, batch_size, hidden_size), dtype='float32'
                 )
                 fetch_list = [static_loss, static_last_hidden, static_last_cell]
-                fetch_list.extend(static_param_name_list)
+                fetch_list.extend(static_params)
                 out = exe.run(
                     base.default_main_program(),
                     feed={
diff --git a/test/legacy_test/test_imperative_ptb_rnn_sorted_gradient.py b/test/legacy_test/test_imperative_ptb_rnn_sorted_gradient.py
@@ -20,7 +20,32 @@
 
 import paddle
 from paddle import base
-from paddle.base import core, framework
+from paddle.autograd.backward_utils import ValueDict
+from paddle.base import core
+
+
+def create_parameter_mapping(startup_program, main_program):
+    startup_params = {}
+    main_params = {}
+    parameter_mapping = ValueDict()
+    for op in startup_program.global_block().ops:
+        if op.name() == "builtin.set_parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.operand(0).source()
+            startup_params[name] = param
+
+    for op in main_program.global_block().ops:
+        if op.name() == "builtin.parameter":
+            name = op.attrs()["parameter_name"]
+            param = op.result(0)
+            main_params[name] = param
+
+    assert len(startup_params) == len(main_params)
+    for name, startup_param in startup_params.items():
+        assert name in main_params
+        main_param = main_params[name]
+        parameter_mapping[main_param] = startup_param
+    return parameter_mapping
 
 
 class TestDygraphPtbRnnSortGradient(unittest.TestCase):
@@ -41,7 +66,13 @@ def ptb_rnn_sort_gradient_cpu_float32(self, is_sparse):
         with base.dygraph.guard():
             base.set_flags({'FLAGS_sort_sum_gradient': True})
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
 
             # TODO: marsyang1993 Change seed to
             ptb_model = PtbModel(
@@ -96,7 +127,13 @@ def ptb_rnn_sort_gradient_cpu_float32(self, is_sparse):
 
         with new_program_scope():
             paddle.seed(seed)
-            paddle.framework.random._manual_program_seed(seed)
+            if paddle.framework.use_pir_api():
+                with paddle.pir_utils.OldIrGuard():
+                    # Note: dygraph use self.main_program.global_block().create_parameter(), it's need manual seed to old Program
+                    paddle.framework.random._manual_program_seed(seed)
+                paddle.framework.random._manual_program_seed(seed)
+            else:
+                paddle.framework.random._manual_program_seed(seed)
 
             ptb_model = PtbModel(
                 hidden_size=hidden_size,
@@ -116,17 +153,18 @@ def ptb_rnn_sort_gradient_cpu_float32(self, is_sparse):
             x = paddle.static.data(
                 name="x", shape=[-1, num_steps, 1], dtype='int64'
             )
-            x.desc.set_need_check_feed(False)
             y = paddle.static.data(name="y", shape=[-1, 1], dtype='float32')
-            y.desc.set_need_check_feed(False)
             init_hidden = paddle.static.data(
                 name="init_hidden", shape=[-1, 1], dtype='float32'
             )
-            init_hidden.desc.set_need_check_feed(False)
             init_cell = paddle.static.data(
                 name="init_cell", shape=[-1, 1], dtype='float32'
             )
-            init_cell.desc.set_need_check_feed(False)
+            if not paddle.framework.use_pir_api():
+                x.desc.set_need_check_feed(False)
+                y.desc.set_need_check_feed(False)
+                init_hidden.desc.set_need_check_feed(False)
+                init_cell.desc.set_need_check_feed(False)
 
             static_loss, static_last_hidden, static_last_cell = ptb_model(
                 x, y, init_hidden, init_cell
@@ -135,15 +173,28 @@ def ptb_rnn_sort_gradient_cpu_float32(self, is_sparse):
             static_param_updated = {}
             static_param_init = {}
             static_param_name_list = []
+            static_params = []
             for param in ptb_model.parameters():
                 static_param_name_list.append(param.name)
+                static_params.append(param)
+            if paddle.framework.use_pir_api():
+                parameter_mapping = create_parameter_mapping(
+                    paddle.static.default_startup_program(),
+                    paddle.static.default_main_program(),
+                )
+                startup_params = [
+                    parameter_mapping[param] for param in static_params
+                ]
+            else:
+                startup_params = static_params
 
             out = exe.run(
-                framework.default_startup_program(),
-                fetch_list=static_param_name_list,
+                paddle.static.default_startup_program(),
+                fetch_list=startup_params,
             )
-            for i in range(len(static_param_name_list)):
-                static_param_init[static_param_name_list[i]] = out[i]
+            for i in range(len(static_params)):
+                param_name = static_param_name_list[i]
+                static_param_init[param_name] = out[i]
             static_loss_value = None
             static_last_cell_value = None
             static_last_hidden_value = None
@@ -159,7 +210,7 @@ def ptb_rnn_sort_gradient_cpu_float32(self, is_sparse):
                     (num_layers, batch_size, hidden_size), dtype='float32'
                 )
                 fetch_list = [static_loss, static_last_hidden, static_last_cell]
-                fetch_list.extend(static_param_name_list)
+                fetch_list.extend(static_params)
                 out = exe.run(
                     base.default_main_program(),
                     feed={
diff --git a/test/legacy_test/test_imperative_recurrent_usage.py b/test/legacy_test/test_imperative_recurrent_usage.py
@@ -81,18 +81,23 @@ def test_recurrent_feed(self):
             in2.stop_gradient = False
             rt1 = RecurrentTest("RecurrentTest")
             static_sum_out, static_out = rt1(in1, in2)
-            base.backward.append_backward(static_sum_out)
+            static_out.persistable = True
             exe = base.Executor(
                 base.CPUPlace()
                 if not core.is_compiled_with_cuda()
                 else base.CUDAPlace(0)
             )
 
-            static_dout = (
-                base.default_main_program()
-                .block(0)
-                ._find_var_recursive(static_out.name + "@GRAD")
-            )
+            if paddle.framework.use_pir_api():
+                grad_list = paddle.static.append_backward(static_sum_out)
+                _, static_dout = grad_list[-1]
+            else:
+                base.backward.append_backward(static_sum_out)
+                static_dout = (
+                    base.default_main_program()
+                    .block(0)
+                    ._find_var_recursive(static_out.name + "@GRAD")
+                )
             fetch_list = [static_sum_out, static_out, static_dout]
             for i in range(3):
                 out = exe.run(
diff --git a/test/legacy_test/test_install_check.py b/test/legacy_test/test_install_check.py
