[Kunlun] PR2: Support MultiDevicePass and BKCL in parallel executor (#29574)
liuyuhui
4 years ago
committed by
GitHub
parent
0b74428db8
commit
4427df37cf
@ -0,0 +1,131 @@
|
||||
// Copyright (c) 2020 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
|
||||
|
||||
#include "xpu/bkcl.h"
|
||||
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/fluid/framework/details/op_handle_base.h"
|
||||
#include "paddle/fluid/framework/lod_tensor.h"
|
||||
#include "paddle/fluid/framework/scope.h"
|
||||
#include "paddle/fluid/platform/bkcl_helper.h"
|
||||
|
||||
DECLARE_bool(sync_bkcl_allreduce);
|
||||
|
||||
namespace paddle {
|
||||
namespace framework {
|
||||
namespace details {
|
||||
|
||||
class BKCLOpHandleBase : public OpHandleBase {
|
||||
public:
|
||||
BKCLOpHandleBase(ir::Node* node, const std::vector<platform::Place>& places,
|
||||
const platform::BKCLCommunicator* bkcl_ctxs)
|
||||
: OpHandleBase(node), places_(places), bkcl_ctxs_(bkcl_ctxs) {
|
||||
if (bkcl_ctxs == nullptr) {
|
||||
return;
|
||||
}
|
||||
// init device context
|
||||
auto default_bkcl_ctxs = bkcl_ctxs_->DefaultFlatCtx();
|
||||
for (auto& p : places_) {
|
||||
this->SetDeviceContext(p, default_bkcl_ctxs->DevCtx(p));
|
||||
}
|
||||
}
|
||||
|
||||
virtual ~BKCLOpHandleBase() {}
|
||||
|
||||
void SetRunEnv(int run_order, bool use_hierarchical_allreduce) {
|
||||
PADDLE_ENFORCE_GE(
|
||||
run_order, 0,
|
||||
platform::errors::InvalidArgument(
|
||||
"The argument run_order must be >= 0, but got %d.", run_order));
|
||||
PADDLE_ENFORCE_NE(use_hierarchical_allreduce, true,
|
||||
platform::errors::Unimplemented(
|
||||
"xpu doesn't support hierarchical_allreduce"));
|
||||
|
||||
run_order_ = run_order;
|
||||
use_hierarchical_allreduce_ = use_hierarchical_allreduce;
|
||||
|
||||
VLOG(10) << "SetRunEnv "
|
||||
<< " run_order:" << run_order
|
||||
<< ", use_hierarchical_allreduce:" << use_hierarchical_allreduce;
|
||||
|
||||
if (bkcl_ctxs_ == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (!use_hierarchical_allreduce_) {
|
||||
auto ctxs = bkcl_ctxs_->GetFlatCtx(run_order);
|
||||
for (auto& p : places_) {
|
||||
this->SetDeviceContext(p, ctxs->DevCtx(p));
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
void FlatBKCLAllReduce(platform::Place place, const void* sendbuff,
|
||||
void* recvbuff, size_t count, BKCLDataType datatype,
|
||||
BKCLOp op) {
|
||||
PADDLE_ENFORCE_GE(
|
||||
run_order_, 0,
|
||||
platform::errors::InvalidArgument(
|
||||
"The argument run_order_ must be >= 0, but got %d.", run_order_));
|
||||
auto flat_bkcl_ctxs = bkcl_ctxs_->GetFlatCtx(run_order_);
|
||||
int dev_id = BOOST_GET_CONST(platform::XPUPlace, place).device;
|
||||
auto& bkcl_ctx = flat_bkcl_ctxs->at(dev_id);
|
||||
auto comm = bkcl_ctx.comm_;
|
||||
|
||||
VLOG(10) << "before all reduce buffer:" << sendbuff << ", numel:" << count
|
||||
<< ", dev_id:" << dev_id << ", dtype:" << datatype
|
||||
<< ", place:" << place;
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
bkcl_all_reduce(comm, sendbuff, recvbuff, count, datatype, op, NULL),
|
||||
BKCL_SUCCESS,
|
||||
platform::errors::PreconditionNotMet("bckl all reduce failed"));
|
||||
}
|
||||
|
||||
void BKCLAllReduce(platform::Place place, const void* sendbuff,
|
||||
void* recvbuff, size_t count, BKCLDataType datatype,
|
||||
BKCLOp op) {
|
||||
PADDLE_ENFORCE_GE(
|
||||
run_order_, 0,
|
||||
platform::errors::InvalidArgument(
|
||||
"The argument run_order_ must be >= 0, but got %d.", run_order_));
|
||||
PADDLE_ENFORCE_EQ(use_hierarchical_allreduce_, false,
|
||||
platform::errors::Unimplemented(
|
||||
"xpu doesn't support hierarchical all reduce"));
|
||||
if (!use_hierarchical_allreduce_) {
|
||||
FlatBKCLAllReduce(place, sendbuff, recvbuff, count, datatype, op);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
protected:
|
||||
std::vector<platform::Place> places_;
|
||||
const platform::BKCLCommunicator* bkcl_ctxs_{nullptr};
|
||||
// When multi trainer call collective function, they need run the same order.
|
||||
// Or the program will hang.So we use allreduce_deps_pass to set this
|
||||
// run_order_.
|
||||
int run_order_{0};
|
||||
// Use 2d allreduce or not.
|
||||
bool use_hierarchical_allreduce_{false};
|
||||
};
|
||||
|
||||
} // namespace details
|
||||
} // namespace framework
|
||||
} // namespace paddle
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue