From 979af475514b7f8f7249c7fabccd965e213b7bfa Mon Sep 17 00:00:00 2001
From: zhaoyingli <86812880+zhaoyinglia@users.noreply.github.com>
Date: Thu, 20 Oct 2022 14:14:20 +0800
Subject: [PATCH] [AutoParallel] fix fp16 for subblock (#47189)

* [AutoParallel] fix fp16 for subblock

* fix engine

* fix comment
---
 python/paddle/distributed/auto_parallel/engine.py      | 6 +++---
 python/paddle/distributed/passes/auto_parallel_fp16.py | 3 ++-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/python/paddle/distributed/auto_parallel/engine.py b/python/paddle/distributed/auto_parallel/engine.py
index a2e1477f88..2eaefee6cd 100644
--- a/python/paddle/distributed/auto_parallel/engine.py
+++ b/python/paddle/distributed/auto_parallel/engine.py
@@ -240,8 +240,8 @@ class Engine:
                 else:
                     specs.append(spec.batch(batch_size))
             elif isinstance(item, (Variable, core.VarBase, core.eager.Tensor)):
-                _adjust_item_spec(num_shards, spec)
                 spec = InputSpec.from_tensor(item, name)
+                _adjust_item_spec(num_shards, spec)
                 if batch_size is None:
                     specs.append(spec)
                 else:
@@ -1508,10 +1508,10 @@ class Engine:
             strict (bool, optional): Whether to skip the loading of mismatch
                 parameter or raise an error when mismatch happens (not found
                 the parameter in file storing model states of or receives a
-                mismatch shape). Default: False.
+                mismatch shape). Default: True.
             load_optimizer (bool, optional): If True, the stored optimizer
                 states is restored. Otherwise, the optimizer states is initialized
-                from scratch. Default: False.
+                from scratch. Default: True.
 
         Returns:
             None
diff --git a/python/paddle/distributed/passes/auto_parallel_fp16.py b/python/paddle/distributed/passes/auto_parallel_fp16.py
index 541901f0c7..35a4831af3 100644
--- a/python/paddle/distributed/passes/auto_parallel_fp16.py
+++ b/python/paddle/distributed/passes/auto_parallel_fp16.py
@@ -181,7 +181,8 @@ class FP16State(object):
         try:
             var = block.var(var_name)
         except ValueError as e:
-            var = self.program.global_block().var(var_name)
+            var = block._var_recursive(var_name)
+            # var = self.program.global_block().var(var_name)
 
         # NOTE(JZ-LIANG) "array_" is a hack to adopt for ernie3.0 inference, since there is
         # a trick which make the LOD_TENSOR_ARRAY to the float32 in while block to reset the LOD_TENSOR_ARRAY
-- 
GitLab