support shard master weight in stage1. (#67748)

winter-wang · web-flow · commit e96eb549f18c · 2024-08-28T18:00:19.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -37,9 +37,6 @@
 from paddle.distributed.auto_parallel.interface import (
     shard_tensor as shard_tensor_static,
 )
-from paddle.distributed.auto_parallel.placement_type import (
-    to_placements,
-)
 from paddle.distributed.auto_parallel.process_mesh import ProcessMesh
 from paddle.distributed.auto_parallel.static.completion import (
     mark_as_sharding_propagation_skip_op,
@@ -60,7 +57,12 @@
 )
 from paddle.optimizer import Optimizer
 
-from .placement_type import check_placements_equal, get_shard_spec
+from .placement_type import (
+    check_placements_equal,
+    get_shard_spec,
+    to_dim_map,
+    to_placements,
+)
 from .random import determinate_rng, rng_state
 
 if TYPE_CHECKING:
@@ -1035,8 +1037,13 @@ def _set_and_check_sharding_prop_from_param(self):
     def _shard_accumulator(self, param):
         target_name = param.name
         if param.name in self._inner_opt._master_weights.keys():
-            target_name = self._inner_opt._master_weights[param.name].name
-
+            master_weight = self._inner_opt._master_weights[param.name]
+            target_name = master_weight.name
+            # shard the master weight
+            if self._shard_fn is not None:
+                self._inner_opt._master_weights[param.name] = (
+                    self._shard_fn.shard_master_weight(param, master_weight)
+                )
         # shard the accumulators
         for key in self._inner_opt._accumulators.keys():
             accumulator = self._inner_opt._accumulators[key][target_name]
@@ -1205,6 +1212,37 @@ def __init__(self, mesh):
     def _set_sharding_mesh_axis(self, sharding_mesh_axis):
         self._sharding_mesh_axis = sharding_mesh_axis
 
+    def shard_master_weight(
+        self, param: Tensor, master_weight: Tensor
+    ) -> Tensor:
+        if param.is_dist():
+            if isinstance(master_weight, pir.Value):
+                data_op = master_weight.get_defining_op()
+                assert (
+                    data_op.name() == "pd_op.data"
+                ), "The master weight must be a result of data op."
+                placements = get_placement_with_sharding(
+                    param, self._sharding_mesh_axis
+                )
+                dim_map, partial_status = to_dim_map(
+                    placements, len(master_weight.shape)
+                )
+                dist_attr = (
+                    paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                        param.process_mesh, dim_map, partial_status
+                    )
+                )
+                dist_type = paddle.base.libpaddle.pir.cvt_to_dist_type(
+                    master_weight.type(), dist_attr
+                )
+                master_weight.set_type(dist_type)
+                data_op.dist_attr = (
+                    paddle.base.libpaddle.pir.create_op_dist_attribute(
+                        param.process_mesh, [], [dist_attr]
+                    )
+                )
+        return master_weight
+
 
 class ShardingStage1(_ShardingStageBase):
     """
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -26,7 +26,7 @@
 
 import paddle
 import paddle.distributed.auto_parallel.static.utils as auto_utils
-from paddle import static, utils
+from paddle import pir, static, utils
 from paddle.base.executor import _to_name_str
 from paddle.distributed import fleet
 from paddle.distributed.passes.pass_base import new_pass
@@ -1299,6 +1299,7 @@ def _initialize(self, mode, init_parameters=True):
             self.program_helper.init_pir(
                 self._pir_dist_main_progs[mode], self._place
             )
+            changed_ouput_op_list = []
             if self._executor is None:
                 self._executor = paddle.static.Executor(self._place)
                 startup_prog = self._startup_progs[mode].clone()
@@ -1335,18 +1336,33 @@ def _initialize(self, mode, init_parameters=True):
                         if local_shape != global_shape:
                             src_value = op.operand_source(0)
                             assert src_value.shape == global_shape
-                            initial_op = src_value.get_defining_op()
-                            if initial_op.name() == "pd_op.full":
-                                initial_op.set_int_array_attr(
-                                    "shape", local_shape
+                            dst_dist_attr = name_map_value[var_name].dist_attr()
+                            if not src_value.is_dist():
+                                src_dist_attr = paddle.base.libpaddle.pir.create_tensor_dist_attribute(
+                                    dst_dist_attr.process_mesh,
+                                    [-1] * len(src_value.shape),
+                                    {},
                                 )
                                 src_value.set_type(
-                                    name_map_value[var_name].type()
+                                    paddle.base.libpaddle.pir.cvt_to_dist_type(
+                                        src_value.type(), src_dist_attr
+                                    )
                                 )
-                            # initial_op.name() == "pd_op.cast": # master_weight
-
+                            pir.set_insertion_point_after(
+                                src_value.get_defining_op()
+                            )
+                            reshard_var = paddle._C_ops.reshard_v2(
+                                src_value, dst_dist_attr
+                            )
+                            if src_value.persistable:
+                                src_value.persistable = False
+                                changed_ouput_op_list.append(op)
+                            op.operand(0).set_source(reshard_var)
                 for del_op in del_ops:
                     del_op.erase()
+                apply_reshard_pass(startup_prog)
+                for op in changed_ouput_op_list:
+                    op.operand_source(0).persistable = True
                 self._executor.run(startup_prog)
                 if self._job_plan is not None:
                     # pipeline scheduling should be enabled after running
@@ -2115,7 +2131,7 @@ def run(
             loss_value = program_for_executor.get_output_value_by_name(
                 self._loss_names[0]
             )
-            if paddle.pir.is_fake_value(loss_value):
+            if pir.is_fake_value(loss_value):
                 no_fetch = True
                 fetch_names = []
             else:
