From ac2d6350075da6d789b824195ed0a021dd88a478 Mon Sep 17 00:00:00 2001 From: Yusheng Su Date: Thu, 2 May 2024 02:11:38 +0400 Subject: [PATCH 01/29] add required package (requests) into requirements.txt --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 220ccbd..44496a3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ numpy torch tiktoken transformers +requests From 80072802d7ef9a42bad60e83904b2305784325b1 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Mon, 29 Apr 2024 21:46:35 -0700 Subject: [PATCH 02/29] 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 03/29] 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 04/29] 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 05/29] 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 06/29] 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 07/29] 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)); From 4ffcf5b90adae5c2c056fd62effc1aba3c1c63ca Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Mon, 29 Apr 2024 19:58:16 -0700 Subject: [PATCH 08/29] Update residual_forward.cu --- dev/cuda/residual_forward.cu | 88 ++++++++++++++++++++++++++++++------ 1 file changed, 74 insertions(+), 14 deletions(-) diff --git a/dev/cuda/residual_forward.cu b/dev/cuda/residual_forward.cu index 71c14ea..75c9cc9 100644 --- a/dev/cuda/residual_forward.cu +++ b/dev/cuda/residual_forward.cu @@ -6,6 +6,8 @@ nvcc -O3 --use_fast_math residual_forward.cu -o residual_forward version 1 is naive port from CPU code to kernel ./residual_forward 1 +version 2 packs input into 128 bit memory reads +./residual_forward 2 */ #include @@ -13,6 +15,21 @@ version 1 is naive port from CPU code to kernel #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 lol @@ -26,33 +43,60 @@ void residual_forward_cpu(float* out, const float* inp1, const float* inp2, int // GPU kernels // elementwise ops are nice and ez -__global__ void residual_forward_kernel(float* out, const float* inp1, const float* inp2, int N) { +__global__ void residual_forward_kernel1(float* out, const float* inp1, const float* inp2, int N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { out[idx] = inp1[idx] + inp2[idx]; } } +__global__ void residual_forward_kernel2(floatX* out, const floatX* inp1, const floatX* inp2, int N) { + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + if (idx < N) { + x128 packed_out; + x128 packed_inp1 = load128cs(inp1 + idx); + x128 packed_inp2 = load128cs(inp2 + idx); + for (int k = 0; k < packed_inp1.size; ++k) + { + packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); + } + store128(out + idx, packed_out); + } +} + // ---------------------------------------------------------------------------- // kernel launcher void residual_forward1(float* out, const float* inp1, const float* inp2, int N, const int block_size) { const int grid_size = ceil_div(N, block_size); - residual_forward_kernel<<>>(out, inp1, inp2, N); + residual_forward_kernel1<<>>(out, inp1, inp2, N); + cudaCheck(cudaGetLastError()); +} + +void residual_forward2(floatX* out, const floatX* inp1, const floatX* inp2, int N, const int block_size) { + const int grid_size = ceil_div(N, block_size)/x128::size; + residual_forward_kernel2<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } // kernel version dispatch void residual_forward(int kernel_num, - float* out, - const float* inp1, - const float* inp2, + floatX* out, + const floatX* inp1, + const floatX* inp2, int N, int block_size) { switch (kernel_num) { +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) case 1: residual_forward1(out, inp1, inp2, N, block_size); break; +#endif +#if defined(ENABLE_BF16) + case 2: + residual_forward2(out, inp1, inp2, N, block_size); + break; +#endif default: printf("Invalid kernel number\n"); exit(1); @@ -76,15 +120,24 @@ int main(int argc, char **argv) { float* inp1 = make_random_float(B * T * C); float* inp2 = make_random_float(B * T * C); + // create X host memory of random numbers + floatX* inp1X = (floatX*)malloc(B * T * C * sizeof(float)); + floatX* inp2X = (floatX*)malloc(B * T * C * sizeof(float)); + + for (int i = 0; i < B * T * C; i++) { + inp1X[i] = (floatX)inp1[i]; + inp2X[i] = (floatX)inp2[i]; + } + // move to GPU - float* d_out; - float* d_inp1; - float* d_inp2; - cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float))); - cudaCheck(cudaMalloc(&d_inp1, B * T * C * sizeof(float))); - cudaCheck(cudaMalloc(&d_inp2, B * T * C * sizeof(float))); - cudaCheck(cudaMemcpy(d_inp1, inp1, B * T * C * sizeof(float), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_inp2, inp2, B * T * C * sizeof(float), cudaMemcpyHostToDevice)); + floatX* d_out; + floatX* d_inp1; + floatX* d_inp2; + cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_inp1, B * T * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_inp2, B * T * C * sizeof(floatX))); + cudaCheck(cudaMemcpy(d_inp1, inp1X, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(d_inp2, inp2X, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); // read kernel_num from command line int kernel_num = 1; @@ -104,7 +157,12 @@ int main(int argc, char **argv) { int block_size = block_sizes[j]; printf("Checking block size %d.\n", block_size); residual_forward(kernel_num, d_out, d_inp1, d_inp2, B * T * C, block_size); +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) validate_result(d_out, out, "out", B * T * C, 1e-5f); +#endif +#if defined(ENABLE_BF16) + validate_result(d_out, out, "out", B * T * C, 1e-2f); +#endif } printf("All results match. Starting benchmarks.\n\n"); @@ -130,9 +188,11 @@ int main(int argc, char **argv) { free(out); free(inp1); free(inp2); + free(inp1X); + free(inp2X); cudaCheck(cudaFree(d_out)); cudaCheck(cudaFree(d_inp1)); cudaCheck(cudaFree(d_inp2)); return 0; -} \ No newline at end of file +} From 32df703a0541239ece0f19bf1a1e12e9e6147ee0 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Mon, 29 Apr 2024 20:02:25 -0700 Subject: [PATCH 09/29] Update train_gpt2.cu --- train_gpt2.cu | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8412658..8a38044 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -935,9 +935,17 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons } __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2, int N) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; if (idx < N) { - out[idx] = (floatX)((float)__ldcs(&inp1[idx]) + (float)__ldcs(&inp2[idx])); + x128 packed_out; + x128 packed_inp1 = load128cs(inp1 + idx); + x128 packed_inp2 = load128cs(inp2 + idx); + #pragma unroll packed_inp1.size + for (int k = 0; k < packed_inp1.size; ++k) + { + packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); + } + store128(out + idx, packed_out); } } @@ -1481,7 +1489,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); + const int grid_size = CEIL_DIV(N/x128::size, block_size); residual_forward_kernel<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } From 52b4dd02a1beb392323fc7d9a2645c836f80f997 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 04:47:57 +0000 Subject: [PATCH 10/29] update residual forward to allow kernels to use both types --- dev/cuda/residual_forward.cu | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/dev/cuda/residual_forward.cu b/dev/cuda/residual_forward.cu index 75c9cc9..509d6d1 100644 --- a/dev/cuda/residual_forward.cu +++ b/dev/cuda/residual_forward.cu @@ -16,7 +16,7 @@ version 2 packs input into 128 bit memory reads #include "common.h" // turn on bf16 as default, done up here for now -//#define ENABLE_BF16 +#define ENABLE_BF16 #if defined(ENABLE_BF16) typedef __nv_bfloat16 floatX; @@ -43,10 +43,10 @@ void residual_forward_cpu(float* out, const float* inp1, const float* inp2, int // GPU kernels // elementwise ops are nice and ez -__global__ void residual_forward_kernel1(float* out, const float* inp1, const float* inp2, int N) { +__global__ void residual_forward_kernel1(floatX* out, const floatX* inp1, const floatX* inp2, int N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { - out[idx] = inp1[idx] + inp2[idx]; + out[idx] = (floatX)((float)inp1[idx] + (float)inp2[idx]); } } @@ -67,14 +67,14 @@ __global__ void residual_forward_kernel2(floatX* out, const floatX* inp1, const // ---------------------------------------------------------------------------- // kernel launcher -void residual_forward1(float* out, const float* inp1, const float* inp2, int N, const int block_size) { +void residual_forward1(floatX* out, const floatX* inp1, const floatX* inp2, int N, const int block_size) { const int grid_size = ceil_div(N, block_size); residual_forward_kernel1<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } void residual_forward2(floatX* out, const floatX* inp1, const floatX* inp2, int N, const int block_size) { - const int grid_size = ceil_div(N, block_size)/x128::size; + const int grid_size = ceil_div(N, (int)(block_size * x128::size)); residual_forward_kernel2<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } @@ -87,16 +87,12 @@ void residual_forward(int kernel_num, int N, int block_size) { switch (kernel_num) { -#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) case 1: residual_forward1(out, inp1, inp2, N, block_size); break; -#endif -#if defined(ENABLE_BF16) case 2: residual_forward2(out, inp1, inp2, N, block_size); break; -#endif default: printf("Invalid kernel number\n"); exit(1); @@ -158,11 +154,11 @@ int main(int argc, char **argv) { printf("Checking block size %d.\n", block_size); residual_forward(kernel_num, d_out, d_inp1, d_inp2, B * T * C, block_size); #if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) - validate_result(d_out, out, "out", B * T * C, 1e-5f); -#endif -#if defined(ENABLE_BF16) - validate_result(d_out, out, "out", B * T * C, 1e-2f); + float tol = 1e-5; +#else + float tol = 1e-2f; #endif + validate_result(d_out, out, "out", B * T * C, tol); } printf("All results match. Starting benchmarks.\n\n"); From 4698164e3eba918efd8589f5fe53738675759819 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Wed, 1 May 2024 04:49:23 +0000 Subject: [PATCH 11/29] update train to multiply block size --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8a38044..2525bb8 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1489,7 +1489,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { const int block_size = 256; - const int grid_size = CEIL_DIV(N/x128::size, block_size); + const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); residual_forward_kernel<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } From 7203c7875f095b7a6ece41501b0e0270db8fad10 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Thu, 2 May 2024 00:14:41 +0000 Subject: [PATCH 12/29] remove int cast --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 2525bb8..4edd043 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1489,7 +1489,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { const int block_size = 256; - const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); + const int grid_size = CEIL_DIV(N, block_size * x128::size); residual_forward_kernel<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } From d0e5fea2dd0da96c0d5d75dc0f1c79200e2decf5 Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Thu, 2 May 2024 00:17:41 +0000 Subject: [PATCH 13/29] update residual_forward to use util functions --- dev/cuda/residual_forward.cu | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/dev/cuda/residual_forward.cu b/dev/cuda/residual_forward.cu index 509d6d1..41b3e31 100644 --- a/dev/cuda/residual_forward.cu +++ b/dev/cuda/residual_forward.cu @@ -102,29 +102,17 @@ void residual_forward(int kernel_num, // ---------------------------------------------------------------------------- int main(int argc, char **argv) { - srand(0); + set_main(); int B = 8; int T = 1024; int C = 768; - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - // create host memory of random numbers float* out = (float*)malloc(B * T * C * sizeof(float)); float* inp1 = make_random_float(B * T * C); float* inp2 = make_random_float(B * T * C); - - // create X host memory of random numbers - floatX* inp1X = (floatX*)malloc(B * T * C * sizeof(float)); - floatX* inp2X = (floatX*)malloc(B * T * C * sizeof(float)); - - for (int i = 0; i < B * T * C; i++) { - inp1X[i] = (floatX)inp1[i]; - inp2X[i] = (floatX)inp2[i]; - } - + // move to GPU floatX* d_out; floatX* d_inp1; @@ -132,8 +120,8 @@ int main(int argc, char **argv) { cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_inp1, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_inp2, B * T * C * sizeof(floatX))); - cudaCheck(cudaMemcpy(d_inp1, inp1X, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_inp2, inp2X, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); + cudaCheck(memcpy_convert(d_inp1, inp1, B * T * C)); + cudaCheck(memcpy_convert(d_inp2, inp2, B * T * C)); // read kernel_num from command line int kernel_num = 1; @@ -184,8 +172,6 @@ int main(int argc, char **argv) { free(out); free(inp1); free(inp2); - free(inp1X); - free(inp2X); cudaCheck(cudaFree(d_out)); cudaCheck(cudaFree(d_inp1)); cudaCheck(cudaFree(d_inp2)); From d6d0d50627f5202f20f12e9dcdbcff75a33fe65c Mon Sep 17 00:00:00 2001 From: Jane Illarionova Date: Thu, 2 May 2024 00:18:55 +0000 Subject: [PATCH 14/29] fix typo --- dev/cuda/residual_forward.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/residual_forward.cu b/dev/cuda/residual_forward.cu index 41b3e31..bbbcde2 100644 --- a/dev/cuda/residual_forward.cu +++ b/dev/cuda/residual_forward.cu @@ -102,7 +102,7 @@ void residual_forward(int kernel_num, // ---------------------------------------------------------------------------- int main(int argc, char **argv) { - set_main(); + setup_main(); int B = 8; int T = 1024; From 665d0a4f895e81d709ee5ba51f5d710715580d5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CChristopher?= <“christopher.paul.dryden@gmail.com”> Date: Thu, 2 May 2024 02:09:57 +0000 Subject: [PATCH 15/29] Added example of removing cooperative groups --- dev/cuda/layernorm_backward.cu | 117 +++++++++++++++++++++++++++++++++ train_gpt2.cu | 34 +++++----- 2 files changed, 136 insertions(+), 15 deletions(-) diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index 4e95cab..dffd1b1 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -129,6 +129,13 @@ void layernorm_backward_cpu(float* dinp, float* dweight, float* dbias, // GPU kernels // GPU helper functions for atomicAdd on smaller than 32-bit types +__device__ floatX warpReduceSum(floatX val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} + #ifdef ENABLE_BF16 __device__ void atomicAddX(__nv_bfloat16* addr, __nv_bfloat16 val) { uintptr_t ptr_val = reinterpret_cast(addr); @@ -655,6 +662,99 @@ __global__ void layernorm_backward_kernel6(Tdinp* dinp, Tparams* dweight, Tparam } } + +// Same as kernel 6 but without cooperative groups or templates +__global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, + const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, + int B, int T, int C) { + extern __shared__ float shared[]; // size = 2 * C + 1 + int warpId = threadIdx.x / warpSize; // warp index within a block + int warpsInBlock = blockDim.x / warpSize; + int base_idx = blockIdx.x * warpsInBlock + warpId; + int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp + int warps_in_grid = gridDim.x * warpsInBlock; + + // the first half of shared memory is bias, second is weight + float* dbias_shared = shared; + float* dweight_shared = shared + C; + + // init shared memory to zero + #pragma unroll 4 + for(int i = threadIdx.x; i < C; i+= blockDim.x){ + dbias_shared[i] = 0.0f; + dweight_shared[i] = 0.0f; + } + uint *tmp_flag = (uint*)(shared + C*2); + __syncthreads(); + + for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { + int b = idx / T; + int t = idx % T; + + const floatX* dout_bt = dout + b * T * C + t * C; + const floatX* inp_bt = inp + b * T * C + t * C; + floatX* dinp_bt = dinp + b * T * C + t * C; + const float mean_bt = (float)mean[b * T + t]; + const float rstd_bt = (float)rstd[b * T + t]; + + // first: two reduce operations + float dnorm_mean = 0.0f; + float dnorm_norm_mean = 0.0f; + for (int i = warpThreadIdx; i < C; i += warpSize) { + float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; + float dnorm_i = (float)weight[i] * (float)dout_bt[i]; + dnorm_mean += dnorm_i; + dnorm_norm_mean += dnorm_i * norm_bti; + } + dnorm_mean = warpReduceSum(dnorm_mean); + dnorm_norm_mean = warpReduceSum(dnorm_norm_mean); + + dnorm_mean = dnorm_mean / C; + dnorm_norm_mean = dnorm_norm_mean / C; + + // now iterate again and accumulate all the gradients + for (int i = warpThreadIdx; i < C; i += warpSize) { + float dout_i = (float)__ldcs(&dout_bt[i]); + float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; + float dnorm_i = (float)weight[i] * dout_i; + // gradient contribution to bias + atomicAdd(&dbias_shared[i], dout_i); + // gradient contribution to weight + atomicAdd(&dweight_shared[i], norm_bti * dout_i); + // gradient contribution to input + float dval = 0.0f; + dval += dnorm_i; // term 1 + dval -= dnorm_mean; // term 2 + dval -= norm_bti * dnorm_norm_mean; // term 3 + dval *= rstd_bt; // final scale + dinp_bt[i] = (floatX)((float)dinp_bt[i] + dval); + } + } + + // Accumulate into a FP32 scratchpad + // BF16 atomics are potentially much slower... and this is more precise! + __syncthreads(); + float* scratch_dbias = scratch; + float* scratch_dweight = scratch + C; + uint* scratchFlag = (uint*)(scratch + (2 * C)); + for(int i = threadIdx.x; i < C; i+= blockDim.x) { + atomicAdd(&scratch_dbias[i], dbias_shared[i]); + atomicAdd(&scratch_dweight[i], dweight_shared[i]); + } + __syncthreads(); + if (threadIdx.x == 0) { + *tmp_flag = atomicAdd(scratchFlag, 1); + } + __syncthreads(); + if (*tmp_flag == gridDim.x-1) { + for(int i = threadIdx.x; i < C; i+= blockDim.x) { + // todo - potentially do stochastic rounding here as well + dbias[i] = (floatX)scratch_dbias[i]; + dweight[i] = (floatX)scratch_dweight[i]; + } + } +} + // ---------------------------------------------------------------------------- // kernel launchers @@ -718,6 +818,20 @@ void layernorm_backward6(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* s layernorm_backward_kernel6<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } +template +void layernorm_backward7(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, + const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, + int B, int T, int C, int block_size) { + const int grid_size = (1024/block_size) * cuda_num_SMs; + size_t shared_mem_size = (2 * C + 1) * sizeof(float); + + // Including this as part of the timing until we can parallelise it + // It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams + cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float)); + + layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); +} + // kernel version dispatch void layernorm_backward(int kernel_num, floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, @@ -747,6 +861,9 @@ void layernorm_backward(int kernel_num, case 6: layernorm_backward6(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; + case 7: + layernorm_backward7(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); diff --git a/train_gpt2.cu b/train_gpt2.cu index ec0c125..828831c 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -125,6 +125,13 @@ int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per namespace cg = cooperative_groups; +__device__ floatX warpReduceSum(floatX val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} + // convenience macro for calculating grid/block dimensions for kernels #define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) @@ -743,17 +750,15 @@ __global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* dout, } } -// single FP32 scratchpad shared by all the threadblocks (based on kernels 3 & 5) -__global__ void layernorm_backward_kernel6(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, +__global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C + 1 - - namespace cg = cooperative_groups; - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); - + int warpId = threadIdx.x / warpSize; // warp index within a block + int warpsInBlock = blockDim.x / warpSize; //number of warps in block + int baseIdx = blockIdx.x * warpsInBlock + warpId; + int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp + int warpsInGrid = gridDim.x * warpsInBlock; // the first half of shared memory is bias, second is weight float* dbias_shared = shared; @@ -768,8 +773,7 @@ __global__ void layernorm_backward_kernel6(floatX* dinp, floatX* dweight, floatX uint *tmp_flag = (uint*)(shared + C*2); __syncthreads(); - int warps_in_grid = gridDim.x * warp.meta_group_size(); - for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { + for (int idx = baseIdx; idx < B * T; idx += warpsInGrid) { int b = idx / T; int t = idx % T; @@ -782,19 +786,19 @@ __global__ void layernorm_backward_kernel6(floatX* dinp, floatX* dweight, floatX // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { + for (int i = warpThreadIdx; i < C; i += warpSize) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } - dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus{}); - dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus{}); + dnorm_mean = warpReduceSum(dnorm_mean); + dnorm_norm_mean = warpReduceSum(dnorm_norm_mean); dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; // now iterate again and accumulate all the gradients - for (int i = warp.thread_rank(); i < C; i += warp.size()) { + for (int i = warpThreadIdx; i < C; i += warpSize) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; @@ -1246,7 +1250,7 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr const int grid_size = 1 * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); cudaMemset(scratch, 0, (2 * C + 1) * sizeof(float)); // todo - memset in parallel with previous kernels using streams - layernorm_backward_kernel6<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); + layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); } From 398cdaf4cdc95f18a7cbe6ee08173b1079603eca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CChristopher?= <“christopher.paul.dryden@gmail.com”> Date: Thu, 2 May 2024 02:33:51 +0000 Subject: [PATCH 16/29] Removed cooperative groups and added packed128 to fused classifier and prepare softmax blockwide --- dev/cuda/classifier_fused.cu | 130 +++++++++++++++++++++++++++++++++++ train_gpt2.cu | 83 ++++++++++++++-------- 2 files changed, 186 insertions(+), 27 deletions(-) diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index 904e8f0..f03d2ce 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -83,6 +83,25 @@ void crossentropy_softmax_backward_cpu(float* dlogits, } } +// ---------------------------------------------------- +// Kernel Utils + +// warp-level reduction for summing values +__device__ float warpReduceSum(float val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} + +// warp-level reduction for finding the maximum value +__device__ float warpReduceMax(float val) { + for (int offset = 16; offset > 0; offset /= 2) { + val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + } + return val; +} + // ---------------------------------------------------------------------------- // GPU kernels @@ -362,6 +381,105 @@ __global__ void fused_classifier_kernel3(float* dlogits, float* losses, float* p } } +__device__ SoftmaxParams prepare_softmax_blockwide2(int idx, const float* inp, int V, int P) { + // one row of inp, i.e. inp[idx, :] of shape (V,) + + const float* x = inp + idx * P; + float thread_maxval = -INFINITY; + float thread_sumval = 0.0f; + // do the loop in reverse to maximise probability of L2 cache hits + // so even small L2s get some hits on the 2nd read of the same thread + for (int i = (V+3)/4 + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { + f128 packed_x = load128cs(x + i * f128::size); // load and do not keep in cache + for(int k = 0; k < packed_x.size; ++k) { + if (i*4+k >= V) { // bounds checking against real V + continue; + } + float v = (float)packed_x[k]; + float old_maxval = thread_maxval; + thread_maxval = fmaxf(thread_maxval, v); + thread_sumval *= expf((old_maxval - thread_maxval)); + thread_sumval += expf(v - thread_maxval); + } + } + // two reductions of up to 1024 threads: + // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) + // this results in much cleaner assembly than a multi-warp cg::reduce + __shared__ float shared_maxval[32]; + __shared__ float shared_sumval[32]; + int num_warps = blockDim.x / 32; + int warp_id = threadIdx.x / 32; + int lane_id = threadIdx.x % 32; + + // reduce maxval within each warp + float warp_maxval = warpReduceMax(thread_maxval); + // thread 0 in each warp writes to shared memory + if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; } + __syncthreads(); + // each thread now loads the maxval across previous warps + // if the thread is "out of range" of data, use -FLT_MAX as the maxval + warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; + // now reduce the maxval among the warp threads + float block_maxval = warpReduceMax(warp_maxval); + // each thread uses maxval to scale sumval to avoid numerical instability / overflow + thread_sumval *= expf(thread_maxval - block_maxval); + // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory + float warp_sumval = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus{}); + + if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; } + __syncthreads(); + // same strategy, now reduce sumval across warps + warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f; + float block_sumval = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus{}); + // return the softmax parameters + return SoftmaxParams{1.f / block_sumval, block_maxval}; +} + +// same as 2 but not using float4 +__global__ void fused_classifier_kernel4(float* dlogits, float* losses, float* probs, + const float* logits, const float* dlosses, const int* targets, + int B, int T, int V, int P) { + int idx = blockIdx.x; + int ix = targets[idx]; + + // softmax (reading B * T * V, same logits read again below, hopefully still in cache) + SoftmaxParams sp = prepare_softmax_blockwide2(idx, logits, V, P); + + // calculate the probability needed for the loss and update (single-threaded) + if(threadIdx.x == 0) { + float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale; + losses[idx] = -logf(prob); + } + + // very sensible default for dlosses is 1/(B*T), which is the uniform loss + float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B*T); + // calculate the gradients directly, saves bandwidth from probs during training + // but also supports writing probs for inference-only and debugging + const float* logits_vec = logits + idx * P; + for (int i = threadIdx.x; i < (V+f128::size-1)/f128::size; i += blockDim.x) { + // this is the 2nd read of logits after the one in prepare_softmax2 + // this data will never be needed again, so we reduce cache persistence + f128 packed_logits_vec = load128cs(logits_vec + i * f128::size); // load and do not keep in cache + f128 packed_probs; + f128 packed_dlogits; + for(int k = 0; k < packed_logits_vec.size; ++k) { + int element = i*packed_logits_vec.size + k; + if (element >= V) { // bounds checking against real V + continue; + } + float v = packed_logits_vec[k]; + float prob = expf(v - sp.Offset) * sp.Scale; + packed_probs[k] = prob; + float indicator = (element == ix) ? 1.0f : 0.0f; + packed_dlogits[k] = (prob - indicator) * dloss; + } + store128(dlogits + idx * P + i * packed_logits_vec.size, packed_dlogits); + if (probs != NULL) { + store128(probs + idx * P + i * packed_logits_vec.size, packed_probs); + } + } +} + // ---------------------------------------------------------------------------- // kernel launcher @@ -395,6 +513,15 @@ void fused_classifier3(float* dlogits, float* losses, cudaCheck(cudaGetLastError()); } +void fused_classifier4(float* dlogits, float* losses, + const float* logits, const float* dlosses, const int* targets, + int B, int T, int V, int P, int block_size) { + const int N = B * T; + const int grid_size = N; + fused_classifier_kernel4<<>>(dlogits, losses, NULL, logits, dlosses, targets, B, T, V, P); + cudaCheck(cudaGetLastError()); +} + void fused_classifier(int kernel_num, float* dlogits, float* losses, const float* logits, const float* dlosses, const int* targets, int B, int T, int V, int P, int block_size) { @@ -408,6 +535,9 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses, case 3: fused_classifier3(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size); break; + case 4: + fused_classifier4(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); diff --git a/train_gpt2.cu b/train_gpt2.cu index 8412658..d2c529e 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -213,6 +213,22 @@ __device__ void atomicAddX(float* addr, float val) { atomicAdd(addr, val); } +// warp-level reduction for summing values +__device__ float warpReduceSum(float val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} + +// warp-level reduction for finding the maximum value +__device__ float warpReduceMax(float val) { + for (int offset = 16; offset > 0; offset /= 2) { + val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset)); + } + return val; +} + // ---------------------------------------------------------------------------- // Packed128 data structure, which forces the compiler to use 128-bit loads/stores // in GPUs that support (the LDG.128 and STS.128 instructions) @@ -1195,8 +1211,7 @@ struct SoftmaxParams { float Offset; }; -__device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp, - int idx, const floatX* inp, int V, int P) { +__device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, int V, int P) { // same but not float4 // one row of inp, i.e. inp[idx, :] of shape (V,) @@ -1205,14 +1220,19 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til float thread_sumval = 0.0f; // do the loop in reverse to maximise probability of L2 cache hits // so even small L2s get some hits on the 2nd read of the same thread - for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { - float v = (float)x[i]; - float old_maxval = thread_maxval; - thread_maxval = fmaxf(thread_maxval, v); - thread_sumval *= expf((old_maxval - thread_maxval)); - thread_sumval += expf(v - thread_maxval); + for (int i = (V+x128::size-1)/x128::size + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { + x128 packed_x = load128cs(x + i * x128::size); // load and do not keep in cache + for(int k = 0; k < packed_x.size; ++k) { + if (i*x128::size+k >= V) { // bounds checking against real V + continue; + } + float v = (float)packed_x[k]; + float old_maxval = thread_maxval; + thread_maxval = fmaxf(thread_maxval, v); + thread_sumval *= expf((old_maxval - thread_maxval)); + thread_sumval += expf(v - thread_maxval); + } } - // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) // this results in much cleaner assembly than a multi-warp cg::reduce @@ -1223,7 +1243,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til int lane_id = threadIdx.x % 32; // reduce maxval within each warp - float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater{}); + float warp_maxval = warpReduceMax(thread_maxval); // thread 0 in each warp writes to shared memory if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; } __syncthreads(); @@ -1231,16 +1251,17 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til // if the thread is "out of range" of data, use -FLT_MAX as the maxval warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; // now reduce the maxval among the warp threads - float block_maxval = cg::reduce(warp, warp_maxval, cg::greater{}); + float block_maxval = warpReduceMax(warp_maxval); // each thread uses maxval to scale sumval to avoid numerical instability / overflow thread_sumval *= expf(thread_maxval - block_maxval); // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory - float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus{}); + float warp_sumval = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus{}); + if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; } __syncthreads(); // same strategy, now reduce sumval across warps warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f; - float block_sumval = cg::reduce(warp, warp_sumval, cg::plus{}); + float block_sumval = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus{}); // return the softmax parameters return SoftmaxParams{1.f / block_sumval, block_maxval}; } @@ -1250,19 +1271,16 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* probs, const floatX* dlosses, const int* targets, int B, int T, int V, int P) { - namespace cg = cooperative_groups; - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); int idx = blockIdx.x; int ix = targets[idx]; // softmax (reading B * T * V, same logits read again below, hopefully still in cache) - SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P); + SoftmaxParams sp = prepare_softmax_blockwide(idx, logits, V, P); // calculate the probability needed for the loss and update (single-threaded) if(threadIdx.x == 0) { float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale; - losses[idx] = (floatX)(-logf(prob)); + losses[idx] = (floatX)(-logf(prob)); } // very sensible default for dlosses is 1/(B*T), which is the uniform loss @@ -1270,18 +1288,29 @@ __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* // calculate the gradients directly, saves bandwidth from probs during training // but also supports writing probs for inference-only and debugging const floatX* logits_vec = logits + idx * P; - // note that we use the padded dimension P to access data, but we only ever - // modify the elements up to V, ignoring the padded dimensions and leaving them at 0 - for (int i = threadIdx.x; i < V; i += blockDim.x) { + for (int i = threadIdx.x; i < (V+x128::size-1)/x128::size; i += blockDim.x) { // this is the 2nd read of logits after the one in prepare_softmax2 // this data will never be needed again, so we reduce cache persistence - float v = (float)__ldcs(&logits_vec[i]); - float prob = expf(v - sp.Offset) * sp.Scale; - if (probs != NULL) { - probs[idx * P + i] = (floatX)prob; + x128 packed_logits_vec = load128cs(logits_vec + i * x128::size); // load and do not keep in cache + x128 packed_probs; + x128 packed_logits; + for(int k = 0; k < packed_logits_vec.size; ++k) { + int element = i*packed_logits_vec.size + k; + if (element >= V) { // bounds checking against real V + continue; + } + float v = (float)packed_logits_vec[k]; + float prob = expf(v - sp.Offset) * sp.Scale; + packed_probs[k] = (floatX)prob; + float indicator = (element == ix) ? 1.0f : 0.0f; + packed_logits[k] = (floatX)((prob - indicator) * dloss); + } + if (logits != NULL){ + store128(logits + idx * P + i * packed_logits_vec.size, packed_logits); + } + if (probs != NULL) { + store128(probs + idx * P + i * packed_logits_vec.size, packed_probs); } - float indicator = (i == ix) ? 1.0f : 0.0f; - logits[idx * P + i] = (floatX)((prob - indicator) * dloss); } } From bcf7d4fe93bed7b47dd80e94f0e52465db2a374d Mon Sep 17 00:00:00 2001 From: Peter Zhizhin Date: Sun, 28 Apr 2024 00:12:31 +0000 Subject: [PATCH 17/29] Add NSight Compute ranges, use CUDA events for timings --- train_gpt2.cu | 72 ++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 63 insertions(+), 9 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8412658..6dee6d8 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -33,6 +33,8 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), -a 1 is "overfit single batch", -x 10 is 10 iterations, and -f 0 disables tf32 */ +#include + #include #include #include @@ -45,11 +47,14 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), #include // GPU / CUDA related #include +#include #include #include #include #include #include +#include + // Multi-GPU related #ifdef MULTI_GPU #include @@ -128,6 +133,18 @@ static void* cudnn_workspace = NULL; // ---------------------------------------------------------------------------- // CUDA utils +// Profiler utils +class NvtxRange { + public: + NvtxRange(const char* s) { nvtxRangePush(s); } + NvtxRange(const std::string& base_str, int number) { + std::string range_string = base_str + " " + std::to_string(number); + nvtxRangePush(range_string.c_str()); + } + ~NvtxRange() { nvtxRangePop(); } +}; +#define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__) + // cuBLAS workspace. Hardcoding to 32MiB but only Hopper needs 32, for others 4 is OK static size_t cublaslt_workspace_size = 32 * 1024 * 1024; static void* cublaslt_workspace = NULL; @@ -1297,6 +1314,7 @@ __global__ void copy_and_cast_kernel(float* dst, const floatX* src, size_t n) { void encoder_forward(floatX* out, int* inp, floatX* wte, floatX* wpe, int B, int T, int C) { + NVTX_RANGE_FN(); const int N = B * T * C; const int block_size = 256; const int grid_size = CEIL_DIV(N, block_size); @@ -1307,6 +1325,7 @@ void encoder_forward(floatX* out, void encoder_backward(floatX* dwte, floatX* dwpe, const floatX* dout, const int* inp, int B, int T, int C) { + NVTX_RANGE_FN(); const int N = B * T * C; const int block_size = 256; const int grid_size = CEIL_DIV(N, block_size); @@ -1317,6 +1336,7 @@ void encoder_backward(floatX* dwte, floatX* dwpe, void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, floatX* weight, floatX* bias, int B, int T, int C) { + NVTX_RANGE_FN(); const int block_size = 512; const int N = B * T; const int grid_size = CEIL_DIV(N * 32, block_size); @@ -1330,6 +1350,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, void matmul_forward_cublaslt(floatX* out, floatX* inp, floatX* weight, floatX* bias, int B, int T, int C, int OC) { + NVTX_RANGE_FN(); int has_bias = (bias != NULL); // check bias alignment @@ -1409,6 +1430,7 @@ void matmul_forward_cublaslt(floatX* out, void attention_forward(floatX* out, floatX* qkvr, floatX* att, floatX* inp, int B, int T, int C, int NH) { + NVTX_RANGE_FN(); // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. // Its contents will be overwritten by this function. const int block_size = 256; @@ -1480,6 +1502,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, } void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { + NVTX_RANGE_FN(); const int block_size = 256; const int grid_size = CEIL_DIV(N, block_size); residual_forward_kernel<<>>(out, inp1, inp2, N); @@ -1487,6 +1510,7 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { } void gelu_forward(floatX* out, const floatX* inp, int N) { + NVTX_RANGE_FN(); const int block_size = 512; const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_forward_kernel2<<>>(out, inp, N); @@ -1494,6 +1518,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) { + NVTX_RANGE_FN(); const int block_size = 128; const int grid_size = CEIL_DIV(N, block_size); gelu_backward_kernel<<>>(dinp, inp, dout, N); @@ -1503,6 +1528,7 @@ void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const in void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, floatX* dout, floatX* inp, floatX* weight, int B, int T, int C, int OC) { + NVTX_RANGE_FN(); float one = 1.0f; float zero = 0.0f; // backward to input, uses = in the backward pass (set the gradient) @@ -1525,6 +1551,7 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { + NVTX_RANGE_FN(); const int block_size = 1024; const int grid_size = 1 * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); @@ -1539,6 +1566,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da const floatX* dout, const floatX* qkvr, const floatX* att, int B, int T, int C, int NH) { + NVTX_RANGE_FN(); const int block_size = 256; int HS = C / NH; // head size @@ -1599,6 +1627,7 @@ template void fused_classifier3(Type* logits, Type* losses, const Type* dlosses, const int* targets, int B, int T, int V, int P) { + NVTX_RANGE_FN(); const int block_size = 1024; const int N = B * T; const int grid_size = N; @@ -1928,6 +1957,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { } void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { + NVTX_RANGE_FN(); // targets are optional and could be NULL // in this function we must be careful and use size_t instead of int, otherwise // we could overflow int. E.g. l * B * NH * T * T overflows int at B 16. @@ -1993,6 +2023,7 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] for (int l = 0; l < L; l++) { + NvtxRange layer_range("Layer", l); residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; @@ -2057,6 +2088,7 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { // also forward the cross-entropy loss function if we have the targets if (targets != NULL) { + NvtxRange classifier_and_loss_range("classifier_and_loss"); // fused classifier: does the forward pass and first part of the backward pass // we're passing dlosses = NULL, which will default them to 1.0f/(B*T), i.e. uniform loss fused_classifier3(acts.output, acts.losses, (floatX*)NULL, model->targets, B, T, V, Vp); @@ -2067,7 +2099,6 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { for (int i=0; icpu_losses[i]); } mean_loss /= B*T; model->mean_loss = mean_loss; - } else { // if we don't have targets, we don't have loss model->mean_loss = -1.0f; @@ -2075,11 +2106,13 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { } void gpt2_zero_grad(GPT2 *model) { + NVTX_RANGE_FN(); if (model->grads_acts_memory != NULL) { cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(floatX))); } if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); } } void gpt2_backward(GPT2 *model) { + NVTX_RANGE_FN(); // double check we forwarded previously, with targets if (model->mean_loss == -1.0f) { printf("Error: must forward with targets before backward\n"); @@ -2136,6 +2169,8 @@ void gpt2_backward(GPT2 *model) { // now backward all the layers for (int l = L-1; l >= 0; l--) { + NvtxRange layer_range("Layer", l); + residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; // get the pointers of the weights for this layer @@ -2223,6 +2258,7 @@ float multi_gpu_cpu_float_mean(float value, const MultiGpuConfig* multi_gpu_conf // Averages out the loss and gradients across all GPUs. No-op when multi-GPU is disabled. // todo - this version only works if all the parameters are the same size (floatX) void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { + NVTX_RANGE_FN(); // Average all losses. model->accumulated_mean_loss = multi_gpu_cpu_float_mean(model->mean_loss, multi_gpu_config); #ifdef MULTI_GPU @@ -2237,6 +2273,7 @@ void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { } void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { + NVTX_RANGE_FN(); // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html // lazily allocate the memory for m_memory and v_memory @@ -2342,6 +2379,7 @@ void dataloader_reset(DataLoader *loader) { } void dataloader_next_batch(DataLoader *loader) { + NVTX_RANGE_FN(); size_t B = loader->B; size_t T = loader->T; // if we are at the end of the file, loop back to the beginning @@ -2582,13 +2620,19 @@ int main(int argc, char *argv[]) { float* cpu_logits = (float*)mallocCheck(model.config.vocab_size * sizeof(float)); // train - struct timespec start, end; + cudaEvent_t start, end; + cudaCheck(cudaEventCreate(&start)); + cudaCheck(cudaEventCreate(&end)); + cudaCheck(cudaProfilerStart()); double total_sum_iteration_time_s = 0.0; for (int step = 0; step <= train_num_batches; step++) { + NvtxRange step_range("Train step", step); + int last_step = step == train_num_batches; // once in a while estimate the validation loss if (step % val_loss_every == 0 || last_step) { + NvtxRange validation_range("validation"); float val_loss = 0.0f; dataloader_reset(&val_loader); for (int i = 0; i < val_num_batches; i++) { @@ -2604,6 +2648,7 @@ int main(int argc, char *argv[]) { // once in a while do model inference to print generated text if (multi_gpu_config.process_rank == 0 && (step > 0 && (step % sample_every) == 0 || last_step)) { + NvtxRange generation_range("generation"); // fill up gen_tokens with the <|endoftext|> token, which kicks off the generation int eot_token = tokenizer.eot_token; for(int i = 0; i < B * T; ++i) { @@ -2612,6 +2657,7 @@ int main(int argc, char *argv[]) { // now sample from the model autoregressively printf("generating:\n---\n"); for (int t = 1; t < genT; t++) { + NvtxRange generation_range("Generation step", t); // note that inference is very wasteful here because for each token // we re-calculate the forward pass for all of (B,T) positions from scratch // but the inference here is just for sanity checking anyway @@ -2652,11 +2698,12 @@ int main(int argc, char *argv[]) { if (last_step) { break; } // do a training step - clock_gettime(CLOCK_MONOTONIC, &start); + cudaEventRecord(start); if (overfit_single_batch == 0 || (step == 0 && overfit_single_batch == 1)) { // if we're overfitting a single batch, we'll only call this at step = 0 dataloader_next_batch(&train_loader); } + dataloader_next_batch(&train_loader); gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); gpt2_zero_grad(&model); gpt2_backward(&model); @@ -2664,22 +2711,29 @@ int main(int argc, char *argv[]) { gpt2_multi_gpu_accumulate(&model, &multi_gpu_config); } gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); - cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings - clock_gettime(CLOCK_MONOTONIC, &end); - double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; + + cudaEventRecord(end); + float time_elapsed_ms; + cudaCheck(cudaEventSynchronize(end)); // wait for the end event to finish to get correct timings + cudaCheck(cudaEventElapsedTime(&time_elapsed_ms, start, end)); if (step > 0) { // consider the first batch to be a warmup (e.g. cuBLAS/cuDNN initialisation) - total_sum_iteration_time_s += time_elapsed_s; + total_sum_iteration_time_s += time_elapsed_ms / 1000.0; } - int tokens_per_second = multi_gpu_config.num_processes * (B * T) / time_elapsed_s; + int tokens_per_second = multi_gpu_config.num_processes * (B * T) / time_elapsed_ms * 1000.0; float accumulated_loss = multi_gpu_config.num_processes == 1 ? model.mean_loss : model.accumulated_mean_loss; - printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, accumulated_loss, time_elapsed_s * 1000, tokens_per_second); + printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, accumulated_loss, time_elapsed_ms, tokens_per_second); logger_log_train(&logger, step, model.mean_loss); + } // add a total average, for optimizations that are only mild improvements (excluding 1st batch as warmup) printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / (train_num_batches-1) * 1000); + cudaProfilerStop(); + // free and destroy everything + cudaCheck(cudaEventDestroy(end)); + cudaCheck(cudaEventDestroy(start)); dataloader_free(&train_loader); dataloader_free(&val_loader); tokenizer_free(&tokenizer); From 2202c9a51bd1f9f4ed4dd495688207950226169f Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 18:09:27 +0000 Subject: [PATCH 18/29] add kernel 4 to docs. have to improve these docs more and document them better --- dev/cuda/classifier_fused.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index f03d2ce..df68941 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -10,6 +10,7 @@ nvcc -O3 --use_fast_math classifier_fused.cu -o classifier_fused ./classifier_fused 1 ./classifier_fused 2 ./classifier_fused 3 +./classifier_fused 4 */ #include @@ -448,7 +449,7 @@ __global__ void fused_classifier_kernel4(float* dlogits, float* losses, float* p // calculate the probability needed for the loss and update (single-threaded) if(threadIdx.x == 0) { float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale; - losses[idx] = -logf(prob); + losses[idx] = -logf(prob); } // very sensible default for dlosses is 1/(B*T), which is the uniform loss From d1771a7b5923b2f5b9475f24585aa410c36138be Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 19:01:26 +0000 Subject: [PATCH 19/29] quick fix master. there is some weirdness here in the adam update that is todo to understand better, for now i just want the master to be ok --- train_gpt2.cu | 56 +++++++++++++++++++++++++-------------------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index d2c529e..e378699 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1175,35 +1175,35 @@ __device__ inline float lerp(float start, float end, float weight) { // Termplate type T instead of floatx template -__global__ void adamw_kernel3(Tp* params_memory, float* master_params, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, +__global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay, unsigned int seed) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= num_parameters) return; // guard - float grad = (float)grads_memory[i]; - float m = m_memory[i]; - float v = v_memory[i]; - // update the first moment (momentum) - m = lerp(grad, m, beta1); - m_memory[i] = m; - // update the second moment (RMSprop) - v = lerp(grad * grad, v, beta2); - v_memory[i] = v; - m /= beta1_correction; // m_hat - v /= beta2_correction; // v_hat - // update the parameters (weight/bias) - float old_param = master_params != NULL ? master_params[i] : (float)params_memory[i]; - float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param)); - // if we have master parameters, directly update the two weight copies - if (master_params != NULL) { - params_memory[i] = (floatX)param; // low-precision copy, for use in the forward pass - master_params[i] = param; // float copy, for use in the next parameter update - } else { - // without a master copy of params in float, do a direct update in low precision - // and use stochastic rounding to mitigate loss of training stability - unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); - stochastic_rounding(param, ¶ms_memory[i], random); - } + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_parameters) return; // guard + // get the gradient, m, and v for this parameter + float grad = (float)grads_memory[i]; + float m = m_memory[i]; + float v = v_memory[i]; + // update the first moment (momentum) + m = lerp(grad, m, beta1); + m_memory[i] = m; + // update the second moment (RMSprop) + v = lerp(grad * grad, v, beta2); + v_memory[i] = v; + m /= beta1_correction; // m_hat + v /= beta2_correction; // v_hat + // fetch the old value of this parameter as a float, from either source + float old_param = (master_params_memory != NULL) ? master_params_memory[i] : (float)params_memory[i]; + // update this parameter + float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param)); + // update our low precision version of the parameters using stochastic rounding + // this will be used in the next forward pass + // TODO: simply doing `params_memory[i] = (floatX)param;` breaks everything (why?) + unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); + stochastic_rounding(param, ¶ms_memory[i], random); + // write the full, float version of the param into our master copy, if we maintain one + // this will be used in the next update + if (master_params_memory != NULL) { master_params_memory[i] = param; } } struct SoftmaxParams { @@ -1280,7 +1280,7 @@ __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* // calculate the probability needed for the loss and update (single-threaded) if(threadIdx.x == 0) { float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale; - losses[idx] = (floatX)(-logf(prob)); + losses[idx] = (floatX)(-logf(prob)); } // very sensible default for dlosses is 1/(B*T), which is the uniform loss From 8ccd05ec2953383f99c1aa8874c058ed07c884c7 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 19:09:53 +0000 Subject: [PATCH 20/29] group warp reduce ops together --- train_gpt2.cu | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 341aa2f..d322e7c 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -140,13 +140,6 @@ int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per namespace cg = cooperative_groups; -__device__ floatX warpReduceSum(floatX val) { - for (int offset = 16; offset > 0; offset /= 2) { - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - } - return val; -} - // convenience macro for calculating grid/block dimensions for kernels #define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) @@ -220,6 +213,13 @@ __device__ void atomicAddX(float* addr, float val) { atomicAdd(addr, val); } +__device__ floatX warpReduceSum(floatX val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} + // warp-level reduction for summing values __device__ float warpReduceSum(float val) { for (int offset = 16; offset > 0; offset /= 2) { From 01df41fbba703f951269dec3aefade5633f53c8f Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 19:24:04 +0000 Subject: [PATCH 21/29] careful to guard the overloaded warpReduceSum and small fixes --- train_gpt2.cu | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index f293b2b..00b7669 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -213,13 +213,6 @@ __device__ void atomicAddX(float* addr, float val) { atomicAdd(addr, val); } -__device__ floatX warpReduceSum(floatX val) { - for (int offset = 16; offset > 0; offset /= 2) { - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - } - return val; -} - // warp-level reduction for summing values __device__ float warpReduceSum(float val) { for (int offset = 16; offset > 0; offset /= 2) { @@ -236,6 +229,15 @@ __device__ float warpReduceMax(float val) { return val; } +#if defined(ENABLE_BF16) || defined(ENABLE_FP16) +__device__ floatX warpReduceSum(floatX val) { + for (int offset = 16; offset > 0; offset /= 2) { + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + } + return val; +} +#endif + // ---------------------------------------------------------------------------- // Packed128 data structure, which forces the compiler to use 128-bit loads/stores // in GPUs that support (the LDG.128 and STS.128 instructions) From 8a9510f32b415501563f978e064080e9abcb4296 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 19:48:56 +0000 Subject: [PATCH 22/29] very minor stylistic preference let's keep the code style consistent --- train_gpt2.cu | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index ac94698..aecdde1 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -966,8 +966,7 @@ __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2, x128 packed_inp1 = load128cs(inp1 + idx); x128 packed_inp2 = load128cs(inp2 + idx); #pragma unroll packed_inp1.size - for (int k = 0; k < packed_inp1.size; ++k) - { + for (int k = 0; k < packed_inp1.size; k++) { packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); } store128(out + idx, packed_out); From 41a0789a44a0e1ac3ad2d18327b4952a439e52b6 Mon Sep 17 00:00:00 2001 From: Peter Zhizhin Date: Thu, 2 May 2024 20:30:08 +0000 Subject: [PATCH 23/29] Added NVTX ranges for FlashAttention --- train_gpt2.cu | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 6dee6d8..333055c 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -639,6 +639,7 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) float* stats, // output for backward pass: (B, NH, T) floatX* inp, // input: (B, T, 3, NH, HS) QKV int B, int T, int NH, int C) { + NVTX_RANGE_FN(); int HS = C / NH; // number of features per head bool is_inference_only = (stats == nullptr); @@ -680,6 +681,7 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) void attention_backward_cudnn(floatX* dqkvr, // output floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs int B, int T, int NH, int C) { + NVTX_RANGE_FN(); int HS = C / NH; // number of features per head // Get graph and tensors from cache (or generate it on first use) @@ -2725,11 +2727,14 @@ int main(int argc, char *argv[]) { printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, accumulated_loss, time_elapsed_ms, tokens_per_second); logger_log_train(&logger, step, model.mean_loss); + if (step == 10) { + cudaProfilerStop(); + } + } // add a total average, for optimizations that are only mild improvements (excluding 1st batch as warmup) printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / (train_num_batches-1) * 1000); - cudaProfilerStop(); // free and destroy everything cudaCheck(cudaEventDestroy(end)); From 9da2729c5a0813f00b30de55168adb224da75075 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 20:40:05 +0000 Subject: [PATCH 24/29] update encoder_forward with the latest and greatest Packed128 and incorporate into train_gpt2cu --- dev/cuda/encoder_forward.cu | 121 ++++++++++++++++++++++-------------- train_gpt2.cu | 29 +++++---- 2 files changed, 91 insertions(+), 59 deletions(-) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index f56a9c8..b978b35 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -20,6 +20,22 @@ version 3 is like version 2 but uses float4 reads/writes #include "common.h" #include +// 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 @@ -44,8 +60,8 @@ void encoder_forward_cpu(float* out, // GPU kernels // naive implementation into kernel, parallelize over B,T, loop over C -__global__ void encoder_forward_kernel1(float* out, - const int* inp, const float* wte, const float* wpe, +__global__ void encoder_forward_kernel1(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int N = B * T; @@ -53,19 +69,19 @@ __global__ void encoder_forward_kernel1(float* out, if (idx < N) { int b = idx / T; int t = idx % T; - float* out_bt = out + b * T * C + t * C; + floatX* out_bt = out + b * T * C + t * C; int ix = inp[b * T + t]; - const float* wte_ix = wte + ix * C; - const float* wpe_t = wpe + t * C; + const floatX* wte_ix = wte + ix * C; + const floatX* wpe_t = wpe + t * C; for (int i = 0; i < C; i++) { - out_bt[i] = wte_ix[i] + wpe_t[i]; + out_bt[i] = (floatX)((float)wte_ix[i] + (float)wpe_t[i]); } } } // optimized implementation: parallelize over all of B,T,C -__global__ void encoder_forward_kernel2(float* out, - const int* inp, const float* wte, const float* wpe, +__global__ void encoder_forward_kernel2(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int N = B * T * C; @@ -78,40 +94,46 @@ __global__ void encoder_forward_kernel2(float* out, int ix = inp[b * T + t]; - float* out_btc = out + b * T * C + t * C + c; - const float* wte_ix = wte + ix * C + c; - const float* wpe_tc = wpe + t * C + c; - *out_btc = *wte_ix + *wpe_tc; + floatX* out_btc = out + b * T * C + t * C + c; + const floatX* wte_ix = wte + ix * C + c; + const floatX* wpe_tc = wpe + t * C + c; + *out_btc = (floatX)((float)*wte_ix + (float)*wpe_tc); } } -__device__ inline float4 add_float4(const float4& a, const float4& b) { - return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); -} - -// use of float4 leads to using 128-bit LDG / STG instructions in SASS, -// very helpful in memory-bound kernels like encoder_forward -__global__ void encoder_forward_kernel3(float4* out, - const int* inp, const float4* wte, const float4* wpe, +__global__ void encoder_forward_kernel3(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C) { - int C4 = C / 4; - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int N = B * T * C4; + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + int N = B * T * C; if (idx < N) { - int bt = idx / C4; + int bt = idx / C; int b = bt / T; int t = bt % T; - int c4 = idx % C4; + int c = idx % C; + int ix = inp[b * T + t]; - out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]); + + floatX* out_btc = out + b * T * C + t * C + c; + const floatX* wte_ix = wte + ix * C + c; + const floatX* wpe_tc = wpe + t * C + c; + + x128 packed_out; + x128 wte = load128cs(wte_ix); + x128 wpe = load128cs(wpe_tc); + #pragma unroll wte.size + for (int k = 0; k < wte.size; k++) { + packed_out[k] = (floatX)((float)wte[k] + (float)wpe[k]); + } + store128(out_btc, packed_out); } } // ---------------------------------------------------------------------------- // kernel launcher -void encoder_forward1(float* out, - const int* inp, const float* wte, const float* wpe, +void encoder_forward1(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C, const int block_size) { const int N = B * T; @@ -120,8 +142,8 @@ void encoder_forward1(float* out, cudaCheck(cudaGetLastError()); } -void encoder_forward2(float* out, - const int* inp, const float* wte, const float* wpe, +void encoder_forward2(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C, const int block_size) { const int N = B * T * C; @@ -130,21 +152,20 @@ void encoder_forward2(float* out, cudaCheck(cudaGetLastError()); } -void encoder_forward3(float* out, - const int* inp, const float* wte, const float* wpe, +void encoder_forward3(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C, const int block_size) { - assert(C % 4 == 0); const int N = B * T * C; - const int grid_size = ceil_div(N / 4, block_size); - encoder_forward_kernel3<<>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C); + const int grid_size = ceil_div(N, (int)(block_size * x128::size)); + encoder_forward_kernel3<<>>(out, inp, wte, wpe, B, T, C); cudaCheck(cudaGetLastError()); } // kernel version dispatch void encoder_forward(int kernel_num, - float* out, - const int* inp, const float* wte, const float* wpe, + floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C, const int block_size) { switch (kernel_num) { @@ -166,7 +187,7 @@ void encoder_forward(int kernel_num, // ---------------------------------------------------------------------------- int main(int argc, char **argv) { - srand(0); + setup_main(); int B = 8; int T = 1024; @@ -183,17 +204,17 @@ int main(int argc, char **argv) { float* wpe = make_random_float(T * C); // move to GPU - float* d_out; + floatX* d_out; int* d_inp; - float* d_wte; - float* d_wpe; - cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float))); + floatX* d_wte; + floatX* d_wpe; + cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_inp, B * T * sizeof(int))); - cudaCheck(cudaMalloc(&d_wte, V * C * sizeof(float))); - cudaCheck(cudaMalloc(&d_wpe, T * C * sizeof(float))); + cudaCheck(cudaMalloc(&d_wte, V * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_wpe, T * C * sizeof(floatX))); cudaCheck(cudaMemcpy(d_inp, inp, B * T * sizeof(int), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_wte, wte, V * C * sizeof(float), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_wpe, wpe, T * C * sizeof(float), cudaMemcpyHostToDevice)); + cudaCheck(memcpy_convert(d_wte, wte, V * C)); + cudaCheck(memcpy_convert(d_wpe, wpe, T * C)); // read kernel_num from command line int kernel_num = 2; @@ -205,7 +226,6 @@ int main(int argc, char **argv) { // first check the correctness of the kernel encoder_forward_cpu(out, inp, wte, wpe, B, T, C); - // time the kernel at different block sizes int block_sizes[] = {32, 64, 128, 256, 512, 1024}; @@ -213,7 +233,12 @@ int main(int argc, char **argv) { int block_size = block_sizes[j]; printf("Checking block size %d.\n", block_size); encoder_forward(kernel_num, d_out, d_inp, d_wte, d_wpe, B, T, C, block_size); - validate_result(d_out, out, "out", B * T * C, 1e-5f); +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) + float tol = 1e-5; +#else + float tol = 1e-2f; +#endif + validate_result(d_out, out, "out", B * T * C, tol); } printf("All results match. Starting benchmarks.\n\n"); diff --git a/train_gpt2.cu b/train_gpt2.cu index aecdde1..263ba47 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -732,12 +732,11 @@ void attention_backward_cudnn(floatX* dqkvr, // ---------------------------------------------------------------------------- // all the kernels -__global__ void encoder_forward_kernel2(floatX* out, - int* inp, floatX* wte, floatX* wpe, +__global__ void encoder_forward_kernel3(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; int N = B * T * C; - if (idx < N) { int bt = idx / C; int b = bt / T; @@ -747,9 +746,17 @@ __global__ void encoder_forward_kernel2(floatX* out, int ix = inp[b * T + t]; floatX* out_btc = out + b * T * C + t * C + c; - floatX* wte_ix = wte + ix * C + c; - floatX* wpe_tc = wpe + t * C + c; - *out_btc = (floatX)((float)*wte_ix + (float)*wpe_tc); + const floatX* wte_ix = wte + ix * C + c; + const floatX* wpe_tc = wpe + t * C + c; + + x128 packed_out; + x128 wte = load128cs(wte_ix); + x128 wpe = load128cs(wpe_tc); + #pragma unroll wte.size + for (int k = 0; k < wte.size; k++) { + packed_out[k] = (floatX)((float)wte[k] + (float)wpe[k]); + } + store128(out_btc, packed_out); } } @@ -1344,12 +1351,12 @@ __global__ void copy_and_cast_kernel(float* dst, const floatX* src, size_t n) { // kernel launchers void encoder_forward(floatX* out, - int* inp, floatX* wte, floatX* wpe, + const int* inp, const floatX* wte, const floatX* wpe, int B, int T, int C) { - const int N = B * T * C; const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); - encoder_forward_kernel2<<>>(out, inp, wte, wpe, B, T, C); + const int N = B * T * C; + const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); + encoder_forward_kernel3<<>>(out, inp, wte, wpe, B, T, C); cudaCheck(cudaGetLastError()); } From 50714d22548ffa460b1ec8fc06d8e47e968d3769 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 21:06:11 +0000 Subject: [PATCH 25/29] pragma unroll quick fix --- dev/cuda/encoder_forward.cu | 2 +- train_gpt2.cu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dev/cuda/encoder_forward.cu b/dev/cuda/encoder_forward.cu index b978b35..16df62f 100644 --- a/dev/cuda/encoder_forward.cu +++ b/dev/cuda/encoder_forward.cu @@ -121,7 +121,7 @@ __global__ void encoder_forward_kernel3(floatX* out, x128 packed_out; x128 wte = load128cs(wte_ix); x128 wpe = load128cs(wpe_tc); - #pragma unroll wte.size + #pragma unroll for (int k = 0; k < wte.size; k++) { packed_out[k] = (floatX)((float)wte[k] + (float)wpe[k]); } diff --git a/train_gpt2.cu b/train_gpt2.cu index 263ba47..f7be4fa 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -752,7 +752,7 @@ __global__ void encoder_forward_kernel3(floatX* out, x128 packed_out; x128 wte = load128cs(wte_ix); x128 wpe = load128cs(wpe_tc); - #pragma unroll wte.size + #pragma unroll for (int k = 0; k < wte.size; k++) { packed_out[k] = (floatX)((float)wte[k] + (float)wpe[k]); } @@ -972,7 +972,7 @@ __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2, x128 packed_out; x128 packed_inp1 = load128cs(inp1 + idx); x128 packed_inp2 = load128cs(inp2 + idx); - #pragma unroll packed_inp1.size + #pragma unroll for (int k = 0; k < packed_inp1.size; k++) { packed_out[k] = (floatX)((float)packed_inp1[k] + (float)packed_inp2[k]); } From 6e48501d60f3e33dedcb0685347e3fb99fa76ba0 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 01:05:47 +0300 Subject: [PATCH 26/29] bias backward kernel that will use all available threads --- dev/cuda/matmul_backward_bias.cu | 35 ++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 7feab39..4ff2945 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -166,6 +166,19 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in } } +__global__ void matmul_backward_bias_kernel5(float* dbias, const float* dout, int B, int T, int OC) { + int oc = blockIdx.x * blockDim.x + threadIdx.x; + if(oc >= OC) return; + float sum = 0.0; + // grid-wide loop for maximum parallelism + for (int i = blockIdx.y; i < B * T; i += gridDim.y) { + sum += dout[i * OC + oc]; + } + // and atomcially add everything together. atomics within one block are conflict-free! + atomicAdd(dbias + oc, sum); +} + + // ---------------------------------------------------------------------------- // kernel launcher @@ -202,6 +215,14 @@ void matmul_backward_bias4(float* dinp, float* dweight, float* dbias, matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); } +void matmul_backward_bias5(float* dinp, float* dweight, float* dbias, + float* dout, float* inp, float* weight, float* ones, + int B, int T, int C, int OC, int block_size) { + const int grid_size_x = ceil_div(OC, block_size); + const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / block_size); + matmul_backward_bias_kernel5<<>>(dbias, dout, B, T, OC); +} + void matmul_backward_bias(int kernel_num, float* dinp, float* dweight, float* dbias, float* dout, float* inp, float* weight, float* ones, @@ -219,6 +240,9 @@ void matmul_backward_bias(int kernel_num, case 4: matmul_backward_bias4(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); break; + case 5: + matmul_backward_bias5(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); @@ -228,20 +252,13 @@ void matmul_backward_bias(int kernel_num, // ---------------------------------------------------------------------------- int main(int argc, char **argv) { - srand(0); + setup_main(); int B = 8; int T = 1024; int C = 768; int OC = 768 * 4; // expansion of 4, e.g. in the MLP - // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); - printf("Device %d: %s\n", deviceIdx, deviceProp.name); - // read kernel_num from command line int kernel_num = 1; if (argc > 1) { @@ -280,7 +297,7 @@ int main(int argc, char **argv) { // memset the bias to zero cudaCheck(cudaMemset(d_dbias, 0, OC * sizeof(float))); // calculate the GPU version - matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, 128); + matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, block_size); // compare printf("Checking correctness...\n"); validate_result(d_dbias, dbias, "dbias", OC, 5e-3f); From 937fbd8e365406924b55da5f837e2208fc83c723 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 22:17:22 +0000 Subject: [PATCH 27/29] small doc fixes --- dev/cuda/matmul_backward_bias.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 4ff2945..7aef545 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -8,6 +8,7 @@ nvcc -O3 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias ./matmul_backward_bias 2 ./matmul_backward_bias 3 ./matmul_backward_bias 4 +./matmul_backward_bias 5 ncu: sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1 @@ -174,7 +175,7 @@ __global__ void matmul_backward_bias_kernel5(float* dbias, const float* dout, in for (int i = blockIdx.y; i < B * T; i += gridDim.y) { sum += dout[i * OC + oc]; } - // and atomcially add everything together. atomics within one block are conflict-free! + // and atomically add everything together. atomics within one block are conflict-free! atomicAdd(dbias + oc, sum); } From 6ebef46f832b4e55b46237c4d06c2597050819ae Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 23:23:28 +0000 Subject: [PATCH 28/29] ugh didn't notice this tiny rebasing mistake, introduced a bug. good candidate for a CI that we can overfit a single batch in the train_gpt2.cu script and get the exact same numbers as we expect in the test_gpt2.cu file --- train_gpt2.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index a4a6566..c52a994 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2761,7 +2761,6 @@ int main(int argc, char *argv[]) { // if we're overfitting a single batch, we'll only call this at step = 0 dataloader_next_batch(&train_loader); } - dataloader_next_batch(&train_loader); gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); gpt2_zero_grad(&model); gpt2_backward(&model); From 776541dac7774ed3dcc07451853967ee3cc0c769 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 2 May 2024 23:32:20 +0000 Subject: [PATCH 29/29] v1 of the new matmul backward bias kernel --- train_gpt2.cu | 84 ++++++++++++++++++++++----------------------------- 1 file changed, 36 insertions(+), 48 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index c52a994..084a7b8 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -27,8 +27,8 @@ mpirun -np 4 ./train_gpt2cu -b 8 -v 200 -s 200 -i data/TinyStories If you'd like to see train_gpt2.cu produce identical results to `python train_gpt2.py`, you can run it like this: -make train_gpt2cu PRECISION=FP32 -./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -f 0 +make train_gpt2cu && ./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -f 0 +make train_gpt2cu PRECISION=FP32 && ./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -f 0 This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), -a 1 is "overfit single batch", -x 10 is 10 iterations, and -f 0 disables tf32 */ @@ -154,6 +154,7 @@ cublasLtHandle_t cublaslt_handle; int cuda_arch_major = 0; int cuda_arch_minor = 0; int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per SM +int cuda_threads_per_SM = 0; namespace cg = cooperative_groups; @@ -1037,44 +1038,19 @@ __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floa } } -// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to: -// dbias = dout.sum((0,1)) -// the idea is to employ one block to reduce along several columns, -// where each block has a width of 32 columns to ensure coalesced access. -// at the end we accumulate the reductions performed by the warps in each block via shared memory -__global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* dout, int B, int T, int OC) { - // this kernel is launched with 1D grid_dim of OC/32 - // for example let's say block_size is 128 - extern __shared__ float smem[]; // of size block_size (128) - const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3 - const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31 - const int tl = blockIdx.x * warpSize; // pointer to the start column for this block - const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4 - - // pointer to the start of the column for one lane of threads - // so e.g. 4 threads (of the same lane_id) will reduce this one column - const floatX* dout_col = dout + tl + lane_id; - - // column reductions by looping through the rows - // each of the 4 threads offsets by its warp_id and then skips by vstep - // together these 4 threads cover all B*T rows of this (lane_id) column - // importantly, consecutive threads (in threadId) are processing adjacent columns, - // leading to a coalesced memory access pattern - float dout_sum = 0.0f; - for (int row = warp_id; row < B * T; row += vstep) { - dout_sum += (float)dout_col[row * OC]; - } - smem[lane_id + warp_id * warpSize] = dout_sum; - __syncthreads(); - - // warp_id 0 reduces the shared memory column-wise, linearly - dout_sum = 0.0f; - if (warp_id == 0) { - for (int j = 0; j < vstep; j++) { - dout_sum += smem[lane_id + j * warpSize]; - } - dbias[tl + lane_id] = (floatX)dout_sum; +__global__ void matmul_backward_bias_kernel5(float* dbias, const floatX* dout, int B, int T, int OC) { + // note: this kernel reads in floatX, but it writes to float! + // this is because we're using atomics, which are super slow in < fp32 precision on < H100 GPUs + // so the trick is do fp32 atomics to a buffer, and then copy_and_cast the result to floatX + int oc = blockIdx.x * blockDim.x + threadIdx.x; + if(oc >= OC) return; + float sum = 0.0; + // grid-wide loop for maximum parallelism + for (int i = blockIdx.y; i < B * T; i += gridDim.y) { + sum += (float)dout[i * OC + oc]; } + // and atomically add everything together. atomics within one block are conflict-free! + atomicAdd(dbias + oc, sum); } __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, @@ -1366,6 +1342,12 @@ __global__ void copy_and_cast_kernel(float* dst, const floatX* src, size_t n) { if (i < n) { dst[i] = (float)src[i]; } } +__global__ void cast_and_add_kernel(floatX* dst, const float* src, size_t n) { + // used only for matmul_backward_bias kernel, a little bit embarassing TODO delete later + const size_t i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { dst[i] += (floatX)src[i]; } // have to += because dbias is a paramater +} + // ---------------------------------------------------------------------------- // kernel launchers @@ -1585,6 +1567,7 @@ void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const in void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, floatX* dout, floatX* inp, floatX* weight, + float* dbias_buffer, int B, int T, int C, int OC) { NVTX_RANGE_FN(); float one = 1.0f; @@ -1599,9 +1582,13 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, dweight, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); // backward to bias, if given, does a += if (dbias != NULL) { - const int block_size = 1024; - const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work - matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); + const int block_size = 128; + const int grid_size_x = CEIL_DIV(OC, block_size); + const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / block_size); + cudaMemset(dbias_buffer, 0, OC * sizeof(float)); + matmul_backward_bias_kernel5<<>>(dbias_buffer, dout, B, T, OC); + cudaCheck(cudaGetLastError()); + cast_and_add_kernel<<>>(dbias, dbias_buffer, OC); cudaCheck(cudaGetLastError()); } } @@ -2212,14 +2199,14 @@ void gpt2_backward(GPT2 *model) { GradActTensors grads_acts = model->grads_acts; // re-use the output buffer of the forward pass as a scratchpad during backward pass - float* scratchF = (float*)acts.output; + float* scratchF = (float*)acts.output; // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) // this was done in the fused classifier kernel as last step of forward pass // technically that is a small, inline backward() pass of calculating // total, final loss as the mean over all losses over all (B,T) positions in the batch // next: backward the classifier matmul - matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, B, T, C, Vp); + matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp); // backward the final layernorm floatX* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 floatX* dresidual = (floatX*)grads_acts.residual3; // the main buffer holding the gradient in the backward pass @@ -2273,12 +2260,12 @@ void gpt2_backward(GPT2 *model) { floatX* dl_bt4c = (floatX*)grads_acts.bt4c; // backprop this layer - matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, B, T, 4*C, C); + matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C); gelu_backward(dl_bt4c, l_fch, dl_bt4c, B*T*4*C); - matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, B, T, C, 4 * C); + matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, scratchF, B, T, C, 4 * C); // layernorm backward does += to the dresidual, so it correctly accumulates grad from the MLP block above layernorm_backward(dresidual, dl_ln2w, dl_ln2b, scratchF, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); - matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, B, T, C, C); + matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, scratchF, B, T, C, C); #ifdef ENABLE_CUDNN float* l_att = (float*)acts.att + l * B * NH * T; // cuDNN needs a smaller FP32 tensor @@ -2294,7 +2281,7 @@ void gpt2_backward(GPT2 *model) { #endif // QKV parameter gradients - matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, B, T, C, 3 * C); + matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, scratchF, B, T, C, 3 * C); // layernorm backward does += to dresidual, so it correctly accumulates gradient for the Attention block above layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); } @@ -2596,6 +2583,7 @@ int main(int argc, char *argv[]) { cudaDeviceProp deviceProp; cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx); cuda_num_SMs = deviceProp.multiProcessorCount; + cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; cuda_arch_major = deviceProp.major; cuda_arch_minor = deviceProp.minor;