[Dy2St] Skip prune parameter in dy2st and save (#68442)

SigureMo · web-flow · commit d5fd91c5205b · 2024-09-27T02:45:47.000-07:00
* [Dy2St] Skip prune parameter in dy2st and save

* update ut name
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -819,10 +819,8 @@ def _verify_program(self, main_program, outputs):
         Verify that the program parameter is initialized, prune some unused params,
         and remove redundant op callstack.
         """
-        # 1. Check all params from main program can be found in self._params
+        # Check all params from main program can be found in self._params
         self._check_params_all_inited(main_program)
-        # 2. Prune the parameters not used anywhere in the program.
-        self._prune_unused_params(main_program, outputs)
 
         return main_program
 
@@ -1050,29 +1048,6 @@ def _append_backward_desc(self, train_runnable_program: RunnableProgram):
             (backward_start_op_index, backward_end_op_index),
         )
 
-    def _prune_unused_params(self, program, outputs):
-        """
-        Prune the parameters not used anywhere in the program.
-        The `@to_static` may only decorated a sub function which
-        contains some unused parameters created in `__init__`.
-        So prune these parameters to avoid unnecessary operations in
-        `run_program_op`.
-        """
-        required_params = []
-        required_param_values = []
-        block = program.global_block()
-        for param, param_value in zip(self._params, self._param_values):
-            if not param_value.use_empty() or any(
-                out.is_same(param_value) for out in outputs
-            ):
-                required_params.append(param)
-                required_param_values.append(param_value)
-            else:
-                # in pir, we need remove the get_parameter op for unused parameters.
-                block.remove_op(param_value.get_defining_op())
-        self._params = required_params
-        self._param_values = required_param_values
-
     def _prepare_attributes(self):
         attrs = [
             'forward_program',
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -172,11 +172,15 @@ def pir_prune_with_input(program, feed_vars, target_vars):
 
     total_ops = program.global_block().ops
     intersection_op_flags = [True] * len(total_ops)
+    skip_prune_ops = ["builtin.parameter"]
 
     # from output to input
     target_vars_ = ValueSet(target_vars)
     for i, op in reversed(list(enumerate(total_ops))):
-        if some_in_set(get_real_op_outputs(op), target_vars_):
+        if (
+            some_in_set(get_real_op_outputs(op), target_vars_)
+            or op.name() in skip_prune_ops
+        ):
             for operand in get_real_op_inputs(op):
                 target_vars_.add(operand)
         else:
diff --git a/test/legacy_test/test_jit_save_load.py b/test/legacy_test/test_jit_save_load.py
@@ -2246,5 +2246,52 @@ def test_strided_buffer(self):
         np.testing.assert_allclose(layer(x).numpy(), loaded_layer(x).numpy())
 
 
+class LayerWithUnusedBuffer(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.linear = paddle.nn.Linear(7, 10)
+        self.register_buffer("buffer", paddle.randn([5, 1]))
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class TestLayerWithUnusedBuffer(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def check_program_has_buffer(self, program, buffer_shape):
+        for op in program.global_block().ops:
+            if (
+                op.name() == "builtin.parameter"
+                and op.result(0).shape == buffer_shape
+            ):
+                return True
+        return False
+
+    def test_layer_with_unused_buffer(self):
+        layer = LayerWithUnusedBuffer()
+        save_dir = os.path.join(
+            self.temp_dir.name, "test_layer_with_unused_buffer"
+        )
+        path = save_dir + "/model"
+        paddle.jit.save(
+            layer=layer,
+            path=path,
+            input_spec=[InputSpec([5, 7], dtype="float32")],
+        )
+
+        loaded_layer = paddle.jit.load(path)
+        x = paddle.rand([5, 7]).astype('float32')
+        self.assertTrue(
+            self.check_program_has_buffer(
+                loaded_layer.program(), layer.buffer.shape
+            )
+        )
+
+
 if __name__ == '__main__':
     unittest.main()
