wehub-resource-sync
162 lines
7.2 KiB
C++
162 lines
7.2 KiB
C++
// Copyright (c) 2025 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/phi/api/backward/backward_api.h"
|
|
#include "paddle/phi/api/include/api.h"
|
|
#include "paddle/phi/backends/all_context.h"
|
|
#include "paddle/phi/backends/device_manager.h"
|
|
#include "paddle/phi/core/distributed/collective/process_group.h"
|
|
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
|
#include "paddle/phi/core/distributed/xccl_comm_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/kernels/funcs/axis_utils.h"
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void CSoftmaxWithEntropyKernel(const Context& dev_ctx,
|
|
const DenseTensor& logits_in,
|
|
const DenseTensor& label_in,
|
|
int64_t ignore_index,
|
|
int rank,
|
|
int nranks,
|
|
DenseTensor* softmax,
|
|
DenseTensor* loss) {
|
|
auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
|
|
dev_ctx.GetCommContext());
|
|
PADDLE_ENFORCE_NE(comm,
|
|
nullptr,
|
|
common::errors::Unavailable(
|
|
"XCCLCommContext is nullptr, collective op should "
|
|
"has ring_id attr."));
|
|
|
|
const DenseTensor* logits = &logits_in;
|
|
const DenseTensor* labels = &label_in;
|
|
auto softmax_dims = softmax->dims();
|
|
auto loss_dims = loss->dims();
|
|
|
|
const int axis = logits->dims().size() - 1;
|
|
const int N = funcs::SizeToAxis(axis, logits->dims());
|
|
const int D = funcs::SizeFromAxis(axis, logits->dims());
|
|
|
|
auto logits_2d = std::make_shared<DenseTensor>();
|
|
auto labels_1d = std::make_shared<DenseTensor>();
|
|
logits_2d->ShareDataWith(*logits).Resize({N, D});
|
|
labels_1d->ShareDataWith(*labels).Resize({N});
|
|
paddle::Tensor logits_2d_tensor(logits_2d), labels_1d_tensor(labels_1d);
|
|
|
|
// step 1, obtain logit_max
|
|
auto logits_2d_max_tensor = logits_2d_tensor.max({1}, true);
|
|
auto logits_2d_max =
|
|
reinterpret_cast<DenseTensor*>(logits_2d_max_tensor.impl().get());
|
|
auto& stream = *dev_ctx.GetStream();
|
|
phi::DeviceManager::CCLAllReduce(dev_ctx.GetPlace().GetDeviceType(),
|
|
logits_2d_max->data<float>(),
|
|
logits_2d_max->data<float>(),
|
|
logits_2d_max->numel(),
|
|
logits_2d_max->dtype(),
|
|
phi::ccl::CCLReduceOp::MAX,
|
|
comm->GetXcclComm(),
|
|
stream.raw_stream());
|
|
|
|
// step 2, obtain logit - logit_max
|
|
auto logits_2d_sub_max = paddle::experimental::clip(
|
|
logits_2d_tensor - logits_2d_max_tensor, -64., 0.);
|
|
|
|
// step 3, obtain predict target
|
|
const int start_index = rank * D;
|
|
auto start_index_tensor =
|
|
paddle::experimental::full_like(labels_1d_tensor,
|
|
start_index,
|
|
labels_1d_tensor.dtype(),
|
|
labels_1d_tensor.place());
|
|
auto end_index_tensor =
|
|
paddle::experimental::full_like(labels_1d_tensor,
|
|
start_index + D,
|
|
labels_1d_tensor.dtype(),
|
|
labels_1d_tensor.place());
|
|
auto labels_1d_mask = paddle::experimental::logical_and(
|
|
labels_1d_tensor.greater_equal(start_index_tensor),
|
|
labels_1d_tensor.less_than(end_index_tensor));
|
|
auto real_label_tensor = (labels_1d_tensor - start_index_tensor)
|
|
.multiply(paddle::experimental::cast(
|
|
labels_1d_mask, labels_1d_tensor.dtype()));
|
|
|
|
auto predicted_logits_tensor =
|
|
logits_2d_sub_max
|
|
.multiply(paddle::experimental::cast(
|
|
paddle::experimental::one_hot(real_label_tensor, D),
|
|
logits_2d_sub_max.dtype()))
|
|
.sum({1}, logits_2d_sub_max.dtype(), false)
|
|
.multiply(paddle::experimental::cast(labels_1d_mask,
|
|
logits_2d_sub_max.dtype()));
|
|
|
|
auto predicted_logits =
|
|
reinterpret_cast<DenseTensor*>(predicted_logits_tensor.impl().get());
|
|
phi::DeviceManager::CCLAllReduce(dev_ctx.GetPlace().GetDeviceType(),
|
|
predicted_logits->data<float>(),
|
|
predicted_logits->data<float>(),
|
|
predicted_logits->numel(),
|
|
predicted_logits->dtype(),
|
|
phi::ccl::CCLReduceOp::SUM,
|
|
comm->GetXcclComm(),
|
|
stream.raw_stream());
|
|
|
|
// step 4, obtain exp(logit)
|
|
auto softmax_2d_tensor = logits_2d_sub_max.exp();
|
|
|
|
// step 5, obtain sum_exp_logits
|
|
auto sum_exp_logits_tensor =
|
|
softmax_2d_tensor.sum({1}, softmax_2d_tensor.dtype(), false);
|
|
|
|
auto sum_exp_logits =
|
|
reinterpret_cast<DenseTensor*>(sum_exp_logits_tensor.impl().get());
|
|
phi::DeviceManager::CCLAllReduce(dev_ctx.GetPlace().GetDeviceType(),
|
|
sum_exp_logits->data<float>(),
|
|
sum_exp_logits->data<float>(),
|
|
sum_exp_logits->numel(),
|
|
sum_exp_logits->dtype(),
|
|
phi::ccl::CCLReduceOp::SUM,
|
|
comm->GetXcclComm(),
|
|
stream.raw_stream());
|
|
|
|
auto softmax_out = softmax_2d_tensor.divide(
|
|
paddle::experimental::reshape(sum_exp_logits_tensor, {N, 1}));
|
|
auto labels_1d_not_equal_ignore = labels_1d_tensor.not_equal(
|
|
paddle::experimental::full_like(labels_1d_tensor,
|
|
ignore_index,
|
|
labels_1d_tensor.dtype(),
|
|
labels_1d_tensor.place()));
|
|
auto loss_out =
|
|
(sum_exp_logits_tensor.log() - predicted_logits_tensor)
|
|
.multiply(paddle::experimental::cast(labels_1d_not_equal_ignore,
|
|
sum_exp_logits_tensor.dtype()));
|
|
softmax
|
|
->ShareDataWith(*reinterpret_cast<DenseTensor*>(softmax_out.impl().get()))
|
|
.Resize(softmax_dims);
|
|
loss->ShareDataWith(*reinterpret_cast<DenseTensor*>(loss_out.impl().get()))
|
|
.Resize(loss_dims);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(c_softmax_with_cross_entropy,
|
|
Custom,
|
|
ALL_LAYOUT,
|
|
phi::CSoftmaxWithEntropyKernel,
|
|
float,
|
|
double,
|
|
phi::float16) {}
|
|
#endif
|