nccl_op.cc 9.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
L
Luo Tao 已提交
2 3 4 5 6 7 8 9 10 11 12 13

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. */
D
Dong Zhihong 已提交
14

Y
Yi Wang 已提交
15 16
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/nccl/nccl_gpu_common.h"
Y
Yang Yang 已提交
17
#include "paddle/fluid/operators/nccl/nccl_gpu_common.h"
D
dongzhihong 已提交
18 19 20 21

namespace paddle {
namespace operators {

Y
Yang Yang 已提交
22 23
static constexpr char kParallelScopes[] = "parallel_scopes";

D
Dong Zhihong 已提交
24
// NCCLinitOp
25
class NCCLInitOp : public framework::OperatorBase {
D
Dong Zhihong 已提交
26
 public:
27 28 29 30 31
  NCCLInitOp(const std::string &type, const framework::VariableNameMap &inputs,
             const framework::VariableNameMap &outputs,
             const framework::AttributeMap &attrs)
      : OperatorBase(type, inputs, outputs, attrs) {}

32 33 34
 private:
  void RunImpl(const framework::Scope &scope,
               const platform::Place &place) const override {
Y
Yang Yang 已提交
35 36 37
    PADDLE_ENFORCE_NOT_NULL(scope.FindVar(Input(kParallelScopes)),
                            "Can not find variable '%s' in the scope.",
                            kParallelScopes);
38 39 40
    const auto &name = Output("Communicator");
    PADDLE_ENFORCE_NOT_NULL(scope.FindVar(name),
                            "Can not find variable '%s' in the scope.", name);
Y
Yang Yang 已提交
41 42 43 44 45 46 47
    // A parallel do may not use all the gpus. For example, the batch size is 7
    // in the last batch while we have 8 gpu. In this case, parallel_do will
    // create 7 parallel scopes, so should ncclInitOp create 7 gpu peers
    auto &parallel_scopes = scope.FindVar(Input(kParallelScopes))
                                ->Get<std::vector<framework::Scope *>>();
    std::vector<int> gpus(parallel_scopes.size());
    for (int i = 0; i < static_cast<int>(parallel_scopes.size()); ++i) {
Y
Yang Yang 已提交
48 49 50
      gpus[i] = i;
    }
    PADDLE_ENFORCE(!gpus.empty(), "NCCL init with 0 gpus.");
D
dzhwinter 已提交
51 52 53 54 55

    if (scope.FindVar(name) == nullptr) {
      PADDLE_THROW("Output(Communicator) is needed for ncclInit operator.");
    }

56 57 58
    platform::Communicator *comm =
        scope.FindVar(name)->GetMutable<platform::Communicator>();
    comm->InitAll(gpus);
D
Dong Zhihong 已提交
59 60 61
  }
};

Y
Yang Yang 已提交
62 63 64 65 66 67
class NCCLInitOpVarTypeInference : public framework::VarTypeInference {
 public:
  void operator()(const framework::OpDesc &op_desc,
                  framework::BlockDesc *block) const override {
    auto out_var_name = op_desc.Output("Communicator").front();
    auto &out_var = block->FindRecursiveOrCreateVar(out_var_name);
T
typhoonzero 已提交
68
    auto var_type = framework::proto::VarType::RAW;
Y
Yang Yang 已提交
69 70 71 72 73 74 75 76 77
    out_var.SetType(var_type);
  }
};

class NCCLInitOpShapeInference : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext *ctx) const override {}
};

D
Dong Zhihong 已提交
78 79
class NCCLInitOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
80
  NCCLInitOpMaker(OpProto *proto, OpAttrChecker *op_checker)
D
Dong Zhihong 已提交
81
      : OpProtoAndCheckerMaker(proto, op_checker) {
Y
Yang Yang 已提交
82
    AddInput(kParallelScopes, "The working place of parallel do.");
D
Dong Zhihong 已提交
83 84 85
    AddOutput("Communicator",
              "Create Communicator for communicating between gpus");
    AddComment(R"DOC(
K
kexinzhao 已提交
86 87 88 89 90
NCCLInit Operator.

Create communicator.

)DOC");
D
Dong Zhihong 已提交
91 92 93 94
  }
};

// AllReduceOp
D
dzhwinter 已提交
95
class NCCLAllReduceOp : public framework::OperatorWithKernel {
D
dongzhihong 已提交
96 97 98 99
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
D
Dong Zhihong 已提交
100 101 102
  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   " Input(X) of AllReduce op input should not be NULL");
