fix bugs

Author: JiayiFeng

paddle/fluid/framework/details/data_balance_op_handle.cc

@@ -20,10 +20,24 @@ namespace paddle {
 namespace framework {
 namespace details {
 
+#ifdef PADDLE_WITH_CUDA
+DataBalanceOpHandle::DataBalanceOpHandle(
+    const std::vector<Scope *> &local_scopes,
+    const std::vector<platform::Place> &places,
+    const platform::NCCLContextMap *ctxs)
+    : local_scopes_(local_scopes), places_(places) {
+  if (ctxs) {
+    for (auto &p : places_) {
+      this->dev_ctxes_[p] = ctxs->DevCtx(p);
+    }
+  }
+}
+#else
 DataBalanceOpHandle::DataBalanceOpHandle(
     const std::vector<Scope *> &local_scopes,
     const std::vector<platform::Place> &places)
     : local_scopes_(local_scopes), places_(places) {}
+#endif
 
 std::string DataBalanceOpHandle::Name() const { return "data balance"; }
 
@@ -104,6 +118,7 @@ void DataBalanceOpHandle::RunImpl() {
     }
   }
   const auto &balance_plan = GetBalancePlan(device_sizes);
+
   for (const auto &trans : balance_plan) {
     for (int data_idx = 0; data_idx < data_num; ++data_idx) {
       LoDTensor *src_tensor = lod_tensors[data_idx][trans[0]];

paddle/fluid/framework/details/data_balance_op_handle.h

@@ -19,15 +19,24 @@
 #include "paddle/fluid/framework/details/op_handle_base.h"
 #include "paddle/fluid/framework/lod_tensor.h"
 #include "paddle/fluid/framework/scope.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/fluid/platform/nccl_helper.h"
+#endif
 
 namespace paddle {
 namespace framework {
 namespace details {
 
 struct DataBalanceOpHandle : public OpHandleBase {
  public:
+#ifdef PADDLE_WITH_CUDA
   DataBalanceOpHandle(const std::vector<Scope *> &local_scopes,
-                      const std::vector<platform::Place> &places);
+                      const std::vector<platform::Place> &places,
+                      const platform::NCCLContextMap *ctxs);
+#else
+  DataBalanceOpHandle(const std::vector<Scope *> &local_scopes,
+                      const std::vector<platform::Place> *places)
+#endif
 
   std::string Name() const override;
 

paddle/fluid/framework/details/multi_devices_graph_builder.cc

@@ -368,7 +368,12 @@ void MultiDevSSAGraphBuilder::InsertAllReduceOp(SSAGraph *result,
 
 void MultiDevSSAGraphBuilder::InsertDataBalanceOp(
     SSAGraph *result, const std::vector<std::string> &datas) const {
+#ifdef PADDLE_WITH_CUDA
+  result->ops_.emplace_back(
+      new DataBalanceOpHandle(local_scopes_, places_, nccl_ctxs_));
+#else
   result->ops_.emplace_back(new DataBalanceOpHandle(local_scopes_, places_));
+#endif
   auto *op_handle = result->ops_.back().get();
   for (size_t i = 0; i < places_.size(); ++i) {
     auto &p = places_[i];

paddle/fluid/framework/details/op_handle_base.cc

@@ -60,6 +60,7 @@ void OpHandleBase::RecordWaitEventOnCtx(platform::DeviceContext *waited_ctx) {
 #ifdef PADDLE_WITH_CUDA
   if (platform::is_cpu_place(waited_ctx->GetPlace()) || events_.empty()) {
     for (auto &dev_ctx : dev_ctxes_) {
+      PADDLE_ENFORCE_NOT_NULL(dev_ctx.second);
       dev_ctx.second->Wait();
     }
   } else {