multi_devices_graph_builder.cc 20.0 KB
Newer Older
Y
Yu Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
//   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.
C
chengduoZH 已提交
14
#include <algorithm>
Y
Yancey1989 已提交
15
#include <fstream>
C
chengduoZH 已提交
16
#include <string>
C
chengduoZH 已提交
17
#include <utility>
C
chengduoZH 已提交
18 19
#include <vector>

20
#include "paddle/fluid/framework/details/all_reduce_op_handle.h"
C
chengduoZH 已提交
21
#include "paddle/fluid/framework/details/broadcast_op_handle.h"
Y
Yu Yang 已提交
22
#include "paddle/fluid/framework/details/computation_op_handle.h"
C
chengduoZH 已提交
23
#include "paddle/fluid/framework/details/multi_devices_graph_builder.h"
C
chengduoZH 已提交
24
#include "paddle/fluid/framework/details/reduce_op_handle.h"
Y
Yancey1989 已提交
25
#include "paddle/fluid/framework/details/rpc_op_handle.h"
Y
Yu Yang 已提交
26
#include "paddle/fluid/framework/details/scale_loss_grad_op_handle.h"
Y
Fix bug  
yuyang18 已提交
27
#include "paddle/fluid/framework/op_info.h"
Y
Yu Yang 已提交
28
#include "paddle/fluid/framework/scope.h"
Y
Yu Yang 已提交
29