D
Dong Zhihong 已提交
103 104 105
    PADDLE_ENFORCE(
        ctx->HasInput("Communicator"),
        " Input(Communicator) of AllReduce op input should not be NULL");
D
Dong Zhihong 已提交
106
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
Y
Yang Yang 已提交
107
                   " Output(Out) of AllReduce op output should not be NULL");
D
dongzhihong 已提交
108

D
Dong Zhihong 已提交
109
    auto x_dims = ctx->GetInputsDim("X");
D
dongzhihong 已提交
110

D
Dong Zhihong 已提交
111 112 113 114
    std::string reduction = ctx->Attrs().Get<std::string>("reduction");
    PADDLE_ENFORCE((reduction == "ncclSum" || reduction == "ncclProd" ||
                    reduction == "ncclMin" || reduction == "ncclMax"),
                   "invalid reduction.");
D
dongzhihong 已提交
115

D
Dong Zhihong 已提交
116 117
    ctx->SetOutputsDim("Out", x_dims);
    ctx->ShareLoD("X", /*->*/ "Out");
D
dzhwinter 已提交
118 119 120
  }
};

D
Dong Zhihong 已提交
121 122 123 124 125 126 127 128 129 130 131 132 133 134
// ReduceOp
class NCCLReduceOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   " Input(X) of Reduce op input should not be NULL");
    PADDLE_ENFORCE(
        ctx->HasInput("Communicator"),
        " Input(Communicator) of Reduce op input should not be NULL");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   " Input(X) of Reduce op input should not be NULL");
D
Dong Zhihong 已提交
135

D
Dong Zhihong 已提交
136 137 138 139 140
    std::string reduction = ctx->Attrs().Get<std::string>("reduction");
    PADDLE_ENFORCE((reduction == "ncclSum" || reduction == "ncclProd" ||
                    reduction == "ncclMin" || reduction == "ncclMax"),
                   "invalid reduction.");

D
Dong Zhihong 已提交
141 142 143
    auto x_dims = ctx->GetInputsDim("X");
    ctx->SetOutputsDim("Out", x_dims);
    ctx->ShareLoD("X", /*->*/ "Out");
D
Dong Zhihong 已提交
144 145 146
  }
};

D
Dong Zhihong 已提交
147 148
// BcastOp
class NCCLBcastOp : public framework::OperatorWithKernel {
D
Dong Zhihong 已提交
149 150 151 152 153 154 155 156 157 158 159
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   " Input(X) of Bcast op input should not be NULL");
    PADDLE_ENFORCE(ctx->HasInput("Communicator"),
                   " Input(Communicator) of Bcast op input should not be NULL");
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   " Output(Out) of Bcast op output should not be NULL");
D
Dong Zhihong 已提交
160

D
Dong Zhihong 已提交
161
    int root = ctx->Attrs().Get<int>("root");
D
Dong Zhihong 已提交
162
    PADDLE_ENFORCE(root != platform::kInvalidGPUId, "Bcast root must be set.");
D
Dong Zhihong 已提交
163

D
Dong Zhihong 已提交
164 165 166
    auto x_dims = ctx->GetInputsDim("X");
    ctx->SetOutputsDim("Out", x_dims);
    ctx->ShareLoD("X", /*->*/ "Out");
D
Dong Zhihong 已提交
167 168 169
  }
};

D
Dong Zhihong 已提交
170
// AllreduceOp
D
dzhwinter 已提交
171
class NCCLAllReduceOpMaker : public framework::OpProtoAndCheckerMaker {
D
Dong Zhihong 已提交
172
 public:
173
  NCCLAllReduceOpMaker(OpProto *proto, OpAttrChecker *op_checker)
D
Dong Zhihong 已提交
174
      : OpProtoAndCheckerMaker(proto, op_checker) {
D
dzhwinter 已提交
175
    AddInput("X", "The input of AllReduce op");
D
Dong Zhihong 已提交
176
    AddInput("Communicator", "Communicator for communicating between gpus");
D
dzhwinter 已提交
177
    AddOutput("Out", "The output of AllReduce op");
D
Dong Zhihong 已提交
178
    AddAttr<std::string>("reduction",
K
kexinzhao 已提交
179
                         "(string, default 'ncclSum') "
D
Dong Zhihong 已提交
180 181
                         "{'ncclMin', 'ncclMax', 'ncclProd', 'ncclSum'}.")
        .SetDefault("ncclSum");
D
dzhwinter 已提交
182
    AddComment(R"DOC(
K
kexinzhao 已提交
183 184 185 186 187
NCCLAllReduce Operator.

AllReduce the input tensors.

)DOC");
D
dzhwinter 已提交
188
  }
D
dongzhihong 已提交
189
};
D
dzhwinter 已提交
190

