From 2287da01207d8d4d92eaeb12cee0f2093dd46a6a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 12 May 2024 19:42:32 +0300 Subject: [PATCH] enable bf16 --- dev/cuda/matmul_backward_bias.cu | 89 ++++++++++++++------------------ 1 file changed, 40 insertions(+), 49 deletions(-) diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 233b7a1..024eca4 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -23,7 +23,7 @@ sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1 #include #include -//#define ENABLE_BF16 +#define ENABLE_BF16 #include "common.h" // ---------------------------------------------------------------------------- @@ -49,16 +49,16 @@ void matmul_backward_bias_cpu(float* dinp, float* dweight, float* dbias, float* dbias_buffer; -__global__ void matmul_backward_bias_kernel1(float* dbias, const float* dout, int B, int T, int OC) { +__global__ void matmul_backward_bias_kernel1(floatX* dbias, const floatX* dout, int B, int T, int OC) { extern __shared__ float shared[]; int o = blockIdx.x; // range [0, OC) int tid = threadIdx.x; // range [0, block_size) int block_size = blockDim.x; - const float* x = dout + o; + const floatX* x = dout + o; // thread coarsening float sum = 0.0; for (int i = tid; i < B * T; i += block_size) { - sum += x[i * OC]; + sum += (float)x[i * OC]; } shared[tid] = sum; __syncthreads(); @@ -71,12 +71,12 @@ __global__ void matmul_backward_bias_kernel1(float* dbias, const float* dout, in } // write the final result (at thread 0) to global memory if (tid == 0) { - dbias[o] += shared[0]; + dbias[o] = (float)dbias[o] + shared[0]; } } // cooperative groups solution, one warp per output channel -__global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, int B, int T, int OC) { +__global__ void matmul_backward_bias_kernel2(floatX* dbias, const floatX* dout, int B, int T, int OC) { // dout is (B, T, OC), dbias is (OC) // e.g. if block_size = 128, then we have 4 warps per block, each in charge of one output channel namespace cg = cooperative_groups; @@ -89,7 +89,7 @@ __global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, in // first, thread coarsening to sum reduce the problem size from B*T to 32 float sum = 0.0f; for(int i = warp.thread_rank(); i < BT; i += warp.size()) { - sum += dout[i * OC + idx]; + sum += (float)dout[i * OC + idx]; } // now do a warp-level reduce to get the sum across the 32 threads in this warp sum = cg::reduce(warp, sum, cg::plus{}); @@ -99,7 +99,7 @@ __global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, in } } -__global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, int B, int T, int OC) { +__global__ void matmul_backward_bias_kernel3(floatX* dbias, const floatX* dout, int B, int T, int OC) { // dout is (B, T, OC), dbias is (OC) // in this version of the kernel the entire block of block_size is dedicated to one output channel namespace cg = cooperative_groups; @@ -114,7 +114,7 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in // round 1: thread coarsening to reduce the problem size from B*T to 32 float thread_sum = 0.0f; for(int i = threadIdx.x; i < BT; i += blockDim.x) { - thread_sum += dout[i * OC + idx]; + thread_sum += (float)dout[i * OC + idx]; } // now do a warp-level reduce to get the sum across the 32 threads in each warp float warp_sum = cg::reduce(warp, thread_sum, cg::plus{}); @@ -136,7 +136,7 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in // 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(float* dbias, const float* dout, int B, int T, int OC) { +__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) @@ -147,7 +147,7 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in // 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 float* dout_col = dout + tl + lane_id; + 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 @@ -156,7 +156,7 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in // leading to a coalesced memory access pattern float dout_sum = 0.0f; for (int row = warp_id; row < B * T; row += vstep) { - dout_sum += dout_col[row * OC]; + dout_sum += (float)dout_col[row * OC]; } smem[lane_id + warp_id * warpSize] = dout_sum; __syncthreads(); @@ -171,13 +171,13 @@ __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) { +__global__ void matmul_backward_bias_kernel5(floatX* dbias, const floatX* 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]; + sum += (float)dout[i * OC + oc]; } // and atomically add everything together. atomics within one block are conflict-free! atomicAdd(dbias + oc, sum); @@ -248,8 +248,7 @@ __global__ void matmul_backward_bias_kernel7(float* dbias, const floatX* dout, i // kernel launcher // version1: simple cuBLAS calls -void matmul_backward_bias1(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias1(floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { dim3 block_dim(block_size); dim3 grid_dim(OC); @@ -257,42 +256,37 @@ void matmul_backward_bias1(float* dinp, float* dweight, float* dbias, matmul_backward_bias_kernel1<<>>(dbias, dout, B, T, OC); } -void matmul_backward_bias2(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias2(floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { // block_size 512 seems best const int grid_size = ceil_div(OC * 32, block_size); matmul_backward_bias_kernel2<<>>(dbias, dout, B, T, OC); } -void matmul_backward_bias3(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias3(floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { // block_size 256 seems best matmul_backward_bias_kernel3<<>>(dbias, dout, B, T, OC); } -void matmul_backward_bias4(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias4(floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { assert(OC % 32 == 0); // OC must be divisible by 32 for this kernel const int grid_size = OC / 32; 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, +void matmul_backward_bias5(floatX* dbias, floatX* dout, 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_bias7(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias7(floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { - if(block_size < 128) { - block_size = 128; + if(block_size < 256) { + block_size = 256; } // Each warp is responsible for 32 * "x128::size" = 256 OCs at BF16 (OC must be a multiple of 256!) // Block size is 512 threads (16 warps) and we reduce those 16 values into 1 at the end @@ -315,28 +309,26 @@ void matmul_backward_bias7(float* dinp, float* dweight, float* dbias, cast_and_add_kernel<<>>(dbias, dbias_buffer, OC); } -void matmul_backward_bias(int kernel_num, - float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, float* ones, +void matmul_backward_bias(int kernel_num, floatX* dbias, floatX* dout, int B, int T, int C, int OC, int block_size) { switch (kernel_num) { case 1: - matmul_backward_bias1(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + matmul_backward_bias1(dbias, dout, B, T, C, OC, block_size); break; case 2: - matmul_backward_bias2(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + matmul_backward_bias2(dbias, dout, B, T, C, OC, block_size); break; case 3: - matmul_backward_bias3(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + matmul_backward_bias3(dbias, dout, B, T, C, OC, block_size); break; case 4: - matmul_backward_bias4(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + matmul_backward_bias4(dbias, dout, 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); + matmul_backward_bias5(dbias, dout, B, T, C, OC, block_size); break; case 7: - matmul_backward_bias7(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + matmul_backward_bias7(dbias, dout, B, T, C, OC, block_size); break; default: printf("Invalid kernel number\n"); @@ -366,13 +358,13 @@ int main(int argc, char **argv) { float* dout = make_random_float(B * T * OC); // move to GPU - float* d_dbias; - float* d_dout; - cudaCheck(cudaMalloc(&d_dbias, OC * sizeof(float))); - cudaCheck(cudaMalloc(&d_dout, B * T * OC * sizeof(float))); + floatX* d_dbias; + floatX* d_dout; + cudaCheck(cudaMalloc(&d_dbias, OC * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_dout, B * T * OC * sizeof(floatX))); cudaCheck(cudaMalloc(&dbias_buffer, OC * sizeof(float))); - cudaCheck(cudaMemcpy(d_dbias, dbias, OC * sizeof(float), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_dout, dout, B * T * OC * sizeof(float), cudaMemcpyHostToDevice)); + cudaCheck(memcpy_convert(d_dbias, dbias, OC)); + cudaCheck(memcpy_convert(d_dout, dout, B * T * OC)); // ncu debugging / profiling, do a single call // int block_size_debug; @@ -391,23 +383,22 @@ int main(int argc, char **argv) { for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; // memset the bias to zero - cudaCheck(cudaMemset(d_dbias, 0, OC * sizeof(float))); + cudaCheck(cudaMemset(d_dbias, 0, OC * sizeof(floatX))); // calculate the GPU version - matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, block_size); + matmul_backward_bias(kernel_num, d_dbias, d_dout, B, T, C, OC, block_size); // compare printf("Checking correctness...\n"); - validate_result(d_dbias, dbias, "dbias", OC, 5e-3f); + float tol = std::is_same_v ? 5e-3f : 1.0f; + validate_result(d_dbias, dbias, "dbias", OC, tol); printf("All results match for block_size=%d.\n\n", block_size); } // now benchmark the kernel for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; - float *d_dinp, *d_dweight, *d_inp, *d_weight, *d_ones; int repeat_times = 2000; float elapsed_time = benchmark_kernel(repeat_times, matmul_backward_bias, kernel_num, - d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight, d_ones, - B, T, C, OC, block_size); + d_dbias, d_dout, B, T, C, OC, block_size); printf("block_size %d time %.4f ms\n", block_size, elapsed_time); }