PaddlePaddle
diff --git a/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+24-19Lines changed: 24 additions & 19 deletions b/‎python/paddle/distributed/auto_parallel/api.py
Copy file name to clipboard
+24-19Lines changed: 24 additions & 19 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/utils.py
Copy file name to clipboard
+40-17Lines changed: 40 additions & 17 deletions b/‎python/paddle/distributed/auto_parallel/static/utils.py
Copy file name to clipboard
+40-17Lines changed: 40 additions & 17 deletions
diff --git a/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎test/auto_parallel/hybrid_strategy/CMakeLists.txt
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
@@ -2242,10 +2242,19 @@ def state_dict(self, mode="all"):
                 'all' : The return value contains the variable in the network and optimizer.
                 Default: 'all'
         """
-        local_state_dict = self.dist_main_program(
-            mode=self._engine._mode
-        ).state_dict(mode)
+
+        if use_pir_api():
+            scope = paddle.static.global_scope()
+            local_state_dict = self._engine._pir_dist_main_progs[
+                self._engine._mode
+            ].state_dict(mode, scope)
+        else:
+            local_state_dict = self.dist_main_program(
+                mode=self._engine._mode
+            ).state_dict(mode)
+
         dist_state_dict = self._build_distributed_state_dict(local_state_dict)
+
         mapping_names = [
             (
                 self._parameter_to_structured_name[k]
@@ -2264,10 +2273,18 @@ def _build_distributed_state_dict(self, local_state_dict):
         Args:
             local_state_dict(Dict[str, libpaddle.Tensor]): The state dict from program.
         """
-        dist_main_program = self.dist_main_program(mode=self._engine._mode)
-        dist_context = self._engine._dist_contexts[self._mode]
-        # Dict[var.name, Dict["process_shape": process_mesh.shape, "process_group": process_mesh.process_ids, "dims_mapping": dims_mapping]]
-        dist_attrs = get_dist_attr(dist_main_program, dist_context)
+        if use_pir_api():
+            dist_main_program = self._engine._pir_dist_main_progs[
+                self._engine._mode
+            ]
+            dist_attrs = get_dist_attr(dist_main_program)
+
+        else:
+            dist_main_program = self.dist_main_program(mode=self._engine._mode)
+            # Dict[var.name, Dict["process_shape": process_mesh.shape, "process_group": process_mesh.process_ids, "dims_mapping": dims_mapping]]
+            dist_attrs = get_dist_attr(
+                dist_main_program, self._engine._dist_contexts[self._mode]
+            )
 
         def build_distributed_tensor(local_tensor, dist_attr):
             assert isinstance(
@@ -2282,18 +2299,6 @@ def build_distributed_tensor(local_tensor, dist_attr):
                 dist_attr["dims_mapping"]
             ), f"local tensor shape {local_tensor.shape} not equal to dims_mapping shape {dist_attr['dims_mapping']}."
             global_shape = local_tensor.shape
-            for i, dim in enumerate(dist_attr["dims_mapping"]):
-                assert dim >= -1 and dim < len(
-                    local_tensor.shape
-                ), f"dim {dim} out of range."
-                if dim == -1:
-                    continue
-                elif dim >= 0:
-                    global_shape[i] = (
-                        dist_attr["process_shape"][dim] * local_tensor.shape[i]
-                    )
-                else:
-                    raise ValueError(f"dim {dim} is not supported.")
             mesh = ProcessMesh(
                 np.array(dist_attr["process_group"]).reshape(
                     dist_attr["process_shape"]
 
@@ -22,6 +22,7 @@
 import numpy as np
 
 import paddle
+from paddle.base.framework import use_pir_api
 from paddle.base.wrapped_decorator import wrap_decorator
 from paddle.framework import core
 from paddle.framework.io_utils import is_belong_to_optimizer, is_parameter
@@ -830,24 +831,46 @@ def get_dist_attr(program, dist_context=None):
     Args:
         program(Program): main program for training
     """
-    from .dist_context import get_default_distributed_context
-
-    assert isinstance(program, paddle.static.Program)
-    if dist_context is None:
-        dist_context = get_default_distributed_context()
     dist_attr = {}
-    for var in program.list_vars():
-        if is_parameter(var) or is_belong_to_optimizer(var):
-            tensor_dist_attr = dist_context.get_tensor_dist_attr_for_program(
-                var
-            )
-            process_mesh = tensor_dist_attr.process_mesh
-            dims_mapping = tensor_dist_attr.dims_mapping
-            dist_attr[var.name] = {
-                "process_shape": process_mesh.shape,
-                "process_group": process_mesh.process_ids,
-                "dims_mapping": dims_mapping,
-            }
+    if use_pir_api():
+        ops = program.global_block().ops
+        for op in ops:
+            if op.name() == "builtin.parameter" or (
+                op.name() == "pd_op.data"
+                and op.has_attr("persistable")
+                and op.attrs()["persistable"]
+            ):
+                op_dist_attr = op.dist_attr
+                var_dist_attr = op_dist_attr.result(0).as_tensor_dist_attr()
+                var_name = (
+                    op.str_attr("parameter_name")
+                    if op.name() == "builtin.parameter"
+                    else op.str_attr("name")
+                )
+                process_mesh = var_dist_attr.process_mesh
+                dist_attr[var_name] = {
+                    "process_shape": process_mesh.shape,
+                    "process_group": process_mesh.process_ids,
+                    "dims_mapping": var_dist_attr.dims_mapping,
+                }
+    else:
+        from .dist_context import get_default_distributed_context
+
+        assert isinstance(program, paddle.static.Program)
+        if dist_context is None:
+            dist_context = get_default_distributed_context()
+        for var in program.list_vars():
+            if is_parameter(var) or is_belong_to_optimizer(var):
+                tensor_dist_attr = (
+                    dist_context.get_tensor_dist_attr_for_program(var)
+                )
+                process_mesh = tensor_dist_attr.process_mesh
+                dims_mapping = tensor_dist_attr.dims_mapping
+                dist_attr[var.name] = {
+                    "process_shape": process_mesh.shape,
+                    "process_group": process_mesh.process_ids,
+                    "dims_mapping": dims_mapping,
+                }
     return dist_attr
 
 
 
@@ -118,3 +118,15 @@ if((WITH_GPU) AND (LINUX))
   set_tests_properties(test_semi_auto_llama_acc_align
                        PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
 endif()
+
+if((WITH_GPU) AND (LINUX))
+  py_test_modules(
+    test_semi_auto_llama_save_load
+    MODULES
+    test_semi_auto_llama_save_load
+    ENVS
+    "http_proxy=;https_proxy=;PYTHONPATH=../..:${PADDLE_BINARY_DIR}/python;FLAGS_enable_pir_api=1"
+  )
+  set_tests_properties(test_semi_auto_llama_save_load
+                       PROPERTIES TIMEOUT "180" LABELS "RUN_TYPE=HYBRID")
+endif()