D
Dong Zhihong 已提交
191 192
// ReduceOp
class NCCLReduceOpMaker : public framework::OpProtoAndCheckerMaker {
D
Dong Zhihong 已提交
193
 public:
194
  NCCLReduceOpMaker(OpProto *proto, OpAttrChecker *op_checker)
D
Dong Zhihong 已提交
195
      : OpProtoAndCheckerMaker(proto, op_checker) {
D
Dong Zhihong 已提交
196
    AddInput("X", "The input of Reduce op");
D
Dong Zhihong 已提交
197
    AddInput("Communicator", "Communicator for communicating between gpus");
D
Dong Zhihong 已提交
198
    AddOutput("Out", "The output of Reduce op");
D
Dong Zhihong 已提交
199
    AddAttr<std::string>("reduction",
K
kexinzhao 已提交
200
                         "(string, default 'ncclSum') "
D
Dong Zhihong 已提交
201 202
                         "{'ncclMin', 'ncclMax', 'ncclProd', 'ncclSum'}.")
        .SetDefault("ncclSum");
D
Dong Zhihong 已提交
203
    AddAttr<int>("root",
K
kexinzhao 已提交
204 205 206
                 "(int, default kInvalidGPUId) "
                 "Root gpu of the parameter. If not, "
                 "set(platform::kInvalidGPUId). Hashed by name.")
D
Dong Zhihong 已提交
207
        .SetDefault(platform::kInvalidGPUId);
D
Dong Zhihong 已提交
208
    AddComment(R"DOC(
K
kexinzhao 已提交
209 210 211 212 213
NCCLReduce Operator.

Reduce the tensors.

)DOC");
D
Dong Zhihong 已提交
214 215 216
  }
};

D
Dong Zhihong 已提交
217
// BcastOp
D
Dong Zhihong 已提交
218
class NCCLBcastOpMaker : public framework::OpProtoAndCheckerMaker {
D
Dong Zhihong 已提交
219
 public:
220
  NCCLBcastOpMaker(OpProto *proto, OpAttrChecker *op_checker)
D
Dong Zhihong 已提交
221
      : OpProtoAndCheckerMaker(proto, op_checker) {
D
Dong Zhihong 已提交
222
    AddInput("X", "The input of BcastSend op");
D
Dong Zhihong 已提交
223
    AddInput("Communicator", "Communicator for communicating between gpus");
D
Dong Zhihong 已提交
224
    AddOutput("Out", "The output of Bcast");
D
Dong Zhihong 已提交
225
    AddAttr<int>("root",
K
kexinzhao 已提交
226 227 228
                 "(int, default kInvalidGPUId) "
                 "Root gpu of the parameter. If not, "
                 "set(platform::kInvalidGPUId). Hashed by name.")
D
Dong Zhihong 已提交
229
        .SetDefault(platform::kInvalidGPUId);
D
Dong Zhihong 已提交
230
    AddComment(R"DOC(
K
kexinzhao 已提交
231 232 233 234 235
NCCLBcast Operator.

Bcast the tensors.

)DOC");
D
Dong Zhihong 已提交
236 237
  }
};
D
dzhwinter 已提交
238

D
Dong Zhihong 已提交
239 240 241 242
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
243
REGISTER_OPERATOR(ncclInit, ops::NCCLInitOp,
Y
Yang Yang 已提交
244 245 246
                  paddle::framework::EmptyGradOpMaker, ops::NCCLInitOpMaker,
                  ops::NCCLInitOpVarTypeInference,
                  ops::NCCLInitOpShapeInference);
247

D
Dong Zhihong 已提交
248 249
REGISTER_OP_WITHOUT_GRADIENT(ncclAllReduce, ops::NCCLAllReduceOp,
                             ops::NCCLAllReduceOpMaker);
D
Dong Zhihong 已提交
250 251
REGISTER_OP_WITHOUT_GRADIENT(ncclBcast, ops::NCCLBcastOp,
                             ops::NCCLBcastOpMaker);
D
Dong Zhihong 已提交
252 253
REGISTER_OP_WITHOUT_GRADIENT(ncclReduce, ops::NCCLReduceOp,
                             ops::NCCLReduceOpMaker);