From 16123c3af8ec6e649415dbe4b518ed5ba7db4c23 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 19 Apr 2024 10:56:34 +0300 Subject: [PATCH 1/4] towards an even better backward attention kernel --- dev/cuda/attention_backward.cu | 115 ++++++++++++++++++++++++++++++--- dev/cuda/common.h | 2 +- train_gpt2.cu | 110 +++++++++++-------------------- 3 files changed, 146 insertions(+), 81 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index db386bf..3cdc679 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -560,6 +560,78 @@ __global__ void softmax_autoregressive_backward_kernel5(float* __restrict__ dpre } } + +// I want `BlockSize` to be statically known to the compiler, thus we get a template here. +// This kernel takes a step back, and looks at the original CPU code again. We have some simple outer loops +// That are independent, (b, t, h), and then the inner loops over (t2, t3) where we're combining elements -- this is +// where we can reuse data and be more efficient +// => handle b, t, h through block indices; each block does all the work for the (t2, t3) loop cooperatively. +// Now we have two nested loops, and in the inner instruction, we combine indexing from both => this calls for +// loop tiling, and lifting some of the memory ops out of the loop. +// We're in luck here; if we tile so that t3 is the outer loop, we can get a sinlge write op per result, AND also cache +// the t2-indexed part of the computation, which is the problematic one because it contains a multiplication that now we +// do not have to repeat over and over. +// => do an outer t3 loop where each thread gets one t3 index. Then, do an outer t2 loop in steps of BlockSize, and +// prepare BlockSize many elements for the inner loop. Here, each thread calculates one element and stores it in shmem. +// Then, in the inner t2 loop, each thread reads *all* the elements previously stored and does its computations. +// This way, we do 3*BlockSize loads, but BlockSize^2 computation steps => This kernel is now entirely compute bound. +// To fix up the compute issues, as above, we replace ifs in memory reading with min, and also split the inner loop +// into a large region where we don't have to calculate the indicator, and a small, costly region where we do. +template +__global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_kernel6(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(); + __shared__ float att_bth_s[BlockSize]; + + int idx = blockIdx.y; + int t = blockIdx.x; + + att += idx * T * T; + datt += idx * T * T; + dpreatt += idx * T * T; + + int hs = C / NH; // head size + float scale = 1.0f / sqrtf(hs); + const float* att_bth = att + t * T; + const float* datt_bth = datt + t * T; + float* dpreatt_bth = dpreatt + t * T; + + int block_steps = ceil_div(t+1, BlockSize); + // very important: This loop condition needs to be the same for all threads. + // even if a thread later on is not going to do any work, it needs to participate in the + // data loading process! + for (int t3f = 0; t3f < block_steps; ++t3f) { + int t3 = t3f * BlockSize + block.thread_rank(); + float acc = 0.f; + float at3 = att_bth[t3]; + for (int t2b = 0; t2b <= t; t2b += BlockSize) { + int end = min(t + 1 - t2b, BlockSize); + block.sync(); + { + int t2i = block.thread_rank(); + int t2 = min(t, t2b + t2i); + att_bth_s[t2i] = att_bth[t2] * datt_bth[t2]; + } + + block.sync(); + if(t3f * BlockSize == t2b) { + for (int t2i = 0; t2i < end; t2i++) { + int t2 = t2b + t2i; + float indicator = t2 == t3 ? 1.0f : 0.0f; + acc += att_bth_s[t2i] * (indicator - at3); + } + } else { + for (int t2i = 0; t2i < end; t2i++) { + int t2 = t2b + t2i; + acc += att_bth_s[t2i] * (0.f - at3); + } + } + } + dpreatt_bth[t3] += scale * acc; + } +} + // ---------------------------------------------------------------------------- // kernel launchers @@ -644,6 +716,29 @@ void launch_softmax_5(float* dpreatt, float* datt, const float* att, int B, int softmax_autoregressive_backward_kernel5<<>>(dpreatt, datt, att, B, T, C, NH); } +void launch_softmax_6(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + switch(block_size) { + case 32: + softmax_autoregressive_backward_kernel6<32><<>>(dpreatt, datt, att, B, T, C, NH); + break; + case 64: + softmax_autoregressive_backward_kernel6<64><<>>(dpreatt, datt, att, B, T, C, NH); + break; + case 128: + softmax_autoregressive_backward_kernel6<128><<>>(dpreatt, datt, att, B, T, C, NH); + break; + case 256: + softmax_autoregressive_backward_kernel6<256><<>>(dpreatt, datt, att, B, T, C, NH); + break; + case 512: + softmax_autoregressive_backward_kernel6<512><<>>(dpreatt, datt, att, B, T, C, NH); + break; + case 1024: + softmax_autoregressive_backward_kernel6<1024><<>>(dpreatt, datt, att, B, T, C, NH); + break; + } +} + // 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 @@ -671,7 +766,7 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, unpermute_kernel_backward<<>>(dvaccum, dout, B, T, NH, HS); // backward into datt - cublasSgemmStridedBatched(cublas_handle, + cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, @@ -679,10 +774,10 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, dvaccum, HS, T * HS, &beta, datt, T, T * T, - B * NH); + B * NH)); // backward into dv - cublasSgemmStridedBatched(cublas_handle, + cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, @@ -690,14 +785,14 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, att, T, T * T, &beta, dv, HS, T * HS, - B * NH); + B * NH)); // backward into preatt softmax_autoregressive_backward(dpreatt, datt, att, B, T, C, NH, block_size); cudaCheck(cudaGetLastError()); // backward into q - cublasSgemmStridedBatched(cublas_handle, + cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, @@ -705,9 +800,9 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, dpreatt, T, T * T, &beta, dq, HS, T * HS, - B * NH); + B * NH)); // backward into k - cublasSgemmStridedBatched(cublas_handle, + cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha, @@ -715,7 +810,7 @@ void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, dpreatt, T, T * T, &beta, dk, HS, T * HS, - B * NH); + B * NH)); // backward into inp num_blocks = ceil_div(B * NH * T * HS, block_size); @@ -750,6 +845,10 @@ void attention_backward(int kernel_num, attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, launch_softmax_5, block_size); break; + case 6: + attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, + launch_softmax_6, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 2a6c6ce..6723a0f 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -5,7 +5,7 @@ template -T ceil_div(T dividend, T divisor) { +__host__ __device__ T ceil_div(T dividend, T divisor) { return (dividend + divisor-1) / divisor; } diff --git a/train_gpt2.cu b/train_gpt2.cu index 90af1b7..3539795 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -668,88 +668,54 @@ __global__ void layernorm_backward_kernel(float* dinp, float* dweight, float* db // 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) { - constexpr int UNROLL = 8; + constexpr const int BlockSize = 64; cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - int t3 = UNROLL * (blockIdx.x * warp.meta_group_size() + warp.meta_group_rank()); + __shared__ float att_bth_s[BlockSize]; - int idx = blockIdx.y * T * T; - if (t3 >= T) { return; } + int idx = blockIdx.y; + int t = blockIdx.x; + + att += idx * T * T; + datt += idx * T * T; + dpreatt += idx * T * T; int hs = C / NH; // head size float scale = 1.0f / sqrtf(hs); - for (int t = t3; t < T; t++) { - float result[UNROLL] = {}; - 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_bth = att + t * T; + const float* datt_bth = datt + t * T; + float* dpreatt_bth = dpreatt + t * T; - float att_at_t3[UNROLL]; - for(int k = 0; k < UNROLL; ++k) { - // if t < t3+k, we're out of bounds. - // in that case, we don't care what we read, because later on, - // we won't write the corresponding result. So just clip to - // make sure this is a valid (in-bounds) memory access. - att_at_t3[k] = att_bth[min(t, t3 + k)]; - } - - // the code below is actually just a for loop; except, - // we have to do something special in one iteration in - // the middle, and an if turned out to have significant - // performance impact. - // so we split the loop in three parts. Ugly, but effective. - - // the beginning/end loop does the same thing, so we write the code - // just once in a lambda. In this step, we're guaranteed that - // indicator == 0 - auto loop_step = [&](int t2){ - float p = att_bth[t2] * datt_bth[t2]; - for (int k = 0; k < UNROLL; ++k) { - result[k] -= p * att_at_t3[k]; - } - }; - - // Now the actual loop. - { - // declare the loop iterator. Needs to be kept across the - // three different parts, so it's not a local variable in - // the for loop. - int t2 = warp.thread_rank(); - - // first part, as long as t2 < t3, indicator == 0 - for (; t2 < t3; t2 += warp.size()) { - loop_step(t2); + int block_steps = CEIL_DIV(t+1, BlockSize); + for (int t3f = 0; t3f < block_steps; ++t3f) { + int t3 = t3f * BlockSize + block.thread_rank(); + float acc = 0.f; + float at3 = att_bth[t3]; + for (int t2b = 0; t2b <= t; t2b += BlockSize) { + int end = min(t + 1 - t2b, BlockSize); + block.sync(); + { + int t2i = block.thread_rank(); + int t2 = min(t, t2b + t2i); + att_bth_s[t2i] = att_bth[t2] * datt_bth[t2]; } - // because k <= warp.size() (==32), the event that t3+k == t2 - // has to happen at this particular step. - static_assert(UNROLL <= 32, "UNROLL is too large, this won't produce correct results."); - if (t2 <= t) { - float att_t2 = att_bth[t2]; - float datt_t2 = datt_bth[t2]; - float p = att_t2 * datt_t2; - for (int k = 0; k < UNROLL; ++k) { - float indicator = t2 == (t3 + k) ? 1.0f : 0.0f; - result[k] += p * (indicator - att_at_t3[k]); + block.sync(); + if(t3f * BlockSize == t2b) { + for (int t2i = 0; t2i < end; t2i++) { + int t2 = t2b + t2i; + float indicator = t2 == t3 ? 1.0f : 0.0f; + float local_derivative = att_bth_s[t2i] * (indicator - at3); + acc += local_derivative; + } + } else { + for (int t2i = 0; t2i < end; t2i++) { + int t2 = t2b + t2i; + float local_derivative = att_bth_s[t2i] * (0.f - at3); + acc += local_derivative; } - t2 += warp.size(); - } - - // rest of the loop, indicator == 0 again - for (; t2 <= t; t2 += warp.size()) { - loop_step(t2); } } - - for(int k = 0; k < UNROLL; ++k) { - result[k] = cg::reduce(warp, result[k], cg::plus()); - } - - // when storing, we need to check that this is actually a valid result. - // here, warp.thread_rank() corresponds to `k` in the previous loops. - if (warp.thread_rank() < UNROLL && t >= t3 + warp.thread_rank()) { - dpreatt_bth[t3 + warp.thread_rank()] = scale * result[warp.thread_rank()]; - } + dpreatt_bth[t3] = scale * acc; } } @@ -1182,7 +1148,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, From 02aa8e47dc1775b61824144dbd3f14f5fcd5f078 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 19 Apr 2024 17:36:18 +0300 Subject: [PATCH 2/4] further simplification of the loop --- dev/cuda/attention_backward.cu | 97 ++++++++++++++++++++++++++++------ 1 file changed, 82 insertions(+), 15 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index 3cdc679..3a48672 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -623,15 +623,65 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker } } else { for (int t2i = 0; t2i < end; t2i++) { - int t2 = t2b + t2i; acc += att_bth_s[t2i] * (0.f - at3); } } } - dpreatt_bth[t3] += scale * acc; + dpreatt_bth[t3] = scale * acc; } } +template +__global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_kernel7(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(); + __shared__ float att_bth_s[BlockSize]; + + int idx = blockIdx.y; + int t = blockIdx.x; + + att += idx * T * T; + datt += idx * T * T; + dpreatt += idx * T * T; + + int hs = C / NH; // head size + float scale = 1.0f / sqrtf(hs); + const float* att_bth = att + t * T; + const float* datt_bth = datt + t * T; + float* dpreatt_bth = dpreatt + t * T; + + int block_steps = ceil_div(t+1, BlockSize); + // very important: This loop condition needs to be the same for all threads. + // even if a thread later on is not going to do any work, it needs to participate in the + // data loading process! + for (int t3f = 0; t3f < block_steps; ++t3f) { + int t3 = t3f * BlockSize + block.thread_rank(); + float acc = 0.f; + float at3 = att_bth[t3]; + for (int t2b = 0; t2b <= t; t2b += BlockSize) { + int end = min(t + 1 - t2b, BlockSize); + block.sync(); + { + int t2i = block.thread_rank(); + int t2 = min(t, t2b + t2i); + att_bth_s[t2i] = att_bth[t2] * datt_bth[t2]; + } + + block.sync(); + float sub = 0.0; + for (int t2i = 0; t2i < end; t2i++) { + sub += att_bth_s[t2i]; + } + acc -= sub * at3; + } + float at_t2_eq_t3 = att_bth[t3] * datt_bth[t3]; + acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3)); + dpreatt_bth[t3] = scale * acc; + } +} + + // ---------------------------------------------------------------------------- // kernel launchers @@ -716,29 +766,42 @@ void launch_softmax_5(float* dpreatt, float* datt, const float* att, int B, int softmax_autoregressive_backward_kernel5<<>>(dpreatt, datt, att, B, T, C, NH); } -void launch_softmax_6(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { +template +void dispatch_launch(Launcher&& launch, int block_size) { switch(block_size) { case 32: - softmax_autoregressive_backward_kernel6<32><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); case 64: - softmax_autoregressive_backward_kernel6<64><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); case 128: - softmax_autoregressive_backward_kernel6<128><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); case 256: - softmax_autoregressive_backward_kernel6<256><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); case 512: - softmax_autoregressive_backward_kernel6<512><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); case 1024: - softmax_autoregressive_backward_kernel6<1024><<>>(dpreatt, datt, att, B, T, C, NH); - break; + return launch(std::integral_constant{}); + default: + assert(false && "Invalid block size"); } } +void launch_softmax_6(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + auto launch = [&](auto int_const) { + softmax_autoregressive_backward_kernel6<<>>(dpreatt, datt, att, B, T, C, NH); + }; + dispatch_launch(launch, block_size); +} + +void launch_softmax_7(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) { + auto launch = [&](auto int_const) { + constexpr int block_size = int_const.value; + softmax_autoregressive_backward_kernel7<<>> + (dpreatt, datt, att, B, T, C, NH); + }; + dispatch_launch(launch, block_size); +} + // 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 @@ -849,6 +912,10 @@ void attention_backward(int kernel_num, attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, launch_softmax_6, block_size); break; + case 7: + attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH, + launch_softmax_7, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); From 4294820a8a675e9222f5641170831a4c3362becc Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 19 Apr 2024 18:31:26 +0300 Subject: [PATCH 3/4] don't do stupid redundant work; the inner loop is just a block-level sum!!!! --- dev/cuda/attention_backward.cu | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index 3a48672..0e16be3 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -28,6 +28,7 @@ OMP_NUM_THREADS=32 ./attention_backward 5 #include #include #include +#include #include "common.h" // ---------------------------------------------------------------------------- @@ -636,7 +637,9 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker 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); __shared__ float att_bth_s[BlockSize]; + __shared__ float att_bth_prefix[32]; int idx = blockIdx.y; int t = blockIdx.x; @@ -651,30 +654,28 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker const float* datt_bth = datt + t * T; float* dpreatt_bth = dpreatt + t * T; + if(warp.meta_group_rank() == 0) { + att_bth_prefix[warp.thread_rank()] = 0; + } + int block_steps = ceil_div(t+1, BlockSize); // very important: This loop condition needs to be the same for all threads. // even if a thread later on is not going to do any work, it needs to participate in the // data loading process! for (int t3f = 0; t3f < block_steps; ++t3f) { int t3 = t3f * BlockSize + block.thread_rank(); - float acc = 0.f; - float at3 = att_bth[t3]; - for (int t2b = 0; t2b <= t; t2b += BlockSize) { - int end = min(t + 1 - t2b, BlockSize); - block.sync(); - { - int t2i = block.thread_rank(); - int t2 = min(t, t2b + t2i); - att_bth_s[t2i] = att_bth[t2] * datt_bth[t2]; - } - block.sync(); - float sub = 0.0; - for (int t2i = 0; t2i < end; t2i++) { - sub += att_bth_s[t2i]; - } - acc -= sub * at3; + float at3 = att_bth[t3]; + float local_sum = 0; + for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + local_sum += att_bth[t2] * datt_bth[t2]; } + block.sync(); + att_bth_prefix[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); + block.sync(); + local_sum = cg::reduce(warp, att_bth_prefix[warp.thread_rank()], cg::plus{}); + + float acc = -local_sum * at3; float at_t2_eq_t3 = att_bth[t3] * datt_bth[t3]; acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3)); dpreatt_bth[t3] = scale * acc; From 50e105a2504b17360995178735cfe5a6cc84da6a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 19 Apr 2024 19:03:40 +0300 Subject: [PATCH 4/4] further cleanup and ability to handle arbitrary sequence lengths --- dev/cuda/attention_backward.cu | 17 ++++++----- train_gpt2.cu | 56 ++++++++++++++++------------------ 2 files changed, 36 insertions(+), 37 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index 0e16be3..7a53a22 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -638,8 +638,7 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker 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 att_bth_s[BlockSize]; - __shared__ float att_bth_prefix[32]; + __shared__ float block_acc[32]; int idx = blockIdx.y; int t = blockIdx.x; @@ -655,7 +654,7 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker float* dpreatt_bth = dpreatt + t * T; if(warp.meta_group_rank() == 0) { - att_bth_prefix[warp.thread_rank()] = 0; + block_acc[warp.thread_rank()] = 0; } int block_steps = ceil_div(t+1, BlockSize); @@ -665,20 +664,22 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker for (int t3f = 0; t3f < block_steps; ++t3f) { int t3 = t3f * BlockSize + block.thread_rank(); - float at3 = att_bth[t3]; + float at3 = att_bth[min(t, t3)]; float local_sum = 0; for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { local_sum += att_bth[t2] * datt_bth[t2]; } block.sync(); - att_bth_prefix[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); + block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); block.sync(); - local_sum = cg::reduce(warp, att_bth_prefix[warp.thread_rank()], cg::plus{}); + local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{}); float acc = -local_sum * at3; - float at_t2_eq_t3 = att_bth[t3] * datt_bth[t3]; + float at_t2_eq_t3 = at3 * datt_bth[min(t, t3)]; acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3)); - dpreatt_bth[t3] = scale * acc; + if(t3 <= t) { + dpreatt_bth[t3] = scale * acc; + } } } diff --git a/train_gpt2.cu b/train_gpt2.cu index 3539795..0ebc948 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -668,9 +668,10 @@ __global__ void layernorm_backward_kernel(float* dinp, float* dweight, float* db // 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) { - constexpr const int BlockSize = 64; + constexpr const int BlockSize = 256; cg::thread_block block = cg::this_thread_block(); - __shared__ float att_bth_s[BlockSize]; + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + __shared__ float block_acc[32]; int idx = blockIdx.y; int t = blockIdx.x; @@ -685,41 +686,38 @@ __global__ void softmax_autoregressive_backward_kernel(float* dpreatt, const flo const float* datt_bth = datt + t * T; float* dpreatt_bth = dpreatt + t * T; + if(warp.meta_group_rank() == 0) { + block_acc[warp.thread_rank()] = 0; + } + int block_steps = CEIL_DIV(t+1, BlockSize); + // very important: This loop condition needs to be the same for all threads. + // even if a thread later on is not going to do any work, it needs to participate in the + // data loading process! for (int t3f = 0; t3f < block_steps; ++t3f) { int t3 = t3f * BlockSize + block.thread_rank(); - float acc = 0.f; - float at3 = att_bth[t3]; - for (int t2b = 0; t2b <= t; t2b += BlockSize) { - int end = min(t + 1 - t2b, BlockSize); - block.sync(); - { - int t2i = block.thread_rank(); - int t2 = min(t, t2b + t2i); - att_bth_s[t2i] = att_bth[t2] * datt_bth[t2]; - } - block.sync(); - if(t3f * BlockSize == t2b) { - for (int t2i = 0; t2i < end; t2i++) { - int t2 = t2b + t2i; - float indicator = t2 == t3 ? 1.0f : 0.0f; - float local_derivative = att_bth_s[t2i] * (indicator - at3); - acc += local_derivative; - } - } else { - for (int t2i = 0; t2i < end; t2i++) { - int t2 = t2b + t2i; - float local_derivative = att_bth_s[t2i] * (0.f - at3); - acc += local_derivative; - } - } + float at3 = att_bth[min(t, t3)]; + float local_sum = 0; + for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + local_sum += att_bth[t2] * datt_bth[t2]; + } + block.sync(); + block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); + block.sync(); + local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{}); + + float acc = -local_sum * at3; + float at_t2_eq_t3 = at3 * datt_bth[min(t, t3)]; + acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3)); + if(t3 <= t) { + dpreatt_bth[t3] = scale * acc; } - dpreatt_bth[t3] = scale * acc; } } + // 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) { @@ -1148,7 +1146,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,