char-fork/llama.cpp

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

Aman GuptaCUDA: add dynamic shared mem to softmax, refactor general usage (#14497)55c2646b4

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
5.9 KiB177 linesraw
1#include "common.cuh"
2#include "cross-entropy-loss.cuh"
3#include "sum.cuh"
4
5#include <cmath>
6#include <cstdint>
7
8template <bool use_shared>
9static __global__ void cross_entropy_loss_f32(
10        const float * __restrict__ logits, const float * __restrict__ labels, float * __restrict__ dst, const int nclasses, const int k) {
11    extern __shared__ float tmp[];
12
13    logits += int64_t(blockIdx.x)*nclasses;
14    labels += int64_t(blockIdx.x)*nclasses;
15
16    // Find maximum for softmax:
17    float max_logit = -INFINITY;
18    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
19        const float val = logits[i];
20        max_logit = fmaxf(max_logit, val);
21
22        if (use_shared) {
23            tmp[i] = val;
24        }
25    }
26    max_logit = warp_reduce_max(max_logit);
27
28    // Calculate log(softmax(logits)) which is just logits - max:
29    float sum = 0.0f;
30    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
31        const float logit_i = use_shared ? tmp[i] : logits[i];
32        sum += expf(logit_i - max_logit);
33    }
34    sum = warp_reduce_sum(sum);
35    sum = logf(sum);
36
37    // log(exp(logits - max) / sum) = (logits - max) - log(sum)
38    float loss = 0.0f;
39    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
40        const float logit_i = use_shared ? tmp[i] : logits[i];
41        loss += (logit_i - max_logit - sum) * labels[i];
42    }
43    loss = -warp_reduce_sum(loss) / (float)k;
44
45    if (threadIdx.x != 0) {
46        return;
47    }
48
49    dst[blockIdx.x] = loss;
50}
51
52template <bool use_shared>
53static __global__ void cross_entropy_loss_back_f32(
54        const float * __restrict__ grad, const float * __restrict__ logits, const float * __restrict__ labels,
55        float * __restrict__ dst, const int nclasses) {
56    extern __shared__ float tmp[];
57
58    logits += int64_t(blockIdx.x)*nclasses;
59    labels += int64_t(blockIdx.x)*nclasses;
60    dst    += int64_t(blockIdx.x)*nclasses;
61
62    float maxval = -INFINITY;
63    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
64        const float val = logits[i];
65        maxval = fmaxf(maxval, val);
66
67        if (use_shared) {
68            tmp[i] = val;
69        }
70    }
71    maxval = warp_reduce_max(maxval);
72
73    float sum = 0.0f;
74    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
75        const float val = expf((use_shared ? tmp[i] : logits[i]) - maxval);
76        sum += val;
77
78        if (use_shared) {
79            tmp[i] = val;
80        } else {
81            dst[i] = val;
82        }
83    }
84    sum = warp_reduce_sum(sum);
85    const float sm_scale = 1.0f/sum;
86
87    const float d_by_nrows = *grad/gridDim.x;
88    for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) {
89        const float val = use_shared ? tmp[i] : dst[i];
90        dst[i] = (val*sm_scale - labels[i])*d_by_nrows;
91    }
92}
93
94void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
95    const ggml_tensor * src0 = dst->src[0];
96    const ggml_tensor * src1 = dst->src[1];
97
98    GGML_ASSERT(src0->type == GGML_TYPE_F32);
99    GGML_ASSERT(src1->type == GGML_TYPE_F32);
100    GGML_ASSERT( dst->type == GGML_TYPE_F32);
101
102    GGML_ASSERT(ggml_is_contiguous(src0));
103    GGML_ASSERT(ggml_is_contiguous(src1));
104    GGML_ASSERT(ggml_is_contiguous(dst));
105
106    const int64_t ne00  = src0->ne[0];
107    const int64_t nrows = ggml_nrows(src0);
108
109    const float * src0_d = (const float *) src0->data;
110    const float * src1_d = (const float *) src1->data;
111    float       * dst_d  = (float       *) dst->data;
112
113    ggml_cuda_pool & pool = ctx.pool();
114    cudaStream_t stream = ctx.stream();
115
116    const dim3 blocks_dim(WARP_SIZE, 1, 1);
117    const dim3 blocks_num(nrows, 1, 1);
118    const size_t nbytes_shared = ne00*sizeof(float);
119
120    const int    id    = ggml_cuda_get_device();
121    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
122
123    ggml_cuda_pool_alloc<float> dst_tmp(pool, blocks_num.x);
124
125    if (nbytes_shared <= smpbo) {
126        CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32<true>), smpbo);
127        cross_entropy_loss_f32<true><<<blocks_num, blocks_dim, nbytes_shared, stream>>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows);
128    } else {
129        cross_entropy_loss_f32<false><<<blocks_num, blocks_dim, 0, stream>>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows);
130    }
131    CUDA_CHECK(cudaGetLastError());
132
133    // Combine results from individual blocks:
134    sum_f32_cuda(pool, dst_tmp.ptr, dst_d, blocks_num.x, stream);
135}
136
137void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
138    const ggml_tensor * grad  = dst->src[0];
139    const ggml_tensor * src0f = dst->src[1];
140    const ggml_tensor * src1f = dst->src[2];
141
142    GGML_ASSERT(src0f->type == GGML_TYPE_F32);
143    GGML_ASSERT(src1f->type == GGML_TYPE_F32);
144    GGML_ASSERT( grad->type == GGML_TYPE_F32);
145    GGML_ASSERT(  dst->type == GGML_TYPE_F32);
146
147    GGML_ASSERT(ggml_is_scalar(grad));
148    GGML_ASSERT(ggml_is_contiguous(src0f));
149    GGML_ASSERT(ggml_is_contiguous(src1f));
150    GGML_ASSERT(ggml_is_contiguous(dst));
151    GGML_ASSERT(ggml_are_same_shape(src0f, src1f));
152    GGML_ASSERT(ggml_are_same_shape(src0f, dst));
153
154    const int64_t ne00  = src0f->ne[0];
155    const int64_t nrows = ggml_nrows(src0f);
156
157    const float * grad_d  = (const float *) grad->data;
158    const float * src0f_d = (const float *) src0f->data;
159    const float * src1f_d = (const float *) src1f->data;
160    float       * dst_d   = (float       *) dst->data;
161
162    cudaStream_t stream = ctx.stream();
163
164    const dim3 blocks_dim(WARP_SIZE, 1, 1);
165    const dim3 blocks_num(nrows, 1, 1);
166    const size_t nbytes_shared = ne00*sizeof(float);
167
168    const int    id    = ggml_cuda_get_device();
169    const size_t smpbo = ggml_cuda_info().devices[id].smpbo;
170
171    if (nbytes_shared <= smpbo) {
172        CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32<true>), smpbo);
173        cross_entropy_loss_back_f32<true><<<blocks_num, blocks_dim, nbytes_shared, stream>>>(grad_d, src0f_d, src1f_d, dst_d, ne00);
174    } else {
175        cross_entropy_loss_back_f32<false><<<blocks_num, blocks_dim, 0, stream>>>(grad_d, src0f_d, src1f_d, dst_d, ne00);
176    }
177}