paddlepaddle--paddle/paddle/phi/kernels/cpu/sgd_kernel.cc

// 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/cpu/cpu_context.h"
#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/eigen/common.h"
#include "paddle/phi/kernels/funcs/jit/kernels.h"

namespace phi {

template <typename T>
void sgd_dense_param_dense_grad_impl(const DenseTensor& param,
                                     const DenseTensor& learning_rate,
                                     const DenseTensor& grad,
                                     DenseTensor* param_out) {
  const auto sz = param_out->numel();
  jit::sgd_attr_t attr(1, sz, 1, sz, 1);
  const T* lr = learning_rate.data<T>();
  const T* param_data = param.data<T>();
  const T* grad_data = grad.data<T>();
  int64_t rows_idx = 0;
  T* out_data = param_out->data<T>();

  auto sgd = jit::KernelFuncs<jit::SgdTuple<T>, CPUPlace>::Cache().At(attr);
  sgd(lr, param_data, grad_data, &rows_idx, out_data, &attr);
}

template <>
void sgd_dense_param_dense_grad_impl<bfloat16>(const DenseTensor& param,
                                               const DenseTensor& learning_rate,
                                               const DenseTensor& grad,
                                               DenseTensor* param_out) {
  auto p = EigenVector<bfloat16>::Flatten(param);
  auto g = EigenVector<bfloat16>::Flatten(grad);
  auto o = EigenVector<bfloat16>::Flatten(*param_out);
  const auto* lr = learning_rate.data<float>();

  o = p - static_cast<bfloat16>(lr[0]) * g;
}

template <typename T>
void sgd_dense_param_dense_grad_mixed_impl(const DenseTensor& param,
                                           const DenseTensor& learning_rate,
                                           const DenseTensor& grad,
                                           DenseTensor* param_out) {
  const T* param_data = param.data<T>();
  const float* grad_data = grad.data<float>();
  const float* lr_ptr = learning_rate.data<float>();

  float lr = lr_ptr[0];
  T* out_data = param_out->data<T>();
  int64_t numel = param.numel();

  for (int64_t i = 0; i < numel; ++i) {
    float p = static_cast<float>(param_data[i]);
    float g = grad_data[i];

    p = p - lr * g;

    out_data[i] = static_cast<T>(p);
  }
}

template <typename T>
void sgd_dense_param_sparse_grad_impl(const DenseTensor& param,
                                      const DenseTensor& learning_rate,
                                      const SelectedRows& grad,
                                      DenseTensor* param_out) {
  const auto& grad_value = grad.value();
  const auto& grad_rows = grad.rows();
  const T* param_data = param.data<T>();
  const T* grad_data = grad_value.data<T>();
  const T* lr = learning_rate.data<T>();
  const int64_t* rows_data = grad_rows.data();
  T* out_data = param_out->data<T>();

  jit::sgd_attr_t attr;
  attr.param_height = param_out->dims()[0];
  attr.param_width = param_out->numel() / attr.param_height;
  attr.grad_height =
      static_cast<int>(grad_rows.size());  // note: it is not grad->height()
  attr.grad_width = grad_value.numel() / attr.grad_height;
  attr.selected_rows_size = static_cast<int>(grad_rows.size());

  auto sgd = jit::KernelFuncs<jit::SgdTuple<T>, CPUPlace>::Cache().At(attr);
  sgd(lr, param_data, grad_data, rows_data, out_data, &attr);
}

template <>
void sgd_dense_param_sparse_grad_impl<bfloat16>(
    const DenseTensor& param,
    const DenseTensor& learning_rate,
    const SelectedRows& grad,
    DenseTensor* param_out) {
  const auto& grad_value = grad.value();
  const auto& grad_rows = grad.rows();
  const auto grad_height = grad.height();
  const int64_t grad_val_height = static_cast<int64_t>(grad_rows.size());
  const auto grad_width = grad_value.numel() / grad_val_height;

  const auto* grad_data = grad_value.data<bfloat16>();
  auto* out_data = param_out->data<bfloat16>();
  const auto* lr = learning_rate.data<float>();

  for (size_t i = 0; i < grad_rows.size(); ++i) {
    PADDLE_ENFORCE_LT(
        grad_rows[i],
        grad_height,
        common::errors::OutOfRange(
            "Grad rows index value should be less than grad height."
            "Got [%s], but expected less than [%s]",
            grad_rows[i],
            grad_height));
    const int64_t row = grad_rows[i];
    for (int64_t j = 0; j < grad_width; ++j) {
      out_data[row * grad_width + j] -=
          static_cast<bfloat16>(lr[0]) * grad_data[i * grad_width + j];
    }
  }
}

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 UNUSED,
                    bool multi_precision UNUSED,
                    DenseTensor* param_out,
                    DenseTensor* master_param_out UNUSED) {
  dev_ctx.template Alloc<T>(param_out);
  if (grad.dtype() == DataType::FLOAT32 && param.dtype() != DataType::FLOAT32) {
    sgd_dense_param_dense_grad_mixed_impl<T>(
        param, learning_rate, grad, param_out);
  } else {
    sgd_dense_param_dense_grad_impl<T>(param, learning_rate, grad, param_out);
  }
}

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
                                       UNUSED,
                                   bool multi_precision UNUSED,
                                   DenseTensor* param_out,
                                   DenseTensor* master_param_out UNUSED) {
  dev_ctx.template Alloc<T>(param_out);
  sgd_dense_param_sparse_grad_impl<T>(param, learning_rate, grad, param_out);
}

