From 994bea93ca34f6628f9fd5a3c94353a205314336 Mon Sep 17 00:00:00 2001
From: zhaoyingli <86812880+zhaoyinglia@users.noreply.github.com>
Date: Thu, 22 Sep 2022 10:04:25 +0800
Subject: [PATCH] [Auto Parallel] fix lazyinit (#46355)

---
 python/paddle/distributed/auto_parallel/helper.py    |  1 +
 .../distributed/auto_parallel/parallelizer_v2.py     | 12 ++++++++----
 python/paddle/fluid/lazy_init.py                     |  2 +-
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/python/paddle/distributed/auto_parallel/helper.py b/python/paddle/distributed/auto_parallel/helper.py
index 6bc177efc9..4a3a1ab5e1 100644
--- a/python/paddle/distributed/auto_parallel/helper.py
+++ b/python/paddle/distributed/auto_parallel/helper.py
@@ -341,6 +341,7 @@ class ProgramHelper(object):
         try:
             return self.proxy_layer.startup_program
         except Exception as err:
+            self._logger.warning("`lazy init` failed.")
             if isinstance(err, AssertionError):
                 return self.concrete_program.startup_program
             raise err
diff --git a/python/paddle/distributed/auto_parallel/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
index b83a19b512..a1dd58fef7 100644
--- a/python/paddle/distributed/auto_parallel/parallelizer_v2.py
+++ b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
@@ -105,9 +105,13 @@ class Parallelizer:
                 format(time.time() - time0, self._mode))
         else:
             # Apply pre optimization passes
-            # self._apply_pre_optimization(serial_main_program,
-            #                              serial_startup_program, None, None,
-            #                              None)
+            time0 = time.time()
+            self._apply_pre_optimization(serial_main_program,
+                                         serial_startup_program, None, None,
+                                         None)
+            self._logger.info(
+                "within parallel apply_pre_optimization time: {}, mode {}".
+                format(time.time() - time0, self._mode))
             # Do logical partition
             time0 = time.time()
             partitioner = Partitioner(self._dist_context, rank)
@@ -176,7 +180,7 @@ class Parallelizer:
         # apply amp pass
         # FIXME we disenable amp for eval since it has a little bug with
         # eval program and which will be fixed in future
-        if self._mode == 'train' and self._strategy.amp.enable:
+        if self._strategy.amp.enable:
             config = copy.deepcopy(self._strategy.amp.to_dict())
             config["dist_context"] = self._dist_context
             config["params_grads"] = params_grads
diff --git a/python/paddle/fluid/lazy_init.py b/python/paddle/fluid/lazy_init.py
index c44ed7f118..8fc175ae3e 100644
--- a/python/paddle/fluid/lazy_init.py
+++ b/python/paddle/fluid/lazy_init.py
@@ -36,7 +36,7 @@ class LazyInitHelper(object):
         """
         if self._state:
             return
-        assert framework.in_dygraph_mode(
+        assert framework._non_static_mode(
         ), "LazyInit.enable() is only available in dygraph mode."
         self._state = True
 
-- 
GitLab