From 3ba2cdc7232f8df796d5fcc8774c5c4ea0b5dacb Mon Sep 17 00:00:00 2001 From: ademeure Date: Thu, 18 Apr 2024 03:41:27 +0100 Subject: [PATCH] fixes + bounds checking --- dev/cuda/classifier_fused.cu | 83 +++++++++++++++++++++--------------- 1 file changed, 49 insertions(+), 34 deletions(-) diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index 6f9391c..f00ee56 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -154,33 +154,31 @@ __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) { +__device__ SoftmaxParams prepare_softmax_blockwide(cg::thread_block_tile<32>& warp, + int idx, const float* inp, int V, int P) { // 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); + const float4* x_vec4 = reinterpret_cast(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/4) + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) { - float4 v = x_vec4[i]; + for (int i = (V+3)/4 + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) { + float4 v4 = x_vec4[i]; #pragma unroll for(int k = 0; k < 4; ++k) { + float v = (i*4+k < V) ? vec_at(v4, k) : 0.f; // bounds checking against real V float old_maxval = thread_maxval; - thread_maxval = fmaxf(thread_maxval, vec_at(v, k)); + thread_maxval = fmaxf(thread_maxval, vec_at(v4, k)); thread_sumval *= expf((old_maxval - thread_maxval)); - thread_sumval += expf(vec_at(v, k) - thread_maxval); + thread_sumval += expf(vec_at(v4, 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 + // 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; @@ -203,31 +201,50 @@ __device__ SoftmaxParams prepare_softmax2(cg::thread_block_tile<32>& warp, return SoftmaxParams{1.f / block_sumval, block_maxval}; } -__global__ void fused_classifier_kernel2(float* dlogits, float* losses, +// Fused forward and backward pass for classifier including softmax, and logit gradients +// Writes to both probs (only for debugging) and dlogits (only for training) are optional +// N.B.: We may want to reuse the logits memory for dlogits, so they should *not* be __restrict__! +__global__ void fused_classifier_kernel2(float* dlogits, float* losses, float* probs, const float* logits, const float* dlosses, const int* targets, - int B, int T, int V) { + 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]; - float dloss = dlosses[idx]; - auto sp = prepare_softmax2(warp, idx, logits, V); + // softmax (reading B * T * V, same logits read again below, hopefully still in cache) + auto sp = prepare_softmax_blockwide(warp, idx, logits, V); - // calculate the probability needed for the loss and update. - // single-threaded + // 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; + float prob = expf(logits[idx * P + 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; + // calculate the gradients directly, saves bandwidth from probs during training + // but also supports writing probs for inference-only and debugging + float dloss = dlosses ? dlosses[idx] : 0.f; + const float4* logits_vec4 = reinterpret_cast(logits + idx * P); + for (int i = threadIdx.x; i < (V+3)/4; 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 + float4 v4 = __ldcs(&logits_vec4[i]); + #pragma unroll + for(int k = 0; k < 4; ++k) { + int element = i*4 + k; + float prob = expf(vec_at(v4, k) - sp.Offset) * sp.Scale; + prob = (element < V) ? prob : 0.f; // bounds checking against real V + + // this kernel is DRAM limited so cost of inner branch is ~zero + if (probs) { + probs[idx * P + element] = prob; + } + if (dlogits) { + float indicator = element == ix ? 1.0f : 0.0f; + dlogits[idx * P + element] = (prob - indicator) * dloss; + } + } } } @@ -239,16 +256,16 @@ void fused_classifier1(float* dlogits, float* losses, int B, int T, int V, int P, 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); + fused_classifier_kernel1<<>>(dlogits, losses, logits, dlosses, targets, B, T, V, P); 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) { + int B, int T, int V, int P, 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, NULL, logits, dlosses, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } @@ -260,11 +277,7 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses, fused_classifier1(dlogits, losses, logits, dlosses, targets, B, T, V, P, 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); + fused_classifier2(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size); break; default: printf("Invalid kernel number\n"); @@ -280,6 +293,8 @@ int main(int argc, char **argv) { int B = 8; int T = 1024; int V = 50257; + // padded size + int P = (V + 63) & ~63; int deviceIdx = 0; cudaCheck(cudaSetDevice(deviceIdx));