Skip to content

[fluid_ops] clean margin_cross_entropy_kernel ProcessGroup #70421

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 25, 2024
Merged

[fluid_ops] clean margin_cross_entropy_kernel ProcessGroup #70421

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 15 additions & 69 deletions paddle/phi/kernels/gpu/margin_cross_entropy_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,39 +132,21 @@ void MarginCrossEntropyKernel(const Context& dev_ctx,
const auto& place = dev_ctx.GetPlace(); // old code

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
const auto& comm_context_manager =
phi::distributed::CommContextManager::GetInstance();
phi::distributed::NCCLCommContext* comm_ctx = nullptr;
distributed::ProcessGroup* pg = nullptr;
gpuStream_t stream;
if (nranks > 1) {
auto map = distributed::ProcessGroupMapFromGid::getInstance();
if (map->has(ring_id)) {
// Use ProcessGroup
pg = map->get(ring_id);
} else {
PADDLE_ENFORCE_EQ(comm_context_manager.Has(std::to_string(ring_id)),
true,
common::errors::InvalidArgument(
"You choose to use new communication library by "
"setting environment "
"variable FLAGS_dynamic_static_unified_comm True. "
"But ring_id(%d) is "
"not found in comm_context_manager.",
std::to_string(ring_id)));
comm_ctx = static_cast<phi::distributed::NCCLCommContext*>(
comm_context_manager.Get(std::to_string(ring_id)));
PADDLE_ENFORCE_NE(comm_ctx,
nullptr,
common::errors::Unavailable(
"NCCLCommContext is nullptr, collective op should "
"has ring_id attr."));

// use global calculate stream
stream = static_cast<GPUContext*>(
phi::DeviceContextPool::Instance().Get(place))
->stream();
}
comm_ctx = static_cast<phi::distributed::NCCLCommContext*>(
dev_ctx.GetCommContext());
PADDLE_ENFORCE_NE(comm_ctx,
nullptr,
common::errors::Unavailable(
"NCCLCommContext is nullptr, collective op should "
"has ring_id attr."));

// use global calculate stream
stream =
static_cast<GPUContext*>(phi::DeviceContextPool::Instance().Get(place))
->stream();
}
#endif

Expand Down Expand Up @@ -261,19 +243,7 @@ void MarginCrossEntropyKernel(const Context& dev_ctx,

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
if (pg) {
std::vector<phi::DenseTensor> in_tensor;
std::vector<phi::DenseTensor> out_tensor;
in_tensor.push_back(logits_max);
out_tensor.push_back(logits_max);

distributed::AllreduceOptions opts;
opts.reduce_op = distributed::ReduceOp::MAX;
auto task = pg->AllReduce(in_tensor, out_tensor, opts);
task->Wait();
} else {
comm_ctx->AllReduce(&logits_max, logits_max, ncclMax, stream);
}
comm_ctx->AllReduce(&logits_max, logits_max, ncclMax, stream);
}
#endif

Expand All @@ -295,19 +265,7 @@ void MarginCrossEntropyKernel(const Context& dev_ctx,

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
if (pg) {
std::vector<phi::DenseTensor> in_tensor;
std::vector<phi::DenseTensor> out_tensor;
in_tensor.push_back(sum_exp_logits);
out_tensor.push_back(sum_exp_logits);

distributed::AllreduceOptions opts;
opts.reduce_op = distributed::ReduceOp::SUM;
auto task = pg->AllReduce(in_tensor, out_tensor, opts);
task->Wait();
} else {
comm_ctx->AllReduce(&sum_exp_logits, sum_exp_logits, ncclSum, stream);
}
comm_ctx->AllReduce(&sum_exp_logits, sum_exp_logits, ncclSum, stream);
}
#endif

Expand Down Expand Up @@ -346,19 +304,7 @@ void MarginCrossEntropyKernel(const Context& dev_ctx,

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
if (pg) {
std::vector<phi::DenseTensor> in_tensor;
std::vector<phi::DenseTensor> out_tensor;
in_tensor.push_back(*loss);
out_tensor.push_back(*loss);

distributed::AllreduceOptions opts;
opts.reduce_op = distributed::ReduceOp::SUM;
auto task = pg->AllReduce(in_tensor, out_tensor, opts);
task->Wait();
} else {
comm_ctx->AllReduce(loss, *loss, ncclSum, stream);
}
comm_ctx->AllReduce(loss, *loss, ncclSum, stream);
}
#endif
}
Expand Down