fix shard_optimizer bug (#68781)

icpcccpc · web-flow · commit 39a0db67a6bd · 2024-11-01T19:39:58.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1043,6 +1043,10 @@ def _set_and_check_sharding_prop_from_param(self):
         ):
             self._sharding_degree = self._shard_fn._mesh.get_dim_size(0)
             self._sharding_mesh_axis = 0
+        elif (self._shard_fn._mesh is not None) and (
+            'dp' in self._shard_fn._mesh.dim_names
+        ):
+            self._sharding_degree = self._shard_fn._mesh.get_dim_size('dp')
         else:
             param_list = self._inner_opt._parameter_list
             for param in param_list:
