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

separating the loops leads to yet more improvements

Browse source
This commit is contained in:
Erik Schultheis 2024-04-19 20:47:19 +03:00
parent 09a117a664
commit b2e6b6867f
2 changed files with 25 additions and 41 deletions
Download patch file Download diff file Expand all files Collapse all files

36
dev/cuda/attention_backward.cu
View file

@ -633,8 +633,8 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker
}
template<int BlockSize>
__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) {
__global__ void 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();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
@ -657,29 +657,21 @@ __global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_ker
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 local_sum = 0;
for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
local_sum += att_bth[t2] * datt_bth[t2];
}
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<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
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[min(t, t3)];
float at_t2_eq_t3 = at3 * datt_bth[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;
}
}
@ -930,7 +922,7 @@ int main(int argc, char **argv) {
srand(0); // reproducibility
// hyperparameters
int B = 8;
int B = 4;
int T = 1024;
int C = 768;
int NH = 12;

30
train_gpt2.cu
View file

@ -690,29 +690,21 @@ __global__ void softmax_autoregressive_backward_kernel(float* dpreatt, const flo
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 local_sum = 0;
for(int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
local_sum += att_bth[t2] * datt_bth[t2];
}
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<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
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[min(t, t3)];
float at_t2_eq_t3 = at3 * datt_bth[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;
}
}