communicator.h 6.7 KB
Newer Older
Q
Qiao Longfei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* Copyright (c) 2019 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. */

#pragma once

17
#include <atomic>
Q
Qiao Longfei 已提交
18 19 20
#include <deque>
#include <memory>
#include <string>
Q
Qiao Longfei 已提交
21 22
#include <unordered_map>
#include <utility>
Q
Qiao Longfei 已提交
23 24
#include <vector>

Q
Qiao Longfei 已提交
25 26
#include <ThreadPool.h>

Q
Qiao Longfei 已提交
27 28 29
#include "paddle/fluid/framework/scope.h"
#include "paddle/fluid/framework/variable.h"
#include "paddle/fluid/operators/distributed/rpc_common.h"
Q
Qiao Longfei 已提交
30 31
#include "paddle/fluid/operators/math/math_function.h"
#include "paddle/fluid/operators/math/selected_rows_functor.h"
Q
Qiao Longfei 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/place.h"

namespace paddle {
namespace operators {
namespace distributed {

using Scope = framework::Scope;
using Variable = framework::Variable;

template <typename T>
class BlockingQueue {
 public:
  explicit BlockingQueue(size_t capacity) : capacity_(capacity) {
    PADDLE_ENFORCE_GT(capacity_, 0, "The capacity must be greater than 0.");
  }

  bool Push(const T& elem) {
Q
Qiao Longfei 已提交
51 52 53 54 55 56 57
    {
      std::unique_lock<std::mutex> lock(mutex_);
      cv_.wait(lock, [&] { return queue_.size() < capacity_; });
      PADDLE_ENFORCE_LT(queue_.size(), capacity_);
      queue_.push_back(elem);
    }
    cv_.notify_one();
Q
Qiao Longfei 已提交
58 59 60 61
    return true;
  }

  bool Push(T&& elem) {
Q
Qiao Longfei 已提交
62 63 64 65 66 67 68
    {
      std::unique_lock<std::mutex> lock(mutex_);
      cv_.wait(lock, [&] { return queue_.size() < capacity_; });
      PADDLE_ENFORCE_LT(queue_.size(), capacity_);
      queue_.emplace_back(std::move(elem));
    }
    cv_.notify_one();
Q
Qiao Longfei 已提交
69 70 71 72 73
    return true;
  }

  T Pop() {
    std::unique_lock<std::mutex> lock(mutex_);
Q
Qiao Longfei 已提交
74
    cv_.wait(lock, [=] { return !queue_.empty(); });
Q
Qiao Longfei 已提交
75 76
    T rc(std::move(queue_.front()));
    queue_.pop_front();
Q
Qiao Longfei 已提交
77
    cv_.notify_one();
Q
Qiao Longfei 已提交
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
    return rc;
  }

  size_t Cap() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return capacity_;
  }

  size_t Size() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return queue_.size();
  }

 private:
  const size_t capacity_;
  std::deque<T> queue_;

