Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
#include <ATen/ATen.h>
2
#include <ATen/AccumulateType.h>
3
#include <ATen/cuda/CUDAContext.h>
4
#include <ATen/cuda/Exceptions.h>
5
// Another possibility:
6
// #include <torch/all.h>
7

8
#include <assert.h>
9

10
#include "type_shim.h"
11
#include "multi_tensor_apply.cuh"
12

13
#define BLOCK_SIZE 512
14
#define ILP 4
15

16
typedef enum{
17
  MOMENT_MODE_0   =0, // Novograd paper mode, momentum caculation with denom then decay inside
18
  MOMENT_MODE_1   =1  // Decoupled weight decay mode
19
} momentMode_t;
20

21
void multi_tensor_norm_out_cuda(
22
  int chunk_size,
23
  at::Tensor noop_flag,
24
  std::vector<std::vector<at::Tensor>> tensor_lists,
25
  at::Tensor out,
26
  const float alpha,
27
  const float beta,
28
  const int norm_type);
29

30
using MATH_T = float;
31

32
template<typename T>
33
struct NovoGradFunctor
34
{
35
   __device__ __forceinline__ void operator()(
36
    int chunk_size,
37
    volatile int* noop_gmem,
38
    TensorListMetadata<3>& tl,
39
    const float beta1,
40
    const float beta2,
41
    const float beta3,
42
    const float beta1_correction,
43
    const float beta2_correction,
44
    const float epsilon,
45
    const float lr,
46
    momentMode_t m_mode,
47
    const float decay,
48
    const float* per_tensor_grad_norm)
49
  {
50
    // I'd like this kernel to propagate infs/nans.
51
    // if(*noop_gmem == 1)
52
    //   return;
53

54
    int tensor_loc = tl.block_to_tensor[blockIdx.x];
55
    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
56
    int chunk_idx = tl.block_to_chunk[blockIdx.x];
57
    int n = tl.sizes[tensor_loc];
58

59
    float grad_norm = per_tensor_grad_norm[tensor_num];
60

61
    T* g = (T*)tl.addresses[0][tensor_loc];
62
    g += chunk_idx*chunk_size;
63

64
    T* p = (T*)tl.addresses[1][tensor_loc];
65
    p += chunk_idx*chunk_size;
66

67
    T* m = (T*)tl.addresses[2][tensor_loc];
68
    m += chunk_idx*chunk_size;
69

70
    n -= chunk_idx*chunk_size;
71

72
    // see note in multi_tensor_scale_kernel.cu
73
    for(int i_start = 0;
74
            i_start < n && i_start < chunk_size;
75
            i_start += blockDim.x*ILP)
76
    {
77
      MATH_T r_g[ILP];
78
      MATH_T r_p[ILP];
79
      MATH_T r_m[ILP];
80
#pragma unroll
81
      for(int ii = 0; ii < ILP; ii++)
82
      {
83
        int i = i_start + threadIdx.x + ii*blockDim.x;
84
        if(i < n && i < chunk_size)
85
        {
86
          r_g[ii] = g[i];
87
          r_p[ii] = p[i];
88
          r_m[ii] = m[i];
89
        } else {
90
          r_g[ii] = MATH_T(0);
91
          r_p[ii] = MATH_T(0);
92
          r_m[ii] = MATH_T(0);
93
        }
94
      }
95
#pragma unroll
96
      for(int ii = 0; ii < ILP; ii++)
97
      {
98
        if (m_mode == MOMENT_MODE_0) {
99
          MATH_T next_v_unbiased = grad_norm / beta2_correction;
100
          MATH_T denom = next_v_unbiased + epsilon;
101
          r_g[ii] = (r_g[ii] / denom) + (decay * r_p[ii]);
102
          r_m[ii] = beta1 * r_m[ii] + beta3 * r_g[ii];
103
          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
104
          r_p[ii] = r_p[ii] - (lr * next_m_unbiased);
105
        }
106
        else {
107
          r_m[ii] = beta1 * r_m[ii] + beta3 * r_g[ii];
108
          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
109
          MATH_T next_v_unbiased = grad_norm / beta2_correction;
110
          MATH_T denom = next_v_unbiased + epsilon;
111
          MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]);
112
          r_p[ii] = r_p[ii] - (lr * update);
113
        }
114
      }
115
#pragma unroll
116
      for(int ii = 0; ii < ILP; ii++)
117
      {
118
        int i = i_start + threadIdx.x + ii*blockDim.x;
119
        if(i < n && i < chunk_size)
120
        {
121
          p[i] = r_p[ii];
122
          m[i] = r_m[ii];
123
        }
124
      }
125
    }
126
  }
127
};
128

129
void multi_tensor_novograd_cuda(
130
  int chunk_size,
131
  at::Tensor noop_flag,
132
  std::vector<std::vector<at::Tensor>> tensor_lists,
133
  at::Tensor grad_norms,
134
  const float lr,
135
  const float beta1,
136
  const float beta2,
137
  const float epsilon,
138
  const int step,
139
  const int bias_correction,
140
  const float weight_decay,
141
  const int grad_averaging,
142
  const int moment_mode,
143
  const int norm_type)
144
{
145
  using namespace at;
146

147
  // Handle bias correction mode
148
  float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
149
  if (bias_correction == 1) {
150
    bias_correction1 = 1 - std::pow(beta1, step);
151
    bias_correction2 = std::sqrt(1 - std::pow(beta2, step));
152
  }
153

154
  // Handle grad averaging mode
155
  float beta3 = 1;
156
  if (grad_averaging == 1) beta3 = 1 - beta1;
157

158
  std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(), tensor_lists.begin()+1);
159

160
  // Compute and update grad norm
161
  // Here use a per tensor norm, and blend new norm(n) and old norm(gn) by
162
  // L-2: gn = sqrt(a * gn^2 + b * n^2)
163
  // L-inf: gn = a * gn + b * n
164
  multi_tensor_norm_out_cuda(chunk_size, noop_flag, grad_list, grad_norms, beta2, (1.0f - beta2), norm_type);
165

166
  // Assume single type across p,g,m1,m2 now
167
  DISPATCH_DOUBLE_FLOAT_AND_HALF(
168
    tensor_lists[0][0].scalar_type(), 0, "novograd",
169
    multi_tensor_apply<3>(
170
      BLOCK_SIZE,
171
      chunk_size,
172
      noop_flag,
173
      tensor_lists,
174
      NovoGradFunctor<scalar_t_0>(),
175
      beta1,
176
      beta2,
177
      beta3, // 1-beta1 or 1 depends on averaging mode
178
      bias_correction1,
179
      bias_correction2,
180
      epsilon,
181
      lr,
182
      (momentMode_t) moment_mode,
183
      weight_decay,
184
      grad_norms.DATA_PTR<float>()); )
185

186
  AT_CUDA_CHECK(cudaGetLastError());
187

188
}
189

190