gen_nccl_id_op.cc 4.9 KB
Newer Older
T
typhoonzero 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
/* Copyright (c) 2016 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 <nccl.h>
#include <stdint.h>
#include <ostream>
#include <string>

#include "paddle/fluid/framework/executor.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/threadpool.h"
G
gongweibao 已提交
24
#include "paddle/fluid/operators/detail/macros.h"
25
#include "paddle/fluid/operators/detail/request_handler_impl.h"
Y
update  
yi.wu 已提交
26
#include "paddle/fluid/platform/nccl_helper.h"
T
typhoonzero 已提交
27 28 29 30 31 32 33 34 35 36 37 38 39 40

namespace paddle {
namespace operators {

class GenNCCLIdOp : public framework::OperatorBase {
 public:
  GenNCCLIdOp(const std::string& type, const framework::VariableNameMap& inputs,
              const framework::VariableNameMap& outputs,
              const framework::AttributeMap& attrs)
      : OperatorBase(type, inputs, outputs, attrs) {}

  void RunImpl(const framework::Scope& scope,
               const platform::Place& dev_place) const override {
    platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
T
typhoonzero 已提交
41 42
    // put nccl id in CPUPlace
    auto& dev_ctx = *pool.Get(platform::CPUPlace());
T
typhoonzero 已提交
43 44 45 46 47 48 49 50 51 52 53 54 55
    int trainer_id = Attr<int>("trainer_id");
    framework::Scope& local_scope = scope.NewScope();

    if (trainer_id == 0) {
      GenerateAndSend(&local_scope, dev_ctx);
    } else {
      GetIdByServer(&local_scope, dev_ctx);
    }
  }

 private:
  void GenerateAndSend(framework::Scope* scope,
                       const platform::DeviceContext& dev_ctx) const {
T
typhoonzero 已提交
56
    auto var = scope->FindVar(NCCL_ID_VARNAME);
T
typhoonzero 已提交
57 58
    PADDLE_ENFORCE_NOT_NULL(var);
    auto id = var->GetMutable<ncclUniqueId>();
T
typhoonzero 已提交
59
    PADDLE_ENFORCE(platform::dynload::ncclGetUniqueId(id));
T
typhoonzero 已提交
60 61 62

    std::vector<std::string> endpoint_list =
        Attr<std::vector<std::string>>("endpoint_list");
G
gongweibao 已提交
63 64
    detail::RPCClient* client = detail::RPCClient::GetInstance<RPCCLIENT_T>();

T
typhoonzero 已提交
65
    for (auto& ep : endpoint_list) {
T
typhoonzero 已提交
66
      VLOG(3) << "sending nccl id to " << ep;
G
gongweibao 已提交
67
      client->AsyncSendVar(ep, dev_ctx, *scope, NCCL_ID_VARNAME);
T
typhoonzero 已提交
68
    }
G
gongweibao 已提交
69
    client->Wait();
Y
yi.wu 已提交
70 71 72 73
    for (auto& ep : endpoint_list) {
      client->AsyncSendBatchBarrier(ep);
    }
    client->Wait();
T
typhoonzero 已提交
74
    VLOG(3) << "sending completed...";
T
typhoonzero 已提交
75 76 77 78 79
  }

  void GetIdByServer(framework::Scope* scope,
                     const platform::DeviceContext& dev_ctx) const {
    std::string endpoint = Attr<std::string>("endpoint");
T
typhoonzero 已提交
80 81 82
    // NOTE: Can not use unique_ptr here because the default
    // deleter will call GRPC Server's base class's dtor and
    // that will cause a wired crash.
83
    detail::RequestSendHandler rpc_h(true);
G
gongweibao 已提交
84 85 86 87 88
    std::unique_ptr<detail::RPCServer> rpc_service(
        new RPCSERVER_T(endpoint, 1));

    rpc_service->RegisterRPC(detail::kRequestSend, &rpc_h);
    rpc_h.SetRPCServer(rpc_service.get());
89

T
typhoonzero 已提交
90 91
    framework::ProgramDesc empty_program;
    framework::Executor executor(dev_ctx.GetPlace());
92 93 94 95
    rpc_h.SetScope(scope);
    rpc_h.SetDevCtx(&dev_ctx);
    rpc_h.SetProgram(&empty_program);
    rpc_h.SetExecutor(&executor);
T
typhoonzero 已提交
96

97
    std::thread server_thread(
G
gongweibao 已提交
98 99 100
        std::bind(&detail::RPCServer::StartServer, rpc_service.get()));

    rpc_service->SetCond(detail::kRequestSend);
T
typhoonzero 已提交
101
    VLOG(3) << "start getting nccl id from trainer 0...";
G
gongweibao 已提交
102
    rpc_service->WaitBarrier(detail::kRequestSend);
T
typhoonzero 已提交
103
    VLOG(3) << "got nccl id and stop server...";
G
gongweibao 已提交
104
    rpc_service->ShutDown();
T
typhoonzero 已提交
105
    VLOG(3) << "rpc server stopped";
106
    server_thread.join();
T
typhoonzero 已提交
107 108 109 110 111
  }
};

class GenNCCLIdOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
T
typhoonzero 已提交
112
  void Make() override {
T
typhoonzero 已提交
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
    AddOutput("NCCLID", "Raw variable contains a NCCL UniqueId instaces.");
    AddComment(R"DOC(
GenNCCLId operator

For trainer 0: generate a new UniqueId and send it to all the other trainers.
For trainer 1~n: start a gRPC server to get the UniqueId, once got, stop the server.
)DOC");
    AddAttr<std::string>("endpoint",
                         "(string), e.g. 127.0.0.1:6175 "
                         "current listen endpoint");
    AddAttr<std::vector<std::string>>(
        "endpoint_list",
        "['trainer1_ip:port', 'trainer2_ip:port', ...] "
        "list of trainer endpoints start from trainer 1")
        .SetDefault({});
    AddAttr<int>("trainer_id",
                 "(int default 0) "
                 "The index of the trainer in distributed training.")
        .SetDefault(0);
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

T
typhoonzero 已提交
140
REGISTER_OPERATOR(gen_nccl_id, ops::GenNCCLIdOp, ops::GenNCCLIdOpMaker);