make gelu constant be a #define, speeds up the kernel by 1% or so

This commit is contained in:
Andrej Karpathy 2024-04-11 15:49:01 +00:00
parent a08c11b60e
commit f26cf00a61
3 changed files with 10 additions and 11 deletions

View file

@ -38,12 +38,13 @@ void cudaCheck(cudaError_t error, const char *file, int line) {
// ----------------------------------------------------------------------------
// CPU code reference
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
void gelu_forward_cpu(float* out, float* inp, int N) {
float s = sqrtf(2.0f / M_PI);
for (int i = 0; i < N; i++) {
float x = inp[i];
float cube = 0.044715f * x * x * x;
out[i] = 0.5f * x * (1.0f + tanhf(s * (x + cube)));
out[i] = 0.5f * x * (1.0f + tanhf(GELU_SCALING_FACTOR * (x + cube)));
}
}
@ -53,11 +54,10 @@ void gelu_forward_cpu(float* out, float* inp, int N) {
// elementwise ops are nice and ez
__global__ void gelu_kernel(float* out, const float* inp, int N) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
float s = sqrtf(2.0f / M_PI);
if (i < N) {
float xi = inp[i];
float cube = 0.044715f * xi * xi * xi;
out[i] = 0.5f * xi * (1.0f + tanhf(s * (xi + cube)));
out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
}
}

View file

@ -335,25 +335,24 @@ void attention_backward(float* dinp, float* dpreatt, float* datt,
}
}
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
void gelu_forward(float* out, float* inp, int N) {
float s = sqrtf(2.0f / M_PI);
for (int i = 0; i < N; i++) {
float x = inp[i];
float cube = 0.044715f * x * x * x;
out[i] = 0.5f * x * (1.0f + tanhf(s * (x + cube)));
out[i] = 0.5f * x * (1.0f + tanhf(GELU_SCALING_FACTOR * (x + cube)));
}
}
void gelu_backward(float* dinp, float* inp, float* dout, int N) {
float s = sqrtf(2.0f / M_PI);
for (int i = 0; i < N; i++) {
float x = inp[i];
float cube = 0.044715f * x * x * x;
float tanh_arg = s * (x + cube);
float tanh_arg = GELU_SCALING_FACTOR * (x + cube);
float tanh_out = tanhf(tanh_arg);
float coshf_out = coshf(tanh_arg);
float sech_out = 1.0f / (coshf_out * coshf_out);
float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * s * (1.0f + 3.0f * 0.044715f * x * x);
float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x);
dinp[i] += local_grad * dout[i];
}
}

View file

@ -305,13 +305,13 @@ __global__ void residual_forward_kernel(float* out, float* inp1, float* inp2, in
}
}
#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI)
__global__ void gelu_kernel(float* out, const float* inp, int N) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
float s = sqrtf(2.0f / M_PI);
if (i < N) {
float xi = inp[i];
float cube = 0.044715f * xi * xi * xi;
out[i] = 0.5f * xi * (1.0f + tanhf(s * (xi + cube)));
out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)));
}
}