[Auto Parallel] fix auto_parallel api in static_mode (#68203)

zhangyuqin1998 · web-flow · commit d60119924f00 · 2024-09-14T14:52:56.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -117,6 +117,19 @@ class _Config(TypedDict, total=False):
 # Part1: Shard attributes related APIs
 
 
+def _to_lodtensor(tensor: paddle.Tensor):
+    lodtensor = core.LoDTensor()
+    if tensor.is_dist():
+        if tensor._is_initialized():
+            lodtensor._share_data_with(tensor._local_value().get_tensor())
+        else:
+            lodtensor = None
+    else:
+        lodtensor._share_data_with(tensor.get_tensor())
+
+    return lodtensor
+
+
 class DistAttr(core.TensorDistAttr):
     """
     DistAttr specifies how tensors are distributed or sliced on ProcessMesh.
@@ -2263,20 +2276,6 @@ def _make_feeds(self, data_list):
                 f"The model takes {feed_name_list} as input"
             )
 
-        def _to_lodtensor(tensor: paddle.Tensor):
-            lodtensor = core.LoDTensor()
-            if tensor.is_dist():
-                if tensor._is_initialized():
-                    lodtensor._share_data_with(
-                        tensor._local_value().get_tensor()
-                    )
-                else:
-                    lodtensor = None
-            else:
-                lodtensor._share_data_with(tensor.get_tensor())
-
-            return lodtensor
-
         feed_list = []
         no_data_ids = []
         # If the feed_var is None, its feed_name should be deleted.
@@ -2503,8 +2502,10 @@ def set_state_dict(self, state_dict: dict[str, Tensor]) -> None:
                 if k in self._structured_to_parameter_name
                 else k
             )
-            local_state_dict[param_name] = v._local_value()
-        dist_main_program.set_state_dict(local_state_dict)
+            local_state_dict[param_name] = _to_lodtensor(v._local_value())
+        dist_main_program.set_state_dict(
+            local_state_dict, paddle.static.global_scope()
+        )
 
 
 def to_static(
