File size: 5,477 Bytes
4d5ab1c | 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 | /*
AdamW kernel
*/
// llmc internal imports
#include "cuda_common.h"
#include "cuda_utils.cuh"
// ----------------------------------------------------------------------------
// CUDA kernels
// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation).
// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda
__device__ float lerp(float start, float end, float weight) {
return fma(weight, end, fma(-weight, start, start));
}
template <typename Tp, typename Tg>
__device__ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay,
float grad_scale, unsigned int seed) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= num_parameters) { return; } // guard
// get the gradient, m, and v for this parameter
float grad = grad_scale * (float)grads_memory[idx];
float m = m_memory[idx];
float v = v_memory[idx];
// update the first moment (momentum)
m = lerp(grad, m, beta1);
m_memory[idx] = m;
// update the second moment (RMSprop)
v = lerp(grad * grad, v, beta2);
v_memory[idx] = v;
m /= beta1_correction; // m_hat
v /= beta2_correction; // v_hat
// fetch the old value of this parameter as a float, from either source
float old_param = (master_params_memory != NULL) ? master_params_memory[idx] : (float)params_memory[idx];
// update this parameter
float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param));
// update our low precision version of the parameters using stochastic rounding
// this will be used in the next forward pass
stochastic_rounding(param, ¶ms_memory[idx], seed);
// write the full, float version of the param into our master copy, if we maintain one
// this will be used in the next update
if (master_params_memory != NULL) { master_params_memory[idx] = param; }
}
template <typename Tp, typename Tg>
__global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride,
float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay,
float grad_scale, unsigned int seed) {
adamw_update(params_memory + blockIdx.y * w_stride,
master_params_memory ? master_params_memory + blockIdx.y * s_stride : NULL,
grads_memory + blockIdx.y * g_stride,
m_memory + blockIdx.y * s_stride,
v_memory + blockIdx.y * s_stride,
num_parameters, learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, grad_scale,
seed
);
}
template <typename Tp>
__global__ void init_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters,
ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= num_parameters) { return; }
params_memory += blockIdx.y * w_stride; // adjust for layer offset
master_params_memory += blockIdx.y * s_stride;
stochastic_rounding(master_params_memory[idx], ¶ms_memory[idx], seed);
}
template <typename Tp, typename Tg>
void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride, int num_slices, float learning_rate, float beta1, float beta2, int t, float eps, float weight_decay,
float grad_scale, unsigned int seed, cudaStream_t stream) {
// AdamW update
int block_size = 512;
int num_blocks = CEIL_DIV(num_parameters, block_size);
float beta1_correction = 1.0f - powf(beta1, t);
float beta2_correction = 1.0f - powf(beta2, t);
adamw_kernel3<<<dim3(num_blocks, num_slices), block_size, 0, stream>>>(params_memory, master_params_memory, grads_memory,
m_memory, v_memory, num_parameters, w_stride, g_stride, s_stride,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay,
grad_scale, seed);
cudaCheck(cudaGetLastError());
}
template <typename Tp>
void init_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters,
ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream) {
int block_size = 512; // must match block size of adamw_update so that RNG also matches
int num_blocks = CEIL_DIV(num_parameters, block_size);
init_from_master_kernel<<<dim3(num_blocks, num_slices), block_size, 0, stream>>>
(params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed);
cudaCheck(cudaGetLastError());
}
|