fix bug in pir mode fp16. (#66674)

winter-wang · web-flow · commit 4d135d27a9e6 · 2024-07-29T16:47:01.000+08:00
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
@@ -203,38 +203,52 @@ bool HasDistInput(const std::vector<pir::Value>& inputs,
   return false;
 }
 
+// return the tensor dist attribute of converted value.
+pir::Attribute CvtValueToDist(pir::Value value,
+                              ProcessMeshAttribute mesh_attr) {
+  if (!value || !value.type()) return nullptr;
+  PADDLE_ENFORCE_EQ(
+      value.type().isa<DistTypeInterface>(),
+      false,
+      common::errors::InvalidArgument("Can't convert value type to dist. "
+                                      "Because it is already a dist type"));
+  auto dense_type = value.type().dyn_cast<pir::DenseTensorType>();
+  if (!dense_type) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Currently only support convert dense_tensor_type to dist type."));
+  }
+  auto ctx = pir::IrContext::Instance();
+  auto dist_type = DistDenseTensorType::get(ctx, dense_type, mesh_attr);
+  value.set_type(dist_type);
+  if (auto op = value.defining_op()) {
+    PADDLE_ENFORCE_EQ(
+        op->num_results(),
+        1u,
+        common::errors::InvalidArgument(
+            "Currently only allowed add dist attribue for operation with "
+            "single output. The current op is %s",
+            op->name()));
+    std::vector<pir::Attribute> dist_operand_attrs;
+    for (auto pre_value : op->operands_source()) {
+      dist_operand_attrs.push_back(CvtValueToDist(pre_value, mesh_attr));
+    }
+    op->set_attribute(
+        kAttrOpDistAttr,
+        OperationDistAttribute::get(ctx,
+                                    mesh_attr,
+                                    dist_operand_attrs,
+                                    {dist_type.tensor_dist_attr()}));
+  }
+  return dist_type.tensor_dist_attr();
+}
+
 void CvtAllInputsToDist(const std::vector<pir::Value>& inputs,
                         ProcessMeshAttribute mesh_attr) {
   for (auto value : inputs) {
     if (auto type = value.type()) {
       if (type.isa<DistTypeInterface>() || type.isa<pir::VectorType>())
         continue;
-      auto dense_type = type.dyn_cast<pir::DenseTensorType>();
-      if (!dense_type) {
-        PADDLE_THROW(common::errors::Unimplemented(
-            "Currently only support convert dense_tensor_type to dist type."));
-      }
-      auto ctx = pir::IrContext::Instance();
-      auto dist_type = DistDenseTensorType::get(ctx, dense_type, mesh_attr);
-      value.set_type(dist_type);
-      if (auto define_op = value.defining_op()) {
-        if (define_op->num_operands() != 0u) {
-          PADDLE_THROW(common::errors::InvalidArgument(
-              "Currently only allowed add dist attribue for leaf nodes "
-              "operation. The current op is %s",
-              define_op->name()));
-        }
-        if (define_op->num_results() != 1u) {
-          PADDLE_THROW(common::errors::InvalidArgument(
-              "Currently only allowed add dist attribue for operation with "
-              "single output. The current op is %s",
-              define_op->name()));
-        }
-        define_op->set_attribute(
-            kAttrOpDistAttr,
-            OperationDistAttribute::get(
-                ctx, mesh_attr, {}, {dist_type.tensor_dist_attr()}));
-      }
+      CvtValueToDist(value, mesh_attr);
     }
   }
 }
diff --git a/paddle/phi/infermeta/spmd_rules/optimizer.cc b/paddle/phi/infermeta/spmd_rules/optimizer.cc
@@ -98,11 +98,7 @@ SpmdInfo AdamInferSpmdDynamic(const DistMetaTensor& param,
         errors::InvalidArgument(
             "skip_update should be replicated, but got shard on mesh %d.",
             skip_update_dist_attr.dims_mapping()[0]));
-    PADDLE_ENFORCE_EQ(
-        skip_update_dist_attr.partial_status().size(),
-        0,
-        errors::InvalidArgument("skip_update should be replicated, but got "
-                                "patial status not empty"));
+    skip_update_dist_attr.clean_partial_status();
     if (skip_update_dist_attr.process_mesh().ndim() > 1 &&
         phi::distributed::IsSubMesh(skip_update_dist_attr.process_mesh(),
                                     param_dist_attr.process_mesh())) {
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -296,7 +296,12 @@ def remove_other_rank_input_output_pass(dist_program):
 
 
 # Note: this is the pass in the dense program
-comm_ops = ["pd_op.c_allreduce_sum", "pd_op.all_gather"]
+comm_ops = [
+    "pd_op.c_allreduce_sum",
+    "pd_op.all_gather",
+    "pd_op.c_allreduce_max",
+    "pd_op.c_reducescatter",
+]
 
 
 def remove_unuseful_comm_op_pass(program):
