multi_devices_graph_builder.cc 7.3 KB
Newer Older
Y
Yu Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
//   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"
T
wip  
typhoonzero 已提交
18
#include "paddle/fluid/framework/details/send_op_handle.h"
Y
Yu Yang 已提交
19
#include "paddle/fluid/framework/scope.h"
Y
Yu Yang 已提交
20 21 22 23

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

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

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

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

T
wip  
typhoonzero 已提交
58 59
void MultiDevSSAGraphBuilder::CreateOpHandleIOs(SSAGraph *result, OpDesc *op,
                                                const platform::Place &p,
T
finish  
typhoonzero 已提交
60
                                                const size_t &i) const {
T
wip  
typhoonzero 已提交
61
  auto *op_handle = result->ops_.back().get();
T
typhoonzero 已提交
62 63
  op_handle->dev_ctxes_[p] = const_cast<platform::DeviceContext *>(
      platform::DeviceContextPool::Instance().Get(p));
T
wip  
typhoonzero 已提交
64 65 66 67 68 69 70 71

  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);
  }

T
finish  
typhoonzero 已提交
72 73 74 75
  var_names = op->OutputArgumentNames();

  for (auto &each_var_name : var_names) {
    CreateOpOutput(result, op_handle, each_var_name, p, i);
T
wip  
typhoonzero 已提交
76 77 78
  }
}

Y
Yu Yang 已提交
79 80 81
std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
    const ProgramDesc &program) const {
  auto graph = new SSAGraph();
Y
Yu Yang 已提交
82
  SSAGraph &result = *graph;
C
chengduoZH 已提交
83
  std::unordered_set<std::string> og_has_been_broadcast;
Y
Yu Yang 已提交
84 85 86 87 88

  // 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 已提交
89 90 91 92 93 94 95 96 97 98 99 100

  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;
      }
    }

T
wip  
typhoonzero 已提交
101 102
    // append send op if program is distributed trainer main program.
    // always use the first device
T
typhoonzero 已提交
103
    if (!is_forwarding && op->Type() == "send") {
T
wip  
typhoonzero 已提交
104 105
      auto &p = places_[0];
      auto *s = local_scopes_[0];
T
typhoonzero 已提交
106
      // FIXME(wuyi): send op always copy from GPU 0
T
finish  
typhoonzero 已提交
107
      result.ops_.emplace_back(new SendOpHandle(*op, s, p));
T
typhoonzero 已提交
108 109
      // Create inputs for output on original place and no ssa output
      // is created for send op.
T
finish  
typhoonzero 已提交
110
      CreateOpHandleIOs(&result, op, p, 0);
T
wip  
typhoonzero 已提交
111 112 113
      continue;
    }

Y
Yu Yang 已提交
114 115 116 117 118 119
    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();
T
wip  
typhoonzero 已提交
120
      CreateOpHandleIOs(&result, op, p, i);
Y
Yu Yang 已提交
121

T
wip  
typhoonzero 已提交
122
      auto var_names = op->OutputArgumentNames();
Y
Yu Yang 已提交
123 124 125

      if (is_forwarding) {
        if (var_names.size() == 1 && var_names[0] == loss_var_name_) {
Y
Yu Yang 已提交
126 127 128 129 130 131 132 133
// 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 已提交
134
          op_handle = new ScaleLossGradOpHandle(local_scopes_.size(), s, p,
Y
Yu Yang 已提交
135
                                                communication_dev_ctx);
Y
Yu Yang 已提交
136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155
          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 已提交
156 157 158 159
      // 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 已提交
160
      for (auto &og : var_names) {
C
chengduoZH 已提交
161
        if (grad_names_.count(og) != 0 &&
C
chengduoZH 已提交
162 163 164
            og_has_been_broadcast.count(og) == 0) {  // is param grad
                                                     // Insert NCCL AllReduce Op
          og_has_been_broadcast.insert(og);
Y
Yu Yang 已提交
165
#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
166 167 168 169 170 171 172 173 174 175 176
          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 已提交
177 178
            auto &prev_grad = vars[vars.size() - 1];
            op_handle->AddInput(prev_grad.get());
Y
Yu Yang 已提交
179

Y
Yu Yang 已提交
180 181 182 183 184
            vars.emplace_back(new VarHandle);
            auto &var = vars.back();
            var->place_ = p;
            var->name_ = og;
            var->version_ = vars.size() - 1;
Y
Yu Yang 已提交
185

Y
Yu Yang 已提交
186
            op_handle->AddOutput(var.get());
Y
Yu Yang 已提交
187
          }
Y
Yu Yang 已提交
188 189 190
#else
          PADDLE_ENFORCE("Not implemented");
#endif
Y
Yu Yang 已提交
191 192 193 194 195 196 197 198 199 200
        }
      }
    }
  }

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

Y
Yu Yang 已提交
202 203 204 205 206
  /*
   * Only variables should be the leaves of graph.
   */
  AddOutputToLeafOps(&result);

Y
Yu Yang 已提交
207 208 209 210 211 212
  if (VLOG_IS_ON(10)) {
    std::ostringstream sout;
    PrintGraphviz(*graph, sout);
    VLOG(10) << sout.str();
  }

Y
Yu Yang 已提交
213
  return std::unique_ptr<SSAGraph>(graph);
Y
Yu Yang 已提交
214
}  // namespace details
Y
Yu Yang 已提交
215 216 217
}  // namespace details
}  // namespace framework
}  // namespace paddle