support sharding of master weights in dynamic mode

PaddlePaddle · Mar 9, 2025 · 0735da2 · 0735da2
1 parent bcfa081
commit 0735da2
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1167,6 +1167,8 @@ def _shard_accumulator(self, param):
                 self._inner_opt._master_weights[param.name] = (
                     self._shard_fn.shard_master_weight(param, master_weight)
                 )
+                self._inner_opt._master_weights[param.name].name = target_name
+
         # shard the accumulators
         for key in self._inner_opt._accumulators.keys():
             accumulator = self._inner_opt._accumulators[key][target_name]
@@ -1414,14 +1416,12 @@ def shard_master_weight(
         self, param: Tensor, master_weight: Tensor
     ) -> Tensor:
         if param.is_dist():
+            placements = get_placement_with_sharding(param, self._sharding_axis)
             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_axis
-                )
                 dim_map, partial_status = to_dim_map(
                     placements, len(master_weight.shape)
                 )
@@ -1439,6 +1439,13 @@ def shard_master_weight(
                         param.process_mesh, [], [dist_attr]
                     )
                 )
+
+            if paddle.in_dynamic_mode() and master_weight.is_dist():
+                master_weight = reshard(
+                    master_weight,
+                    mesh=param.process_mesh,
+                    placements=placements,
+                )
         return master_weight