wehub-resource-sync
925 lines
38 KiB
C++
925 lines
38 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/phi/kernels/adamw_kernel.h"
|
|
|
|
#include <vector>
|
|
|
|
#include "glog/logging.h"
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/backends/xpu/xpu_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
|
|
#include "paddle/phi/common/amp_type_traits.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename Context>
|
|
float GetAbsMax(const Context& dev_ctx,
|
|
const float* input,
|
|
float* buffer_xpu,
|
|
int64_t numel) {
|
|
int max_ptr_size = backends::xpu::get_xpu_max_ptr_size(-1);
|
|
std::vector<float> buffer_cpu(max_ptr_size);
|
|
// int findmax(Context* xpu_ctx, const T* x, float* maxptr, int64_t len);
|
|
int r = xpu::findmax<float>(dev_ctx.x_context(), input, buffer_xpu, numel);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax");
|
|
memory_utils::Copy(CPUPlace(),
|
|
static_cast<void*>(buffer_cpu.data()),
|
|
dev_ctx.GetPlace(),
|
|
static_cast<void*>(buffer_xpu),
|
|
sizeof(float) * max_ptr_size);
|
|
return *std::max_element(buffer_cpu.begin(), buffer_cpu.end());
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void AdamwDenseKernelKL3(const Context& dev_ctx,
|
|
const DenseTensor& param,
|
|
const DenseTensor& grad,
|
|
const DenseTensor& learning_rate,
|
|
const DenseTensor& moment1,
|
|
const DenseTensor& moment2,
|
|
const DenseTensor& beta1_pow,
|
|
const DenseTensor& beta2_pow,
|
|
const optional<DenseTensor>& master_param,
|
|
const optional<DenseTensor>& skip_update,
|
|
const Scalar& beta1,
|
|
const Scalar& beta2,
|
|
const Scalar& epsilon,
|
|
float lr_ratio,
|
|
float coeff,
|
|
bool with_decay,
|
|
bool lazy_mode,
|
|
int64_t min_row_size_to_use_multithread,
|
|
bool multi_precision,
|
|
bool use_global_beta_pow,
|
|
DenseTensor* param_out,
|
|
DenseTensor* moment1_out,
|
|
DenseTensor* moment2_out,
|
|
DenseTensor* beta1_pow_out,
|
|
DenseTensor* beta2_pow_out,
|
|
DenseTensor* master_param_outs) {
|
|
// TODO(houj04):
|
|
// 当KL3稳定以后，并且不需要支持KL1和KL2的时候，拿这里的AdamwDenseKernelKL3替换掉AdamwDenseKernel
|
|
using MT = typename phi::dtype::MPTypeTrait<T>::Type;
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
|
|
const auto grad_type = grad.dtype();
|
|
|
|
VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
|
|
|
|
MT coeff_ = static_cast<MT>(coeff);
|
|
MT lr_ratio_ = static_cast<MT>(lr_ratio);
|
|
|
|
bool skip_update_ = false;
|
|
if (skip_update.is_initialized()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
skip_update->numel(),
|
|
1,
|
|
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
|
skip_update->numel()));
|
|
std::vector<bool> skip_update_vec;
|
|
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
|
skip_update_ = skip_update_vec[0];
|
|
}
|
|
|
|
// skip_update=true, just copy input to output
|
|
if (skip_update_) {
|
|
VLOG(4) << "Adamw skip update";
|
|
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
|
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
|
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
|
if (!use_global_beta_pow) {
|
|
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
|
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// if with_decay = false, coeff = 0
|
|
if (!with_decay) {
|
|
coeff_ = static_cast<MT>(0.0);
|
|
}
|
|
|
|
MT beta1_ = beta1.to<MT>();
|
|
MT beta2_ = beta2.to<MT>();
|
|
MT epsilon_ = epsilon.to<MT>();
|
|
VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel()
|
|
<< "beta2_pow.numel() : " << beta2_pow.numel();
|
|
VLOG(3) << "param.numel(): " << param.numel();
|
|
PADDLE_ENFORCE_EQ(
|
|
beta1_pow_out->numel(),
|
|
1,
|
|
errors::InvalidArgument("beta1 pow output size should be 1, but received "
|
|
"value is:%d.",
|
|
beta1_pow_out->numel()));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
beta2_pow_out->numel(),
|
|
1,
|
|
errors::InvalidArgument("beta2 pow output size should be 1, but received "
|
|
"value is:%d.",
|
|
beta2_pow_out->numel()));
|
|
|
|
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_outs) : nullptr;
|
|
|
|
// check moment_dtype
|
|
auto moment1_dtype = moment1.dtype();
|
|
auto moment2_dtype = moment2.dtype();
|
|
PADDLE_ENFORCE_EQ(moment1_dtype,
|
|
moment1_out->dtype(),
|
|
errors::InvalidArgument(
|
|
"moment1.dtype does not match moment1_out->dtype"));
|
|
PADDLE_ENFORCE_EQ(moment2_dtype,
|
|
moment2_out->dtype(),
|
|
errors::InvalidArgument(
|
|
"moment2.dtype does not match moment2_out->dtype"));
|
|
PADDLE_ENFORCE_EQ(
|
|
moment1_dtype,
|
|
moment2_dtype,
|
|
errors::InvalidArgument("moment1.dtype does not match moment2.dtype"));
|
|
|
|
bool moment_in_fp16 = false;
|
|
if (moment1_dtype == DataType::FLOAT16) {
|
|
moment_in_fp16 = true;
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(
|
|
moment1_dtype,
|
|
DataType::FLOAT32,
|
|
errors::InvalidArgument("moment1.dtype is neither fp32 nor fp16"));
|
|
}
|
|
|
|
float* moment1_input_for_xdnn = nullptr;
|
|
float* moment2_input_for_xdnn = nullptr;
|
|
float* moment1_output_for_xdnn = nullptr;
|
|
float* moment2_output_for_xdnn = nullptr;
|
|
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
if (moment_in_fp16) {
|
|
// allocate temp buffer on XPU
|
|
moment1_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_input_for_xdnn);
|
|
moment2_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_input_for_xdnn);
|
|
moment1_output_for_xdnn =
|
|
RAII_GUARD.alloc_l3_or_gm<float>(moment1_out->numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_output_for_xdnn);
|
|
moment2_output_for_xdnn =
|
|
RAII_GUARD.alloc_l3_or_gm<float>(moment2_out->numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_output_for_xdnn);
|
|
|
|
int r = 0;
|
|
using XPUType16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
|
|
// cast moment1 and moment2, from fp16 to fp32
|
|
// int cast(Context* xpu_ctx, const TX* x, TY* y, int64_t len);
|
|
r = xpu::cast<XPUType16, float>(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType16*>(
|
|
moment1.template data<phi::float16>()),
|
|
moment1_input_for_xdnn,
|
|
moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1 from fp16 to float");
|
|
r = xpu::cast<XPUType16, float>(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType16*>(
|
|
moment2.template data<phi::float16>()),
|
|
moment2_input_for_xdnn,
|
|
moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2 from fp16 to float");
|
|
|
|
// acquire xpu_scale_value
|
|
float moment1_scale_value = XPUStorageProperties::default_xpu_scale_value;
|
|
if (moment1.storage_properties_initialized()) {
|
|
moment1_scale_value =
|
|
moment1.storage_properties<XPUStorageProperties>().xpu_scale_value;
|
|
}
|
|
float moment2_scale_value = XPUStorageProperties::default_xpu_scale_value;
|
|
if (moment2.storage_properties_initialized()) {
|
|
moment2_scale_value =
|
|
moment2.storage_properties<XPUStorageProperties>().xpu_scale_value;
|
|
}
|
|
|
|
// de-scale using scale_value
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
if (moment1_scale_value > 0) {
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment1_input_for_xdnn,
|
|
moment1_input_for_xdnn,
|
|
moment1.numel(),
|
|
false,
|
|
1.0f / moment1_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment1");
|
|
}
|
|
if (moment2_scale_value > 0) {
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment2_input_for_xdnn,
|
|
moment2_input_for_xdnn,
|
|
moment2.numel(),
|
|
false,
|
|
1.0f / moment2_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment2");
|
|
}
|
|
}
|
|
|
|
// learning_rate may be float64 (get_lr_dtype returns float64 for all
|
|
// platforms), but XPU kernels only support float32 (MT). Cast if needed.
|
|
const MT* lr_for_xdnn = nullptr;
|
|
MT* lr_cast_buf = nullptr;
|
|
if (learning_rate.dtype() == DataType::FLOAT64) {
|
|
lr_cast_buf = RAII_GUARD.alloc_l3_or_gm<MT>(learning_rate.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(lr_cast_buf);
|
|
int r = xpu::cast<double, MT>(dev_ctx.x_context(),
|
|
learning_rate.template data<double>(),
|
|
lr_cast_buf,
|
|
learning_rate.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast lr from float64 to MT");
|
|
lr_for_xdnn = lr_cast_buf;
|
|
} else {
|
|
lr_for_xdnn = learning_rate.data<MT>();
|
|
}
|
|
|
|
// template <typename T, typename TG, typename MT> DLL_EXPORT int
|
|
// adamw(Context* xpu_ctx, MT beta1, MT beta2, MT epsilon, MT coeff, MT
|
|
// lr_ratio, const MT* beta1_pow, MT beta1_pow_scalar, const MT* beta2_pow, MT
|
|
// beta2_pow_scalar, const MT* moment1, MT* moment1_out, const MT* moment2,
|
|
// MT* moment2_out, const MT* lr, const TG* grad, const T* param, T*
|
|
// param_out, const MT* master_param, MT* master_param_out, int64_t n);
|
|
if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
|
|
// Compute with betapow in REG
|
|
if (grad_type == DataType::FLOAT32) {
|
|
int r = xpu::adamw<XPUType, float, MT>(
|
|
dev_ctx.x_context(),
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff_,
|
|
lr_ratio_,
|
|
nullptr, // beta1_pow
|
|
*beta1_pow.data<MT>(), // beta1_pow_scalar
|
|
nullptr, // beta2_pow
|
|
*beta2_pow.data<MT>(), // beta2_pow_scalar
|
|
moment_in_fp16 ? moment1_input_for_xdnn : moment1.template data<MT>(),
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment1_out),
|
|
moment_in_fp16 ? moment2_input_for_xdnn : moment2.template data<MT>(),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment2_out),
|
|
lr_for_xdnn,
|
|
grad.data<float>(),
|
|
reinterpret_cast<const XPUType*>(param.data<T>()),
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
master_in_data,
|
|
master_out_data,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
} else {
|
|
int r = xpu::adamw<XPUType, XPUType, MT>(
|
|
dev_ctx.x_context(),
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff_,
|
|
lr_ratio_,
|
|
nullptr, // beta1_pow
|
|
*beta1_pow.data<MT>(), // beta1_pow_scalar
|
|
nullptr, // beta2_pow
|
|
*beta2_pow.data<MT>(), // beta2_pow_scalar
|
|
moment_in_fp16 ? moment1_input_for_xdnn : moment1.template data<MT>(),
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment1_out),
|
|
moment_in_fp16 ? moment2_input_for_xdnn : moment2.template data<MT>(),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment2_out),
|
|
lr_for_xdnn,
|
|
reinterpret_cast<const XPUType*>(grad.data<T>()),
|
|
reinterpret_cast<const XPUType*>(param.data<T>()),
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
master_in_data,
|
|
master_out_data,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
}
|
|
if (!use_global_beta_pow) {
|
|
// Cpu update
|
|
dev_ctx.template HostAlloc<MT>(beta1_pow_out)[0] =
|
|
beta1_ * beta1_pow.data<MT>()[0];
|
|
dev_ctx.template HostAlloc<MT>(beta2_pow_out)[0] =
|
|
beta2_ * beta2_pow.data<MT>()[0];
|
|
}
|
|
} else {
|
|
if (grad_type == DataType::FLOAT32) {
|
|
int r = xpu::adamw<XPUType, float, MT>(
|
|
dev_ctx.x_context(),
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff_,
|
|
lr_ratio_,
|
|
beta1_pow.data<MT>(), // beta1_pow
|
|
0.0f, // beta1_pow_scalar
|
|
beta2_pow.data<MT>(), // beta2_pow
|
|
0.0f, // beta2_pow_scalar
|
|
moment_in_fp16 ? moment1_input_for_xdnn : moment1.template data<MT>(),
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment1_out),
|
|
moment_in_fp16 ? moment2_input_for_xdnn : moment2.template data<MT>(),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment2_out),
|
|
lr_for_xdnn,
|
|
grad.data<float>(),
|
|
reinterpret_cast<const XPUType*>(param.data<T>()),
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
master_in_data,
|
|
master_out_data,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
} else {
|
|
int r = xpu::adamw<XPUType, XPUType, MT>(
|
|
dev_ctx.x_context(),
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff_,
|
|
lr_ratio_,
|
|
beta1_pow.data<MT>(), // beta1_pow
|
|
0.0f, // beta1_pow_scalar
|
|
beta2_pow.data<MT>(), // beta2_pow
|
|
0.0f, // beta2_pow_scalar
|
|
moment_in_fp16 ? moment1_input_for_xdnn : moment1.template data<MT>(),
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment1_out),
|
|
moment_in_fp16 ? moment2_input_for_xdnn : moment2.template data<MT>(),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<MT>(moment2_out),
|
|
lr_for_xdnn,
|
|
reinterpret_cast<const XPUType*>(grad.data<T>()),
|
|
reinterpret_cast<const XPUType*>(param.data<T>()),
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
master_in_data,
|
|
master_out_data,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
}
|
|
if (!use_global_beta_pow) {
|
|
// Update with xpu
|
|
int r = xpu::scale(dev_ctx.x_context(),
|
|
beta1_pow.data<MT>(),
|
|
dev_ctx.template Alloc<MT>(beta1_pow_out),
|
|
beta1_pow.numel(),
|
|
false,
|
|
beta1_,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
|
|
r = xpu::scale(dev_ctx.x_context(),
|
|
beta2_pow.data<MT>(),
|
|
dev_ctx.template Alloc<MT>(beta2_pow_out),
|
|
beta2_pow.numel(),
|
|
false,
|
|
beta2_,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
|
|
}
|
|
}
|
|
|
|
if (moment_in_fp16) {
|
|
int r = 0;
|
|
using XPUType16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
|
|
// findmax and calculate scale_value for moment1 and moment2
|
|
int max_ptr_size = backends::xpu::get_xpu_max_ptr_size(-1);
|
|
float* buffer_for_findmax = RAII_GUARD.alloc_l3_or_gm<float>(max_ptr_size);
|
|
|
|
// for moment1
|
|
float moment1_max = GetAbsMax<Context>(dev_ctx,
|
|
moment1_output_for_xdnn,
|
|
buffer_for_findmax,
|
|
moment1_out->numel());
|
|
float moment1_scale_value = 65504.0f / moment1_max / 2.0f;
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment1_output_for_xdnn,
|
|
moment1_output_for_xdnn,
|
|
moment1_out->numel(),
|
|
false,
|
|
moment1_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "scale before convert to fp16, for moment1_output_for_xdnn");
|
|
// write to moment1_out
|
|
std::unique_ptr<phi::StorageProperties> moment1_out_sp =
|
|
std::make_unique<phi::XPUStorageProperties>(moment1_scale_value);
|
|
moment1_out->set_storage_properties(std::move(moment1_out_sp));
|
|
|
|
// for moment2
|
|
float moment2_max_ = GetAbsMax<Context>(dev_ctx,
|
|
moment2_output_for_xdnn,
|
|
buffer_for_findmax,
|
|
moment2_out->numel());
|
|
float moment2_scale_value = 65504.0f / moment2_max_ / 2.0f;
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment2_output_for_xdnn,
|
|
moment2_output_for_xdnn,
|
|
moment2_out->numel(),
|
|
false,
|
|
moment2_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "scale before convert to fp16, for moment2_output_for_xdnn");
|
|
// write to moment2_out
|
|
std::unique_ptr<phi::StorageProperties> moment2_out_sp =
|
|
std::make_unique<phi::XPUStorageProperties>(moment2_scale_value);
|
|
moment2_out->set_storage_properties(std::move(moment2_out_sp));
|
|
|
|
// cast moment1 and moment2 output, from fp32 to fp16
|
|
// int cast(Context* xpu_ctx, const TX* x, TY* y, int64_t len);
|
|
r = xpu::cast<float, XPUType16>(
|
|
dev_ctx.x_context(),
|
|
moment1_output_for_xdnn,
|
|
reinterpret_cast<XPUType16*>(
|
|
dev_ctx.template Alloc<phi::float16>(moment1_out)),
|
|
moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1_out from float to fp16");
|
|
r = xpu::cast<float, XPUType16>(
|
|
dev_ctx.x_context(),
|
|
moment2_output_for_xdnn,
|
|
reinterpret_cast<XPUType16*>(
|
|
dev_ctx.template Alloc<phi::float16>(moment2_out)),
|
|
moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2_out from float to fp16");
|
|
}
|
|
return;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void AdamwDenseKernel(const Context& dev_ctx,
|
|
const DenseTensor& param,
|
|
const DenseTensor& grad,
|
|
const DenseTensor& learning_rate,
|
|
const DenseTensor& moment1,
|
|
const DenseTensor& moment2,
|
|
const optional<DenseTensor>& moment2_max, // UNUSED
|
|
const DenseTensor& beta1_pow,
|
|
const DenseTensor& beta2_pow,
|
|
const optional<DenseTensor>& master_param,
|
|
const optional<DenseTensor>& skip_update,
|
|
const Scalar& beta1,
|
|
const Scalar& beta2,
|
|
const Scalar& epsilon,
|
|
double lr_ratio,
|
|
double coeff,
|
|
bool with_decay,
|
|
bool lazy_mode,
|
|
int64_t min_row_size_to_use_multithread,
|
|
bool multi_precision,
|
|
bool use_global_beta_pow,
|
|
bool amsgrad, // UNUSED
|
|
DenseTensor* param_out,
|
|
DenseTensor* moment1_out,
|
|
DenseTensor* moment2_out,
|
|
DenseTensor* moment2_max_out, // UNUSED
|
|
DenseTensor* beta1_pow_out,
|
|
DenseTensor* beta2_pow_out,
|
|
DenseTensor* master_param_outs) {
|
|
PADDLE_ENFORCE_NE(
|
|
amsgrad,
|
|
true,
|
|
common::errors::Unimplemented("Operation amsgrad is not supported yet."));
|
|
|
|
auto dev_version =
|
|
backends::xpu::get_xpu_version(dev_ctx.GetPlace().GetDeviceId());
|
|
if (dev_version == backends::xpu::XPUVersion::XPU3) {
|
|
AdamwDenseKernelKL3<T, Context>(dev_ctx,
|
|
param,
|
|
grad,
|
|
learning_rate,
|
|
moment1,
|
|
moment2,
|
|
beta1_pow,
|
|
beta2_pow,
|
|
master_param,
|
|
skip_update,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
lr_ratio,
|
|
coeff,
|
|
with_decay,
|
|
lazy_mode,
|
|
min_row_size_to_use_multithread,
|
|
multi_precision,
|
|
use_global_beta_pow,
|
|
param_out,
|
|
moment1_out,
|
|
moment2_out,
|
|
beta1_pow_out,
|
|
beta2_pow_out,
|
|
master_param_outs);
|
|
return;
|
|
}
|
|
|
|
// check moment_dtype
|
|
auto moment1_dtype = moment1.dtype();
|
|
auto moment2_dtype = moment2.dtype();
|
|
PADDLE_ENFORCE_EQ(moment1_dtype,
|
|
moment1_out->dtype(),
|
|
errors::InvalidArgument(
|
|
"moment1.dtype does not match moment1_out->dtype"));
|
|
PADDLE_ENFORCE_EQ(moment2_dtype,
|
|
moment2_out->dtype(),
|
|
errors::InvalidArgument(
|
|
"moment2.dtype does not match moment2_out->dtype"));
|
|
PADDLE_ENFORCE_EQ(
|
|
moment1_dtype,
|
|
moment2_dtype,
|
|
errors::InvalidArgument("moment1.dtype does not match moment2.dtype"));
|
|
|
|
bool moment_in_fp16 = false;
|
|
if (moment1_dtype == DataType::FLOAT16) {
|
|
moment_in_fp16 = true;
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(
|
|
moment1_dtype,
|
|
DataType::FLOAT32,
|
|
errors::InvalidArgument("moment1.dtype is neither fp32 nor fp16"));
|
|
}
|
|
|
|
float* moment1_input_for_xdnn = nullptr;
|
|
float* moment2_input_for_xdnn = nullptr;
|
|
float* moment1_output_for_xdnn = nullptr;
|
|
float* moment2_output_for_xdnn = nullptr;
|
|
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
if (moment_in_fp16) {
|
|
// allocate temp buffer on XPU
|
|
moment1_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_input_for_xdnn);
|
|
moment2_input_for_xdnn = RAII_GUARD.alloc_l3_or_gm<float>(moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_input_for_xdnn);
|
|
moment1_output_for_xdnn =
|
|
RAII_GUARD.alloc_l3_or_gm<float>(moment1_out->numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment1_output_for_xdnn);
|
|
moment2_output_for_xdnn =
|
|
RAII_GUARD.alloc_l3_or_gm<float>(moment2_out->numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(moment2_output_for_xdnn);
|
|
|
|
int r = 0;
|
|
using XPUType16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
|
|
// cast moment1 and moment2, from fp16 to fp32
|
|
// int cast(Context* xpu_ctx, const TX* x, TY* y, int64_t len);
|
|
r = xpu::cast<XPUType16, float>(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType16*>(
|
|
moment1.template data<phi::float16>()),
|
|
moment1_input_for_xdnn,
|
|
moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1 from fp16 to float");
|
|
r = xpu::cast<XPUType16, float>(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType16*>(
|
|
moment2.template data<phi::float16>()),
|
|
moment2_input_for_xdnn,
|
|
moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2 from fp16 to float");
|
|
|
|
// acquire xpu_scale_value
|
|
float moment1_scale_value = XPUStorageProperties::default_xpu_scale_value;
|
|
if (moment1.storage_properties_initialized()) {
|
|
moment1_scale_value =
|
|
moment1.storage_properties<XPUStorageProperties>().xpu_scale_value;
|
|
}
|
|
float moment2_scale_value = XPUStorageProperties::default_xpu_scale_value;
|
|
if (moment2.storage_properties_initialized()) {
|
|
moment2_scale_value =
|
|
moment2.storage_properties<XPUStorageProperties>().xpu_scale_value;
|
|
}
|
|
|
|
// de-scale using scale_value
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
if (moment1_scale_value > 0) {
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment1_input_for_xdnn,
|
|
moment1_input_for_xdnn,
|
|
moment1.numel(),
|
|
false,
|
|
1.0f / moment1_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment1");
|
|
}
|
|
if (moment2_scale_value > 0) {
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment2_input_for_xdnn,
|
|
moment2_input_for_xdnn,
|
|
moment2.numel(),
|
|
false,
|
|
1.0f / moment2_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "de-scale for moment2");
|
|
}
|
|
}
|
|
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
bool skip_update_ = false;
|
|
if (skip_update.is_initialized()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
skip_update->numel(),
|
|
1,
|
|
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
|
skip_update->numel()));
|
|
std::vector<bool> skip_update_vec;
|
|
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
|
skip_update_ = skip_update_vec[0];
|
|
}
|
|
if (skip_update_) {
|
|
VLOG(4) << "Adamw skip update";
|
|
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
|
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
|
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
|
if (!use_global_beta_pow) {
|
|
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
|
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
|
}
|
|
return;
|
|
}
|
|
|
|
auto beta1_ = beta1.to<float>();
|
|
auto beta2_ = beta2.to<float>();
|
|
auto epsilon_ = epsilon.to<float>();
|
|
|
|
const float* beta1_pow_ptr = beta1_pow.template data<float>();
|
|
const float* beta2_pow_ptr = beta2_pow.template data<float>();
|
|
DenseTensor xpu_beta1_pow;
|
|
DenseTensor xpu_beta2_pow;
|
|
if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
|
|
Copy(dev_ctx, beta1_pow, dev_ctx.GetPlace(), false, &xpu_beta1_pow);
|
|
Copy(dev_ctx, beta2_pow, dev_ctx.GetPlace(), false, &xpu_beta2_pow);
|
|
dev_ctx.Wait();
|
|
beta1_pow_ptr = xpu_beta1_pow.template data<float>();
|
|
beta2_pow_ptr = xpu_beta2_pow.template data<float>();
|
|
}
|
|
if (!with_decay) {
|
|
coeff = static_cast<float>(0.0);
|
|
}
|
|
|
|
float* new_lr = RAII_GUARD.alloc_l3_or_gm<float>(learning_rate.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(new_lr);
|
|
int r = 0;
|
|
// learning_rate may be float64 (get_lr_dtype returns float64 for all
|
|
// platforms), cast to float32 for XPU kernels which only support float.
|
|
if (learning_rate.dtype() == DataType::FLOAT64) {
|
|
float* lr_fp32 = RAII_GUARD.alloc_l3_or_gm<float>(learning_rate.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(lr_fp32);
|
|
r = xpu::cast<double, float>(dev_ctx.x_context(),
|
|
learning_rate.template data<double>(),
|
|
lr_fp32,
|
|
learning_rate.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast lr from float64 to float32");
|
|
r = xpu::scale(dev_ctx.x_context(),
|
|
lr_fp32,
|
|
new_lr,
|
|
learning_rate.numel(),
|
|
false,
|
|
static_cast<float>(lr_ratio),
|
|
0.0f);
|
|
} else {
|
|
r = xpu::scale(dev_ctx.x_context(),
|
|
learning_rate.template data<float>(),
|
|
new_lr,
|
|
learning_rate.numel(),
|
|
false,
|
|
static_cast<float>(lr_ratio),
|
|
0.0f);
|
|
}
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
|
|
|
|
if (multi_precision) {
|
|
const float* master_param_in_data = master_param->data<float>();
|
|
float* master_param_out_data =
|
|
dev_ctx.template Alloc<float>(master_param_outs);
|
|
// convert grad to float if necessary
|
|
float* grad_fp32 = nullptr;
|
|
const auto grad_type = grad.dtype();
|
|
if (grad_type != DataType::FLOAT32) {
|
|
grad_fp32 = RAII_GUARD.alloc_l3_or_gm<float>(grad.numel());
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(grad_fp32);
|
|
// int cast(Context* xpu_ctx, const TX* x, TY* y, int64_t len);
|
|
int r = xpu::cast<XPUType, float>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(grad.template data<T>()),
|
|
grad_fp32,
|
|
grad.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
|
|
}
|
|
// int adamw(Context* xpu_ctx, const T* g, const float* mom1, const float*
|
|
// mom2, const T* param, const float* beta1_pow, const float* beta2_pow,
|
|
// const float* lr, float* moment1_out, float* moment2_out, T* param_out,
|
|
// float beta1, float beta2, float epsilon, float coeff, int64_t n);
|
|
r = xpu::adamw<float>(
|
|
dev_ctx.x_context(),
|
|
(grad_type == DataType::FLOAT32) ? grad.data<float>() : grad_fp32,
|
|
moment_in_fp16 ? moment1_input_for_xdnn
|
|
: moment1.template data<float>(),
|
|
moment_in_fp16 ? moment2_input_for_xdnn
|
|
: moment2.template data<float>(),
|
|
master_param_in_data,
|
|
beta1_pow_ptr,
|
|
beta2_pow_ptr,
|
|
new_lr,
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<float>(moment1_out),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<float>(moment2_out),
|
|
master_param_out_data,
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
// convert master_param_out(fp32) to param_out(T)
|
|
r = xpu::cast<float, XPUType>(
|
|
dev_ctx.x_context(),
|
|
master_param_out_data,
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
param_out->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
|
|
} else {
|
|
// int adamw(Context* xpu_ctx, const T* g, const float* mom1, const float*
|
|
// mom2, const T* param, const float* beta1_pow, const float* beta2_pow,
|
|
// const float* lr, float* moment1_out, float* moment2_out, T* param_out,
|
|
// float beta1, float beta2, float epsilon, float coeff, int64_t n);
|
|
r = xpu::adamw(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(grad.template data<T>()),
|
|
moment_in_fp16 ? moment1_input_for_xdnn
|
|
: moment1.template data<float>(),
|
|
moment_in_fp16 ? moment2_input_for_xdnn
|
|
: moment2.template data<float>(),
|
|
reinterpret_cast<const XPUType*>(param.template data<T>()),
|
|
beta1_pow_ptr,
|
|
beta2_pow_ptr,
|
|
new_lr,
|
|
moment_in_fp16 ? moment1_output_for_xdnn
|
|
: dev_ctx.template Alloc<float>(moment1_out),
|
|
moment_in_fp16 ? moment2_output_for_xdnn
|
|
: dev_ctx.template Alloc<float>(moment2_out),
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(param_out)),
|
|
beta1_,
|
|
beta2_,
|
|
epsilon_,
|
|
coeff,
|
|
param.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "adamw");
|
|
}
|
|
|
|
if (moment_in_fp16) {
|
|
int r = 0;
|
|
using XPUType16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
|
|
// findmax and calculate scale_value for moment1 and moment2
|
|
int max_ptr_size = backends::xpu::get_xpu_max_ptr_size(-1);
|
|
float* buffer_for_findmax = RAII_GUARD.alloc_l3_or_gm<float>(max_ptr_size);
|
|
|
|
// for moment1
|
|
float moment1_max = GetAbsMax<Context>(dev_ctx,
|
|
moment1_output_for_xdnn,
|
|
buffer_for_findmax,
|
|
moment1_out->numel());
|
|
float moment1_scale_value = 65504.0f / moment1_max / 2.0f;
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment1_output_for_xdnn,
|
|
moment1_output_for_xdnn,
|
|
moment1_out->numel(),
|
|
false,
|
|
moment1_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "scale before convert to fp16, for moment1_output_for_xdnn");
|
|
// write to moment1_out
|
|
std::unique_ptr<phi::StorageProperties> moment1_out_sp =
|
|
std::make_unique<phi::XPUStorageProperties>(moment1_scale_value);
|
|
moment1_out->set_storage_properties(std::move(moment1_out_sp));
|
|
|
|
// for moment2
|
|
float moment2_max_ = GetAbsMax<Context>(dev_ctx,
|
|
moment2_output_for_xdnn,
|
|
buffer_for_findmax,
|
|
moment2_out->numel());
|
|
float moment2_scale_value = 65504.0f / moment2_max_ / 2.0f;
|
|
// int scale(Context* xpu_ctx, const T* x, T* y, int64_t len, bool
|
|
// bias_after_scale, float _scale, float _bias);
|
|
r = xpu::scale<float>(dev_ctx.x_context(),
|
|
moment2_output_for_xdnn,
|
|
moment2_output_for_xdnn,
|
|
moment2_out->numel(),
|
|
false,
|
|
moment2_scale_value,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "scale before convert to fp16, for moment2_output_for_xdnn");
|
|
// write to moment2_out
|
|
std::unique_ptr<phi::StorageProperties> moment2_out_sp =
|
|
std::make_unique<phi::XPUStorageProperties>(moment2_scale_value);
|
|
moment2_out->set_storage_properties(std::move(moment2_out_sp));
|
|
|
|
// cast moment1 and moment2 output, from fp32 to fp16
|
|
// int cast(Context* xpu_ctx, const TX* x, TY* y, int64_t len);
|
|
r = xpu::cast<float, XPUType16>(
|
|
dev_ctx.x_context(),
|
|
moment1_output_for_xdnn,
|
|
reinterpret_cast<XPUType16*>(
|
|
dev_ctx.template Alloc<phi::float16>(moment1_out)),
|
|
moment1.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment1_out from float to fp16");
|
|
r = xpu::cast<float, XPUType16>(
|
|
dev_ctx.x_context(),
|
|
moment2_output_for_xdnn,
|
|
reinterpret_cast<XPUType16*>(
|
|
dev_ctx.template Alloc<phi::float16>(moment2_out)),
|
|
moment2.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast moment2_out from float to fp16");
|
|
}
|
|
|
|
if (!use_global_beta_pow) {
|
|
// update in cpu
|
|
if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
|
|
const float* beta1_pow_p = beta1_pow.template data<float>();
|
|
dev_ctx.template HostAlloc<float>(beta1_pow_out)[0] =
|
|
beta1_ * beta1_pow_p[0];
|
|
const float* beta2_pow_p = beta2_pow.template data<float>();
|
|
dev_ctx.template HostAlloc<float>(beta2_pow_out)[0] =
|
|
beta2_ * beta2_pow_p[0];
|
|
xpu_wait(dev_ctx.x_context()->xpu_stream);
|
|
} else { // update in xpu
|
|
float* beta1_pow_out_p = dev_ctx.template Alloc<float>(beta1_pow_out);
|
|
float* beta2_pow_out_p = dev_ctx.template Alloc<float>(beta2_pow_out);
|
|
int r = xpu::scale(dev_ctx.x_context(),
|
|
beta1_pow_ptr,
|
|
beta1_pow_out_p,
|
|
beta1_pow.numel(),
|
|
false,
|
|
beta1_,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
|
|
r = xpu::scale(dev_ctx.x_context(),
|
|
beta2_pow_ptr,
|
|
beta2_pow_out_p,
|
|
beta2_pow.numel(),
|
|
false,
|
|
beta2_,
|
|
0.0f);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(adamw,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::AdamwDenseKernel,
|
|
float,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
// Skip beta1_pow, beta2_pow, skip_update data transform
|
|
kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
|
|
kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
|
|
kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
|
|
|
|
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
|
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
|
|
kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
|
|
}
|
|
kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
|
|
kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
|
|
}
|