wehub-resource-sync
1084 lines
40 KiB
C++
1084 lines
40 KiB
C++
// Copyright (c) 2022 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 "paddle/fluid/distributed/collective/process_group_bkcl.h"
|
|
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/fluid/distributed/collective/bkcl_tools.h"
|
|
#include "paddle/fluid/distributed/collective/common.h"
|
|
#include "paddle/fluid/distributed/collective/process_group_kernel_utils.h"
|
|
#include "paddle/phi/api/lib/utils/allocator.h"
|
|
#include "paddle/phi/core/device_context.h"
|
|
#include "paddle/phi/core/distributed/check/bkcl_dynamic_check.h"
|
|
#include "paddle/phi/core/distributed/check/static_check.h"
|
|
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
|
#include "paddle/phi/core/distributed/utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/memory/allocation/allocator_facade.h"
|
|
#include "paddle/phi/core/memory/memcpy.h"
|
|
#include "paddle/phi/core/platform/device/xpu/bkcl_helper.h"
|
|
#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
|
|
#include "paddle/utils/string/string_helper.h"
|
|
|
|
COMMON_DECLARE_bool(enable_bkcl_dynamic_check);
|
|
|
|
namespace paddle {
|
|
namespace distributed {
|
|
|
|
using phi::distributed::CheckSizeOnEachRank;
|
|
|
|
ProcessGroupBKCL::BKCLTask::BKCLTask(const Place& place,
|
|
int rank,
|
|
CommType comm_type,
|
|
bool sync_op,
|
|
bool use_calc_stream)
|
|
: TaskStream(rank, comm_type, sync_op, use_calc_stream), place_(place) {
|
|
comm_event_ = std::make_shared<XPUEventManager>();
|
|
}
|
|
|
|
ProcessGroupBKCL::BKCLTask::~BKCLTask() {}
|
|
|
|
bool ProcessGroupBKCL::BKCLTask::IsCompleted() {
|
|
LOG_FIRST_N(WARNING, 1) << "XPU do not support event query now.";
|
|
return true;
|
|
}
|
|
|
|
// TODO(sheniang03): Add timeout for wait, now timeout unused
|
|
bool ProcessGroupBKCL::BKCLTask::Wait(std::chrono::milliseconds timeout) {
|
|
const auto* calc_ctx =
|
|
static_cast<XPUContext*>(phi::DeviceContextPool::Instance().Get(place_));
|
|
if (barrier_) {
|
|
// If we use the work to do barrier, we should block cpu
|
|
|
|
// TODO(zhangxiaoci) There is no such function that can sync entire device
|
|
// for xpu (for now), so all we can do is sync whatever stream that we know
|
|
// and hope for the best. Note that for correctness the communication stream
|
|
// needs to be in sync mode.
|
|
phi::backends::xpu::XPUDeviceGuard guard(place_.GetDeviceId());
|
|
xpu_wait();
|
|
calc_ctx->Wait();
|
|
}
|
|
// Warning here when use calc stream but also invoke waiting explicitly.
|
|
if (UseCalcStream()) {
|
|
VLOG(3) << "Warning: The communication is on calc stream, wait here is "
|
|
"useless.";
|
|
return true;
|
|
}
|
|
|
|
comm_event_->Block(*calc_ctx);
|
|
|
|
return true;
|
|
}
|
|
|
|
// Same as Wait
|
|
void ProcessGroupBKCL::BKCLTask::Synchronize() { Wait(kWaitTimeout); }
|
|
|
|
ProcessGroupBKCL::ProcessGroupBKCL(
|
|
const std::shared_ptr<phi::distributed::Store>& store,
|
|
int rank,
|
|
int size,
|
|
int gid)
|
|
: ProcessGroupWithStream(rank, size, gid), store_(store) {}
|
|
|
|
void ProcessGroupBKCL::GroupStart() {
|
|
PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_group_start());
|
|
}
|
|
|
|
void ProcessGroupBKCL::GroupEnd() {
|
|
PADDLE_ENFORCE_BKCL_SUCCESS(bkcl_group_end());
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Recv(
|
|
DenseTensor* tensor,
|
|
int src_rank,
|
|
int64_t offset,
|
|
int64_t numel,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
// numel > 0 indicates the tensor need to be sliced
|
|
DenseTensor partial_tensor;
|
|
if (numel > 0) {
|
|
partial_tensor = GetPartialTensor(*tensor, offset, numel);
|
|
tensor = &partial_tensor;
|
|
}
|
|
|
|
return Point2Point(
|
|
[&](phi::distributed::BKCLCommContext* comm_context,
|
|
XPUStream stream,
|
|
int rank_in_group) {
|
|
VLOG(3) << "bkcl_recv "
|
|
<< "recvbuff: " << tensor->data()
|
|
<< ", count: " << tensor->numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(tensor->dtype()))
|
|
<< ", src_in_group: " << src_rank
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream
|
|
<< ", rank_in_group: " << rank_in_group << ", nranks: " << size_
|
|
<< ", offset: " << offset << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->Recv(tensor, tensor->numel(), rank_in_group, stream);
|
|
},
|
|
src_rank,
|
|
*tensor,
|
|
CommType::RECV,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Send(
|
|
const DenseTensor& tensor,
|
|
int dst_rank,
|
|
int64_t offset,
|
|
int64_t numel,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(tensor);
|
|
// numel > 0 indicates the tensor need to be sliced
|
|
const DenseTensor& tensor_maybe_partial =
|
|
numel > 0 ? GetPartialTensor(tensor, offset, numel) : tensor;
|
|
|
|
return Point2Point(
|
|
[&](phi::distributed::BKCLCommContext* comm_context,
|
|
XPUStream stream,
|
|
int rank_in_group) {
|
|
VLOG(3) << "bkcl_send "
|
|
<< "sendbuff: " << tensor_maybe_partial.data()
|
|
<< ", count: " << tensor_maybe_partial.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(
|
|
phi::ToBKCLDataType(tensor_maybe_partial.dtype()))
|
|
<< ", dst_in_group: " << dst_rank
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream
|
|
<< ", rank_in_group: " << rank_in_group << ", nranks: " << size_
|
|
<< ", offset: " << offset << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->Send(tensor_maybe_partial,
|
|
tensor_maybe_partial.numel(),
|
|
rank_in_group,
|
|
stream);
|
|
},
|
|
dst_rank,
|
|
tensor_maybe_partial,
|
|
CommType::SEND,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroupBKCL::BKCLTask> ProcessGroupBKCL::CreateTask(
|
|
const Place& place,
|
|
int rank,
|
|
CommType comm_type,
|
|
bool is_sync,
|
|
bool use_calc_stream) {
|
|
return std::make_shared<ProcessGroupBKCL::BKCLTask>(
|
|
place, rank, comm_type, is_sync, use_calc_stream);
|
|
}
|
|
|
|
void ProcessGroupBKCL::BroadcastUniqueBKCLID(BKCLUniqueId* bkcl_id) {
|
|
auto key = "ProcessGroupBKCL/bkcl_ids/" + std::to_string(gid_) + "/0";
|
|
if (rank_ == 0) {
|
|
auto id = std::vector<uint8_t>(
|
|
reinterpret_cast<uint8_t*>(bkcl_id),
|
|
reinterpret_cast<uint8_t*>(bkcl_id) + BKCL_UNIQUE_ID_BYTES);
|
|
store_->set(key, id);
|
|
} else {
|
|
const auto& ret = store_->get(key);
|
|
std::memcpy(bkcl_id, ret.data(), ret.size());
|
|
}
|
|
}
|
|
|
|
void ProcessGroupBKCL::CreateBKCLEnvCache(const Place& place,
|
|
const std::string& place_key) {
|
|
phi::backends::xpu::XPUDeviceGuard guard(place.GetDeviceId());
|
|
|
|
VLOG(3) << "init bkcl rank: " << rank_ << ", nranks: " << size_
|
|
<< ", place: " << place_key;
|
|
|
|
int num_ranks = GetSize();
|
|
int rank = GetRank();
|
|
|
|
phi::distributed::CommContextManager::CreateBKCLCommContext(
|
|
store_, std::to_string(gid_), rank_, size_);
|
|
|
|
auto bkcl_comm_ctx = this->GetCommContext();
|
|
VLOG(3) << "Get nccl comm: " << bkcl_comm_ctx->GetBKCLComm()
|
|
<< " for place_key: " << place_key << " on rank_in_group: " << rank
|
|
<< " nranks: " << num_ranks << " gid: " << gid_;
|
|
|
|
// must use phi::XPUContext here to make sure XPUContext::Init() is called
|
|
auto comm_ctx = std::make_unique<phi::XPUContext>(place, true);
|
|
// comm_ctx does not require a pre-allocated GM buffer
|
|
comm_ctx->x_context()->set_option("XPUAPI_DEFAULT_SIZE", "1");
|
|
comm_ctx->SetBkclContext(bkcl_comm_ctx->GetBKCLComm());
|
|
|
|
calc_event_ = std::make_shared<XPUEventManager>();
|
|
auto* calc_ctx = static_cast<phi::XPUContext*>(
|
|
phi::DeviceContextPool::Instance().Get(place));
|
|
calc_ctx->CreateStream();
|
|
|
|
// Note(lijin23): XPU use calc stream for communication now, so we disable the
|
|
// creation of comm stream to reduce the total number of streams used.
|
|
// comm_ctx->CreateStream();
|
|
|
|
place_to_calc_ctx_.emplace(place_key, calc_ctx);
|
|
place_to_comm_ctx_.emplace(place_key, std::move(comm_ctx));
|
|
}
|
|
|
|
void ProcessGroupBKCL::SyncCalcStream(const Place& place) {
|
|
const std::string& key = GetKeyFromPlace(place);
|
|
const auto* calc_ctx = place_to_calc_ctx_[key];
|
|
const auto* comm_ctx = place_to_comm_ctx_[key].get();
|
|
calc_event_->Record(*calc_ctx);
|
|
calc_event_->Block(*comm_ctx);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Collective(
|
|
std::function<void(phi::distributed::BKCLCommContext*, XPUStream)> fn,
|
|
const std::vector<DenseTensor>& tensors,
|
|
CommType op_type,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
if (!use_calc_stream) {
|
|
VLOG(3) << "For XPU, Communication on non-calc stream has minor effect on "
|
|
"performance and might be conflict with streams in calc_ctx, so "
|
|
"we disable it currently.";
|
|
use_calc_stream = true;
|
|
}
|
|
CheckTensorContiguous(tensors);
|
|
|
|
PADDLE_ENFORCE_GT(
|
|
tensors.size(),
|
|
0,
|
|
common::errors::InvalidArgument("Num of tensors must be greater than 0"));
|
|
const auto& place = tensors[0].place();
|
|
const auto& key = GetKeyFromPlace(place);
|
|
|
|
phi::backends::xpu::XPUDeviceGuard xpu_guard(place);
|
|
|
|
if (!calc_event_ ||
|
|
(place_to_comm_ctx_.find(key) == place_to_comm_ctx_.end())) {
|
|
CreateBKCLEnvCache(place, key);
|
|
}
|
|
|
|
if (!use_calc_stream) {
|
|
SyncCalcStream(place);
|
|
}
|
|
|
|
auto task = CreateTask(place, rank_, op_type, sync_op, use_calc_stream);
|
|
|
|
const auto* calc_ctx = place_to_calc_ctx_.at(key);
|
|
const auto& comm_ctx = place_to_comm_ctx_.at(key);
|
|
auto bkcl_stream = use_calc_stream ? calc_ctx->stream() : comm_ctx->stream();
|
|
|
|
auto bkcl_comm_ctx = this->GetCommContext();
|
|
|
|
fn(bkcl_comm_ctx, bkcl_stream);
|
|
|
|
if (!use_calc_stream) {
|
|
PADDLE_ENFORCE_NOT_NULL(comm_ctx.get(),
|
|
common::errors::Fatal("comm context is nullptr."));
|
|
if (!is_coalescing_) {
|
|
task->comm_event_->Record(*comm_ctx.get());
|
|
} else {
|
|
coalescing_place_keys_.push_back(key);
|
|
}
|
|
}
|
|
|
|
if (sync_op) {
|
|
task->Wait();
|
|
}
|
|
|
|
return task;
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Collective(
|
|
std::function<void(phi::distributed::BKCLCommContext*, XPUStream)> fn,
|
|
const DenseTensor& tensor,
|
|
CommType op_type,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
const std::vector<DenseTensor> tensors = {tensor};
|
|
return Collective(fn, tensors, op_type, sync_op, use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Point2Point(
|
|
std::function<void(phi::distributed::BKCLCommContext*, XPUStream, int)> fn,
|
|
int peer,
|
|
const DenseTensor& tensor,
|
|
CommType comm_type,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
if (!use_calc_stream) {
|
|
VLOG(3) << "For XPU, Communication on non-calc stream has minor effect on "
|
|
"performance and might be conflict with streams in calc_ctx, so "
|
|
"we disable it currently.";
|
|
use_calc_stream = true;
|
|
}
|
|
CheckTensorContiguous(tensor);
|
|
const auto& place = tensor.place();
|
|
|
|
int p2p_target_rank = peer;
|
|
std::string key = GetKeyFromPlace(place);
|
|
|
|
phi::backends::xpu::XPUDeviceGuard xpu_guard(place);
|
|
|
|
if (place_to_comm_ctx_.find(key) == place_to_comm_ctx_.end()) {
|
|
CreateBKCLEnvCache(place, key);
|
|
}
|
|
|
|
if (!use_calc_stream) {
|
|
SyncCalcStream(place);
|
|
}
|
|
|
|
auto task = CreateTask(place, rank_, comm_type, sync_op, use_calc_stream);
|
|
const auto* calc_ctx = place_to_calc_ctx_.at(key);
|
|
const auto& comm_ctx = place_to_comm_ctx_.at(key);
|
|
auto bkcl_stream = use_calc_stream ? calc_ctx->stream() : comm_ctx->stream();
|
|
|
|
auto bkcl_comm_ctx = this->GetCommContext();
|
|
fn(bkcl_comm_ctx, bkcl_stream, p2p_target_rank);
|
|
|
|
if (!use_calc_stream) {
|
|
PADDLE_ENFORCE_NOT_NULL(comm_ctx.get(),
|
|
common::errors::Fatal("comm context is nullptr."));
|
|
if (!is_coalescing_) {
|
|
task->comm_event_->Record(*comm_ctx.get());
|
|
} else {
|
|
coalescing_place_keys_.push_back(key);
|
|
}
|
|
}
|
|
|
|
if (sync_op) {
|
|
task->Wait();
|
|
}
|
|
|
|
return task;
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllReduce(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const AllreduceOptions& opts,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
VLOG(3) << "bkcl_all_reduce"
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", redop: " << ToBKCLRedType(opts.reduce_op)
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
|
|
comm_context->AllReduce(
|
|
out_tensor, in_tensor, ToBKCLRedType(opts.reduce_op), stream);
|
|
},
|
|
in_tensor,
|
|
CommType::ALLREDUCE,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllToAll(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const std::vector<int64_t>& out_size_each_rank,
|
|
const std::vector<int64_t>& in_size_each_rank,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
CheckTensorContiguous(*out_tensor);
|
|
|
|
std::vector<int64_t> out_split_sizes;
|
|
std::vector<int64_t> in_split_sizes;
|
|
bool is_equal_split = false;
|
|
if (out_size_each_rank.empty() && in_size_each_rank.empty()) {
|
|
out_split_sizes =
|
|
std::vector<int64_t>(size_, out_tensor->dims()[0] / size_);
|
|
in_split_sizes = std::vector<int64_t>(size_, in_tensor.dims()[0] / size_);
|
|
is_equal_split = true;
|
|
} else {
|
|
out_split_sizes = out_size_each_rank;
|
|
in_split_sizes = in_size_each_rank;
|
|
}
|
|
|
|
const DDim& out_dim = out_tensor->dims();
|
|
const DDim& in_dim = in_tensor.dims();
|
|
CheckSizeOnEachRank(out_dim, out_split_sizes, size_);
|
|
CheckSizeOnEachRank(in_dim, in_split_sizes, size_);
|
|
|
|
// AllToAllUnequalSplit requires allocating temporary memory and must use
|
|
// calc_stream to ensure the correct lifecycle management of the temporary
|
|
// tensor.
|
|
if (!use_calc_stream) {
|
|
VLOG(3) << "For XPU, Communication on non-calc stream has minor effect on "
|
|
"performance and might be conflict with streams in calc_ctx, so "
|
|
"we disable it currently.";
|
|
use_calc_stream = true;
|
|
}
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
if (FLAGS_enable_bkcl_dynamic_check) {
|
|
phi::distributed::BKCLDynamicCheck::CheckShape(
|
|
*out_tensor,
|
|
in_tensor,
|
|
in_split_sizes,
|
|
rank_,
|
|
size_,
|
|
comm_context->GetBKCLComm());
|
|
}
|
|
|
|
VLOG(3) << "[bkcl_all_to_all] "
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", out_split_sizes: "
|
|
<< string::join_strings(out_split_sizes, ',')
|
|
<< ", in_split_sizes: "
|
|
<< string::join_strings(in_split_sizes, ',')
|
|
<< ", is_equal_split: " << is_equal_split
|
|
<< ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
|
|
if (is_equal_split) {
|
|
comm_context->AllToAll(out_tensor, in_tensor, stream);
|
|
} else {
|
|
int64_t in_row_size =
|
|
in_dim[0] == 0 ? 0 : in_tensor.numel() / in_dim[0];
|
|
int64_t out_row_size =
|
|
out_dim[0] == 0 ? 0 : out_tensor->numel() / out_dim[0];
|
|
|
|
int64_t nranks = size_;
|
|
|
|
std::vector<int64_t> in_numel_vec(nranks);
|
|
std::vector<int64_t> in_offset_vec(nranks);
|
|
std::vector<int64_t> out_numel_vec(nranks);
|
|
std::vector<int64_t> out_offset_vec(nranks);
|
|
|
|
int64_t in_offset = 0;
|
|
int64_t out_offset = 0;
|
|
for (int64_t i = 0; i < nranks; i++) {
|
|
int64_t in_numel = in_split_sizes[i] * in_row_size;
|
|
int64_t out_numel = out_split_sizes[i] * out_row_size;
|
|
|
|
in_numel_vec[i] = in_numel;
|
|
in_offset_vec[i] = in_offset;
|
|
in_offset += in_numel;
|
|
|
|
out_numel_vec[i] = out_numel;
|
|
out_offset_vec[i] = out_offset;
|
|
out_offset += out_numel;
|
|
}
|
|
|
|
PADDLE_ENFORCE_GE(
|
|
in_tensor.place().GetDeviceId(),
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"The all_to_all device id must greater or equal than 0."));
|
|
phi::XPUPlace place = in_tensor.place();
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
auto allocator_cpu = std::unique_ptr<phi::Allocator>(
|
|
new paddle::experimental::DefaultAllocator(CPUPlace()));
|
|
#endif
|
|
auto allocator = std::unique_ptr<phi::Allocator>(
|
|
new paddle::experimental::DefaultAllocator(place));
|
|
DenseTensorMeta meta(DataType::INT64, DDim{nranks});
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
DenseTensor in_size_tensor = {allocator_cpu.get(), meta};
|
|
DenseTensor in_offset_tensor = {allocator_cpu.get(), meta};
|
|
DenseTensor out_size_tensor = {allocator_cpu.get(), meta};
|
|
DenseTensor out_offset_tensor = {allocator_cpu.get(), meta};
|
|
#else
|
|
DenseTensor in_size_tensor = {allocator.get(), meta};
|
|
DenseTensor in_offset_tensor = {allocator.get(), meta};
|
|
DenseTensor out_size_tensor = {allocator.get(), meta};
|
|
DenseTensor out_offset_tensor = {allocator.get(), meta};
|
|
#endif
|
|
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
memory::Copy(CPUPlace(),
|
|
in_size_tensor.data(),
|
|
CPUPlace(),
|
|
in_numel_vec.data(),
|
|
in_size_tensor.numel() * sizeof(int64_t));
|
|
memory::Copy(CPUPlace(),
|
|
in_offset_tensor.data(),
|
|
CPUPlace(),
|
|
in_offset_vec.data(),
|
|
in_offset_tensor.numel() * sizeof(int64_t));
|
|
memory::Copy(CPUPlace(),
|
|
out_size_tensor.data(),
|
|
CPUPlace(),
|
|
out_numel_vec.data(),
|
|
out_size_tensor.numel() * sizeof(int64_t));
|
|
memory::Copy(CPUPlace(),
|
|
out_offset_tensor.data(),
|
|
CPUPlace(),
|
|
out_offset_vec.data(),
|
|
out_offset_tensor.numel() * sizeof(int64_t));
|
|
#else
|
|
|
|
memory::Copy(place,
|
|
in_size_tensor.data(),
|
|
CPUPlace(),
|
|
in_numel_vec.data(),
|
|
in_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
in_offset_tensor.data(),
|
|
CPUPlace(),
|
|
in_offset_vec.data(),
|
|
in_offset_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
out_size_tensor.data(),
|
|
CPUPlace(),
|
|
out_numel_vec.data(),
|
|
out_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
out_offset_tensor.data(),
|
|
CPUPlace(),
|
|
out_offset_vec.data(),
|
|
out_offset_tensor.numel() * sizeof(int64_t));
|
|
#endif
|
|
|
|
comm_context->AllToAllUnequalSplit(out_tensor,
|
|
in_tensor,
|
|
out_size_tensor,
|
|
out_offset_tensor,
|
|
in_size_tensor,
|
|
in_offset_tensor,
|
|
stream);
|
|
}
|
|
},
|
|
in_tensor,
|
|
CommType::ALLTOALL,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllToAll(
|
|
std::vector<DenseTensor>* out_tensors,
|
|
const std::vector<DenseTensor>& in_tensors,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensors);
|
|
CheckTensorContiguous(*out_tensors);
|
|
CheckTensorSamePlace(in_tensors);
|
|
CheckTensorSamePlace(*out_tensors);
|
|
phi::distributed::CommStaticCheck::CheckDataType(*out_tensors, in_tensors);
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
out_tensors->size(),
|
|
size_,
|
|
common::errors::InvalidArgument(
|
|
"Number of out tensors[%d] do not match the world size[%d].",
|
|
out_tensors->size(),
|
|
size_));
|
|
PADDLE_ENFORCE_EQ(
|
|
in_tensors.size(),
|
|
size_,
|
|
common::errors::InvalidArgument(
|
|
"Number of in tensors[%d] do not match the world size[%d].",
|
|
in_tensors.size(),
|
|
size_));
|
|
|
|
// AllToAllUnequalSplit requires allocating temporary memory and must use
|
|
// calc_stream to ensure the correct lifecycle management of the temporary
|
|
// tensor.
|
|
if (!use_calc_stream) {
|
|
VLOG(3) << "For XPU, Communication on non-calc stream has minor effect on "
|
|
"performance and might be conflict with streams in calc_ctx, so "
|
|
"we disable it currently.";
|
|
use_calc_stream = true;
|
|
}
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
if (FLAGS_enable_bkcl_dynamic_check) {
|
|
phi::distributed::BKCLDynamicCheck::CheckAlltoAllShape(
|
|
*out_tensors,
|
|
in_tensors,
|
|
rank_,
|
|
size_,
|
|
comm_context->GetBKCLComm());
|
|
}
|
|
|
|
VLOG(3) << "[AllToAll] "
|
|
<< "sendbuff: "
|
|
<< string::join_strings(GetTensorPtrs(in_tensors), ',')
|
|
<< ", recvbuff: "
|
|
<< string::join_strings(GetTensorPtrs(*out_tensors), ',')
|
|
<< ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensors[0].dtype()))
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", out_split_sizes: "
|
|
<< string::join_strings(GetAllToAllSplitSizes(*out_tensors),
|
|
',')
|
|
<< ", in_split_sizes: "
|
|
<< string::join_strings(GetAllToAllSplitSizes(in_tensors), ',')
|
|
<< ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
|
|
int64_t nranks = size_;
|
|
int64_t in_numel_sum = 0;
|
|
int64_t out_numel_sum = 0;
|
|
|
|
for (int64_t i = 0; i < nranks; i++) {
|
|
in_numel_sum += in_tensors[i].numel();
|
|
out_numel_sum += (*out_tensors)[i].numel();
|
|
}
|
|
|
|
std::vector<int64_t> in_numel_vec(nranks);
|
|
std::vector<int64_t> in_offset_vec(nranks);
|
|
std::vector<int64_t> out_numel_vec(nranks);
|
|
std::vector<int64_t> out_offset_vec(nranks);
|
|
|
|
int64_t in_offset = 0;
|
|
int64_t out_offset = 0;
|
|
for (int64_t i = 0; i < nranks; i++) {
|
|
int64_t in_numel = in_tensors[i].numel();
|
|
int64_t out_numel = (*out_tensors)[i].numel();
|
|
|
|
in_numel_vec[i] = in_numel;
|
|
in_offset_vec[i] = in_offset;
|
|
in_offset += in_numel;
|
|
|
|
out_numel_vec[i] = out_numel;
|
|
out_offset_vec[i] = out_offset;
|
|
out_offset += out_numel;
|
|
}
|
|
|
|
PADDLE_ENFORCE_GE(
|
|
in_tensors[0].place().GetDeviceId(),
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"The all_to_all device id must greater or equal than 0."));
|
|
phi::XPUPlace place = in_tensors[0].place();
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
auto allocator_cpu = std::unique_ptr<phi::Allocator>(
|
|
new paddle::experimental::DefaultAllocator(CPUPlace()));
|
|
#endif
|
|
auto allocator = std::unique_ptr<phi::Allocator>(
|
|
new paddle::experimental::DefaultAllocator(place));
|
|
|
|
DenseTensorMeta concated_in_tensor_meta(in_tensors[0].dtype(),
|
|
DDim{in_numel_sum});
|
|
DenseTensorMeta concated_out_tensor_meta((*out_tensors)[0].dtype(),
|
|
DDim{out_numel_sum});
|
|
DenseTensorMeta split_meta(DataType::INT64, DDim{nranks});
|
|
|
|
DenseTensor concated_in_tensor = {allocator.get(),
|
|
concated_in_tensor_meta};
|
|
DenseTensor concated_out_tensor = {allocator.get(),
|
|
concated_out_tensor_meta};
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
DenseTensor in_size_tensor = {allocator_cpu.get(), split_meta};
|
|
DenseTensor in_offset_tensor = {allocator_cpu.get(), split_meta};
|
|
DenseTensor out_size_tensor = {allocator_cpu.get(), split_meta};
|
|
DenseTensor out_offset_tensor = {allocator_cpu.get(), split_meta};
|
|
#else
|
|
DenseTensor in_size_tensor = {allocator.get(), split_meta};
|
|
DenseTensor in_offset_tensor = {allocator.get(), split_meta};
|
|
DenseTensor out_size_tensor = {allocator.get(), split_meta};
|
|
DenseTensor out_offset_tensor = {allocator.get(), split_meta};
|
|
#endif
|
|
|
|
if (in_numel_sum > 0) {
|
|
ConcatTensorByNumel(*GetDeviceContext(place, use_calc_stream),
|
|
in_tensors,
|
|
&concated_in_tensor);
|
|
}
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
memory::Copy(CPUPlace(),
|
|
in_size_tensor.data(),
|
|
CPUPlace(),
|
|
in_numel_vec.data(),
|
|
in_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(CPUPlace(),
|
|
in_offset_tensor.data(),
|
|
CPUPlace(),
|
|
in_offset_vec.data(),
|
|
in_offset_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(CPUPlace(),
|
|
out_size_tensor.data(),
|
|
CPUPlace(),
|
|
out_numel_vec.data(),
|
|
out_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(CPUPlace(),
|
|
out_offset_tensor.data(),
|
|
CPUPlace(),
|
|
out_offset_vec.data(),
|
|
out_offset_tensor.numel() * sizeof(int64_t));
|
|
#else
|
|
memory::Copy(place,
|
|
in_size_tensor.data(),
|
|
CPUPlace(),
|
|
in_numel_vec.data(),
|
|
in_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
in_offset_tensor.data(),
|
|
CPUPlace(),
|
|
in_offset_vec.data(),
|
|
in_offset_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
out_size_tensor.data(),
|
|
CPUPlace(),
|
|
out_numel_vec.data(),
|
|
out_size_tensor.numel() * sizeof(int64_t));
|
|
|
|
memory::Copy(place,
|
|
out_offset_tensor.data(),
|
|
CPUPlace(),
|
|
out_offset_vec.data(),
|
|
out_offset_tensor.numel() * sizeof(int64_t));
|
|
#endif
|
|
|
|
comm_context->AllToAllUnequalSplit(&concated_out_tensor,
|
|
concated_in_tensor,
|
|
out_size_tensor,
|
|
out_offset_tensor,
|
|
in_size_tensor,
|
|
in_offset_tensor,
|
|
stream);
|
|
|
|
if (out_numel_sum > 0) {
|
|
SplitTensorByNumel(*GetDeviceContext(place, use_calc_stream),
|
|
concated_out_tensor,
|
|
out_tensors);
|
|
}
|
|
},
|
|
in_tensors,
|
|
CommType::ALLTOALL,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Broadcast(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const BroadcastOptions& opts,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
CheckTensorContiguous(*out_tensor);
|
|
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
int root = opts.source_rank + opts.source_root;
|
|
|
|
VLOG(3) << "bkcl_broadcast "
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", root: " << root
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->Broadcast(out_tensor, in_tensor, root, stream);
|
|
},
|
|
in_tensor,
|
|
CommType::BROADCAST,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::AllGather(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
int64_t offset,
|
|
int64_t numel,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
|
|
const DenseTensor& in_tensor_maybe_partial =
|
|
numel > 0 ? GetPartialTensor(in_tensor, offset, numel) : in_tensor;
|
|
phi::distributed::CommStaticCheck::GatherLikeShape(*out_tensor,
|
|
in_tensor_maybe_partial,
|
|
/*dst_rank*/ rank_,
|
|
/*cur_rank*/ rank_,
|
|
size_,
|
|
phi::AllocationType::XPU);
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
VLOG(3) << "bkcl_all_gather "
|
|
<< "sendbuff: " << in_tensor_maybe_partial.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor_maybe_partial.numel()
|
|
<< ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", offset: " << offset
|
|
<< ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
|
|
comm_context->AllGather(out_tensor, in_tensor_maybe_partial, stream);
|
|
},
|
|
in_tensor_maybe_partial,
|
|
CommType::ALLGATHER,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Reduce(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const ReduceOptions& opts,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
CheckTensorContiguous(*out_tensor);
|
|
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
VLOG(3) << "bkcl_reduce "
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", redop: "
|
|
<< BKCLRedTypeToString(ToBKCLRedType(opts.reduce_op))
|
|
<< ", root: " << opts.root_rank
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->Reduce(out_tensor,
|
|
in_tensor,
|
|
ToBKCLRedType(opts.reduce_op),
|
|
opts.root_rank,
|
|
stream);
|
|
},
|
|
in_tensor,
|
|
CommType::REDUCE,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::ReduceScatter(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const ReduceScatterOptions& opts,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
CheckTensorContiguous(*out_tensor);
|
|
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
VLOG(3) << "bkcl_reduce_scatter "
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", redop: "
|
|
<< BKCLRedTypeToString(ToBKCLRedType(opts.reduce_op))
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->ReduceScatter(
|
|
out_tensor, in_tensor, ToBKCLRedType(opts.reduce_op), stream);
|
|
},
|
|
in_tensor,
|
|
CommType::REDUCE_SCATTER,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
|
|
#if defined(PADDLE_WITH_FLAGCX)
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Scatter(
|
|
DenseTensor* out_tensor,
|
|
const DenseTensor& in_tensor,
|
|
const ScatterOptions& opts,
|
|
bool sync_op,
|
|
bool use_calc_stream) {
|
|
CheckTensorContiguous(in_tensor);
|
|
CheckTensorContiguous(*out_tensor);
|
|
|
|
phi::distributed::CommStaticCheck::ScatterLikeShape(
|
|
*out_tensor,
|
|
in_tensor,
|
|
/*dst_rank*/ opts.root_rank,
|
|
/*cur_rank*/ rank_,
|
|
size_,
|
|
phi::AllocationType::XPU);
|
|
return Collective(
|
|
[&](phi::distributed::BKCLCommContext* comm_context, XPUStream stream) {
|
|
VLOG(3) << "bkcl_scatter "
|
|
<< "sendbuff: " << in_tensor.data()
|
|
<< ", recvbuff: " << out_tensor->data()
|
|
<< ", count: " << in_tensor.numel() << ", datatype: "
|
|
<< BKCLDTypeToString(phi::ToBKCLDataType(in_tensor.dtype()))
|
|
<< ", bkcl_comm: " << comm_context->GetBKCLComm()
|
|
<< ", stream: " << stream << ", rank_in_group: " << rank_
|
|
<< ", nranks: " << size_ << ", sync_op: " << sync_op
|
|
<< ", use_calc_stream: " << use_calc_stream;
|
|
comm_context->Scatter(out_tensor, in_tensor, opts.root_rank, stream);
|
|
},
|
|
in_tensor,
|
|
CommType::SCATTER,
|
|
sync_op,
|
|
use_calc_stream);
|
|
}
|
|
#endif
|
|
|
|
std::shared_ptr<ProcessGroup::Task> ProcessGroupBKCL::Barrier(
|
|
const BarrierOptions& opts) {
|
|
PADDLE_ENFORCE_GE(opts.device_id,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"The barrier device id must greater or equal than 0."));
|
|
phi::XPUPlace place(opts.device_id);
|
|
auto allocator = std::unique_ptr<phi::Allocator>(
|
|
new paddle::experimental::DefaultAllocator(place));
|
|
DenseTensorMeta meta(DataType::FLOAT32, DDim{1});
|
|
DenseTensor barrier_tensor{allocator.get(), meta};
|
|
|
|
auto task = AllReduce(&barrier_tensor,
|
|
barrier_tensor,
|
|
{},
|
|
/*sync_op*/ true,
|
|
/*use_calc_stream*/ false);
|
|
auto bkcl_task = dynamic_cast<BKCLTask*>(task.get());
|
|
bkcl_task->barrier_ = true;
|
|
return task;
|
|
}
|
|
|
|
phi::DeviceContext* ProcessGroupBKCL::GetDeviceContext(
|
|
const Place& place) const {
|
|
return GetDeviceContext(place, /*use_calc_stream*/ false);
|
|
}
|
|
|
|
phi::DeviceContext* ProcessGroupBKCL::GetDeviceContext(
|
|
const Place& place, bool use_calc_stream) const {
|
|
if (!use_calc_stream) {
|
|
VLOG(3) << "For XPU, Communication on non-calc stream has minor effect on "
|
|
"performance and might be conflict with streams in calc_ctx, so "
|
|
"we disable it currently.";
|
|
use_calc_stream = true;
|
|
}
|
|
const std::string& key = GetKeyFromPlace(place);
|
|
if (use_calc_stream) {
|
|
const auto& iter = place_to_calc_ctx_.find(key);
|
|
PADDLE_ENFORCE_NE(iter,
|
|
place_to_calc_ctx_.end(),
|
|
common::errors::InvalidArgument(
|
|
"Cannot find device context in process group."));
|
|
return iter->second;
|
|
} else {
|
|
const auto& iter = place_to_comm_ctx_.find(key);
|
|
PADDLE_ENFORCE_NE(iter,
|
|
place_to_comm_ctx_.end(),
|
|
common::errors::InvalidArgument(
|
|
"Cannot find device context in process group."));
|
|
return iter->second.get();
|
|
}
|
|
}
|
|
|
|
std::shared_ptr<ProcessGroupBKCL> ProcessGroupBKCL::CreateProcessGroupBKCL(
|
|
const std::shared_ptr<phi::distributed::Store>& store,
|
|
int rank,
|
|
int size,
|
|
int gid) {
|
|
auto process_group =
|
|
std::make_shared<ProcessGroupBKCL>(store, rank, size, gid);
|
|
ProcessGroupIdMap::GetInstance().emplace(gid, process_group);
|
|
return process_group;
|
|
}
|
|
|
|
phi::distributed::BKCLCommContext* ProcessGroupBKCL::GetOrCreateCommContext(
|
|
const Place& place, CommType comm_type) {
|
|
const auto& key = GetKeyFromPlace(place);
|
|
if (place_to_comm_ctx_.find(key) == place_to_comm_ctx_.end()) {
|
|
CreateBKCLEnvCache(place, key);
|
|
}
|
|
return GetCommContext();
|
|
}
|
|
|
|
phi::distributed::BKCLCommContext* ProcessGroupBKCL::GetCommContext() {
|
|
const auto& comm_context_manager =
|
|
phi::distributed::CommContextManager::GetInstance();
|
|
auto comm_context = static_cast<phi::distributed::BKCLCommContext*>(
|
|
comm_context_manager.Get(std::to_string(this->gid_)));
|
|
PADDLE_ENFORCE_NE(comm_context,
|
|
nullptr,
|
|
common::errors::Unavailable("BKCLCommContext is nullptr"));
|
|
return comm_context;
|
|
}
|
|
|
|
void ProcessGroupBKCL::StartCoalescing() {
|
|
PADDLE_ENFORCE_EQ(is_coalescing_,
|
|
false,
|
|
common::errors::PreconditionNotMet(
|
|
"Coalescing is on, please call EndCoalesce."));
|
|
is_coalescing_ = true;
|
|
GroupStart();
|
|
}
|
|
|
|
void ProcessGroupBKCL::EndCoalescing(
|
|
std::optional<std::vector<std::shared_ptr<ProcessGroup::Task>>> tasks_opt) {
|
|
GroupEnd();
|
|
|
|
// NOTE(shenliang03): If using calculate stream, no need to record stream and
|
|
// update task.
|
|
if (!tasks_opt.has_value() | coalescing_place_keys_.empty()) {
|
|
is_coalescing_ = false;
|
|
return;
|
|
}
|
|
|
|
auto& tasks = tasks_opt.value();
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
tasks.size(),
|
|
coalescing_place_keys_.size(),
|
|
common::errors::PreconditionNotMet(
|
|
"Number of tasks[%d] do not match number of collectives[%d].",
|
|
tasks.size(),
|
|
coalescing_place_keys_.size()));
|
|
|
|
for (size_t i = 0; i < tasks.size(); ++i) {
|
|
auto* task = static_cast<ProcessGroupBKCL::BKCLTask*>(tasks[i].get());
|
|
const auto& key = coalescing_place_keys_[i];
|
|
const auto& comm_ctx = place_to_comm_ctx_.at(key);
|
|
task->comm_event_->Record(*comm_ctx.get());
|
|
}
|
|
|
|
is_coalescing_ = false;
|
|
coalescing_place_keys_.clear();
|
|
}
|
|
|
|
} // namespace distributed
|
|
} // namespace paddle
|