remove unuseful c_allgather op for pir autp parallel. (#64465)

winter-wang · web-flow · commit cc38216d6625 · 2024-05-21T14:55:28.000+08:00
diff --git a/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc b/paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
@@ -122,6 +122,22 @@ void VerifyDenseBlock(pir::Block* block) {
   }
 }
 
+void RemoveUnusefulCallgatherOp(pir::Block* block) {
+  std::vector<pir::Operation*> del_ops;
+  for (auto& op : *block) {
+    if (op.isa<CAllgatherOp>()) {
+      auto nrank = op.attribute<pir::Int32Attribute>("nranks").data();
+      if (nrank == 1) {
+        op.result(0).ReplaceAllUsesWith(op.operand_source(0));
+        del_ops.emplace_back(&op);
+      }
+    }
+  }
+  for (auto op : del_ops) {
+    op->Erase();
+  }
+}
+
 std::shared_ptr<pir::Program> DistToDensePass(pir::Program* prog) {
   if (FLAGS_print_ir) {
     VLOG(0) << "IR before DistToDense Pass = " << *prog;
@@ -135,6 +151,7 @@ std::shared_ptr<pir::Program> DistToDensePass(pir::Program* prog) {
   ctx->GetOrRegisterDialect<DistDialect>();
 
   ProcessDistBlock(new_prog->block());
+  RemoveUnusefulCallgatherOp(new_prog->block());
   VLOG(6) << "IR before VerifyDenseBlock Pass = " << *new_prog;
   VerifyDenseBlock(new_prog->block());
 
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -28,7 +28,7 @@
 from paddle.distributed import fleet
 from paddle.framework import (
     IrGraph,
-    _current_expected_place as _get_device,
+    _current_expected_place_ as _get_device,
     core,
     in_dynamic_mode,
 )
@@ -714,6 +714,10 @@ def _prepare_program(self, mode, init_parameters=True):
         # TODO(zhiqiu): fit the processes below for pir
         if self._in_pir_mode:
             self._parallel_pir(mode)
+            # Init comm
+            self._init_comm()
+            # startup program
+            self._initialize(mode, init_parameters)
             self._has_prepared[mode] = True
             return
         # Do the planning process
@@ -1024,22 +1028,35 @@ def _init_comm(self):
                 for process_group in all_process_groups:
                     process_group.instantiate()
 
-    def _init_lr(self):
+    def _init_lr(self, main_program):
         # hack to find learning_rate op
         lr_name = None
-        for op in self.main_program.global_block().ops:
+        for op in main_program.global_block().ops:
             if (
                 op.name() == "pd_op.data"
                 and 'learning_rate' in op.attrs()["name"]
             ):
                 lr_name = op.attrs()["name"]
                 break
+            if (
+                op.name() == "builtin.parameter"
+                and 'learning_rate' in op.attrs()["parameter_name"]
+            ):
+                lr_name = op.attrs()["parameter_name"]
+                break
 
         if lr_name is not None:
             buffer_tensor = global_scope().var(lr_name).get_tensor()
-            buffer_tensor.set(
-                np.float32(self._optimizer._learning_rate), self._place
-            )
+            from paddle.optimizer.lr import LRScheduler
+
+            if isinstance(self._optimizer._learning_rate, float):
+                buffer_tensor.set(
+                    np.float32(self._optimizer._learning_rate), self._place
+                )
+            elif isinstance(self._optimizer._learning_rate, LRScheduler):
+                buffer_tensor.set(
+                    np.float32(self._optimizer._learning_rate()), self._place
+                )
 
     def _initialize(self, mode, init_parameters=True):
         self._place = _get_device()
@@ -1058,10 +1075,12 @@ def _initialize(self, mode, init_parameters=True):
             # 6. vpp init adaption
 
             self.program_helper.init_pir(
-                self._pir_dense_main_progs[mode], self._place
+                self._pir_dist_main_progs[mode], self._place
             )
 
-            self._init_lr()
+            self._init_lr(self._pir_dense_main_progs[mode])
+            if self._executor is None:
+                self._executor = paddle.static.Executor(self._place)
             return
 
         if self._strategy.seed:
diff --git a/python/paddle/distributed/auto_parallel/static/pir_pass.py b/python/paddle/distributed/auto_parallel/static/pir_pass.py
@@ -56,7 +56,7 @@ def reshard_combine_value(op, operand, attr):
     reshard_vars = []
     for inner_operand, inner_attr in zip(combine_op.operands(), array_attr):
         reshard_vars.append(reshard_single_value(op, inner_operand, inner_attr))
-    paddle.pir.set_insertion_point(combine_op)
+    paddle.pir.set_insertion_point(op)
     return paddle._C_ops.builtin_combine(reshard_vars)
 
 
