chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,234 @@
|
||||
// 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/phi/kernels/sgd_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_helper.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/mixed_vector.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename MT, typename GradT>
|
||||
__global__ void SGDKernelMT(const T* param,
|
||||
const GradT* grad,
|
||||
const MT* learning_rate,
|
||||
const int64_t num,
|
||||
T* param_out,
|
||||
const MT* master_param,
|
||||
MT* master_param_out) {
|
||||
MT lr = static_cast<MT>(learning_rate[0]);
|
||||
CUDA_KERNEL_LOOP_TYPE(i, num, int64_t) {
|
||||
MT p_data = master_param ? master_param[i] : static_cast<MT>(param[i]);
|
||||
MT g_data = static_cast<MT>(grad[i]);
|
||||
p_data = p_data - lr * g_data;
|
||||
param_out[i] = static_cast<T>(p_data);
|
||||
if (master_param_out) {
|
||||
master_param_out[i] = p_data;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void SparseSGDFunctorKernel(const T* selected_rows,
|
||||
const int64_t* rows,
|
||||
const MT* learning_rate,
|
||||
T* tensor_out,
|
||||
int64_t row_numel,
|
||||
int64_t limit) {
|
||||
MT lr = learning_rate[0];
|
||||
for (int64_t i = blockIdx.x; i < limit; i += gridDim.x) {
|
||||
const T* selected_rows_ptr = selected_rows + i * row_numel;
|
||||
T* tensor_out_ptr = tensor_out + rows[i] * row_numel;
|
||||
for (int64_t index = threadIdx.x; index < row_numel; index += blockDim.x) {
|
||||
// Since index in rows of SelectedRows can be duplicate, we have to use
|
||||
// Atomic Operation to avoid concurrent write error.
|
||||
CudaAtomicAdd(
|
||||
tensor_out_ptr + index,
|
||||
static_cast<T>(-lr * static_cast<MT>(selected_rows_ptr[index])));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SGDDenseKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& grad,
|
||||
const optional<DenseTensor>& master_param,
|
||||
bool multi_precision,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* master_param_out) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
// do check here
|
||||
// if (multi_precision) {
|
||||
// bool has_master =
|
||||
// dev_ctx.HasInput("MasterParam") &&
|
||||
// dev_ctx.HasOutput("MasterParamOut");
|
||||
|
||||
// }
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;
|
||||
const bool use_float32_grad = grad.dtype() == DataType::FLOAT32;
|
||||
|
||||
int block = 512;
|
||||
int64_t grid_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int grid = std::min((param.numel() + block - 1) / block, grid_max);
|
||||
if (use_float32_grad) {
|
||||
SGDKernelMT<T, MT, float><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
param.data<T>(),
|
||||
grad.data<float>(),
|
||||
learning_rate.data<MT>(),
|
||||
param.numel(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data);
|
||||
} else {
|
||||
SGDKernelMT<T, MT, T><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
param.data<T>(),
|
||||
grad.data<T>(),
|
||||
learning_rate.data<MT>(),
|
||||
param.numel(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SGDDenseParamSparseGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& learning_rate,
|
||||
const SelectedRows& grad,
|
||||
const optional<DenseTensor>& master_param,
|
||||
bool multi_precision,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* master_param_out) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
// do some check here
|
||||
// if (multi_precision) {
|
||||
// bool has_master =
|
||||
// dev_ctx.HasInput("MasterParam") &&
|
||||
// dev_ctx.HasOutput("MasterParamOut");
|
||||
|
||||
// }
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
param.IsSharedBufferWith(*param_out),
|
||||
true,
|
||||
common::errors::InvalidArgument(
|
||||
"The input tensor Param of SgdOp should be equal with ParamOut "
|
||||
"if variable's type is SelectedRows."));
|
||||
|
||||
auto in_height = grad.height();
|
||||
auto out_dims = param_out->dims();
|
||||
PADDLE_ENFORCE_EQ(in_height,
|
||||
out_dims[0],
|
||||
common::errors::InvalidArgument(
|
||||
"The input tensor Grad's height of SgdOp should be "
|
||||
"equal with ParamOut's dims. But received Grad's "
|
||||
"height [%s] and ParamOut's dims [%s]",
|
||||
in_height,
|
||||
out_dims[0]));
|
||||
|
||||
auto& in_value = grad.value();
|
||||
auto& in_rows = grad.rows();
|
||||
|
||||
int64_t in_row_numel = in_value.numel() / in_rows.size();
|
||||
PADDLE_ENFORCE_EQ(in_row_numel,
|
||||
param_out->numel() / in_height,
|
||||
common::errors::InvalidArgument(
|
||||
"The in_row_numel of SgdOp should be equal with "
|
||||
"param_out's numel / in_height."));
|
||||
|
||||
auto* in_data = in_value.data<T>();
|
||||
auto* out_data = param_out->data<T>();
|
||||
|
||||
const int kThreadsPerBlock = 256;
|
||||
int thread_x = kThreadsPerBlock;
|
||||
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
int max_blocks = std::max(max_threads / kThreadsPerBlock, 1);
|
||||
MixVector<int64_t> mixv_in_rows(&in_rows);
|
||||
SparseSGDFunctorKernel<T, MT><<<max_blocks, thread_x, 0, dev_ctx.stream()>>>(
|
||||
in_data,
|
||||
mixv_in_rows.CUDAData(dev_ctx.GetPlace()),
|
||||
learning_rate.data<MT>(),
|
||||
out_data,
|
||||
in_row_numel,
|
||||
in_rows.size());
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SGDSparseParamSparseGradKernel(const Context& dev_ctx,
|
||||
const SelectedRows& param,
|
||||
const DenseTensor& learning_rate,
|
||||
const SelectedRows& grad,
|
||||
const optional<SelectedRows>& master_param,
|
||||
bool multi_precision,
|
||||
SelectedRows* param_out,
|
||||
SelectedRows* master_param_out) {
|
||||
PADDLE_THROW("not impl");
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
PD_REGISTER_KERNEL(sgd,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SGDDenseKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
||||
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PD_REGISTER_KERNEL(
|
||||
sgd, GPU, ALL_LAYOUT, phi::SGDDenseKernel, phi::float16, float, double) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
PD_REGISTER_KERNEL(sgd_dense_param_sparse_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SGDDenseParamSparseGradKernel,
|
||||
phi::float16,
|
||||
float,
|
||||
double) {}
|
||||
|
||||
PD_REGISTER_KERNEL(sgd_sparse_param_sparse_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SGDSparseParamSparseGradKernel,
|
||||
phi::float16,
|
||||
float,
|
||||
double) {}
|
||||
Reference in New Issue
Block a user