  mutable std::mutex mutex_;
Q
Qiao Longfei 已提交
96
  std::condition_variable cv_;
Q
Qiao Longfei 已提交
97 98
};

Q
Qiao Longfei 已提交
99 100 101 102 103 104 105 106 107 108 109 110 111
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;

inline void MergeVars(const std::string& var_name,
                      const std::vector<std::shared_ptr<Variable>>& vars,
                      Scope* scope) {
  PADDLE_ENFORCE(!vars.empty(), "should have value to merge!");
  auto cpu_place = platform::CPUPlace();
  auto& var0 = vars[0];
  auto* out_var = scope->Var(var_name);
  if (var0->IsType<framework::LoDTensor>()) {
    auto dims = var0->Get<framework::LoDTensor>().dims();
Q
Qiao Longfei 已提交
112
    VLOG(3) << "merge " << var_name << " LoDTensor dims " << dims;
Q
Qiao Longfei 已提交
113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157

    // init output tensor
    auto* out_t = out_var->GetMutable<framework::LoDTensor>();
    out_t->mutable_data<float>(dims, cpu_place);

    // check the input dims
    for (auto& var : vars) {
      auto& var_t = var->Get<framework::LoDTensor>();
      PADDLE_ENFORCE_EQ(var_t.dims(), dims, "should have the same dims");
    }

    // set output tensor to 0.
    auto cpu_ctx = paddle::platform::CPUDeviceContext();
    math::SetConstant<paddle::platform::CPUDeviceContext, float>
        constant_functor;
    constant_functor(cpu_ctx, out_t, static_cast<float>(0));

    // sum all vars to out
    auto result = EigenVector<float>::Flatten(*out_t);
    for (auto& var : vars) {
      auto& in_t = var->Get<framework::LoDTensor>();
      auto in = EigenVector<float>::Flatten(in_t);
      result.device(*cpu_ctx.eigen_device()) = result + in;
    }
  } else if (var0->IsType<framework::SelectedRows>()) {
    auto& slr0 = var0->Get<framework::SelectedRows>();
    auto* out_slr = out_var->GetMutable<framework::SelectedRows>();
    out_slr->mutable_rows()->clear();
    out_slr->mutable_value()->mutable_data<float>({{}}, cpu_place);
    std::vector<const paddle::framework::SelectedRows*> inputs;
    inputs.reserve(vars.size());
    for (auto& var : vars) {
      inputs.push_back(&var->Get<framework::SelectedRows>());
    }
    math::scatter::MergeAdd<paddle::platform::CPUDeviceContext, float>
        merge_add;
    auto dev_ctx = paddle::platform::CPUDeviceContext();
    merge_add(dev_ctx, inputs, out_slr, false);
    VLOG(3) << "merge " << var_name << " SelectedRows height: " << slr0.height()
            << " dims: " << slr0.value().dims();
  } else {
    PADDLE_THROW("unsupported var type!");
  }
}

Q
Qiao Longfei 已提交
158 159
using RpcCtxMap = std::unordered_map<std::string, RpcContext>;

Q
Qiao Longfei 已提交
160 161
class Communicator {
 public:
Q
Qiao Longfei 已提交
162
  Communicator(const RpcCtxMap& send_varname_to_ctx,
Q
Qiao Longfei 已提交
163
               const RpcCtxMap& recv_varname_to_ctx, Scope* recv_scope);
Q
Qiao Longfei 已提交
164

Q
Qiao Longfei 已提交
165
  ~Communicator();
Q
Qiao Longfei 已提交
166 167

  void Start();
168
  void Stop();
Q
Qiao Longfei 已提交
169 170 171 172 173

  // send grad
  void Send(const std::string& var_name, const framework::Scope& scope);

 private:
Q
Qiao Longfei 已提交
174 175
  // recv all parameter
  void RecvAll();
Q
Qiao Longfei 已提交
176 177 178
  void SendThread();
  void RecvThread();

Q
Qiao Longfei 已提交
179
  bool running_ = false;
Q
Qiao Longfei 已提交
180 181 182
  std::unordered_map<std::string,
                     std::shared_ptr<BlockingQueue<std::shared_ptr<Variable>>>>
      send_varname_to_queue_;
Q
Qiao Longfei 已提交
183 184
  RpcCtxMap send_varname_to_ctx_;
  RpcCtxMap recv_varname_to_ctx_;
185 186
  std::unique_ptr<std::thread> send_thread_{nullptr};
  std::unique_ptr<std::thread> recv_thread_{nullptr};
Q
Qiao Longfei 已提交
187 188
  Scope* recv_scope_;                  // should be global scope
  std::unique_ptr<Scope> send_scope_;  // an independent scope
Q
Qiao Longfei 已提交
189 190
  std::unique_ptr<::ThreadPool> send_threadpool_{nullptr};
  std::unique_ptr<::ThreadPool> recv_threadpool_{nullptr};
191
  std::atomic_uint grad_num_{0};  // the num of gradient sent since last recv
Q
Qiao Longfei 已提交
192 193 194 195 196 197 198 199 200 201 202

  // the following code is for initialize the commnunicator
 public:
  static void Init(const RpcCtxMap& send_varname_to_ctx,
                   const RpcCtxMap& recv_varname_to_ctx, Scope* recv_scope) {
    if (communicator_ == nullptr) {
      communicator_.reset(new Communicator(send_varname_to_ctx,
                                           recv_varname_to_ctx, recv_scope));
    }
  }

203 204 205 206 207 208 209
  static void Init(const paddle::framework::ProgramDesc& program,
                   Scope* param_scope);

  static Communicator* GetInstance();

  static std::shared_ptr<Communicator> GetInstantcePtr();

Q
Qiao Longfei 已提交
210
 private:
211
  static std::shared_ptr<Communicator> communicator_;
Q
Qiao Longfei 已提交
212 213 214 215 216
};

}  // namespace distributed
}  // namespace operators
}  // namespace paddle