Y
Yu Yang 已提交
30 31 32
namespace paddle {
namespace framework {
namespace details {
Y
Yu Yang 已提交
33 34

#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
35 36 37 38
MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
    const std::vector<platform::Place> &places,
    const std::string &loss_var_name,
    const std::unordered_set<std::string> &params,
C
chengduoZH 已提交
39
    const std::vector<Scope *> &local_scopes,
Y
yuyang18 已提交
40
    platform::NCCLContextMap *nccl_ctxs, const BuildStrategy &strategy)
Y
Yu Yang 已提交
41 42 43
    : loss_var_name_(loss_var_name),
      places_(places),
      local_scopes_(local_scopes),
C
chengduoZH 已提交
44
      nccl_ctxs_(nccl_ctxs),
Y
yuyang18 已提交
45
      strategy_(strategy) {
Y
Yu Yang 已提交
46 47 48 49 50
#else
MultiDevSSAGraphBuilder::MultiDevSSAGraphBuilder(
    const std::vector<platform::Place> &places,
    const std::string &loss_var_name,
    const std::unordered_set<std::string> &params,
Y
yuyang18 已提交
51
    const std::vector<Scope *> &local_scopes, const BuildStrategy &strategy)
Y
Yu Yang 已提交
52 53
    : loss_var_name_(loss_var_name),
      places_(places),
C
chengduoZH 已提交
54
      local_scopes_(local_scopes),
Y
yuyang18 已提交
55
      strategy_(strategy) {
Y
Yu Yang 已提交
56
#endif
Y
Yu Yang 已提交
57 58 59
  for (auto &p : params) {
    grad_names_.insert(GradVarName(p));
  }
Y
Yancey1989 已提交
60
  balance_vars_.resize(places_.size(), 0);
Y
Yu Yang 已提交
61 62
}

Y
Yu Yang 已提交
63 64
void MultiDevSSAGraphBuilder::CreateOpHandleIOs(SSAGraph *result,
                                                const OpDesc &op,
Y
Yu Yang 已提交
65 66
                                                size_t place_id) const {
  auto p = places_[place_id];
T
wip  
typhoonzero 已提交
67
  auto *op_handle = result->ops_.back().get();
X
Xin Pan 已提交
68 69
  op_handle->SetDeviceContext(p,
                              platform::DeviceContextPool::Instance().Get(p));
T
wip  
typhoonzero 已提交
70

Y
Yu Yang 已提交
71 72 73
  for (auto &each_var_name : op.InputArgumentNames()) {
    VarHandle *var =
        CreateOrGetLatestVarHandle(result, each_var_name, p, place_id);
T
wip  
typhoonzero 已提交
74 75 76
    op_handle->AddInput(var);
  }

Y
Yu Yang 已提交
77 78
  for (auto &each_var_name : op.OutputArgumentNames()) {
    CreateOpOutput(result, op_handle, each_var_name, p, place_id);
T
wip  
typhoonzero 已提交
79 80
  }
}
Y
fix pe  
Yancey1989 已提交
81 82 83 84

std::vector<std::string> MultiDevSSAGraphBuilder::FindDistTrainSendVars(
    const ProgramDesc &program) const {
  std::vector<std::string> send_vars;
Y
Yancey1989 已提交
85 86
  // since parameters are all in block 0,
  // it's enough to only scan send ops in block 0
Y
fix pe  
Yancey1989 已提交
87
  for (auto *op : program.Block(0).AllOps()) {
Y
Yancey1989 已提交
88 89
    // TODO(Yancey1989): use a graceful method to find send op,
    // instead of the the hard code string
90
    if (op->Type() == "send") {
Y
fix pe  
Yancey1989 已提交
91 92 93 94 95 96 97 98 99 100 101 102 103
      auto op_vars = op->InputArgumentNames();
      send_vars.reserve(send_vars.size() +
                        std::distance(op_vars.begin(), op_vars.end()));
      send_vars.insert(send_vars.end(), op_vars.begin(), op_vars.end());
    }
  }
  return send_vars;
}

std::vector<std::string> MultiDevSSAGraphBuilder::FindDistTrainRecvVars(
    const ProgramDesc &program) const {
  std::vector<std::string> recv_vars;
  for (auto *op : program.Block(0).AllOps()) {
Y
Yancey1989 已提交
104 105 106
    // TODO(Yancey1989): use a graceful method to find recv op,
    // instead of the hard code string
    if (op->Type() == "recv") {
Y
fix pe  
Yancey1989 已提交
107 108 109 110 111 112 113 114 115 116 117 118 119
      auto op_vars = op->OutputArgumentNames();
      recv_vars.reserve(recv_vars.size() +
                        std::distance(op_vars.begin(), op_vars.end()));
      recv_vars.insert(recv_vars.end(), op_vars.begin(), op_vars.end());
    }
  }
  return recv_vars;
}

bool MultiDevSSAGraphBuilder::IsDistTrainOp(
    const OpDesc &op, const std::vector<std::string> &send_vars,
    const std::vector<std::string> &recv_vars) const {
  if (send_vars.size() == 0 || recv_vars.size() == 0) {
T
typhoonzero 已提交
120 121 122
    return false;
  }

Y
Yu Yang 已提交
123 124 125 126
  /**
   * Check any of opvars contains `.block` and in sendvars
   */
  auto checker = [](const std::vector<std::string> &opvars,
Y
fix pe  
Yancey1989 已提交
127
                    const std::vector<std::string> &rpc_vars) -> bool {
T
typhoonzero 已提交
128
    for (auto &var : opvars) {
Y
Yancey1989 已提交
129 130 131
      // a variable name with the suffix `.block` means it's a splited
      // variable by (DistributeTranspiler)
      // [python/paddle/fluid/transpiler/distribute_transpiler.py]
T
typhoonzero 已提交
132
      if (var.find(".block") != std::string::npos &&
Y
fix pe  
Yancey1989 已提交
133
          std::find(rpc_vars.begin(), rpc_vars.end(), var) != rpc_vars.end()) {
Y
Yu Yang 已提交
134
        return true;
T
typhoonzero 已提交
135 136
      }
    }
Y
Yu Yang 已提交
137
    return false;
T
typhoonzero 已提交
138 139
  };

Y
Yancey1989 已提交
140 141
  return checker(op.OutputArgumentNames(), send_vars) ||
         checker(op.InputArgumentNames(), recv_vars);
T
typhoonzero 已提交
142 143
}

Y
Yancey1989 已提交
144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163
size_t MultiDevSSAGraphBuilder::GetAppropriateDeviceID(
    const std::vector<std::string> &var_names) const {
  int64_t numel_sum = 0;
  for (auto var_name : var_names) {
    auto var_desc = all_vars_.at(var_name);
    PADDLE_ENFORCE_NOT_NULL(var_desc);
    auto dim = framework::make_ddim(var_desc->GetShape());
    int64_t numel = framework::product(dim);
    PADDLE_ENFORCE_GT(numel, 0);
    numel_sum += numel;
  }

  auto smallest =
      std::min_element(std::begin(balance_vars_), std::end(balance_vars_));
  size_t dev_id =
      static_cast<size_t>(std::distance(std::begin(balance_vars_), smallest));
  balance_vars_[dev_id] += numel_sum;
  return dev_id;
}

Y
Yu Yang 已提交
164 165
std::unique_ptr<SSAGraph> MultiDevSSAGraphBuilder::Build(
    const ProgramDesc &program) const {
C
fix ci  
chengduoZH 已提交
166
  for (auto *var : program.Block(0).AllVars()) {
Y
Yancey1989 已提交
167
    all_vars_.emplace(var->Name(), var);
C
fix ci  
chengduoZH 已提交
168
  }
C
chengduoZH 已提交
169

Y
Yu Yang 已提交
170
  auto graph = new SSAGraph();
Y
Yu Yang 已提交
171
  SSAGraph &result = *graph;
C
chengduoZH 已提交
172
  std::unordered_set<std::string> og_has_been_broadcast;
Y
Yu Yang 已提交
173 174 175 176 177

  // 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 已提交
178

Y
fix pe  
Yancey1989 已提交
179 180 181 182
  // find send/recv vars so that we can place the distributed training
  // realted op in the place 0
  auto send_vars = FindDistTrainSendVars(program);
  auto recv_vars = FindDistTrainRecvVars(program);
T
typhoonzero 已提交
183

C
chengduoZH 已提交
184 185 186
  std::vector<std::unordered_set<std::string>> bcast_var_name_set;
  bcast_var_name_set.resize(places_.size());

C
chengduoZH 已提交
187
  size_t cur_device_id = 0;
Y
Yu Yang 已提交
188
  bool is_forwarding = true;
189

Y
Yu Yang 已提交
190
  for (auto *op : program.Block(0).AllOps()) {
Y
Yancey1989 已提交
191 192 193
    if (boost::get<int>(
            op->GetAttr(OpProtoAndCheckerMaker::OpRoleAttrName())) ==
        static_cast<int>(OpRole::kRPC)) {
Y
Yancey1989 已提交
194
      CreateRPCOp(&result, *op);
Y
fix pe  
Yancey1989 已提交
195
    } else if (IsDistTrainOp(*op, send_vars, recv_vars)) {
Y
Yancey1989 已提交
196
      CreateDistTrainOp(&result, *op);
Y
Yu Yang 已提交
197
    } else if (IsScaleLossOp(*op)) {
Y
Yu Yang 已提交
198
      // user can customize loss@grad if not use_default_grad_scale_
Y
yuyang18 已提交
199 200
      if (strategy_.gradient_scale_ !=
          BuildStrategy::GradientScaleStrategy::kCustomized) {
Y
Yu Yang 已提交
201 202
        CreateScaleLossGradOp(&result);
      }
203 204 205 206
      // This assumes the backward generating code will ensure IsScaleLossOp
      // is true only for the op that scale the final scalar loss.
      // It also assumes backward op will always follow the forward op in
      // the block.
Y
Yu Yang 已提交
207
      is_forwarding = false;
Y
Yu Yang 已提交
208
    } else {
209
      int op_dev_id = GetOpDeviceID(*op);
C
chengduoZH 已提交
210 211 212 213 214
      if (op_dev_id == -1) {  // var on all device
        CreateComputationalOps(&result, *op, places_.size());
      } else {
        CreateComputationalOp(&result, *op, op_dev_id);
        for (auto &var_name : op->OutputArgumentNames()) {
215
          var_name_on_devices_.emplace(var_name, op_dev_id);
C
chengduoZH 已提交
216 217
        }
      }
C
chengduoZH 已提交
218
      if (!is_forwarding && places_.size() > 1) {
Y
Yu Yang 已提交
219
        // Currently, we assume that once gradient is generated, it can be
Y
Yu Yang 已提交
220
        // broadcast, and each gradient is only broadcast once.
Y
yuyang18 已提交
221 222 223
        if (static_cast<bool>(boost::get<int>(op->GetAttr(
                                  OpProtoAndCheckerMaker::OpRoleAttrName())) &
                              static_cast<int>(OpRole::kBackward))) {
Y
yuyang18 已提交
224 225 226 227 228 229 230
          try {
            auto backward_vars =
                boost::get<std::vector<std::string>>(op->GetNullableAttr(
                    OpProtoAndCheckerMaker::OpRoleVarAttrName()));

            PADDLE_ENFORCE_EQ(backward_vars.size() % 2, 0);

Y
Fix bug  
yuyang18 已提交
231
            for (size_t i = 0; i < backward_vars.size(); i += 2) {
Y
yuyang18 已提交
232 233
              auto &p_name = backward_vars[i];
              auto &g_name = backward_vars[i + 1];
Y
yuyang18 已提交
234 235
              VLOG(10) << "Bcast " << g_name << " for parameter " << p_name;

Y
yuyang18 已提交
236 237
              switch (strategy_.reduce_) {
                case BuildStrategy::ReduceStrategy::kReduce:
Y
Yancey1989 已提交
238
                  cur_device_id = GetAppropriateDeviceID({g_name});
Y
yuyang18 已提交
239
                  CreateReduceOp(&result, g_name, cur_device_id);
240
                  var_name_on_devices_.emplace(g_name, cur_device_id);
Y
yuyang18 已提交
241 242 243
                  bcast_var_name_set[cur_device_id].emplace(p_name);
                  break;
                case BuildStrategy::ReduceStrategy::kAllReduce:
Y
Yancey1989 已提交
244
                  if (IsSparseGradient(g_name)) {
Y
yuyang18 已提交
245 246 247
                    CreateReduceOp(&result, g_name, 0);
                    CreateBroadcastOp(&result, g_name, 0);
                  } else {
C
chengduoZH 已提交
248
                    InsertAllReduceOp(&result, g_name);
Y
yuyang18 已提交
249 250
                  }
                  break;
251 252 253
                default:
                  LOG(FATAL) << "Unknown reduce strategy ";
                  break;
Y
yuyang18 已提交
254
              }
C
chengduoZH 已提交
255
            }
Y
yuyang18 已提交
256
          } catch (boost::bad_get e) {
Y
Yu Yang 已提交
257 258 259 260 261 262
          }
        }
      }
    }
  }

C
chengduoZH 已提交
263 264 265 266 267 268 269
  // Insert BCast Ops
  for (size_t dev_id = 0; dev_id < bcast_var_name_set.size(); ++dev_id) {
    auto &to_bcast_set = bcast_var_name_set[dev_id];
    for (auto &bcast_name : to_bcast_set) {
      CreateBroadcastOp(&result, bcast_name, dev_id);
    }
  }
Y
Yu Yang 已提交
270 271
  /*
    Dependency graph has been constructed. However, there are still data
272
    hazards need to be handled.
Y
Yu Yang 已提交
273 274
   */
  PolishGraphToSupportDataHazards(&result);
Y
Yu Yang 已提交
275

Y
Yu Yang 已提交
276 277 278 279 280
  /*
   * Only variables should be the leaves of graph.
   */
  AddOutputToLeafOps(&result);

Y
Yu Yang 已提交
281
  return std::unique_ptr<SSAGraph>(graph);
Y
Yu Yang 已提交
282 283
}

Y
Yancey1989 已提交
284 285 286
bool MultiDevSSAGraphBuilder::IsSparseGradient(const std::string &og) const {
  PADDLE_ENFORCE(all_vars_.count(og) != 0);
  if (all_vars_.at(og)->GetType() == proto::VarType::SELECTED_ROWS) {
C
fix ci  
chengduoZH 已提交
287 288 289
    return true;
  }
  return false;
290 291
}

292 293 294 295 296 297 298 299 300 301 302 303 304
void MultiDevSSAGraphBuilder::SetCommunicationContext(
    OpHandleBase *op_handle, const platform::Place &p) const {
#ifdef PADDLE_WITH_CUDA
  if (nccl_ctxs_ == nullptr) {
    op_handle->SetDeviceContext(p,
                                platform::DeviceContextPool::Instance().Get(p));
  }
#else
  op_handle->SetDeviceContext(p,
                              platform::DeviceContextPool::Instance().Get(p));
#endif
}

C
chengduoZH 已提交
305 306
void MultiDevSSAGraphBuilder::CreateBroadcastOp(SSAGraph *result,
                                                const std::string &p_name,
C
chengduoZH 已提交
307
                                                size_t src_dev_id) const {
C
chengduoZH 已提交
308 309 310 311 312 313 314
#ifdef PADDLE_WITH_CUDA
  auto *op_handle = new BroadcastOpHandle(local_scopes_, places_, nccl_ctxs_);
#else
  auto *op_handle = new BroadcastOpHandle(local_scopes_, places_);
#endif

  result->ops_.emplace_back(op_handle);
C
chengduoZH 已提交
315
  auto *in = result->vars_.at(src_dev_id).at(p_name).back().get();
C
chengduoZH 已提交
316 317 318 319
  op_handle->AddInput(in);

  for (size_t i = 0; i < places_.size(); ++i) {
    auto &p = places_[i];
C
chengduoZH 已提交
320 321
    SetCommunicationContext(op_handle, p);
    auto &vars = result->vars_.at(i).at(p_name);
C
chengduoZH 已提交
322 323 324 325 326 327 328 329 330 331 332 333 334 335
    auto *out_var = new VarHandle(vars.size(), i, p_name, p);
    vars.emplace_back(out_var);
    op_handle->AddOutput(out_var);
  }
}

void MultiDevSSAGraphBuilder::CreateComputationalOp(SSAGraph *result,
                                                    const OpDesc &op,
                                                    int dev_id) const {
  result->ops_.emplace_back(
      new ComputationOpHandle(op, local_scopes_[dev_id], places_[dev_id]));
  CreateOpHandleIOs(result, op, dev_id);
}

C
chengduoZH 已提交
336 337
void MultiDevSSAGraphBuilder::InsertAllReduceOp(SSAGraph *result,
                                                const std::string &og) const {
Y
Yu Yang 已提交
338 339
#ifdef PADDLE_WITH_CUDA
  result->ops_.emplace_back(
340
      new AllReduceOpHandle(local_scopes_, places_, nccl_ctxs_));
C
chengduoZH 已提交
341
#else
342
  result->ops_.emplace_back(new AllReduceOpHandle(local_scopes_, places_));
C
chengduoZH 已提交
343
#endif
Y
Yu Yang 已提交
344 345 346 347
  auto *op_handle = result->ops_.back().get();

  for (size_t i = 0; i < places_.size(); ++i) {
    auto &p = places_[i];
C
chengduoZH 已提交
348
    SetCommunicationContext(op_handle, p);
Y
Yu Yang 已提交
349
    auto &vars = result->vars_[i][og];
Y
Yu Yang 已提交
350 351
    PADDLE_ENFORCE(!vars.empty());
    auto &prev_grad = vars.back();
Y
Yu Yang 已提交
352 353
    op_handle->AddInput(prev_grad.get());

354
    auto var = new VarHandle(vars.size(), i, og, p);
Y
Yu Yang 已提交
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371
    vars.emplace_back(var);
    op_handle->AddOutput(var);
  }
}

bool MultiDevSSAGraphBuilder::IsParameterGradientOnce(
    const std::string &og,
    std::unordered_set<std::string> *og_has_been_broadcast) const {
  bool is_pg_once =
      grad_names_.count(og) != 0 && og_has_been_broadcast->count(og) == 0;
  if (is_pg_once) {
    // Insert NCCL AllReduce Op
    og_has_been_broadcast->insert(og);
  }
  return is_pg_once;
}

372
int MultiDevSSAGraphBuilder::GetOpDeviceID(const OpDesc &op) const {
Y
yuyang18 已提交
373
  if (strategy_.reduce_ != BuildStrategy::ReduceStrategy::kReduce) {
C
chengduoZH 已提交
374 375 376
    return -1;
  }

377 378 379 380
  for (auto &varname : op.InputArgumentNames()) {
    int dev_id = GetVarDeviceID(varname);
    if (dev_id != -1) {
      return dev_id;
C
chengduoZH 已提交
381 382
    }
  }
383 384 385 386 387 388
  return -1;
}

int MultiDevSSAGraphBuilder::GetVarDeviceID(const std::string &varname) const {
  auto got = var_name_on_devices_.find(varname);
  return got == var_name_on_devices_.end() ? -1 : got->second;
C
chengduoZH 已提交
389 390
}

Y
Yu Yang 已提交
391 392 393 394
void MultiDevSSAGraphBuilder::CreateScaleLossGradOp(SSAGraph *result) const {
  for (size_t i = 0; i < places_.size(); ++i) {
// Insert ScaleCost OpHandle
#ifdef PADDLE_WITH_CUDA
C
chengduoZH 已提交
395 396 397
    auto *communication_dev_ctx =
        nccl_ctxs_ ? nccl_ctxs_->DevCtx(places_[i])
                   : platform::DeviceContextPool::Instance().Get(places_[i]);
Y
Yu Yang 已提交
398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419
#else
    auto *communication_dev_ctx =
        platform::DeviceContextPool::Instance().Get(platform::CPUPlace());
#endif

    auto *op_handle =
        new ScaleLossGradOpHandle(local_scopes_.size(), local_scopes_[i],
                                  places_[i], communication_dev_ctx);
    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_), places_[i],
                   i);
  }
}

void MultiDevSSAGraphBuilder::CreateComputationalOps(SSAGraph *result,
T
typhoonzero 已提交
420 421 422
                                                     const OpDesc &op,
                                                     size_t num_places) const {
  for (size_t scope_idx = 0; scope_idx < num_places; ++scope_idx) {
Y
Yu Yang 已提交
423 424 425
    auto p = places_[scope_idx];
    auto s = local_scopes_[scope_idx];
    result->ops_.emplace_back(new ComputationOpHandle(op, s, p));
Y
Yu Yang 已提交
426
    CreateOpHandleIOs(result, op, scope_idx);
Y
Yu Yang 已提交
427 428 429
  }
}

C
chengduoZH 已提交
430 431 432
VarHandle *MultiDevSSAGraphBuilder::CreateReduceOp(SSAGraph *result,
                                                   const std::string &og,
                                                   int dst_dev_id) const {
C
chengduoZH 已提交
433 434 435 436 437 438 439 440 441 442
#ifdef PADDLE_WITH_CUDA
  result->ops_.emplace_back(
      new ReduceOpHandle(local_scopes_, places_, nccl_ctxs_));
#else
  result->ops_.emplace_back(new ReduceOpHandle(local_scopes_, places_));
#endif
  auto *op_handle = result->ops_.back().get();

  for (size_t i = 0; i < places_.size(); ++i) {
    auto &p = places_[i];
C
chengduoZH 已提交
443 444
    SetCommunicationContext(op_handle, p);
    auto &vars = result->vars_[i][og];
C
chengduoZH 已提交
445 446 447 448 449
    PADDLE_ENFORCE(!vars.empty());
    auto &prev_grad = vars.back();
    op_handle->AddInput(prev_grad.get());
  }
  auto &vars = result->vars_[dst_dev_id][og];
450
  auto var = new VarHandle(vars.size(), dst_dev_id, og, places_[dst_dev_id]);
C
chengduoZH 已提交
451 452 453 454 455
  vars.emplace_back(var);
  op_handle->AddOutput(var);
  return var;
}

456 457
// Find the first occurence of `prev_op_name` and make current `op` depend
// on it.
Y
fix pe  
Yancey1989 已提交
458 459
void MultiDevSSAGraphBuilder::ConnectOp(SSAGraph *result, OpHandleBase *op,
                                        const std::string &prev_op_name) const {
Y
Yancey1989 已提交
460
  for (auto &prev_op : result->ops_) {
Y
fix pe  
Yancey1989 已提交
461
    if (prev_op->Name() == prev_op_name) {
Y
Yancey1989 已提交
462 463 464
      auto *dep_var = new DummyVarHandle();
      prev_op->AddOutput(dep_var);
      result->dep_vars_.emplace(dep_var);
Y
fix pe  
Yancey1989 已提交
465
      op->AddInput(dep_var);
Y
Yancey1989 已提交
466 467 468 469
    }
  }
}

Y
Yancey1989 已提交
470
void MultiDevSSAGraphBuilder::CreateDistTrainOp(SSAGraph *result,
Y
Yancey1989 已提交
471 472
                                                const OpDesc &op) const {
  int op_dev_id = -1;
Y
yi.wu 已提交
473
  if (op.Type() == "split_byref" || op.Type() == "split_selected_rows") {
Y
Yancey1989 已提交
474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
    op_dev_id = GetVarDeviceID(op.InputArgumentNames()[0]);
    if (strategy_.reduce_ == BuildStrategy::ReduceStrategy::kAllReduce) {
      op_dev_id = GetAppropriateDeviceID(op.InputArgumentNames());
      for (auto &varname : op.InputArgumentNames()) {
        var_name_on_devices_.emplace(varname, op_dev_id);
      }
    }
    for (auto &varname : op.OutputArgumentNames()) {
      var_name_on_devices_.emplace(varname, op_dev_id);
    }
  } else if (op.Type() == "concat") {
    op_dev_id = GetVarDeviceID(op.InputArgumentNames()[0]);
  } else {
    PADDLE_ENFORCE(
        "the distribute training related op should be in [split_byref, "
        "concat].");
  }

  PADDLE_ENFORCE(op_dev_id != -1,
                 "can not find right place for distributed op: %s", op.Type());

  CreateComputationalOp(result, op, op_dev_id);
Y
Yancey1989 已提交
496 497 498 499 500
  if (op.Type() == "concat") {
    ConnectOp(result, result->ops_.back().get(), "fetch_barrier");
  }
}

501
// Create RPC related op handles that connects its in ops and out ops.
Y
Yancey1989 已提交
502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531
void MultiDevSSAGraphBuilder::CreateRPCOp(SSAGraph *result,
                                          const OpDesc &op) const {
  int op_dev_id = -1;
  if (op.Type() == "send") {
    op_dev_id = GetVarDeviceID(op.InputArgumentNames()[0]);
    // the variable name which contains .block means it was splited by
    // split_byref op
    // so that we can balance the variable blocks to all the pserver instances.
    if (strategy_.reduce_ == BuildStrategy::ReduceStrategy::kAllReduce &&
        op.InputArgumentNames()[0].find(".block") == std::string::npos) {
      op_dev_id = GetAppropriateDeviceID(op.InputArgumentNames());
      for (auto &varname : op.InputArgumentNames()) {
        var_name_on_devices_.emplace(varname, op_dev_id);
      }
    }
  } else if (op.Type() == "recv") {
    op_dev_id = GetAppropriateDeviceID(op.OutputArgumentNames());
    for (auto &varname : op.OutputArgumentNames()) {
      var_name_on_devices_.emplace(varname, op_dev_id);
    }
  } else {
    // send_barrier and fetch_barrier op can be scheduled on device 0
    op_dev_id = 0;
  }

  PADDLE_ENFORCE(op_dev_id != -1, "can not find the right place for rpc op: %s",
                 op.Type());

  result->ops_.emplace_back(new RPCOpHandle(op, local_scopes_[op_dev_id],
                                            op.Type(), places_[op_dev_id]));
Y
fix pe  
Yancey1989 已提交
532

Y
Yancey1989 已提交
533
  if (op.Type() == "send_barrier") {
534
    ConnectOp(result, result->ops_.back().get(), "send");
Y
Yancey1989 已提交
535
  } else if (op.Type() == "recv") {
Y
fix pe  
Yancey1989 已提交
536
    ConnectOp(result, result->ops_.back().get(), "send_barrier");
Y
Yancey1989 已提交
537
  } else if (op.Type() == "fetch_barrier") {
Y
fix pe  
Yancey1989 已提交
538
    ConnectOp(result, result->ops_.back().get(), "recv");
539
  } else if (op.Type() == "send") {
Y
Yancey1989 已提交
540 541 542
    // do nothing
  } else {
    PADDLE_THROW(
Y
Yancey1989 已提交
543
        "rpc op should be in ["
544
        "send, send_barrier. recv, fetch_barrier]");
Y
Yancey1989 已提交
545 546
  }

Y
Yancey1989 已提交
547
  CreateOpHandleIOs(result, op, op_dev_id);
Y
Yu Yang 已提交
548 549 550
}

bool MultiDevSSAGraphBuilder::IsScaleLossOp(const OpDesc &op) const {
Y
yuyang18 已提交
551 552
  return boost::get<int>(
             op.GetAttr(OpProtoAndCheckerMaker::OpRoleAttrName())) ==
Y
Fix bug  
yuyang18 已提交
553 554 555
             (static_cast<int>(OpRole::kBackward) |
              static_cast<int>(OpRole::kLoss)) &&
         !loss_var_name_.empty();  // If loss_var is empty. This is test mode
Y
Yu Yang 已提交
556
}
Y
Yu Yang 已提交
557 558 559
}  // namespace details
}  // namespace framework
}  // namespace paddle