From 80072802d7ef9a42bad60e83904b2305784325b1 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Mon, 29 Apr 2024 21:46:35 -0700 Subject: [PATCH 1/6] Create gelu_backward.cu --- dev/cuda/gelu_backward.cu | 226 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 dev/cuda/gelu_backward.cu diff --git a/dev/cuda/gelu_backward.cu b/dev/cuda/gelu_backward.cu new file mode 100644 index 0000000..4489a20 --- /dev/null +++ b/dev/cuda/gelu_backward.cu @@ -0,0 +1,226 @@ +/* +Kernels for gelu backward pass. + +Compile example: +nvcc -O3 --use_fast_math gelu_backward.cu -o gelu_backward + +If encountering "error: identifier "M_PI" is undefined", add the following lines to the top of the file: + +#define _USE_MATH_DEFINES +#include OR #include + +version 1 is naive port from CPU code to kernel +./gelu_backward 1 + +version 2 uses the Packed128 data structure +./gelu_backward 2 +*/ + +#include +#include +#include +#include "common.h" + +// turn on bf16 as default, done up here for now +#define ENABLE_BF16 + +#if defined(ENABLE_BF16) +typedef __nv_bfloat16 floatX; +typedef __nv_bfloat16 floatN; +#elif defined(ENABLE_FP16) +typedef half floatX; +typedef half floatN; +#else +typedef float floatX; +typedef float floatN; +#endif + +typedef Packed128 x128; + +// ---------------------------------------------------------------------------- +// CPU code reference + +#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI) + +void gelu_backward_cpu(float* dinp, const float* inp, const float* dout, const int N) { + for (int i = 0; i < N; i++) { + float x = inp[i]; + float cube = 0.044715f * x * x * x; + 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 * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); + dinp[i] = (floatX)(local_grad * (float)dout[i]); + } +} + +// ---------------------------------------------------------------------------- +// GPU kernels + +// elementwise ops are nice and ez +__global__ void gelu_backward1(float* dinp, const float* inp, const float* dout, int N) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < N) { + float x = inp[i]; + float cube = 0.044715f * x * x * x; + 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 * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); + dinp[i] = (floatX)(local_grad * (float)dout[i]); + } +} + +__global__ void gelu_backward2(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { + int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + if (i < N) { + x128 packed_dinp; + x128 packed_inp = load128cs(inp + i); + x128 packed_dout = load128cs(dout + i); + for (int k = 0; k < packed_inp.size; ++k) { + float x = (float)packed_inp[k]; + float cube = 0.044715f * x * x * x; + 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 * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); + packed_dinp[k] = (floatX)(local_grad * (float)packed_dout[k]); + } + + store128(dinp + i, packed_dinp); + } +} + +// ---------------------------------------------------------------------------- +// kernel launcher + +void gelu_backward1(float* dinp, const float* inp, const float* dout, int N, const int block_size) { + const int grid_size = ceil_div(N, block_size); + gelu_backward1<<>>(dinp, inp, dout, N); + cudaCheck(cudaGetLastError()); +} + +void gelu_backward2(floatX* dinp, const floatX* inp, const floatX* dout, int N, const int block_size) { + const int grid_size = ceil_div(N, block_size)/x128::size; + gelu_backward2<<>>(dinp, inp, dout, N); + cudaCheck(cudaGetLastError()); +} + +// kernel version dispatch +void gelu_backward(int kernel_num, + floatX* dinp, + const floatX* inp, + const floatX* dout, + int B, int T, int C, + int block_size) { + switch (kernel_num) { +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) + case 1: + gelu_backward1(dinp, inp, dout, B * T * C, block_size); + break; +#endif +#if defined(ENABLE_BF16) + case 2: + gelu_backward2(dinp, inp, dout, B * T * C, block_size); + break; +#endif + default: + printf("Invalid kernel number\n"); + exit(1); + } +} + +// ---------------------------------------------------------------------------- + +int main(int argc, char **argv) { + srand(0); + + int B = 8; + int T = 1024; + int C = 768; + + int deviceIdx = 0; + cudaCheck(cudaSetDevice(deviceIdx)); + + // create host memory of random numbers + float* dinp = (float*)malloc(B * T * C * sizeof(float)); + float* inp = make_random_float(B * T * C); + float* dout = make_random_float(B * T * C); + + // read kernel_num from command line + int kernel_num = 1; + if (argc > 1) { + kernel_num = atoi(argv[1]); + } + printf("Using kernel %d\n", kernel_num); + + // first check the correctness of the kernel + gelu_backward_cpu(dinp, inp, dout, B * T * C); + + // move to GPU + floatX* d_dinp; + floatX* d_inp; + floatX* d_dout; + cudaCheck(cudaMalloc(&d_dinp, B * T * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(floatX))); + + floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX)); + floatX* doutX = (floatX*)malloc(B * T * C * sizeof(floatX)); + + for (int i = 0; i < B * T * C; i++) { + inpX[i] = (floatX)inp[i]; + doutX[i] = (floatX)dout[i]; + } + + cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(d_dout, doutX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); + + // time the kernel at different block sizes + int block_sizes[] = {32, 64, 128, 256, 512, 1024}; + for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { + int block_size = block_sizes[j]; + printf("Checking block size %d.\n", block_size); + gelu_backward(kernel_num, d_dinp, d_inp, d_dout, B, T, C, block_size); +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) + validate_result(d_dinp, dinp, "dinp", B * T * C, 1e-5f); +#endif +#if defined(ENABLE_BF16) +#endif + validate_result(d_dinp, dinp, "dinp", B * T * C, 1e-2f); + } + + printf("All results match. Starting benchmarks.\n\n"); + + for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { + int block_size = block_sizes[j]; + + int repeat_times = 1000; + + float elapsed_time = benchmark_kernel(repeat_times, gelu_backward, + kernel_num, d_dinp, d_inp, d_dout, + B, T, C, block_size); + + // napkin math: estimate the memory bandwidth achieved + // for each (B,T,C) output element, we do 1 read and 1 write, 4 bytes each + // and e.g. A100 40GB PCIe is advertised at 1,555GB/s + long memory_ops = B * T * C * 2 * 4; + float memory_bandwidth = memory_ops / elapsed_time / 1e6; + + printf("block_size %4d | time %.4f ms | bandwidth %.2f GB/s\n", block_size, elapsed_time, memory_bandwidth); + } + + // free memory + free(dinp); + free(inp); + free(dout); + free(inpX); + free(doutX); + cudaCheck(cudaFree(d_dinp)); + cudaCheck(cudaFree(d_inp)); + cudaCheck(cudaFree(d_dout)); + return 0; +} From ab2de05a139331febbf35b1479cbd39bacfd81b1 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Mon, 29 Apr 2024 21:49:32 -0700 Subject: [PATCH 2/6] Update train_gpt2.cu --- train_gpt2.cu | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8412658..3802066 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -959,16 +959,23 @@ __global__ void gelu_forward_kernel2(floatX* out, const floatX* inp, int N) { } __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { - int i = blockIdx.x * blockDim.x + threadIdx.x; + int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; if (i < N) { - float x = (float)inp[i]; - float cube = 0.044715f * x * x * x; - 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 * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); - dinp[i] = (floatX)(local_grad * (float)dout[i]); + x128 packed_dinp; + x128 packed_inp = load128cs(inp + i); + x128 packed_dout = load128cs(dout + i); + for (int k = 0; k < packed_inp.size; ++k) { + float x = (float)packed_inp[k]; + float cube = 0.044715f * x * x * x; + 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 * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); + packed_dinp[k] = (floatX)(local_grad * (float)packed_dout[k]); + } + + store128(dinp + i, packed_dinp); } } @@ -1495,7 +1502,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N) { void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { const int block_size = 128; - const int grid_size = CEIL_DIV(N, block_size); + const int grid_size = CEIL_DIV(N/x128::size, block_size); gelu_backward_kernel<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } From 7746217433ebb9b7e067b2b785ff335e1c15576c Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 04:26:19 +0000 Subject: [PATCH 3/6] update gelu backward rto allow all kernels to use both types --- dev/cuda/Makefile | 3 ++- dev/cuda/gelu_backward.cu | 20 ++++++++------------ 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/dev/cuda/Makefile b/dev/cuda/Makefile index af5fded..4ea7637 100644 --- a/dev/cuda/Makefile +++ b/dev/cuda/Makefile @@ -18,7 +18,7 @@ MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux- $(NVCC) $(CFLAGS) $(NVCCFLAGS) $< -o $@ # Build all targets -TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward +TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward all: $(TARGETS) # Individual targets: forward pass @@ -26,6 +26,7 @@ attention_forward: attention_forward.cu classifier_fused: classifier_fused.cu crossentropy_forward: crossentropy_forward.cu encoder_forward: encoder_forward.cu +gelu_backward: gelu_backward.cu gelu_forward: gelu_forward.cu layernorm_forward: layernorm_forward.cu residual_forward: residual_forward.cu diff --git a/dev/cuda/gelu_backward.cu b/dev/cuda/gelu_backward.cu index 4489a20..875e0cc 100644 --- a/dev/cuda/gelu_backward.cu +++ b/dev/cuda/gelu_backward.cu @@ -59,10 +59,10 @@ void gelu_backward_cpu(float* dinp, const float* inp, const float* dout, const i // GPU kernels // elementwise ops are nice and ez -__global__ void gelu_backward1(float* dinp, const float* inp, const float* dout, int N) { +__global__ void gelu_backward1(floatX* dinp, const floatX* inp, const floatX* dout, int N) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < N) { - float x = inp[i]; + float x = (float)inp[i]; float cube = 0.044715f * x * x * x; float tanh_arg = GELU_SCALING_FACTOR * (x + cube); float tanh_out = tanhf(tanh_arg); @@ -97,14 +97,14 @@ __global__ void gelu_backward2(floatX* dinp, const floatX* inp, const floatX* do // ---------------------------------------------------------------------------- // kernel launcher -void gelu_backward1(float* dinp, const float* inp, const float* dout, int N, const int block_size) { +void gelu_backward1(floatX* dinp, const floatX* inp, const floatX* dout, int N, const int block_size) { const int grid_size = ceil_div(N, block_size); gelu_backward1<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } void gelu_backward2(floatX* dinp, const floatX* inp, const floatX* dout, int N, const int block_size) { - const int grid_size = ceil_div(N, block_size)/x128::size; + const int grid_size = ceil_div(N, block_size * x128::size); gelu_backward2<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } @@ -117,16 +117,12 @@ void gelu_backward(int kernel_num, int B, int T, int C, int block_size) { switch (kernel_num) { -#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) case 1: gelu_backward1(dinp, inp, dout, B * T * C, block_size); break; -#endif -#if defined(ENABLE_BF16) case 2: gelu_backward2(dinp, inp, dout, B * T * C, block_size); break; -#endif default: printf("Invalid kernel number\n"); exit(1); @@ -186,11 +182,11 @@ int main(int argc, char **argv) { printf("Checking block size %d.\n", block_size); gelu_backward(kernel_num, d_dinp, d_inp, d_dout, B, T, C, block_size); #if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) - validate_result(d_dinp, dinp, "dinp", B * T * C, 1e-5f); + float tol = 1e-5; +#else + float tol = 1e-2f; #endif -#if defined(ENABLE_BF16) -#endif - validate_result(d_dinp, dinp, "dinp", B * T * C, 1e-2f); + validate_result(d_dinp, dinp, "dinp", B * T * C, tol); } printf("All results match. Starting benchmarks.\n\n"); From e9a80b5d84be027b85aa1e1bee2588294425f796 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 04:55:18 +0000 Subject: [PATCH 4/6] update ceildiv for gelu_backward --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 3802066..77a911d 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1502,7 +1502,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N) { void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { const int block_size = 128; - const int grid_size = CEIL_DIV(N/x128::size, block_size); + const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); gelu_backward_kernel<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } From d98e5ae859c70614f3a742d25bed5482a7d2087b Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 23:36:55 +0000 Subject: [PATCH 5/6] remove int casting --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 77a911d..6887999 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1502,7 +1502,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N) { void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { const int block_size = 128; - const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); + const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_backward_kernel<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } From 3de3c53ba9cd1d5a5df89e312fd67aea98348231 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 23:58:49 +0000 Subject: [PATCH 6/6] update kernel with util functions --- dev/cuda/gelu_backward.cu | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/dev/cuda/gelu_backward.cu b/dev/cuda/gelu_backward.cu index 875e0cc..8c64d7c 100644 --- a/dev/cuda/gelu_backward.cu +++ b/dev/cuda/gelu_backward.cu @@ -132,15 +132,12 @@ void gelu_backward(int kernel_num, // ---------------------------------------------------------------------------- int main(int argc, char **argv) { - srand(0); + setup_main(); int B = 8; int T = 1024; int C = 768; - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - // create host memory of random numbers float* dinp = (float*)malloc(B * T * C * sizeof(float)); float* inp = make_random_float(B * T * C); @@ -164,16 +161,8 @@ int main(int argc, char **argv) { cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(floatX))); - floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX)); - floatX* doutX = (floatX*)malloc(B * T * C * sizeof(floatX)); - - for (int i = 0; i < B * T * C; i++) { - inpX[i] = (floatX)inp[i]; - doutX[i] = (floatX)dout[i]; - } - - cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_dout, doutX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); + cudaCheck(memcpy_convert(d_inp, inp, B * T * C)); + cudaCheck(memcpy_convert(d_dout, dout, B * T * C)); // time the kernel at different block sizes int block_sizes[] = {32, 64, 128, 256, 512, 1024}; @@ -213,8 +202,6 @@ int main(int argc, char **argv) { free(dinp); free(inp); free(dout); - free(inpX); - free(doutX); cudaCheck(cudaFree(d_dinp)); cudaCheck(cudaFree(d_inp)); cudaCheck(cudaFree(d_dout));