Skip to content

Commit

Permalink
fix the bug of margin cross entropy loss for eager mode (PaddlePaddle…
Browse files Browse the repository at this point in the history 
…#43161)
  • Loading branch information
@GuoxiaWang
GuoxiaWang authored Jun 2, 2022
1 parent 85baa3c commit 67163fb
Show file tree
Hide file tree
Showing 3 changed files with 99 additions and 35 deletions.
2 changes: 1 addition & 1 deletion paddle/fluid/eager/api/utils/tensor_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
#pragma once

#include "paddle/fluid/eager/eager_tensor.h"
#include "paddle/phi/api/all.h"
#include "paddle/phi/api/include/tensor.h"

namespace egr {
namespace egr_utils_api {
Expand Down
128 changes: 94 additions & 34 deletions paddle/fluid/operators/margin_cross_entropy_op.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,12 @@ namespace cub = hipcub;
#include "paddle/fluid/operators/reduce_ops/reduce_op.cu.h"
#include "paddle/fluid/operators/reduce_ops/reduce_op.h"
#include "paddle/fluid/string/string_helper.h"
#include "paddle/phi/api/include/tensor.h"
#include "paddle/phi/kernels/funcs/axis_utils.h"
#include "paddle/phi/kernels/funcs/math_function.h"

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
#include "paddle/fluid/distributed/collective/ProcessGroup.h"
#include "paddle/fluid/platform/collective_helper.h"
#include "paddle/fluid/platform/device/gpu/nccl_helper.h"
#endif
Expand Down Expand Up @@ -63,19 +65,34 @@ void GetClassInterval(const gpuStream_t& stream, const platform::Place& place,
framework::TensorFromVector(shard_dim_vec, ctx, &num_classes_per_device);
int* num_classes_per_device_ptr = num_classes_per_device.data<int>();

const auto& comm = platform::NCCLCommContext::Instance().Get(rid, place);
// use global calculate stream
const auto calcu_stream =
static_cast<platform::CUDADeviceContext*>(
platform::DeviceContextPool::Instance().Get(place))
->stream();

PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
num_classes_per_device_ptr, num_classes_per_device_ptr,
num_classes_per_device.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(num_classes_per_device.dtype())),
ncclSum, comm->comm(), calcu_stream));
auto map = distributed::ProcessGroupMapFromGid::getInstance();
if (map->has(rid)) {
// Use ProcessGroup
distributed::ProcessGroup* pg = map->get(rid);
std::vector<phi::DenseTensor> in_tensor;
std::vector<phi::DenseTensor> out_tensor;
in_tensor.push_back(num_classes_per_device);
out_tensor.push_back(num_classes_per_device);

distributed::AllreduceOptions opts;
opts.reduce_op = distributed::ReduceOp::SUM;
auto task = pg->AllReduce(in_tensor, out_tensor, opts);
task->Wait();
} else {
const auto& comm = platform::NCCLCommContext::Instance().Get(rid, place);
// use global calculate stream
const auto calcu_stream =
static_cast<platform::CUDADeviceContext*>(
platform::DeviceContextPool::Instance().Get(place))
->stream();

PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
num_classes_per_device_ptr, num_classes_per_device_ptr,
num_classes_per_device.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(num_classes_per_device.dtype())),
ncclSum, comm->comm(), calcu_stream));
}

auto class_interval_ptr =
class_interval->mutable_data<int>({nranks + 1}, place);
Expand Down Expand Up @@ -228,14 +245,21 @@ class MarginCrossEntropyOpCUDAKernel : public framework::OpKernel<T> {

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
platform::NCCLComm* comm;
distributed::ProcessGroup* pg = nullptr;
gpuStream_t stream;
if (nranks > 1) {
comm = platform::NCCLCommContext::Instance().Get(rid, place);

// use global calculate stream
stream = static_cast<platform::CUDADeviceContext*>(
platform::DeviceContextPool::Instance().Get(place))
->stream();
auto map = distributed::ProcessGroupMapFromGid::getInstance();
if (map->has(rid)) {
// Use ProcessGroup
pg = map->get(rid);
} else {
comm = platform::NCCLCommContext::Instance().Get(rid, place);

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

Expand Down Expand Up @@ -306,11 +330,23 @@ class MarginCrossEntropyOpCUDAKernel : public framework::OpKernel<T> {

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
logits_max_buff, logits_max_buff, logits_max.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(logits_max.dtype())),
ncclMax, comm->comm(), stream));
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 {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
logits_max_buff, logits_max_buff, logits_max.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(logits_max.dtype())),
ncclMax, comm->comm(), stream));
}
}
#endif

Expand All @@ -329,11 +365,23 @@ class MarginCrossEntropyOpCUDAKernel : public framework::OpKernel<T> {

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
sum_exp_logits_buff, sum_exp_logits_buff, sum_exp_logits.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(sum_exp_logits.dtype())),
ncclSum, comm->comm(), stream));
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 {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
sum_exp_logits_buff, sum_exp_logits_buff, sum_exp_logits.numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(sum_exp_logits.dtype())),
ncclSum, comm->comm(), stream));
}
}
#endif

Expand Down Expand Up @@ -363,11 +411,23 @@ class MarginCrossEntropyOpCUDAKernel : public framework::OpKernel<T> {

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
if (nranks > 1) {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
loss_ptr, loss_ptr, loss->numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(loss->dtype())),
ncclSum, comm->comm(), stream));
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 {
PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::ncclAllReduce(
loss_ptr, loss_ptr, loss->numel(),
platform::ToNCCLDataType(
framework::TransToProtoVarType(loss->dtype())),
ncclSum, comm->comm(), stream));
}
}
#endif
}
Expand Down
4 changes: 4 additions & 0 deletions python/paddle/fluid/tests/unittests/test_parallel_margin_cross_entropy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from __future__ import print_function

import os
import unittest
import paddle.fluid as fluid

Expand All @@ -23,7 +24,10 @@
class TestParallelMarginSoftmaxWithCrossEntropy(TestMultipleGpus):
def test_parallel_margin_cross_entropy(self):
self.run_mnist_2gpu('parallel_margin_cross_entropy.py')
self.run_mnist_2gpu(
'parallel_margin_cross_entropy.py', eager_mode=False)


if __name__ == "__main__":
os.environ["FLAGS_enable_eager_mode"] = "1"
unittest.main()

0 comments on commit 67163fb

Please sign in to comment.