From e278112c1949d7f1b2f0f8ececa3e812c354b058 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 18 Apr 2024 00:48:54 +0300 Subject: [PATCH] coalesced reading in attention backward pass --- dev/cuda/attention_backward.cu | 64 ++++++++++++++++++++++++++++++++-- train_gpt2.cu | 30 +++++++++------- 2 files changed, 79 insertions(+), 15 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index 1c461dc..3fc8db7 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -371,6 +371,39 @@ __global__ void softmax_autoregressive_backward_kernel2(float* dpreatt, const fl } } +// parallelize across t,b,h +__global__ void softmax_autoregressive_backward_kernel3(float* dpreatt, const float* datt, const float* att, + int B, int T, int C, int NH) { + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int t3 = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); + + int idx = blockIdx.y * T * T; + if (t3 >= T) { return; } + + int hs = C / NH; // head size + float scale = 1.0f / sqrtf(hs); + for (int t = t3; t < T; t++) { + float result = 0.0; + const float* att_bth = att + idx + t*T; + const float* datt_bth = datt + idx + t*T; + float* dpreatt_bth = dpreatt + idx + t*T; + const float att_at_t3 = att_bth[t3]; + + for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) { + float indicator = t2 == t3 ? 1.0f : 0.0f; + float local_derivative = att_bth[t2] * (indicator - att_at_t3); + result += local_derivative * datt_bth[t2]; + } + + result = cg::reduce(warp, result, cg::plus()); + if(warp.thread_rank() == 0) { + dpreatt_bth[t3] += scale * result; + } + } +} + // ---------------------------------------------------------------------------- // kernel launchers @@ -430,12 +463,29 @@ void attention_forward(float* out, float* vaccum, float* qkvr, float* preatt, fl unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); } +void launch_softmax_1(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + int num_blocks = ceil_div(T, block_size); + softmax_autoregressive_backward_kernel1<<>>(dpreatt, datt, att, B, T, C, NH); +} + +void launch_softmax_2(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + int num_blocks = ceil_div(T, block_size); + softmax_autoregressive_backward_kernel2<<>>(dpreatt, datt, att, B, T, C, NH); +} + +void launch_softmax_3(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + int num_blocks = ceil_div(32*T, block_size); + softmax_autoregressive_backward_kernel3<<>>(dpreatt, datt, att, B, T, C, NH); +} + // the sequence of transformations in this compound op is: // inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C) +template void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, float* dvaccum, const float* dout, const float* inp, const float* qkvr, const float* preatt, const float* att, const float* vaccum, int B, int T, int C, int NH, + SoftmaxKernel softmax_autoregressive_backward, const int block_size) { int HS = C / NH; // head size const float alpha = 1.0f; @@ -477,8 +527,7 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, B * NH); // backward into preatt - num_blocks = ceil_div(T, block_size); - softmax_autoregressive_backward_kernel2<<>>(dpreatt, datt, att, B, T, C, NH); + softmax_autoregressive_backward(dpreatt, datt, att, B, T, C, NH, block_size); // backward into q cublasSgemmStridedBatched(cublas_handle, @@ -515,7 +564,16 @@ void attention_backward(int kernel_num, const int block_size) { switch (kernel_num) { case 1: - attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, block_size); + attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, + launch_softmax_1, block_size); + break; + case 2: + attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, + launch_softmax_2, block_size); + break; + case 3: + attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, + launch_softmax_3, block_size); break; default: printf("Invalid kernel number\n"); diff --git a/train_gpt2.cu b/train_gpt2.cu index 331e418..ad54e99 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -669,32 +669,38 @@ __global__ void setConstant(float* vec, float constant, int N) { // naive kernel to backward through an autoregressive softmax, just to get correctness __global__ void softmax_autoregressive_backward_kernel(float* dpreatt, const float* datt, const float* att, - int B, int T, int C, int NH) { - int t3 = blockIdx.x * blockDim.x + threadIdx.x; - if (t3 >= T) { - return; - } + int B, int T, int C, int NH) { + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int t3 = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); + + int idx = blockIdx.y * T * T; + if (t3 >= T) { return; } int hs = C / NH; // head size float scale = 1.0f / sqrtf(hs); - int idx = blockIdx.y * T * T; - for (int t = t3; t < T; t++) { float result = 0.0; const float* att_bth = att + idx + t*T; const float* datt_bth = datt + idx + t*T; float* dpreatt_bth = dpreatt + idx + t*T; + const float att_at_t3 = att_bth[t3]; - for (int t2 = 0; t2 <= t; t2++) { + for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) { float indicator = t2 == t3 ? 1.0f : 0.0f; - float local_derivative = att_bth[t2] * (indicator - att_bth[t3]); - result += scale * local_derivative * datt_bth[t2]; + float local_derivative = att_bth[t2] * (indicator - att_at_t3); + result += local_derivative * datt_bth[t2]; } - dpreatt_bth[t3] += result; + result = cg::reduce(warp, result, cg::plus()); + if(warp.thread_rank() == 0) { + dpreatt_bth[t3] += scale * result; + } } } + // Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation). // Reference: https://developer.nvidia.com/blog/lerp-faster-cuda __device__ inline float lerp(float start, float end, float weight) { @@ -1035,7 +1041,7 @@ void attention_backward(float* dinp, float* dqkvr, float* dpreatt, float* datt, B * NH); // backward into preatt - softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, NH); + softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, NH); // backward into q cublasSgemmStridedBatched(cublas_handle,