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/distributed/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");
63
    distributed::RPCClient* client =
W
Wu Yi 已提交
64
        distributed::RPCClient::GetInstance<RPCCLIENT_T>(0);
G
gongweibao 已提交
65

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

  void GetIdByServer(framework::Scope* scope,
                     const platform::DeviceContext& dev_ctx) const {
    std::string endpoint = Attr<std::string>("endpoint");
T
typhoonzero 已提交
81 82 83
    // 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.
84 85
    distributed::RequestSendHandler rpc_h(true);
    std::unique_ptr<distributed::RPCServer> rpc_service(
G
gongweibao 已提交
86 87
        new RPCSERVER_T(endpoint, 1));

88
    rpc_service->RegisterRPC(distributed::kRequestSend, &rpc_h);
G
gongweibao 已提交
89
    rpc_h.SetRPCServer(rpc_service.get());
90

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

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

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

class GenNCCLIdOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
T
typhoonzero 已提交
113
  void Make() override {
T
typhoonzero 已提交
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
    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 已提交
141
REGISTER_OPERATOR(gen_nccl_id, ops::GenNCCLIdOp, ops::GenNCCLIdOpMaker);