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