char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Neo Zhang[SYCL] support OP OPT_STEP_ADAMW, OPT_STEP_SGD (#25268)37a215c9e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.7 KiB131 linesraw
1#include "opt-step.hpp"
2
3#define SYCL_OPT_STEP_BLOCK_SIZE 256
4
5template <typename T>
6static void opt_step_adamw_f32_kernel(
7    T * __restrict__ x,
8    const T * __restrict__ g,
9    T * __restrict__ g_m,
10    T * __restrict__ g_v,
11    const T * __restrict__ pars,
12    const int64_t k,
13    const sycl::nd_item<1> & item) {
14
15    const int64_t i = (int64_t) item.get_global_id(0);
16    if (i >= k) {
17        return;
18    }
19
20    const float alpha  = pars[0];
21    const float beta1  = pars[1];
22    const float beta2  = pars[2];
23    const float eps    = pars[3];
24    const float wd     = pars[4];
25    const float beta1h = pars[5];
26    const float beta2h = pars[6];
27
28    const float gi  = g[i];
29    const float gmi = g_m[i] * beta1 +      gi * (1.0f - beta1);
30    const float gvi = g_v[i] * beta2 + gi * gi * (1.0f - beta2);
31
32    g_m[i] = gmi;
33    g_v[i] = gvi;
34
35    const float mh = gmi * beta1h;
36    const float vh = sycl::sqrt(gvi * beta2h) + eps;
37
38    x[i] = x[i] * (1.0f - alpha * wd) - alpha * mh / vh;
39}
40
41template <typename T>
42static void opt_step_sgd_f32_kernel(
43    T * __restrict__ x,
44    const T * __restrict__ g,
45    const T * __restrict__ pars,
46    const int64_t k,
47    const sycl::nd_item<1> & item) {
48
49    const int64_t i = (int64_t) item.get_global_id(0);
50    if (i >= k) {
51        return;
52    }
53
54    x[i] = x[i] * (1.0f - pars[0] * pars[1]) - pars[0] * g[i];
55}
56
57void ggml_sycl_opt_step_adamw(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
58    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/5);
59
60    const ggml_tensor * src0         = dst->src[0];
61    const ggml_tensor * src0_grad    = dst->src[1];
62    const ggml_tensor * src0_grad_m  = dst->src[2];
63    const ggml_tensor * src0_grad_v  = dst->src[3];
64    const ggml_tensor * adamw_params = dst->src[4];
65
66    GGML_ASSERT(src0->type         == GGML_TYPE_F32);
67    GGML_ASSERT(src0_grad->type    == GGML_TYPE_F32);
68    GGML_ASSERT(src0_grad_m->type  == GGML_TYPE_F32);
69    GGML_ASSERT(src0_grad_v->type  == GGML_TYPE_F32);
70    GGML_ASSERT(adamw_params->type == GGML_TYPE_F32);
71    GGML_ASSERT(ggml_is_contiguous(src0));
72    GGML_ASSERT(ggml_is_contiguous(src0_grad));
73    GGML_ASSERT(ggml_is_contiguous(src0_grad_m));
74    GGML_ASSERT(ggml_is_contiguous(src0_grad_v));
75    GGML_ASSERT(ggml_is_contiguous(adamw_params));
76    GGML_ASSERT(ggml_are_same_shape(src0, src0_grad));
77    GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_m));
78    GGML_ASSERT(ggml_are_same_shape(src0, src0_grad_v));
79    GGML_ASSERT(ggml_nelements(adamw_params) == 7);
80
81    dpct::queue_ptr stream = ctx.stream();
82    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
83
84    float       * src0_d         = (float       *) src0->data;
85    const float * src0_grad_d    = (const float *) src0_grad->data;
86    float       * src0_grad_m_d  = (float       *) src0_grad_m->data;
87    float       * src0_grad_v_d  = (float       *) src0_grad_v->data;
88    const float * adamw_params_d = (const float *) adamw_params->data;
89
90    const int64_t ne = ggml_nelements(src0);
91    const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE;
92
93    stream->parallel_for(
94        sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE),
95        [=](sycl::nd_item<1> item) {
96            opt_step_adamw_f32_kernel(src0_d, src0_grad_d, src0_grad_m_d, src0_grad_v_d, adamw_params_d, ne, item);
97        });
98}
99
100void ggml_sycl_opt_step_sgd(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
101    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
102
103    const ggml_tensor * src0       = dst->src[0];
104    const ggml_tensor * src0_grad  = dst->src[1];
105    const ggml_tensor * sgd_params = dst->src[2];
106
107    GGML_ASSERT(src0->type      == GGML_TYPE_F32);
108    GGML_ASSERT(src0_grad->type == GGML_TYPE_F32);
109    GGML_ASSERT(sgd_params->type == GGML_TYPE_F32);
110    GGML_ASSERT(ggml_is_contiguous(src0));
111    GGML_ASSERT(ggml_is_contiguous(src0_grad));
112    GGML_ASSERT(ggml_is_contiguous(sgd_params));
113    GGML_ASSERT(ggml_are_same_shape(src0, src0_grad));
114    GGML_ASSERT(ggml_nelements(sgd_params) == 2);
115
116    dpct::queue_ptr stream = ctx.stream();
117    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
118
119    float       * src0_d      = (float       *) src0->data;
120    const float * src0_grad_d = (const float *) src0_grad->data;
121    const float * sgd_params_d = (const float *) sgd_params->data;
122
123    const int64_t ne = ggml_nelements(src0);
124    const int64_t num_blocks = (ne + SYCL_OPT_STEP_BLOCK_SIZE - 1) / SYCL_OPT_STEP_BLOCK_SIZE;
125
126    stream->parallel_for(
127        sycl::nd_range<1>(num_blocks * SYCL_OPT_STEP_BLOCK_SIZE, SYCL_OPT_STEP_BLOCK_SIZE),
128        [=](sycl::nd_item<1> item) {
129            opt_step_sgd_f32_kernel(src0_d, src0_grad_d, sgd_params_d, ne, item);
130        });
131}