diff --git a/dev/cuda/Makefile b/dev/cuda/Makefile index c741788..14eae20 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_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_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 fused_residual_forward global_norm all: $(TARGETS) # Individual targets: forward pass @@ -48,6 +48,7 @@ matmul_backward: matmul_backward.cu # Update kernels adamw: adamw.cu +global_norm: global_norm.cu # NCCL communication kernels nccl_all_reduce: nccl_all_reduce.cu diff --git a/dev/cuda/global_norm.cu b/dev/cuda/global_norm.cu new file mode 100644 index 0000000..434c343 --- /dev/null +++ b/dev/cuda/global_norm.cu @@ -0,0 +1,199 @@ +/* +Kernels for a global norm. +Global norm in this context means that we want to calculate a single norm cooperatively using all avalailable SMs, instead + of multiple norms that can be handled by separate blocks. + +Compile example: +nvcc -O3 --use_fast_math global_norm.cu -o global_norm + +version 1 uses as few blocks as possible to still fill the GPU, and only does atomic adds in the end +./gelu_forward 1 + +version 2 is the same but with only warp-wide reduction inside the kernel, and more global atomics +./gelu_forward 2 +*/ + +#include "common.h" +#include +#include +#include + +// TODO move this into 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; + +float global_norm_cpu(const float* data, size_t count) { + // accumulate in double so we have an accurate numerical reference + double acc = 0.0; + for(size_t i = 0; i < count; ++i) { + acc += (double)data[i] * (double)data[i]; + } + return (float)acc; +} + + +template +__global__ void norm_kernel1(float* out, const T* data, size_t count) { + // we want as few atomics as possible, so each block tries to do + // the maximum amount of work (so no fixed chunk, but instead iterating + // until we run out of data), and then we reduce inside the block + // and finally have just one atomic per block. + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + + __shared__ float block_result[32]; + + // out will be updated atomically from all thread blocks + size_t index = threadIdx.x + blockDim.x * blockIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + // warp-level reduce + float warp_result = cg::reduce(warp, accumulator, cg::plus{}); + block_result[warp.meta_group_rank()] = warp_result; + block.sync(); + if(warp.meta_group_rank() == 0) { + float gather = warp.thread_rank() < warp.meta_group_size() ? block_result[warp.thread_rank()] : 0.f; + float block_sum = cg::reduce(warp, gather, cg::plus{}); + if(warp.thread_rank() == 0) { + atomicAdd(out, block_sum); + } + } +} + + + +template +__global__ void norm_kernel2(float* out, const T* data, size_t count) { + // no shared memory; but one atomic per warp instead of per block + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + + // out will be updated atomically from all thread blocks + size_t index = threadIdx.x + blockDim.x * blockIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + + // warp-level reduce + float warp_result = cg::reduce(warp, accumulator, cg::plus{}); + // and atomic in global buffer + if(warp.thread_rank() == 0) { + atomicAdd(out, warp_result); + } +} + + + +template +void global_norm1(float* out, const T* values, size_t count, int block_size) { + // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. + // having one block less than possible is a tiny performance hit, having + // one block too many is catastrophic, since it only can start once all the other + // blocks finish. anyway, I think cuda_threads_per_SM should be a multiple of 512 + // on all gpus, so the division really is going to be exact. + const int grid_size = cuda_threads_per_SM * cuda_num_SMs / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel1<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +template +void global_norm2(float* out, const T* values, size_t count, int block_size) { + // ditto + const int grid_size = cuda_threads_per_SM * cuda_num_SMs / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel2<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size) { + switch (kernel_num) { + case 1: + return global_norm1(out, values, count, block_size); + case 2: + return global_norm2(out, values, count, block_size); + } +} + +int main(int argc, const char **argv) { + setup_main(); + + int C = 768; + int L = 12; + + size_t num_params = (size_t)(C * 4*C + C*C) * 2 * L; + + // create host memory of random numbers + float* inp = make_random_float(num_params); + // scale them down + for(size_t i = 0; i < num_params; ++i) { + inp[i] *= 1e-3; + } + + // 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 + float out = global_norm_cpu(inp, num_params); + + // move to GPU + float* d_out; + floatX* d_inp; + cudaCheck(cudaMalloc(&d_out, sizeof(float))); + cudaCheck(cudaMalloc(&d_inp, num_params * sizeof(floatX))); + cudaCheck(memcpy_convert(d_inp, inp, num_params)); + + 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); + cudaCheck(cudaMemset(d_out, 0, sizeof(float))); + global_norm(kernel_num, d_out, d_inp, num_params, block_size); + validate_result(d_out, &out, "out", 1, 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, global_norm, + kernel_num, d_out, d_inp, + num_params, block_size); + size_t memory_ops = num_params * sizeof(floatX); + 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(inp); + cudaCheck(cudaFree(d_out)); + cudaCheck(cudaFree(d_inp)); +} \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index cfec532..cb00628 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1206,8 +1206,6 @@ __global__ void norm_kernel(float* out, const T* data, size_t count) { // the maximum amount of work (so no fixed chunk, but instead iterating // until we run out of data), and then we reduce inside the block // and finally have just one atomic per block. - // TODO write a second version that just spams atomics in dev/cuda, - // often they are surprisingly fast namespace cg = cooperative_groups; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); @@ -1711,7 +1709,7 @@ void fused_classifier(Type* logits, Type* losses, } template -void norm(float* out, const T* values, size_t count) { +void global_norm(float* out, const T* values, size_t count) { const int block_size = 512; // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. // having one block less than possible is a tiny performance hit, having @@ -1720,7 +1718,7 @@ void norm(float* out, const T* values, size_t count) { // on all gpus, so the division really is going to be exact. const int grid_size = cuda_threads_per_SM * cuda_num_SMs / block_size; assert(grid_size > 0); // gives a better error than letting the call below fail - norm_kernel<<>>(out, values, count); + norm_kernel<<>>(out, values, count); cudaCheck(cudaGetLastError()); } @@ -2428,7 +2426,7 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo float* grad_norm = (float*)model->acts.output; // global gradient norm - norm(grad_norm, (floatX*)model->grads_memory, model->num_parameters); + global_norm(grad_norm, (floatX*)model->grads_memory, model->num_parameters); int block_size = 512; int num_blocks = CEIL_DIV(num_parameters, block_size);