all_reduce_op_handle.cc 5.1 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
Yu Yang 已提交
15

16
#include "paddle/fluid/framework/details/all_reduce_op_handle.h"
C
chengduoZH 已提交
17
#include "paddle/fluid/framework/details/container_cast.h"
C
chengduoZH 已提交
18
#include "paddle/fluid/framework/details/reduce_and_gather.h"
C
chengduoZH 已提交
19
#include "paddle/fluid/framework/details/variable_visitor.h"
Y
Stash  
Yu Yang 已提交
20

Y
Yu Yang 已提交
21 22 23
namespace paddle {
namespace framework {
namespace details {
C
chengduoZH 已提交
24 25

#ifdef PADDLE_WITH_CUDA
26 27 28
AllReduceOpHandle::AllReduceOpHandle(const std::vector<Scope *> &local_scopes,
                                     const std::vector<platform::Place> &places,
                                     const platform::NCCLContextMap *ctxs)
Y
Yu Yang 已提交
29
    : local_scopes_(local_scopes), places_(places), nccl_ctxs_(ctxs) {
30
  if (nccl_ctxs_) {
C
chengduoZH 已提交
31 32 33
    for (auto &p : places_) {
      this->dev_ctxes_[p] = nccl_ctxs_->DevCtx(p);
    }
Y
Yu Yang 已提交
34 35
  }
}
C
chengduoZH 已提交
36
#else
37 38
AllReduceOpHandle::AllReduceOpHandle(const std::vector<Scope *> &local_scopes,
                                     const std::vector<platform::Place> &places)
C
chengduoZH 已提交
39 40
    : local_scopes_(local_scopes), places_(places) {}
#endif
Y
Yu Yang 已提交
41

42
void AllReduceOpHandle::RunImpl() {
C
chengduoZH 已提交
43
  if (NoDummyInputSize() == 1) {
Y
Yu Yang 已提交
44 45 46
    return;  // No need to all reduce when GPU count = 1;
  } else {
    // Wait input done
C
chengduoZH 已提交
47
    WaitInputVarGenerated();
C
chengduoZH 已提交
48 49 50 51 52 53 54 55
    auto in_var_handles = DynamicCast<VarHandle>(this->Inputs());
    auto out_var_handles = DynamicCast<VarHandle>(this->Outputs());
    PADDLE_ENFORCE_EQ(
        in_var_handles.size(), places_.size(),
        "The NoDummyInputSize should be equal to the number of places.");
    PADDLE_ENFORCE_EQ(
        in_var_handles.size(), out_var_handles.size(),
        "The NoDummyInputSize and NoDummyOutputSize should be equal.");
Y
Yu Yang 已提交
56

C
chengduoZH 已提交
57
    std::vector<const LoDTensor *> lod_tensors;
Y
Yu Yang 已提交
58 59
    for (size_t i = 0; i < local_scopes_.size(); ++i) {
      auto *s = local_scopes_[i];
Y
Yu Yang 已提交
60
      auto &local_scope = *s->FindVar(kLocalExecScopeName)->Get<Scope *>();
C
chengduoZH 已提交
61 62
      auto &lod_tensor =
          local_scope.FindVar(in_var_handles[i]->name_)->Get<LoDTensor>();
C
chengduoZH 已提交
63
      lod_tensors.emplace_back(&lod_tensor);
C
chengduoZH 已提交
64 65
      PADDLE_ENFORCE_EQ(in_var_handles[i]->name_, out_var_handles[i]->name_,
                        "The name of input and output should be equal.");
Y
Stash  
Yu Yang 已提交
66 67
    }

C
chengduoZH 已提交
68
    if (platform::is_gpu_place(lod_tensors[0]->place())) {
C
chengduoZH 已提交
69
#ifdef PADDLE_WITH_CUDA
C
chengduoZH 已提交
70
      PADDLE_ENFORCE(nccl_ctxs_, "nccl_ctxs should not be nullptr.");
C
chengduoZH 已提交
71 72
      int dtype = -1;
      size_t numel = 0;
Y
Stash  
Yu Yang 已提交
73 74 75
      std::vector<std::function<void()>> all_reduce_calls;
      for (size_t i = 0; i < local_scopes_.size(); ++i) {
        auto &p = places_[i];
C
chengduoZH 已提交
76
        auto &lod_tensor = *lod_tensors[i];
Y
Stash  
Yu Yang 已提交
77
        void *buffer = const_cast<void *>(lod_tensor.data<void>());
Y
Yu Yang 已提交
78

Y
Stash  
Yu Yang 已提交
79 80 81 82 83 84 85 86 87
        if (dtype == -1) {
          dtype = platform::ToNCCLDataType(lod_tensor.type());
        }

        if (numel == 0) {
          numel = static_cast<size_t>(lod_tensor.numel());
        }

        int dev_id = boost::get<platform::CUDAPlace>(p).device;
C
chengduoZH 已提交
88
        auto &nccl_ctx = nccl_ctxs_->at(dev_id);
Y
Stash  
Yu Yang 已提交
89 90 91 92 93 94 95
        auto stream = nccl_ctx.stream();
        auto comm = nccl_ctx.comm_;
        all_reduce_calls.emplace_back([=] {
          PADDLE_ENFORCE(platform::dynload::ncclAllReduce(
              buffer, buffer, numel, static_cast<ncclDataType_t>(dtype),
              ncclSum, comm, stream));
        });
Y
Yu Yang 已提交
96
      }
97 98 99 100 101 102
      this->RunAndRecordEvent([&] {
        platform::NCCLGroupGuard guard;
        for (auto &call : all_reduce_calls) {
          call();
        }
      });
C
chengduoZH 已提交
103 104 105
#else
      PADDLE_THROW("Not compiled with CUDA");
#endif
Y
Stash  
Yu Yang 已提交
106
    } else {  // Special handle CPU only Operator's gradient. Like CRF
Y
Yu Yang 已提交
107 108 109
      auto &trg = *this->local_scopes_[0]
                       ->FindVar(kLocalExecScopeName)
                       ->Get<Scope *>()
110
                       ->FindVar(out_var_handles[0]->name_)
Y
Yu Yang 已提交
111
                       ->GetMutable<framework::LoDTensor>();
Y
Yu Yang 已提交
112

Y
Stash  
Yu Yang 已提交
113 114
      // Reduce All Tensor to trg in CPU
      ReduceLoDTensor func(lod_tensors, &trg);
C
chengduoZH 已提交
115
      VisitDataType(ToDataType(lod_tensors[0]->type()), func);
116

117
      for (size_t i = 1; i < local_scopes_.size(); ++i) {
Y
Yu Yang 已提交
118 119
        auto &scope =
            *local_scopes_[i]->FindVar(kLocalExecScopeName)->Get<Scope *>();
120
        auto &p = places_[i];
121
        auto *var = scope.FindVar(out_var_handles[i]->name_);
122 123 124 125 126 127 128 129
        auto *dev_ctx = dev_ctxes_[p];

        RunAndRecordEvent(p, [&trg, var, dev_ctx, p] {
          auto &tensor_gpu = *var->GetMutable<framework::LoDTensor>();
          auto &tensor_cpu = trg;
          TensorCopy(tensor_cpu, p, *dev_ctx, &tensor_gpu);
        });
      }
Y
Yu Yang 已提交
130 131 132
    }
  }
}
Y
Yu Yang 已提交
133

C
chengduoZH 已提交
134
std::string AllReduceOpHandle::Name() const { return "all_reduce"; }
Y
Yu Yang 已提交
135 136 137
}  // namespace details
}  // namespace framework
}  // namespace paddle