From b225501b7b38e8a2ea3ce4dd46ac3eb0eebbb886 Mon Sep 17 00:00:00 2001 From: ademeure Date: Tue, 16 Apr 2024 14:54:23 +0100 Subject: [PATCH] optimised version of the cool new fused classifier + hopefully bugfixes? --- dev/cuda/classifier_fused.cu | 111 ++++++++++++++++++++++++++++++++--- 1 file changed, 104 insertions(+), 7 deletions(-) diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index f38a7cd..dbebe3f 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 #include #include +#include #include #include #include @@ -72,7 +73,7 @@ void crossentropy_softmax_backward_cpu(float* dlogits, for (int i = 0; i < V; i++) { float p = probs_bt[i]; float indicator = i == ix ? 1.0f : 0.0f; - dlogits_bt[i] += (p - indicator) * dloss; + dlogits_bt[i] = (p - indicator) * dloss; } } } @@ -111,8 +112,7 @@ __device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp, return SoftmaxParams{norm, global_maxval}; } - -__global__ void fused_classifier_kernel(float* dlogits, float* losses, +__global__ void fused_classifier_kernel1(float* dlogits, float* losses, const float* logits, const float* dlosses, const int* targets, int B, int T, int V) { namespace cg = cooperative_groups; @@ -139,15 +139,96 @@ __global__ void fused_classifier_kernel(float* dlogits, float* losses, // calculate all the gradients for (int i = warp.thread_rank(); i < V; i += warp.size()) { - float prob = expf(logits[i] - sp.Offset) * sp.Scale; + float prob = expf(logits[idx * V + i] - sp.Offset) * sp.Scale; float* dlogits_bt = dlogits + b * T * V + t * V; float dloss = dlosses[b * T + t]; int ix = targets[b * T + t]; float p = prob; float indicator = i == ix ? 1.0f : 0.0f; - dlogits_bt[i] += (p - indicator) * dloss; + dlogits_bt[i] = (p - indicator) * dloss; + } +} + + +__device__ float vec_at(const float4& vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__device__ SoftmaxParams prepare_softmax2(cg::thread_block_tile<32>& warp, + int idx, const float* inp, int V) { + // one row of inp, i.e. inp[idx, :] of shape (V,) + // float4 to get 128-bit loads and memory level parallelism + // this is only possible if V is a multiple of 4 + const float4* x_vec4 = reinterpret_cast(inp + idx * V); + + 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/4) + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) { + float4 v = x_vec4[i]; + #pragma unroll + for(int k = 0; k < 4; ++k) { + float old_maxval = thread_maxval; + thread_maxval = fmaxf(thread_maxval, vec_at(v, k)); + thread_sumval *= expf((old_maxval - thread_maxval)); + thread_sumval += expf(vec_at(v, k) - thread_maxval); + } } + // reduction in 2 stages: 1) inside warp, 2) between warps + // this results in much cleaner code than using a multi-warp cg::reduce + // todo benchmark to make sure it's faster, possibly too many reductions + // we could do the 2nd set of reductions per-block rather than per-warp + // but that would require an extra __syncthreads() unfortunately + __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 + shared_maxval[warp_id] = cg::reduce(warp, thread_maxval, cg::greater{}); + __syncthreads(); + float warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; + float block_maxval = cg::reduce(warp, warp_maxval, cg::greater{}); + // use maxval to scale sumval to avoid numerical instability / overflow + thread_sumval *= expf(thread_maxval - block_maxval); + // reduce sumval + shared_sumval[warp_id] = cg::reduce(warp, thread_sumval, cg::plus{}); + __syncthreads(); + float warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.f; + float block_sumval = cg::reduce(warp, warp_sumval, cg::plus{}); + + return SoftmaxParams{1.f / block_sumval, block_maxval}; +} + +__global__ void fused_classifier_kernel2(float* dlogits, float* losses, + const float* logits, const float* dlosses, const int* targets, + int B, int T, int V) { + 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]; + float dloss = dlosses[idx]; + + auto sp = prepare_softmax2(warp, idx, logits, V); + + // calculate the probability needed for the loss and update. + // single-threaded + if(threadIdx.x == 0) { + float prob = expf(logits[idx * V + ix] - sp.Offset) * sp.Scale; + losses[idx] = -logf(prob); + } + + // calculate all the gradients + for (int i = threadIdx.x; i < V; i += blockDim.x) { + float prob = expf(__ldcs(&logits[idx * V + i]) - sp.Offset) * sp.Scale; + float indicator = i == ix ? 1.0f : 0.0f; + dlogits[idx * V + i] = (prob - indicator) * dloss; + } } // ---------------------------------------------------------------------------- @@ -157,8 +238,17 @@ void fused_classifier1(float* dlogits, float* losses, const float* logits, const float* dlosses, const int* targets, int B, int T, int V, int block_size) { const int N = B * T; + const int grid_size = N / (block_size / 32); + fused_classifier_kernel1<<>>(dlogits, losses, logits, dlosses, targets, B, T, V); + cudaCheck(cudaGetLastError()); +} + +void fused_classifier2(float* dlogits, float* losses, + const float* logits, const float* dlosses, const int* targets, + int B, int T, int V, int block_size) { + const int N = B * T; const int grid_size = N; - fused_classifier_kernel<<>>(dlogits, losses, logits, dlosses, targets, B, T, V); + fused_classifier_kernel2<<>>(dlogits, losses, logits, dlosses, targets, B, T, V); cudaCheck(cudaGetLastError()); } @@ -169,6 +259,13 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses, case 1: fused_classifier1(dlogits, losses, logits, dlosses, targets, B, T, V, block_size); break; + case 2: + if((V % 4) != 0) { + printf("V needs to be a multiple of 4 for this kernel to work!\n"); + exit(EXIT_FAILURE); + } + fused_classifier2(dlogits, losses, logits, dlosses, targets, B, T, V, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); @@ -182,7 +279,7 @@ int main(int argc, char **argv) { int B = 8; int T = 1024; - int V = 50257; + int V = 50272; // rounded up from 50257 for alignment purposes int deviceIdx = 0; cudaCheck(cudaSetDevice(deviceIdx));