aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

coalesced reading in attention backward pass

Browse source
This commit is contained in:
Erik Schultheis 2024-04-18 00:48:54 +03:00
parent 4c27dd923b
commit e278112c19
2 changed files with 79 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

64
dev/cuda/attention_backward.cu
View file

@ -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<float>());
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<<<num_blocks, block_size>>>(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<<<dim3(num_blocks, B*NH), block_size>>>(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<<<dim3(num_blocks, B*NH), block_size>>>(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<<<dim3(num_blocks, B*NH), block_size>>>(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<class SoftmaxKernel>
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<<<dim3(num_blocks, B*NH), block_size>>>(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");

30
train_gpt2.cu
View file

@ -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<float>());
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<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
softmax_autoregressive_backward_kernel<<<dim3(CEIL_DIV(B * T * C, block_size/32), B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
// backward into q
cublasSgemmStridedBatched(cublas_handle,