pir support lazy_init (#67467)

wanghuancoder · web-flow · commit f6f447e43b23 · 2024-08-16T10:02:54.000+08:00
diff --git a/python/paddle/base/data_feeder.py b/python/paddle/base/data_feeder.py
@@ -26,6 +26,7 @@
 from ..pir.core import _PADDLE_PIR_DTYPE_2_NUMPY_DTYPE, ParameterMeta
 from . import core
 from .framework import (
+    EagerParamBase,
     Variable,
     _cpu_num,
     _cuda_ids,
@@ -165,9 +166,13 @@ def check_variable_and_dtype(
     input, input_name, expected_dtype, op_name, extra_message=''
 ):
     if in_pir_mode():
-        check_type(
-            input, input_name, (Value, ParameterMeta), op_name, extra_message
-        )
+        from ..nn.initializer.lazy_init import lazy_init_helper
+
+        if lazy_init_helper().state:
+            expected = (Value, ParameterMeta, EagerParamBase)
+        else:
+            expected = (Value, ParameterMeta)
+        check_type(input, input_name, expected, op_name, extra_message)
     else:
         check_type(input, input_name, (Variable, Value), op_name, extra_message)
     check_dtype(input.dtype, input_name, expected_dtype, op_name, extra_message)
@@ -187,14 +192,15 @@ def check_type(input, input_name, expected_type, op_name, extra_message=''):
     # NOTE: `in_to_static_mode` is used to determined whether this op is called under
     # @to_static in transformation from dygraph to static layer. We add Tensor in
     # expected_type to skip checking because Tensor may be created and used in unusual way.
+    from ..nn.initializer.lazy_init import lazy_init_helper
     from .dygraph.base import in_to_static_mode
 
     # Need a better design to be fix this.
     if in_to_static_mode():
         if not isinstance(expected_type, tuple):
             expected_type = (expected_type,)
         expected_type += (core.eager.Tensor,)
-    elif isinstance(input, core.eager.Tensor):
+    elif isinstance(input, core.eager.Tensor) and not lazy_init_helper().state:
         raise TypeError(
             "Please use `with base.dygraph.guard()` as context or `base.enable_dygraph()` to switch to imperative mode firstly. "
             f"Because received '{input_name}' in {op_name} is a imperative Variable."
diff --git a/python/paddle/nn/initializer/normal.py b/python/paddle/nn/initializer/normal.py
@@ -25,6 +25,7 @@
     in_pir_mode,
 )
 from .initializer import Initializer
+from .lazy_init import lazy_init_helper
 
 __all__ = []
 
@@ -235,9 +236,16 @@ def forward(
             The initialization op
         """
         block = self._check_block(block)
-        assert isinstance(
-            var, (framework.Variable, paddle.pir.core.ParameterMeta)
-        )
+        if lazy_init_helper().state:
+            expected = (
+                framework.Variable,
+                paddle.pir.core.ParameterMeta,
+                core.eager.Tensor,
+            )
+        else:
+            expected = (framework.Variable, paddle.pir.core.ParameterMeta)
+
+        assert isinstance(var, expected)
         assert isinstance(block, (framework.Block, pir.Block))
 
         if self._seed == 0:
diff --git a/test/legacy_test/test_lazy_init.py b/test/legacy_test/test_lazy_init.py
@@ -57,16 +57,16 @@ def clear_nameset(self):
         unique_name.dygraph_parameter_name_checker._name_set = set()
 
     def test_wrapper(self):
-        with paddle.pir_utils.OldIrGuard():
-            paddle.disable_static()
-            with LazyGuard():
-                fc = Linear(
-                    10,
-                    10,
-                    weight_attr=self.weight_attr,
-                    bias_attr=self.bias_attr,
-                )
-            program = fc._startup_program()
+        paddle.disable_static()
+        with LazyGuard():
+            fc = Linear(
+                10,
+                10,
+                weight_attr=self.weight_attr,
+                bias_attr=self.bias_attr,
+            )
+        program = fc._startup_program()
+        if not paddle.framework.use_pir_api():
             self.check_program(program)
 
     def check_program(self, program):
@@ -119,18 +119,18 @@ def forward(self, x):
 
 class TestNestModelLazy(TestInitializerBase):
     def test_wrapper(self):
-        with paddle.pir_utils.OldIrGuard():
-            paddle.disable_static()
-            with LazyGuard():
-                base_model = Linear(
-                    10,
-                    10,
-                    weight_attr=self.weight_attr,
-                    bias_attr=self.bias_attr,
-                )
-                nest_model = NestModel(base_model)
-
-            self.check_data(nest_model)
+        paddle.disable_static()
+        with LazyGuard():
+            base_model = Linear(
+                10,
+                10,
+                weight_attr=self.weight_attr,
+                bias_attr=self.bias_attr,
+            )
+            nest_model = NestModel(base_model)
+
+        self.check_data(nest_model)
+        if not paddle.framework.use_pir_api():
             self.check_program(nest_model)
 
     def check_data(self, model):
