multi_devices_graph_builder.cc 6.4 KB
Newer Older
Y
Yu Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "paddle/fluid/framework/details/multi_devices_graph_builder.h"
#include "paddle/fluid/framework/details/computation_op_handle.h"
#include "paddle/fluid/framework/details/scale_loss_grad_op_handle.h"
#include "paddle/fluid/framework/scope.h"
Y
Yu Yang 已提交
19 20 21 22

#ifdef PADDLE_WITH_CUDA
#include "paddle/fluid/framework/details/nccl_all_reduce_op_handle.h"
#endif
Y
Yu Yang 已提交
23

Y
Yu Yang 已提交
24 25 26
#include <string>
#include <vector>

Y
Yu Yang 已提交
27 28 29
namespace paddle {
namespace framework {
namespace details {
Y
Yu Yang 已提交
30 31

#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
32 33 34 35 36 37 38 39 40 41
MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
    const std::vector<platform::Place> &places,
    const std::string &loss_var_name,
    const std::unordered_set<std::string> &params,
    const std::vector<Scope *> &local_scopes,
    platform::NCCLContextMap *nccl_ctxs)
    : loss_var_name_(loss_var_name),
      places_(places),
      local_scopes_(local_scopes),
      nccl_ctxs_(nccl_ctxs) {
Y
Yu Yang 已提交
42 43 44 45 46 47 48 49 50 51
#else
MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
    const std::vector<platform::Place> &places,
    const std::string &loss_var_name,
    const std::unordered_set<std::string> &params,
    const std::vector<Scope *> &local_scopes)
    : loss_var_name_(loss_var_name),
      places_(places),
      local_scopes_(local_scopes) {
#endif
Y
Yu Yang 已提交
52 53 54 55 56
  for (auto &p : params) {
    grad_names_.insert(GradVarName(p));
  }
}

Y
Yu Yang 已提交
57 58 59
std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
    const ProgramDesc &program) const {
  auto graph = new SSAGraph();
Y
Yu Yang 已提交
60
  SSAGraph &result = *graph;
C
chengduoZH 已提交
61
  std::unordered_set<std::string> og_has_been_broadcast;
Y
Yu Yang 已提交
62 63 64 65 66

  // We cannot invoke resize. It is a bug of GCC 4.8
  result.vars_ = std::vector<
      std::unordered_map<std::string, std::vector<std::unique_ptr<VarHandle>>>>(
      places_.size());
Y
Yu Yang 已提交
67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84

  bool is_forwarding = true;
  for (auto *op : program.Block(0).AllOps()) {
    bool change_forward = false;
    if (!is_forwarding) {
      // FIXME(yy): Do not hard code like this
      if (op->OutputArgumentNames().size() == 1 &&
          op->OutputArgumentNames()[0] == GradVarName(loss_var_name_)) {
        continue;  // Drop fill 1. for backward coeff;
      }
    }

    for (size_t i = 0; i < places_.size(); ++i) {
      auto &p = places_[i];
      auto *s = local_scopes_[i];

      result.ops_.emplace_back(new ComputationOpHandle(*op, s, p));
      auto *op_handle = result.ops_.back().get();
Y
Yu Yang 已提交
85
      op_handle->dev_ctxes_[p] = const_cast<platform::DeviceContext *>(
Y
Yu Yang 已提交
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
          platform::DeviceContextPool::Instance().Get(p));

      auto var_names = op->InputArgumentNames();

      for (auto &each_var_name : var_names) {
        VarHandle *var =
            CreateOrGetLatestVarHandle(&result, each_var_name, p, i);
        op_handle->AddInput(var);
      }
      var_names = op->OutputArgumentNames();

      for (auto &each_var_name : var_names) {
        CreateOpOutput(&result, op_handle, each_var_name, p, i);
      }

      if (is_forwarding) {
        if (var_names.size() == 1 && var_names[0] == loss_var_name_) {
Y
Yu Yang 已提交
103 104 105 106 107 108 109 110
// Insert ScaleCost OpHandle
#ifdef PADDLE_WITH_CUDA
          auto *communication_dev_ctx = nccl_ctxs_->DevCtx(p);
#else
          auto *communication_dev_ctx =
              platform::DeviceContextPool::Instance().Get(platform::CPUPlace());
#endif

Y
Yu Yang 已提交
111
          op_handle = new ScaleLossGradOpHandle(local_scopes_.size(), s, p,
Y
Yu Yang 已提交
112
                                                communication_dev_ctx);
Y
Yu Yang 已提交
113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
          result.ops_.emplace_back(op_handle);

          // FIXME: Currently ScaleLossGradOp only use device_count as scale
          // factor. So it does not depend on any other operators.
          // VarHandle *loss = GetVarHandle(loss_var_name, place);
          // loss->pending_ops_.emplace_back(op_handle);
          // op_handle->inputs_.emplace_back(loss);

          CreateOpOutput(&result, op_handle, GradVarName(loss_var_name_), p, i);
          change_forward = true;
        }
      }
    }

    if (change_forward) {
      is_forwarding = false;
    }

    if (!is_forwarding) {
      auto var_names = op->OutputArgumentNames();
C
chengduoZH 已提交
133 134 135 136
      // Currently, we assume that once gradient is generated, it can be
      // broadcast, and each gradient is only broadcast once. But there are no
      // other cases, for example, we need to adjust the gradient according to
      // the input when we get the gradient, which is not considered at present.
Y
Yu Yang 已提交
137
      for (auto &og : var_names) {
C
chengduoZH 已提交
138
        if (grad_names_.count(og) != 0 &&
C
chengduoZH 已提交
139 140 141
            og_has_been_broadcast.count(og) == 0) {  // is param grad
                                                     // Insert NCCL AllReduce Op
          og_has_been_broadcast.insert(og);
Y
Yu Yang 已提交
142
#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
143 144 145 146 147 148 149 150 151 152 153
          result.ops_.emplace_back(
              new NCCLAllReduceOpHandle(local_scopes_, places_, *nccl_ctxs_));
          auto *op_handle = result.ops_.back().get();

          for (size_t i = 0; i < places_.size(); ++i) {
            auto &p = places_[i];
            auto &vars = result.vars_[i][og];

            if (vars.empty()) {  // This device has no data. continue.
              continue;
            }
Y
Yu Yang 已提交
154 155
            auto &prev_grad = vars[vars.size() - 1];
            op_handle->AddInput(prev_grad.get());
Y
Yu Yang 已提交
156

Y
Yu Yang 已提交
157 158 159 160 161
            vars.emplace_back(new VarHandle);
            auto &var = vars.back();
            var->place_ = p;
            var->name_ = og;
            var->version_ = vars.size() - 1;
Y
Yu Yang 已提交
162

Y
Yu Yang 已提交
163
            op_handle->AddOutput(var.get());
Y
Yu Yang 已提交
164
          }
Y
Yu Yang 已提交
165 166 167
#else
          PADDLE_ENFORCE("Not implemented");
#endif
Y
Yu Yang 已提交
168 169 170 171 172 173 174 175 176 177
        }
      }
    }
  }

  /*
    Dependency graph has been constructed. However, there are still data
    harzaeds need to be handled.
   */
  PolishGraphToSupportDataHazards(&result);
Y
Yu Yang 已提交
178

Y
Yu Yang 已提交
179 180 181 182 183
  /*
   * Only variables should be the leaves of graph.
   */
  AddOutputToLeafOps(&result);

Y
Yu Yang 已提交
184 185 186 187 188 189
  if (VLOG_IS_ON(10)) {
    std::ostringstream sout;
    PrintGraphviz(*graph, sout);
    VLOG(10) << sout.str();
  }

Y
Yu Yang 已提交
190
  return std::unique_ptr<SSAGraph>(graph);
Y
Yu Yang 已提交
191
}  // namespace details
Y
Yu Yang 已提交
192 193 194
}  // namespace details
}  // namespace framework
}  // namespace paddle