upgrade matmul_grad inferspmd to support shard optimizer stage3. (#66098)

winter-wang · web-flow · commit e42b101af61d · 2024-07-18T10:44:18.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/matmul.cc b/paddle/phi/infermeta/spmd_rules/matmul.cc
@@ -285,11 +285,12 @@ static bool DistAttrsAreBasicallyEqual(
           in_dist_attr.partial_status() == out_dist_attr.partial_status());
 }
 
-SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
-                             const DistMetaTensor& y,
+SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x_,
+                             const DistMetaTensor& y_,
                              const DistMetaTensor& out_grad,
                              bool trans_x,
                              bool trans_y) {
+  DistMetaTensor x = x_, y = y_;
   auto get_attr = [](const ArgDistAttr& attr) -> const TensorDistAttr& {
     return paddle::get<TensorDistAttr>(attr);
   };
@@ -335,22 +336,18 @@ SpmdInfo MatmulGradInferSpmd(const DistMetaTensor& x,
   auto fwd_spmd_info = MatmulInferSpmd(x, y, trans_x, trans_y);
   auto infer_x_dist_attr = get_attr(fwd_spmd_info.first[0]);
   auto infer_y_dist_attr = get_attr(fwd_spmd_info.first[1]);
-  auto is_dist_attr_equal = [&](const TensorDistAttr& dist_attr,
-                                const ArgDistAttr& arg_dist_attr) -> bool {
-    const auto& infer_dist_attr = get_attr(arg_dist_attr);
+  auto is_dist_attr_not_equal =
+      [&](const TensorDistAttr& dist_attr,
+          const TensorDistAttr& infer_dist_attr) -> bool {
     return (dist_attr.process_mesh() != infer_dist_attr.process_mesh() ||
             dist_attr.dims_mapping() != infer_dist_attr.dims_mapping() ||
             dist_attr.partial_status() != infer_dist_attr.partial_status());
   };
-
-  if (is_dist_attr_equal(x.dist_attr(), fwd_spmd_info.first[0]) ||
-      is_dist_attr_equal(y.dist_attr(), fwd_spmd_info.first[1])) {
-    auto x_r_dist_attr = GetReplicatedDistAttr(x.dist_attr());
-    auto y_r_dist_attr = GetReplicatedDistAttr(y.dist_attr());
-    return {{x_r_dist_attr,
-             y_r_dist_attr,
-             GetReplicatedDistAttr(out_grad.dist_attr())},
-            {x_r_dist_attr, y_r_dist_attr}};
+  if (is_dist_attr_not_equal(x.dist_attr(), infer_x_dist_attr)) {
+    x = DistMetaTensor(x.dims(), infer_x_dist_attr);
+  }
+  if (is_dist_attr_not_equal(y.dist_attr(), infer_y_dist_attr)) {
+    y = DistMetaTensor(y.dims(), infer_y_dist_attr);
   }
 
   SpmdInfo dx_spmd_info;
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -1327,6 +1327,14 @@ def __call__(self, key, param, accumulator):
             # Only deal with momentum in optimizer, beta should be replicated cross param's mesh
             if 'beta' not in key:
                 placements = param.placements
+                if all(
+                    isinstance(placement, dist.Replicate)
+                    for placement in placements
+                ):
+                    placements = get_placement_with_sharding(
+                        param, self._sharding_mesh_axis
+                    )
+
             else:
                 placements = [
                     dist.Replicate()
diff --git a/test/auto_parallel/hybrid_strategy/semi_auto_parallel_simple_net_dp_mp.py b/test/auto_parallel/hybrid_strategy/semi_auto_parallel_simple_net_dp_mp.py
@@ -52,11 +52,11 @@ def test_dp_mp_demo_net(self):
             self.dp_mp_parameters,
         ) = self.run_dynamic(model, shard_input=True)
 
-        self.check_tensor_eq(self.dp_mp_loss, self.base_loss)
+        self.check_tensor_eq(self.dp_mp_loss, self.base_loss, rtol=1e-04)
         for param, param_base in zip(
             self.dp_mp_parameters, self.base_parameters
         ):
-            self.check_tensor_eq(param, param_base)
+            self.check_tensor_eq(param, param_base, atol=1e-06)
             self.check_tensor_eq(param.grad, param_base.grad)
 
         # save load
