From 44b5bb028cce36d5c2cab4fb727edf136ea5fe9f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 19 Apr 2024 22:33:05 +0300 Subject: [PATCH] one more version of attention backward for another few percent perf gains --- dev/cuda/attention_backward.cu | 46 ++++++++++++++++++------------ train_gpt2.cu | 52 ++++++++++++++++++---------------- 2 files changed, 56 insertions(+), 42 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index 2f52919..88ce4ae 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -671,43 +671,53 @@ __global__ void softmax_autoregressive_backward_kernel7(float* dpreatt, const fl } } - +// The slightly less pretty version of kernel 7. Adding in all the dirty tricks that can give us a few more percent +// - streaming memory access instructions +// - reordering blocks to prevent tail effect +// - multiple values of T per block template __global__ void softmax_autoregressive_backward_kernel8(float* dpreatt, const float* datt, const float* att, int B, int T, int C, float scale) { namespace cg = cooperative_groups; + constexpr int T_per_block = 4; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); __shared__ float block_acc[32]; int idx = blockIdx.y; // go through blocks in reverse order, so the slowest block starts first - int t = T - 1 - blockIdx.x; + int t0 = T - 1 - T_per_block*blockIdx.x; att += idx * T * T; datt += idx * T * T; dpreatt += idx * T * T; - const float* att_bth = att + t * T; - const float* datt_bth = datt + t * T; - float* dpreatt_bth = dpreatt + t * T; - - if(warp.meta_group_rank() == 0) { + if (warp.meta_group_rank() == 0) { block_acc[warp.thread_rank()] = 0; } - float local_sum = 0; - for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { - local_sum += att_bth[t2] * datt_bth[t2]; - } + for(int to = 0; to < T_per_block; ++to) { + int t = t0 - to; + if(t < 0) return; + const float* att_bth = att + t * T; + const float* datt_bth = datt + t * T; + float* dpreatt_bth = dpreatt + t * T; - 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 local_sum = 0; + for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + local_sum += att_bth[t2] * datt_bth[t2]; + } - for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { - float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum); - __stcs(dpreatt_bth + t3, scale * acc); + 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{}); + + for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { + // don't touch the cache. Some parts will still be here from the previous loop, and + // we want to exploit those. + float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum); + __stcs(dpreatt_bth + t3, scale * acc); + } } } @@ -839,7 +849,7 @@ void launch_softmax_8(float* dpreatt, float* datt, const float* att, int B, int float scale = 1.0f / sqrtf(hs); auto launch = [&](auto int_const) { constexpr int block_size = int_const.value; - softmax_autoregressive_backward_kernel8<<>> + softmax_autoregressive_backward_kernel8<<>> (dpreatt, datt, att, B, T, C, scale); }; dispatch_launch(launch, block_size); diff --git a/train_gpt2.cu b/train_gpt2.cu index 8f801bb..64d1dcc 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -612,46 +612,48 @@ __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) { + int B, int T, int C, float scale) { constexpr const int BlockSize = 256; + constexpr int T_per_block = 4; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); __shared__ float block_acc[32]; int idx = blockIdx.y; - int t = blockIdx.x; + // go through blocks in reverse order, so the slowest block starts first + int t0 = T - 1 - T_per_block*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; - - if(warp.meta_group_rank() == 0) { + if (warp.meta_group_rank() == 0) { block_acc[warp.thread_rank()] = 0; } - float local_sum = 0; - for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { - local_sum += att_bth[t2] * datt_bth[t2]; - } + for(int to = 0; to < T_per_block; ++to) { + int t = t0 - to; + if(t < 0) return; + const float* att_bth = att + t * T; + const float* datt_bth = datt + t * T; + float* dpreatt_bth = dpreatt + t * T; - 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 local_sum = 0; + for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + local_sum += att_bth[t2] * datt_bth[t2]; + } - for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { - float at3 = att_bth[t3]; - float acc = -local_sum * at3; - float at_t2_eq_t3 = at3 * datt_bth[t3]; - acc += (at_t2_eq_t3 * (1.f - at3) - at_t2_eq_t3 * (0.f - at3)); - dpreatt_bth[t3] = scale * acc; + 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{}); + + for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { + // don't touch the cache. Some parts will still be here from the previous loop, and + // we want to exploit those. + float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum); + __stcs(dpreatt_bth + t3, scale * acc); + } } } @@ -1018,7 +1020,9 @@ void attention_backward(float* dinp, float* dqkvr, float* dpreatt, float* datt, // backward into dv cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, dvaccum, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH)); // backward into preatt - softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, NH); + int hs = C / NH; // head size + float scale = 1.0f / sqrtf(hs); + softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); cudaCheck(cudaGetLastError()); // backward into q cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH));