1#include "ggml-impl.h"
2#include "opt-step-adamw.cuh"
3
4#include <cstdint>
5
6static __global__ void opt_step_adamw_f32(
7 float * __restrict__ x, const float * __restrict__ g, float * __restrict__ g_m, float * __restrict__ g_v,
8 const float * __restrict__ pars, const int64_t k) {
9
10 const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x;
11
12 if (i >= k) {
13 return;
14 }
15
16 const float alpha = pars[0];
17 const float beta1 = pars[1];
18 const float beta2 = pars[2];
19 const float eps = pars[3];
20 const float wd = pars[4];
21 const float beta1h = pars[5];
22 const float beta2h = pars[6];
23
24 const float gi = g[i];
25 const float gmi = g_m[i]*beta1 + gi*(1.0f - beta1);
26 const float gvi = g_v[i]*beta2 + gi*gi*(1.0f - beta2);
27
28 g_m[i] = gmi;
29 g_v[i] = gvi;
30
31 const float mh = gmi*beta1h;
32 const float vh = sqrtf(a: gvi*beta2h) + eps;
33
34 x[i] = x[i]*(1.0f - alpha*wd) - alpha*mh/vh;
35}
36
37static void opt_step_adamw_f32_cuda(
38 float * x, const float * g, float * g_m, float * g_v, const float * pars, const int64_t k, cudaStream_t stream) {
39
40 const dim3 block_dims(CUDA_OPT_STEP_ADAMW_BLOCK_SIZE, 1, 1);
41 const dim3 block_nums((k + CUDA_OPT_STEP_ADAMW_BLOCK_SIZE - 1) / CUDA_OPT_STEP_ADAMW_BLOCK_SIZE, 1, 1);
42 opt_step_adamw_f32<<<gridDim: block_nums, blockDim: block_dims, sharedMem: 0, stream>>>(x, g, g_m, g_v, pars, k);
43}
44
45void ggml_cuda_opt_step_adamw(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
46 const ggml_tensor * src0 = dst->src[0];
47 const ggml_tensor * src0_grad = dst->src[1];
48 const ggml_tensor * src0_grad_m = dst->src[2];
49 const ggml_tensor * src0_grad_v = dst->src[3];
50 const ggml_tensor * adamw_params = dst->src[4];
51
52 GGML_ASSERT(src0->type == GGML_TYPE_F32);
53 GGML_ASSERT(src0_grad->type == GGML_TYPE_F32);
54 GGML_ASSERT(src0_grad_m->type == GGML_TYPE_F32);
55 GGML_ASSERT(src0_grad_v->type == GGML_TYPE_F32);
56 GGML_ASSERT(adamw_params->type == GGML_TYPE_F32);
57 GGML_ASSERT(ggml_is_contiguous(src0));
58 GGML_ASSERT(ggml_is_contiguous(src0_grad));
59 GGML_ASSERT(ggml_is_contiguous(src0_grad_m));
60 GGML_ASSERT(ggml_is_contiguous(src0_grad_v));
61 GGML_ASSERT(ggml_is_contiguous(adamw_params));
62 GGML_ASSERT(ggml_are_same_shape(src0, src0_grad));
63 GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m));
64 GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v));
65 GGML_ASSERT(ggml_nelements(adamw_params) == 7);
66
67 float * src0_d = (float *) src0->data;
68 const float * src0_grad_d = (const float *) src0_grad->data;
69 float * src0_grad_m_d = (float *) src0_grad_m->data;
70 float * src0_grad_v_d = (float *) src0_grad_v->data;
71 const float * adamw_params_d = (const float *) adamw_params->data;
72
73 cudaStream_t stream = ctx.stream();
74
75 const int64_t ne = ggml_nelements(src0);
76
77 opt_step_adamw_f32_cuda(x: src0_d, g: src0_grad_d, g_m: src0_grad_m_d, g_v: src0_grad_v_d, pars: adamw_params_d, k: ne, stream);
78}
79