char-fork/llama.cpp

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

Neo Zhang[SYCL] support OP cross_entropy_loss, cross_entropy_loss_back (#25236)55edb2de4

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
8.8 KiB255 linesraw
1#include "cross_entropy_loss.hpp"
2
3#include <cstdint>
4#include <cmath>
5
6template <bool has_shared>
7static __dpct_inline__ void cross_entropy_loss_f32_kernel(
8        const float * __restrict__ logits,
9        const float * __restrict__ labels,
10        float * __restrict__ row_loss,
11        const int nclasses,
12        const int nrows,
13        float * __restrict__ smem,
14        const sycl::nd_item<3> & item) {
15
16    const int row = item.get_group(2);
17    const int tid = item.get_local_id(2);
18
19    logits += (int64_t) row * nclasses;
20    labels += (int64_t) row * nclasses;
21
22    float max_logit = -INFINITY;
23    for (int i = tid; i < nclasses; i += WARP_SIZE) {
24        const float v = logits[i];
25        max_logit = sycl::fmax(max_logit, v);
26        if (has_shared) {
27            smem[i] = v;
28        }
29    }
30    max_logit = warp_reduce_max<WARP_SIZE>(max_logit);
31
32    float sum_exp = 0.0f;
33    for (int i = tid; i < nclasses; i += WARP_SIZE) {
34        const float v = has_shared ? smem[i] : logits[i];
35        sum_exp += sycl::exp(v - max_logit);
36    }
37    sum_exp = warp_reduce_sum<WARP_SIZE>(sum_exp);
38    const float log_sum = sycl::log(sum_exp);
39
40    float loss = 0.0f;
41    for (int i = tid; i < nclasses; i += WARP_SIZE) {
42        const float v = has_shared ? smem[i] : logits[i];
43        loss += (v - max_logit - log_sum) * labels[i];
44    }
45    loss = -warp_reduce_sum<WARP_SIZE>(loss) / (float) nrows;
46
47    if (tid == 0) {
48        row_loss[row] = loss;
49    }
50}
51
52template <bool has_shared>
53static __dpct_inline__ void cross_entropy_loss_back_f32_kernel(
54        const float * __restrict__ grad,
55        const float * __restrict__ logits,
56        const float * __restrict__ labels,
57        float * __restrict__ dst,
58        const int nclasses,
59        const int nrows,
60        float * __restrict__ smem,
61        const sycl::nd_item<3> & item) {
62
63    const int row = item.get_group(2);
64    const int tid = item.get_local_id(2);
65
66    logits += (int64_t) row * nclasses;
67    labels += (int64_t) row * nclasses;
68    dst    += (int64_t) row * nclasses;
69
70    float max_logit = -INFINITY;
71    for (int i = tid; i < nclasses; i += WARP_SIZE) {
72        const float v = logits[i];
73        max_logit = sycl::fmax(max_logit, v);
74        if (has_shared) {
75            smem[i] = v;
76        }
77    }
78    max_logit = warp_reduce_max<WARP_SIZE>(max_logit);
79
80    float sum_exp = 0.0f;
81    for (int i = tid; i < nclasses; i += WARP_SIZE) {
82        const float v = sycl::exp((has_shared ? smem[i] : logits[i]) - max_logit);
83        sum_exp += v;
84        if (has_shared) {
85            smem[i] = v;
86        } else {
87            dst[i] = v;
88        }
89    }
90    sum_exp = warp_reduce_sum<WARP_SIZE>(sum_exp);
91    const float inv_sum = 1.0f / sum_exp;
92
93    const float d_by_nrows = grad[0] / (float) nrows;
94    for (int i = tid; i < nclasses; i += WARP_SIZE) {
95        const float sm_num = has_shared ? smem[i] : dst[i];
96        dst[i] = (sm_num * inv_sum - labels[i]) * d_by_nrows;
97    }
98}
99
100static void cross_entropy_reduce_rows(
101        ggml_backend_sycl_context & ctx,
102        const float * row_loss,
103        float * dst,
104        const int64_t nrows) {
105    if (nrows == 1) {
106        SYCL_CHECK(CHECK_TRY_ERROR(
107            ctx.stream()->memcpy(dst, row_loss, sizeof(float))));
108        return;
109    }
110
111    ggml_sycl_pool_alloc<float> tmp_alloc(ctx.pool(), nrows);
112    float * tmp = tmp_alloc.get();
113    SYCL_CHECK(CHECK_TRY_ERROR(
114        ctx.stream()->memcpy(tmp, row_loss, nrows * sizeof(float))));
115
116    int64_t cur = nrows;
117    while (cur > 1) {
118        const int64_t out = (cur + WARP_SIZE - 1) / WARP_SIZE;
119        const sycl::range<3> block(1, 1, WARP_SIZE);
120        const sycl::range<3> grid(1, 1, out);
121        ctx.stream()->parallel_for(
122            sycl::nd_range<3>(grid * block, block),
123            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
124                const int row = item.get_group(2);
125                const int tid = item.get_local_id(2);
126                const int64_t i = (int64_t) row * WARP_SIZE + tid;
127                float v = i < cur ? tmp[i] : 0.0f;
128                v = warp_reduce_sum<WARP_SIZE>(v);
129                if (tid == 0) {
130                    tmp[row] = v;
131                }
132            });
133        cur = out;
134    }
135
136    SYCL_CHECK(CHECK_TRY_ERROR(
137        ctx.stream()->memcpy(dst, tmp, sizeof(float))));
138}
139
140void ggml_sycl_cross_entropy_loss(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
141    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
142
143    const ggml_tensor * src0 = dst->src[0];
144    const ggml_tensor * src1 = dst->src[1];
145
146    GGML_ASSERT(src0->type == GGML_TYPE_F32);
147    GGML_ASSERT(src1->type == GGML_TYPE_F32);
148    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
149    GGML_ASSERT(ggml_is_contiguous(src0));
150    GGML_ASSERT(ggml_is_contiguous(src1));
151    GGML_ASSERT(ggml_is_contiguous(dst));
152    GGML_ASSERT(ggml_are_same_shape(src0, src1));
153    GGML_ASSERT(ggml_is_scalar(dst));
154
155    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
156
157    const int64_t nclasses = src0->ne[0];
158    const int64_t nrows = ggml_nrows(src0);
159
160    const float * logits_d = (const float *) src0->data;
161    const float * labels_d = (const float *) src1->data;
162    float * dst_d = (float *) dst->data;
163
164    ggml_sycl_pool_alloc<float> row_loss_alloc(ctx.pool(), nrows);
165    float * row_loss = row_loss_alloc.get();
166
167    const sycl::range<3> block(1, 1, WARP_SIZE);
168    const sycl::range<3> grid(1, 1, nrows);
169    const size_t nbytes_shared = (size_t) nclasses * sizeof(float);
170    const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo;
171
172    if (nbytes_shared <= smpbo) {
173        ctx.stream()->submit([&](sycl::handler & cgh) {
174            sycl::local_accessor<float, 1> smem(sycl::range<1>(nclasses), cgh);
175            cgh.parallel_for(
176                sycl::nd_range<3>(grid * block, block),
177                [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
178                    cross_entropy_loss_f32_kernel<true>(
179                        logits_d, labels_d, row_loss,
180                        (int) nclasses, (int) nrows,
181                        get_pointer(smem), item);
182                });
183        });
184    } else {
185        ctx.stream()->parallel_for(
186            sycl::nd_range<3>(grid * block, block),
187            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
188                cross_entropy_loss_f32_kernel<false>(
189                    logits_d, labels_d, row_loss,
190                    (int) nclasses, (int) nrows,
191                    nullptr, item);
192            });
193    }
194
195    cross_entropy_reduce_rows(ctx, row_loss, dst_d, nrows);
196}
197
198void ggml_sycl_cross_entropy_loss_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
199    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
200
201    const ggml_tensor * grad  = dst->src[0];
202    const ggml_tensor * src0f = dst->src[1];
203    const ggml_tensor * src1f = dst->src[2];
204
205    GGML_ASSERT(grad->type == GGML_TYPE_F32);
206    GGML_ASSERT(src0f->type == GGML_TYPE_F32);
207    GGML_ASSERT(src1f->type == GGML_TYPE_F32);
208    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
209
210    GGML_ASSERT(ggml_is_scalar(grad));
211    GGML_ASSERT(ggml_is_contiguous(grad));
212    GGML_ASSERT(ggml_is_contiguous(src0f));
213    GGML_ASSERT(ggml_is_contiguous(src1f));
214    GGML_ASSERT(ggml_is_contiguous(dst));
215    GGML_ASSERT(ggml_are_same_shape(src0f, src1f));
216    GGML_ASSERT(ggml_are_same_shape(src0f, dst));
217
218    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
219
220    const int64_t nclasses = src0f->ne[0];
221    const int64_t nrows = ggml_nrows(src0f);
222
223    const float * grad_d  = (const float *) grad->data;
224    const float * logits_d = (const float *) src0f->data;
225    const float * labels_d = (const float *) src1f->data;
226    float * dst_d = (float *) dst->data;
227
228    const sycl::range<3> block(1, 1, WARP_SIZE);
229    const sycl::range<3> grid(1, 1, nrows);
230    const size_t nbytes_shared = (size_t) nclasses * sizeof(float);
231    const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo;
232
233    if (nbytes_shared <= smpbo) {
234        ctx.stream()->submit([&](sycl::handler & cgh) {
235            sycl::local_accessor<float, 1> smem(sycl::range<1>(nclasses), cgh);
236            cgh.parallel_for(
237                sycl::nd_range<3>(grid * block, block),
238                [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
239                    cross_entropy_loss_back_f32_kernel<true>(
240                        grad_d, logits_d, labels_d, dst_d,
241                        (int) nclasses, (int) nrows,
242                        get_pointer(smem), item);
243                });
244        });
245    } else {
246        ctx.stream()->parallel_for(
247            sycl::nd_range<3>(grid * block, block),
248            [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
249                cross_entropy_loss_back_f32_kernel<false>(
250                    grad_d, logits_d, labels_d, dst_d,
251                    (int) nclasses, (int) nrows,
252                    nullptr, item);
253            });
254    }
255}