cwenzi/neuroflow-cpp

1#include "neuroflow/swiglu.hpp"
2
3#include <algorithm>
4#include <cmath>
5#include <cstring>
6
7#ifdef USE_CBLAS
8#include <cblas.h>
9#endif
10
11#ifdef USE_CUDA
12#include "cuda_context.hpp"
13#endif
14
15namespace neuroflow {
16
17static Tensor linear_backward_weight(const Tensor& input, const Tensor& output_grad) {
18    size_t in_dim = input.shape_.back();
19    size_t out_dim = output_grad.shape_.back();
20    size_t batch = input.numel() / in_dim;
21
22    Tensor grad({out_dim, in_dim}, QuantType::FP32);
23    float* gp = grad.as_fp32();
24    memset(gp, 0, grad.data_size_);
25
26    const float* ip = input.as_fp32();
27    const float* op = output_grad.as_fp32();
28
29#ifdef USE_CBLAS
30    cblas_sgemm(CblasRowMajor, CblasTrans, CblasNoTrans,
31        static_cast<int>(out_dim), static_cast<int>(in_dim), static_cast<int>(batch),
32        1.0f, op, static_cast<int>(out_dim), ip, static_cast<int>(in_dim),
33        0.0f, gp, static_cast<int>(in_dim));
34#else
35    for (size_t b = 0; b < batch; ++b) {
36        for (size_t o = 0; o < out_dim; ++o) {
37            for (size_t i = 0; i < in_dim; ++i) {
38                gp[o * in_dim + i] += op[b * out_dim + o] * ip[b * in_dim + i];
39            }
40        }
41    }
42#endif
43    return grad;
44}
45
46static Tensor linear_backward_input(const Tensor& output_grad, const Tensor& weight) {
47    size_t out_dim = output_grad.shape_.back();
48    size_t batch = output_grad.numel() / out_dim;
49    size_t in_dim = weight.shape_[1];
50
51    Tensor grad({batch, in_dim}, QuantType::FP32);
52    float* gp = grad.as_fp32();
53    const float* op = output_grad.as_fp32();
54    const float* wp = weight.as_fp32();
55
56#ifdef USE_CBLAS
57    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
58        static_cast<int>(batch), static_cast<int>(in_dim), static_cast<int>(out_dim),
59        1.0f, op, static_cast<int>(out_dim), wp, static_cast<int>(in_dim),
60        0.0f, gp, static_cast<int>(in_dim));
61#else
62    for (size_t b = 0; b < batch; ++b) {
63        for (size_t i = 0; i < in_dim; ++i) {
64            float val = 0.0f;
65            for (size_t o = 0; o < out_dim; ++o) {
66                val += op[b * out_dim + o] * wp[o * in_dim + i];
67            }
68            gp[b * in_dim + i] = val;
69        }
70    }
71#endif
72    return grad;
73}
74
75static Tensor bias_backward(const Tensor& output_grad) {
76    size_t out_dim = output_grad.shape_.back();
77    size_t batch = output_grad.numel() / out_dim;
78
79    Tensor grad({out_dim}, QuantType::FP32);
80    float* gp = grad.as_fp32();
81    const float* op = output_grad.as_fp32();
82
83    memset(gp, 0, grad.data_size_);
84    for (size_t b = 0; b < batch; ++b) {
85        for (size_t o = 0; o < out_dim; ++o) {
86            gp[o] += op[b * out_dim + o];
87        }
88    }
89    return grad;
90}
91
92SwiGLUFFN::SwiGLUFFN(size_t d_model, size_t d_ff)
93    : d_model_(d_model), d_ff_(d_ff > 0 ? d_ff : d_model * 4) {
94    w_gate_ = std::make_shared<Linear>(d_model_, d_ff_, true);
95    w_up_ = std::make_shared<Linear>(d_model_, d_ff_, true);
96    w_down_ = std::make_shared<Linear>(d_ff_, d_model_, true);
97}
98
99Tensor SwiGLUFFN::forward(const Tensor& x) {
100    cache_.input = x.clone();
101
102    cache_.gate_out = w_gate_->forward(x);
103    cache_.up_out = w_up_->forward(x);
104
105    size_t n = cache_.gate_out.numel();
106    cache_.gate_activated = cache_.gate_out.clone();
107
108#ifdef USE_CUDA
109    if (CudaContext::instance().is_available() && cache_.gate_activated.is_on_gpu()) {
110        launch_silu(cache_.gate_activated.as_gpu_fp32(), n, CudaContext::instance().stream());
111        cache_.gate_activated.gpu_dirty_ = true;
112    } else {
113#endif
114    float* ga = cache_.gate_activated.as_fp32();
115    for (size_t i = 0; i < n; ++i) {
116        float v = ga[i];
117        ga[i] = v / (1.0f + expf(-v));
118    }
119#ifdef USE_CUDA
120    }
121#endif
122
123    cache_.multiplied = Tensor(cache_.gate_activated.shape_, QuantType::FP32);
124
125#ifdef USE_CUDA
126    if (CudaContext::instance().is_available() && cache_.gate_activated.is_on_gpu()) {
127        cache_.multiplied.to_gpu();
128        launch_elementwise_mul(cache_.multiplied.as_gpu_fp32(),
129                    cache_.gate_activated.as_gpu_fp32(),
130                    cache_.up_out.as_gpu_fp32(), n,
131                    CudaContext::instance().stream());
132        cache_.multiplied.gpu_dirty_ = true;
133    } else {
134#endif
135    const float* ga_p = cache_.gate_activated.as_fp32();
136    const float* up_p = cache_.up_out.as_fp32();
137    float* mp = cache_.multiplied.as_fp32();
138    for (size_t i = 0; i < n; ++i) {
139        mp[i] = ga_p[i] * up_p[i];
140    }
141#ifdef USE_CUDA
142    }
143#endif
144
145    return w_down_->forward(cache_.multiplied);
146}
147
148SwiGLUFFN::Gradients SwiGLUFFN::backward(const Tensor& output_grad) {
149    Gradients grads;
150
151    grads.w_down_weight_grad = linear_backward_weight(cache_.multiplied, output_grad);
152    grads.w_down_bias_grad = bias_backward(output_grad);
153    Tensor d_multiplied = linear_backward_input(output_grad, w_down_->weight);
154
155    Tensor d_gate_activated(d_multiplied.shape_, QuantType::FP32);
156    Tensor d_up(d_multiplied.shape_, QuantType::FP32);
157    const float* dm = d_multiplied.as_fp32();
158    const float* ga_p = cache_.gate_activated.as_fp32();
159    const float* up_p = cache_.up_out.as_fp32();
160    float* dga = d_gate_activated.as_fp32();
161    float* dup = d_up.as_fp32();
162    size_t n = d_multiplied.numel();
163    for (size_t i = 0; i < n; ++i) {
164        dga[i] = dm[i] * up_p[i];
165        dup[i] = dm[i] * ga_p[i];
166    }
167
168    const float* go_p = cache_.gate_out.as_fp32();
169    Tensor d_gate_out(d_gate_activated.shape_, QuantType::FP32);
170    float* dgo = d_gate_out.as_fp32();
171    for (size_t i = 0; i < n; ++i) {
172        float sig = ga_p[i];
173        dgo[i] = dga[i] * sig * (1.0f + go_p[i] * (1.0f - sig));
174    }
175
176    grads.w_gate_weight_grad = linear_backward_weight(cache_.input, d_gate_out);
177    grads.w_gate_bias_grad = bias_backward(d_gate_out);
178    Tensor d_input_gate = linear_backward_input(d_gate_out, w_gate_->weight);
179
180    grads.w_up_weight_grad = linear_backward_weight(cache_.input, d_up);
181    grads.w_up_bias_grad = bias_backward(d_up);
182    Tensor d_input_up = linear_backward_input(d_up, w_up_->weight);
183
184    grads.input_grad = Tensor(cache_.input.shape_, QuantType::FP32);
185    const float* dig = d_input_gate.as_fp32();
186    const float* diu = d_input_up.as_fp32();
187    float* ig = grads.input_grad.as_fp32();
188    for (size_t i = 0; i < grads.input_grad.numel(); ++i) {
189        ig[i] = dig[i] + diu[i];
190    }
191
192    return grads;
193}
194
195#ifdef USE_CUDA
196__global__ void kernel_silu_impl(float* data, size_t n) {
197    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
198    if (idx >= n) return;
199    float v = data[idx];
200    data[idx] = v / (1.0f + expf(-v));
201}
202
203void launch_silu(float* data, size_t n, cudaStream_t stream) {
204    int block = 256;
205    int grid = (static_cast<int>(n) + block - 1) / block;
206    kernel_silu_impl<<<grid, block, 0, stream>>>(data, n);
207}
208
209__global__ void kernel_elementwise_mul_impl(float* out, const float* a, const float* b, size_t n) {
210    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
211    if (idx >= n) return;
212    out[idx] = a[idx] * b[idx];
213}
214
215void launch_elementwise_mul(float* out, const float* a, const float* b, size_t n, cudaStream_t stream) {
216    int block = 256;
217    int grid = (static_cast<int>(n) + block - 1) / block;
218    kernel_elementwise_mul_impl<<<grid, block, 0, stream>>>(out, a, b, n);
219}
220#endif
221
222} // namespace neuroflow
223