template <typename T, typename Context>
void SGDSparseParamSparseGradKernel(const Context& dev_ctx UNUSED,
                                    const SelectedRows& param,
                                    const DenseTensor& learning_rate,
                                    const SelectedRows& grad,
                                    const optional<SelectedRows>& master_param
                                        UNUSED,
                                    bool multi_precision UNUSED,
                                    SelectedRows* param_out,
                                    SelectedRows* master_param_out UNUSED) {
  // for distributed training, a sparse var may be empty,
  // just skip updating.
  if (grad.rows().empty()) {
    return;
  }

  auto param_row_width = param.value().dims()[1];
  auto grad_row_width = grad.value().dims()[1];
  PADDLE_ENFORCE_EQ(
      param_row_width,
      grad_row_width,
      common::errors::InvalidArgument(
          "The param_row in SgdOP should have the same size with grad_row. "
          "But received param_row's width is [%s], and grad_row's width is "
          "[%s]",
          param_row_width,
          grad_row_width));

  using MT = typename dtype::MPTypeTrait<T>::Type;
  const auto* lr = learning_rate.data<MT>();
  const auto* grad_data = grad.value().data<T>();
  auto* out_data = param_out->mutable_value()->data<T>();
  for (size_t i = 0; i < grad.rows().size(); i++) {
    int64_t id_index = param_out->AutoGrownIndex(grad.rows()[i], false);
    PADDLE_ENFORCE_GE(
        id_index,
        static_cast<int64_t>(0),
        common::errors::InvalidArgument(
            "The id in SgdOp should be >= 0. But received id_index is [%s]",
            id_index));
    for (int64_t j = 0; j < grad_row_width; j++) {
      out_data[id_index * grad_row_width + j] -= static_cast<T>(
          lr[0] * static_cast<MT>(grad_data[i * grad_row_width + j]));
    }
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(
    sgd, CPU, ALL_LAYOUT, phi::SGDDenseKernel, phi::bfloat16, float, double) {}

PD_REGISTER_KERNEL(sgd_dense_param_sparse_grad,
                   CPU,
                   ALL_LAYOUT,
                   phi::SGDDenseParamSparseGradKernel,
                   phi::bfloat16,
                   float,
                   double) {}

PD_REGISTER_KERNEL(sgd_sparse_param_sparse_grad,
                   CPU,
                   ALL_LAYOUT,
                   phi::SGDSparseParamSparseGradKernel,
                   phi::bfloat16,
                   float,
                   double) {}