File size: 5,483 Bytes
146b945 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
#include <ATen/core/Tensor.h>
#include <ATen/native/cuda/ForeachFunctors.cuh>
#include <ATen/native/cuda/MultiTensorApply.cuh>
#include <ATen/native/cuda/Pow.cuh>
#include <utility>
namespace adam_atan2 {
using at::native::kILP;
constexpr int kArgsDepth = 4;
constexpr uint8_t kParamIdx = 0;
constexpr uint8_t kGradIdx = 1;
constexpr uint8_t kExpAvgIdx = 2;
constexpr uint8_t kExpAvgSqIdx = 3;
template <typename T>
__device__ __forceinline__ T lerp(const T v0, const T v1, const T t) {
// NOTE(one): Identical to PyTorch when t < 0.5
// https://github.com/pytorch/pytorch/blob/b7f25226929e70187a9f36c393665abad0b25190/aten/src/ATen/native/Lerp.h#L21
return fma(t, v1, fma(-t, v0, v0));
}
template <typename scalar_type, typename opmath_t>
__device__ __forceinline__ void adam_math(
scalar_type r_args[kArgsDepth][kILP],
const opmath_t &step_size,
const opmath_t &wd_alpha,
const opmath_t &mbeta1,
const opmath_t &mbeta2,
const opmath_t &bias_correction2_sqrt)
{
#pragma unroll
for (int ii = 0; ii < kILP; ii++)
{
// Load values.
opmath_t param = static_cast<opmath_t>(r_args[kParamIdx][ii]);
const opmath_t grad = static_cast<opmath_t>(r_args[kGradIdx][ii]);
opmath_t exp_avg = static_cast<opmath_t>(r_args[kExpAvgIdx][ii]);
opmath_t exp_avg_sq = static_cast<opmath_t>(r_args[kExpAvgSqIdx][ii]);
param *= wd_alpha;
exp_avg = lerp(exp_avg, grad, mbeta1);
exp_avg_sq = lerp(exp_avg_sq, grad * grad, mbeta2);
const opmath_t denom = std::sqrt(exp_avg_sq) / bias_correction2_sqrt;
param -= step_size * std::atan2(exp_avg, denom);
// Store results.
r_args[kParamIdx][ii] = param;
r_args[kExpAvgIdx][ii] = exp_avg;
r_args[kExpAvgSqIdx][ii] = exp_avg_sq;
}
}
template <typename scalar_type>
struct FusedAdamMathFunctor {
using opmath_t = at::opmath_type<scalar_type>;
__device__ __forceinline__ void operator()(
int chunk_size,
at::native::FusedOptimizerTensorListMetadata<kArgsDepth>& tl,
const double& lr,
const double& beta1,
const double& beta2,
const double& weight_decay) {
const auto tensor_loc = tl.block_to_tensor[blockIdx.x];
const auto chunk_idx = tl.block_to_chunk[blockIdx.x];
const auto [step_size, wd_alpha, bias_correction2_sqrt, mbeta1, mbeta2] = [&]() -> std::tuple<opmath_t, opmath_t, opmath_t, opmath_t, opmath_t> {
auto* step_count = reinterpret_cast<const float*>(tl.state_steps_addresses[tensor_loc]);
const auto bias_correction1 = 1 - at::native::pow_(beta1, *step_count);
const auto bias_correction2 = 1 - at::native::pow_(beta2, *step_count);
const auto bias_correction2_sqrt = std::sqrt(bias_correction2);
return {
static_cast<opmath_t>(lr / bias_correction1),
static_cast<opmath_t>(1 - lr * weight_decay),
static_cast<opmath_t>(bias_correction2_sqrt),
static_cast<opmath_t>(1 - beta1),
static_cast<opmath_t>(1 - beta2)
};
}();
scalar_type* args[kArgsDepth];
scalar_type r_args[kArgsDepth][kILP];
const auto n = tl.numel_for_tensor[tensor_loc] - chunk_idx * chunk_size;
const bool all_aligned{
at::native::init_args<kArgsDepth>(args, tl, chunk_idx, chunk_size, tensor_loc)};
if ((n % kILP == 0) && (chunk_size % kILP == 0) && all_aligned) {
for (int64_t i_start = threadIdx.x;
i_start * kILP < n && i_start * kILP < chunk_size;
i_start += blockDim.x) {
#pragma unroll
for (int i = 0; i < kArgsDepth; i++) {
at::native::load_store(r_args[i], args[i], 0, i_start);
}
adam_math(
r_args,
step_size,
wd_alpha,
mbeta1,
mbeta2,
bias_correction2_sqrt);
#pragma unroll
for (int i = 0; i < kArgsDepth; i++) {
if (i != kGradIdx) {
at::native::load_store(args[i], r_args[i], i_start, 0);
}
}
}
} else {
for (int64_t i_start = 0; i_start < n && i_start < chunk_size;
i_start += blockDim.x * kILP) {
at::native::load_args<kArgsDepth>(r_args, args, i_start, chunk_size, n);
adam_math(
r_args,
step_size,
wd_alpha,
mbeta1,
mbeta2,
bias_correction2_sqrt);
#pragma unroll
for (int i = 0; i < kArgsDepth; i++) {
if (i != kGradIdx) {
at::native::store_args(args[i], r_args[i], i_start, chunk_size, n);
}
}
}
}
}
};
void adam_atan2_cuda_impl_(
std::vector<at::Tensor> params,
std::vector<at::Tensor> grads,
std::vector<at::Tensor> exp_avgs,
std::vector<at::Tensor> exp_avg_sqs,
std::vector<at::Tensor> state_steps,
const double lr,
const double beta1,
const double beta2,
const double weight_decay) {
std::vector<std::vector<at::Tensor>> tensor_lists{params, grads, exp_avgs, exp_avg_sqs};
AT_DISPATCH_FLOATING_TYPES_AND2(
at::ScalarType::Half,
at::ScalarType::BFloat16,
params[0].scalar_type(),
"adam_atan2_kernel_cuda",
[&]() {
at::native::multi_tensor_apply_for_fused_optimizer<kArgsDepth>(
tensor_lists,
state_steps,
FusedAdamMathFunctor<scalar_t>(),
lr,
beta1,
beta2,
weight_decay);
});
}
} // namespace adam_atan2 |